Merge branch 'develop' into mg/tensor_indices

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Added (new features/APIs/variables/...)
 - [[PR 998]](https://github.com/parthenon-hpc-lab/parthenon/pull/998) tensor indices added to sparse pack
+- [[PR 999]](https://github.com/parthenon-hpc-lab/parthenon/pull/999) Add a post-initialization hook
 - [[PR 987]](https://github.com/parthenon-hpc-lab/parthenon/pull/987) New tasking infrastructure and capabilities
 - [[PR 969]](https://github.com/parthenon-hpc-lab/parthenon/pull/969) New macro-based auto-naming of profiling regions and kernels
 - [[PR 981]](https://github.com/parthenon-hpc-lab/parthenon/pull/981) Add IndexSplit
@@ -20,6 +21,7 @@
 - [[PR978]](https://github.com/parthenon-hpc-lab/parthenon/pull/978) remove erroneous sparse check
 
 ### Infrastructure (changes irrelevant to downstream codes)
+- [[PR 990]](https://github.com/parthenon-hpc-lab/parthenon/pull/990) Partial refactor of HDF5 I/O code for readability/extendability
 - [[PR 982]](https://github.com/parthenon-hpc-lab/parthenon/pull/982) add some gut check testing for parthenon-VIBE
 
 ### Removed (removing behavior/API/varaibles/...)

doc/sphinx/src/README.rst

@@ -178,6 +178,8 @@ Mesh
 ^^^^
 
 -  ``InitUserMeshData``
+-  ``ProblemGenerator``
+-  ``PostInitialization``
 -  ``PreStepUserWorkInLoop``
 -  ``PostStepUserWorkInLoop``
 -  ``UserWorkAfterLoop``
@@ -188,19 +190,20 @@ MeshBlock
 -  ``InitApplicationMeshBlockData``
 -  ``InitMeshBlockUserData``
 -  ``ProblemGenerator``
+-  ``PostInitialization``
 -  ``UserWorkBeforeOutput``
 
 To redefine these functions, the user sets the respective function
 pointers in the ApplicationInput member app_input of the
 ParthenonManager class prior to calling ``ParthenonInit``. This is
 demonstrated in the ``main()`` functions in the examples.
 
-Note that the ``ProblemGenerator``\ s of ``Mesh`` and ``MeshBlock`` are
-mutually exclusive. Moreover, the ``Mesh`` one requires
-``parthenon/mesh/pack_size=-1`` during initialization, i.e., all blocks
-on a rank need to be in a single pack. This allows to use MPI reductions
-inside the function, for example, to globally normalize quantities. The
-``parthenon/mesh/pack_size=-1`` exists only during problem
+Note that the ``ProblemGenerator``\ s (and ``PostInitialization``\ s) of
+``Mesh`` and ``MeshBlock`` are mutually exclusive. Moreover, the ``Mesh``
+ones requires ``parthenon/mesh/pack_size=-1`` during initialization, i.e.,
+all blocks on a rank need to be in a single pack. This allows to use MPI
+reductions inside the function, for example, to globally normalize quantities.
+The ``parthenon/mesh/pack_size=-1`` exists only during problem
 inititalization, i.e., simulations can be restarted with an arbitrary
 ``pack_size``. For an example of the ``Mesh`` version, see the `Poisson
 example <https://github.com/parthenon-hpc-lab/parthenon/blob/develop/example/poisson/parthenon_app_inputs.cpp>`__.

doc/sphinx/src/interface/state.rst

@@ -113,12 +113,19 @@ several useful features and functions.
   if set (defaults to ``nullptr`` an therefore a no-op) to print
   diagnostics after the time-integration advance
 - ``void UserWorkBeforeLoopMesh(Mesh *, ParameterInput *pin, SimTime
-  &tm)`` performs a per-package, mesh-wide calculation after the mesh
-  has been generated, and problem generators called, but before any
-  time evolution. This work is done both on first initialization and
-  on restart. If you would like to avoid doing the work upon restart,
-  you can check for the const ``is_restart`` member field of the ``Mesh``
-  object.
+  &tm)`` performs a per-package, mesh-wide calculation after (1) the mesh
+  has been generated, (2) problem generators are called, and (3) comms
+  are executed, but before any time evolution. This work is done both on
+  first initialization and on restart. If you would like to avoid doing the
+  work upon restart, you can check for the const ``is_restart`` member
+  field of the ``Mesh`` object.  It is worth making a clear distinction
+  between ``UserWorkBeforeLoopMesh`` and ``ApplicationInput``s
+  ``PostInitialization``.  ``PostInitialization`` is very much so tied to
+  initialization, and will not be called upon restarts.  ``PostInitialization``
+  is also carefully positioned after ``ProblemGenerator`` and before
+  ``PreCommFillDerived`` (and hence communications).  In practice, when
+  additional granularity is required inbetween initialization and communication,
+  ``PostInitialization`` may be the desired hook.
 
 The reasoning for providing ``FillDerived*`` and ``EstimateTimestep*``
 function pointers appropriate for usage with both ``MeshData`` and

doc/sphinx/src/outputs.rst

@@ -46,6 +46,7 @@ look like
 
    <parthenon/output1>
    file_type = hdf5
+   write_xdmf = true # Determines whether xdmf annotations are output
    # nonexistent variables/swarms are ignored
    variables = density, velocity, & # comments are still ok
                energy               # notice the & continuation character

doc/sphinx/src/parthenon_manager.rst

@@ -31,12 +31,16 @@ runtimes. The function
 
 Calls the ``Initialize(ParameterInput *pin)`` function of all packages
 to be utilized and creates the grid hierarchy, including the ``Mesh``
-and ``MeshBlock`` objects, and calls the ``ProblemGenerator``
-initialization routines.
+and ``MeshBlock`` objects, and calls the ``ProblemGenerator`` (and
+``PostInitialization``) routines.
 
 The reason these functions are split out is to enable decisions to be
 made by the application between reading the input deck and setting up
-the grid. For example, a common use-case is:
+the grid. For example, during problem initialization, ``ProblemGenerator``
+may be used to be the user-facing API to describe initial conditions,
+whereas, ``PostInitialization`` could use those user-specified fields
+to sync *all* fields prior to entering communication routines. A common
+use-case is:
 
 .. code:: cpp
 
@@ -53,13 +57,14 @@ the grid. For example, a common use-case is:
   if (manager_status == ParthenonStatus::error) {
     pman.ParthenonFinalize();
     return 1;
-  } 
+  }
 
   // Redefine parthenon defaults
   pman.app_input->ProcessPackages = MyProcessPackages;
   std::string prob = pman.pin->GetString("app", "problem");
   if (prob == "problem1") {
     pman.app_input->ProblemGenerator = Problem1Generator;
+    pman.app_input->PostInitialization = Problem1PostInitialization;
   } else {
     pman.app_input->ProblemGenerator = Problem2Generator;
   }

src/application_input.hpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2020-2022. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
@@ -36,6 +36,8 @@ struct ApplicationInput {
   std::function<void(Mesh *, ParameterInput *)> InitUserMeshData = nullptr;
   std::function<void(Mesh *, ParameterInput *, MeshData<Real> *)> MeshProblemGenerator =
       nullptr;
+  std::function<void(Mesh *, ParameterInput *, MeshData<Real> *)> MeshPostInitialization =
+      nullptr;
 
   std::function<void(Mesh *, ParameterInput *, SimTime &)> PreStepMeshUserWorkInLoop =
       nullptr;
@@ -57,6 +59,7 @@ struct ApplicationInput {
       InitApplicationMeshBlockData = nullptr;
   std::function<void(MeshBlock *, ParameterInput *)> InitMeshBlockUserData = nullptr;
   std::function<void(MeshBlock *, ParameterInput *)> ProblemGenerator = nullptr;
+  std::function<void(MeshBlock *, ParameterInput *)> PostInitialization = nullptr;
   std::function<void(MeshBlock *, ParameterInput *)> MeshBlockUserWorkBeforeOutput =
       nullptr;
 };

src/basic_types.hpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2021-2023. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2021-2024. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
@@ -49,7 +49,7 @@ enum class TaskStatus { complete, incomplete, iterate };
 
 enum class AmrTag : int { derefine = -1, same = 0, refine = 1 };
 enum class RefinementOp_t { Prolongation, Restriction, None };
-
+enum class CellLevel : int { coarse = -1, same = 0, fine = 1 };
 // JMM: Not clear this is the best place for this but it minimizes
 // circular dependency nonsense.
 constexpr int NUM_BNDRY_TYPES = 10;
@@ -152,16 +152,17 @@ inline std::vector<TopologicalElement> GetTopologicalElements(TopologicalType tt
   if (tt == TT::Face) return {TE::F1, TE::F2, TE::F3};
   return {TE::CC};
 }
+
 using TE = TopologicalElement;
 // Returns one if the I coordinate of el is offset from the zone center coordinates,
 // and zero otherwise
-KOKKOS_INLINE_FUNCTION int TopologicalOffsetI(TE el) noexcept {
+inline constexpr int TopologicalOffsetI(TE el) {
   return (el == TE::F1 || el == TE::E2 || el == TE::E3 || el == TE::NN);
 }
-KOKKOS_INLINE_FUNCTION int TopologicalOffsetJ(TE el) noexcept {
+inline constexpr int TopologicalOffsetJ(TE el) {
   return (el == TE::F2 || el == TE::E3 || el == TE::E1 || el == TE::NN);
 }
-KOKKOS_INLINE_FUNCTION int TopologicalOffsetK(TE el) noexcept {
+inline constexpr int TopologicalOffsetK(TE el) {
   return (el == TE::F3 || el == TE::E2 || el == TE::E1 || el == TE::NN);
 }
 

src/coordinates/uniform_cartesian.hpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2023. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2023-2024. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
@@ -164,18 +164,11 @@ class UniformCartesian {
 
   template <int dir, TopologicalElement el>
   KOKKOS_FORCEINLINE_FUNCTION Real X(const int idx) const {
-    using TE = TopologicalElement;
-    [[maybe_unused]] bool constexpr X1EDGE =
-        el == TE::F1 || el == TE::E2 || el == TE::E3 || el == TE::NN;
-    [[maybe_unused]] bool constexpr X2EDGE =
-        el == TE::F2 || el == TE::E3 || el == TE::E1 || el == TE::NN;
-    [[maybe_unused]] bool constexpr X3EDGE =
-        el == TE::F3 || el == TE::E1 || el == TE::E2 || el == TE::NN;
-    if constexpr (dir == X1DIR && X1EDGE) {
+    if constexpr (dir == X1DIR && TopologicalOffsetI(el)) {
       return xmin_[dir - 1] + idx * dx_[dir - 1]; // idx - 1/2
-    } else if constexpr (dir == X2DIR && X2EDGE) {
+    } else if constexpr (dir == X2DIR && TopologicalOffsetJ(el)) {
       return xmin_[dir - 1] + idx * dx_[dir - 1]; // idx - 1/2
-    } else if constexpr (dir == X3DIR && X3EDGE) {
+    } else if constexpr (dir == X3DIR && TopologicalOffsetK(el)) {
       return xmin_[dir - 1] + idx * dx_[dir - 1]; // idx - 1/2
     } else {
       return xmin_[dir - 1] + (idx + 0.5) * dx_[dir - 1]; // idx

src/mesh/domain.hpp

@@ -3,7 +3,7 @@
 // Copyright(C) 2014 James M. Stone <[email protected]> and other code contributors
 // Licensed under the 3-clause BSD License, see LICENSE file for details
 //========================================================================================
-// (C) (or copyright) 2020. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
@@ -23,6 +23,7 @@
 #include <type_traits>
 #include <vector>
 
+#include "basic_types.hpp"
 #include "defs.hpp"
 
 namespace parthenon {

src/mesh/mesh.cpp

@@ -3,7 +3,7 @@
 // Copyright(C) 2014 James M. Stone <[email protected]> and other code contributors
 // Licensed under the 3-clause BSD License, see LICENSE file for details
 //========================================================================================
-// (C) (or copyright) 2020-2023. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
@@ -156,6 +156,9 @@ Mesh::Mesh(ParameterInput *pin, ApplicationInput *app_in, Packages_t &packages,
   if (app_in->MeshProblemGenerator != nullptr) {
     ProblemGenerator = app_in->MeshProblemGenerator;
   }
+  if (app_in->MeshPostInitialization != nullptr) {
+    PostInitialization = app_in->MeshPostInitialization;
+  }
   if (app_in->PreStepMeshUserWorkInLoop != nullptr) {
     PreStepUserWorkInLoop = app_in->PreStepMeshUserWorkInLoop;
   }
@@ -930,6 +933,10 @@ void Mesh::Initialize(bool init_problem, ParameterInput *pin, ApplicationInput *
       PARTHENON_REQUIRE_THROWS(
           !(ProblemGenerator != nullptr && block_list[0]->ProblemGenerator != nullptr),
           "Mesh and MeshBlock ProblemGenerators are defined. Please use only one.");
+      PARTHENON_REQUIRE_THROWS(
+          !(PostInitialization != nullptr &&
+            block_list[0]->PostInitialization != nullptr),
+          "Mesh and MeshBlock PostInitializations are defined. Please use only one.");
 
       // Call Mesh ProblemGenerator
       if (ProblemGenerator != nullptr) {
@@ -946,6 +953,23 @@ void Mesh::Initialize(bool init_problem, ParameterInput *pin, ApplicationInput *
           pmb->ProblemGenerator(pmb.get(), pin);
         }
       }
+
+      // Call Mesh PostInitialization
+      if (PostInitialization != nullptr) {
+        PARTHENON_REQUIRE(num_partitions == 1,
+                          "Mesh PostInitialization requires parthenon/mesh/pack_size=-1 "
+                          "during first initialization.");
+
+        auto &md = mesh_data.GetOrAdd("base", 0);
+        PostInitialization(this, pin, md.get());
+        // Call individual MeshBlock PostInitialization
+      } else {
+        for (int i = 0; i < nmb; ++i) {
+          auto &pmb = block_list[i];
+          pmb->PostInitialization(pmb.get(), pin);
+        }
+      }
+
       std::for_each(block_list.begin(), block_list.end(),
                     [](auto &sp_block) { sp_block->SetAllVariablesToInitialized(); });
     }
@@ -1202,6 +1226,22 @@ int Mesh::GetNumberOfMeshBlockCells() const {
 }
 const RegionSize &Mesh::GetBlockSize() const { return base_block_size; }
 
+const IndexShape &Mesh::GetLeafBlockCellBounds(CellLevel level) const {
+  // TODO(JMM): Luke this is for your Metadata::fine stuff.
+  PARTHENON_DEBUG_REQUIRE(level != CellLevel::fine,
+                          "Currently no access to finer cellbounds");
+  MeshBlock *pmb = block_list[0].get();
+  if (level == CellLevel::same) {
+    return pmb->cellbounds;
+    // TODO(JMM):
+    // } else if (level == CellLevel::fine) {
+    //   return pmb->fine_cellbounds;
+    // }
+  } else { // if (level == CellLevel::coarse) {
+    return pmb->c_cellbounds;
+  }
+}
+
 // Functionality re-used in mesh constructor
 void Mesh::RegisterLoadBalancing_(ParameterInput *pin) {
 #ifdef MPI_PARALLEL // JMM: Not sure this ifdef is needed

src/mesh/mesh.hpp

@@ -3,7 +3,7 @@
 // Copyright(C) 2014 James M. Stone <[email protected]> and other code contributors
 // Licensed under the 3-clause BSD License, see LICENSE file for details
 //========================================================================================
-// (C) (or copyright) 2020-2023. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
@@ -98,6 +98,7 @@ class Mesh {
   int GetNumberOfMeshBlockCells() const;
   const RegionSize &GetBlockSize() const;
   RegionSize GetBlockSize(const LogicalLocation &loc) const;
+  const IndexShape &GetLeafBlockCellBounds(CellLevel level = CellLevel::same) const;
 
   // data
   bool modified;
@@ -160,6 +161,8 @@ class Mesh {
   // defined in either the prob file or default_pgen.cpp in ../pgen/
   std::function<void(Mesh *, ParameterInput *, MeshData<Real> *)> ProblemGenerator =
       nullptr;
+  std::function<void(Mesh *, ParameterInput *, MeshData<Real> *)> PostInitialization =
+      nullptr;
   static void UserWorkAfterLoopDefault(Mesh *mesh, ParameterInput *pin,
                                        SimTime &tm); // called in main loop
   std::function<void(Mesh *, ParameterInput *, SimTime &)> UserWorkAfterLoop =
@@ -194,6 +197,20 @@ class Mesh {
   std::vector<int> GetNbList() const noexcept { return nblist; }
   std::vector<LogicalLocation> GetLocList() const noexcept { return loclist; }
 
+  // TODO(JMM): Put in implementation file?
+  auto GetLevelsAndLogicalLocationsFlat() const noexcept {
+    std::vector<std::int64_t> levels, logicalLocations;
+    levels.reserve(nbtotal);
+    logicalLocations.reserve(nbtotal * 3);
+    for (const auto &loc : loclist) {
+      levels.push_back(loc.level() - GetRootLevel());
+      logicalLocations.push_back(loc.lx1());
+      logicalLocations.push_back(loc.lx2());
+      logicalLocations.push_back(loc.lx3());
+    }
+    return std::make_pair(levels, logicalLocations);
+  }
+
   void OutputMeshStructure(const int dim, const bool dump_mesh_structure = true);
 
   // Ordering here is important to prevent deallocation of pools before boundary

src/mesh/meshblock.cpp

@@ -3,7 +3,7 @@
 // Copyright(C) 2014 James M. Stone <[email protected]> and other code contributors
 // Licensed under the 3-clause BSD License, see LICENSE file for details
 //========================================================================================
-// (C) (or copyright) 2020-2023. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
@@ -115,6 +115,12 @@ void MeshBlock::Initialize(int igid, int ilid, LogicalLocation iloc,
   } else if (app_in->MeshProblemGenerator == nullptr) {
     ProblemGenerator = &ProblemGeneratorDefault;
   }
+  if (app_in->PostInitialization != nullptr) {
+    PostInitialization = app_in->PostInitialization;
+    // Only set default post-init when no mesh post-init is set
+  } else if (app_in->MeshPostInitialization == nullptr) {
+    PostInitialization = &PostInitializationDefault;
+  }
   if (app_in->MeshBlockUserWorkBeforeOutput != nullptr) {
     UserWorkBeforeOutput = app_in->MeshBlockUserWorkBeforeOutput;
   }

src/mesh/meshblock.hpp

@@ -3,7 +3,7 @@
 // Copyright(C) 2014 James M. Stone <[email protected]> and other code contributors
 // Licensed under the 3-clause BSD License, see LICENSE file for details
 //========================================================================================
-// (C) (or copyright) 2020-2023. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved.
 //
 // This program was produced under U.S. Government contract 89233218CNA000001 for Los
 // Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
@@ -440,7 +440,9 @@ class MeshBlock : public std::enable_shared_from_this<MeshBlock> {
 
   // defined in either the prob file or default_pgen.cpp in ../pgen/
   static void ProblemGeneratorDefault(MeshBlock *pmb, ParameterInput *pin);
+  static void PostInitializationDefault(MeshBlock *pmb, ParameterInput *pin);
   std::function<void(MeshBlock *, ParameterInput *)> ProblemGenerator = nullptr;
+  std::function<void(MeshBlock *, ParameterInput *)> PostInitialization = nullptr;
   static pMeshBlockApplicationData_t
   InitApplicationMeshBlockDataDefault(MeshBlock *, ParameterInput *pin);
   std::function<pMeshBlockApplicationData_t(MeshBlock *, ParameterInput *)>