Merge pull request #999 from parthenon-hpc-lab/pdmullen/post-init

Add a post-initialization hook

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[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

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/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/mesh/mesh.cpp

@@ -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(); });
     }

src/mesh/mesh.hpp

@@ -161,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 =

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 *)>

src/pgen/default_pgen.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-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
@@ -110,7 +110,7 @@ void MeshBlock::InitMeshBlockUserDataDefault(MeshBlock *pmb, ParameterInput *pin
 }
 
 //========================================================================================
-//! \fn void MeshBlock::ProblemGeneratorDefault(ParameterInput *pin)
+//! \fn void MeshBlock::ProblemGeneratorDefault(MeshBlock *pmb, ParameterInput *pin)
 //  \brief Should be used to set initial conditions.
 //========================================================================================
 
@@ -120,6 +120,13 @@ void MeshBlock::ProblemGeneratorDefault(MeshBlock *pmb, ParameterInput *pin) {
   return;
 }
 
+//========================================================================================
+//! \fn void MeshBlock::PostInitializationDefault(MeshBlock *pmb, ParameterInput *pin)
+//  \brief Should be used to perform post initialization ops.
+//========================================================================================
+
+void MeshBlock::PostInitializationDefault(MeshBlock *pmb, ParameterInput *pin) { return; }
+
 //========================================================================================
 //! \fn void MeshBlock::UserWorkBeforeOutputDefault(MeshBlock *pmb, ParameterInput *pin)
 //  \brief Function called before generating output files