Allow for user overridable boundary conditions

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[PR 1314]](https://github.com/parthenon-hpc-lab/parthenon/pull/1314) Add option of user specified BCs in AddBoundaryExchangeTasks
 - [[PR 1244]](https://github.com/parthenon-hpc-lab/parthenon/pull/1244) Add TaskCollection timeout capability
 - [[PR 1311]](https://github.com/parthenon-hpc-lab/parthenon/pull/1311) Magnitude refinement criteria, per-package PostInitialize function hooks, mergescv and prettyparams utilities
 - [[PR 1142]](https://github.com/parthenon-hpc-lab/parthenon/pull/1142) Unify par_dispatch, par_for_outer & par_for_inner overloads

doc/sphinx/src/boundary_conditions.rst

@@ -145,3 +145,17 @@ for a more complete example):
     pkg->UserBoundaryFunctions[BF::inner_x2].push_back(GetMyBC<X2DIR, BCSide::Inner>());
     ...
   }
+
+User override of boundary conditions in `AddBoundaryExchangeTasks`.
+-------------------------------------------------------------------
+
+Sometimes it is desirable to apply different boundary conditions when communicating 
+on different containers (e.g. when one container represents the value of a variable
+and another container represents a Newton-Raphson correction to that variable). To
+easily allow this, `AddBoundaryExchangeTasks` can take a final argument with type 
+`std::function<TaskStatus(std::shared_ptr<MeshData<Real>> &, bool)>` or with type 
+`std::function<TaskID(TaskID, TaskList *, std::shared_ptr<MeshData<Real>>, bool)>`.
+When this argument is specified, all other boundary conditions will be ignored and
+the passed functions will be used as boundary conditions instead. The final boolean
+arguments of the functions specify if the boundary condition needs to be applied on
+the coarse buffer.

src/bvals/comms/boundary_communication.cpp

@@ -168,6 +168,8 @@ template TaskStatus
 SendBoundBufs<BoundaryType::gmg_prolongate_send>(std::shared_ptr<MeshData<Real>> &);
 template TaskStatus
 SendBoundBufs<BoundaryType::flxcor_send>(std::shared_ptr<MeshData<Real>> &);
+template TaskStatus
+SendBoundBufs<BoundaryType::gmg_same>(std::shared_ptr<MeshData<Real>> &);
 
 template <BoundaryType bound_type>
 TaskStatus StartReceiveBoundBufs(std::shared_ptr<MeshData<Real>> &md) {
@@ -196,6 +198,8 @@ template TaskStatus StartReceiveBoundBufs<BoundaryType::gmg_prolongate_recv>(
     std::shared_ptr<MeshData<Real>> &);
 template TaskStatus
 StartReceiveBoundBufs<BoundaryType::flxcor_recv>(std::shared_ptr<MeshData<Real>> &);
+template TaskStatus
+StartReceiveBoundBufs<BoundaryType::gmg_same>(std::shared_ptr<MeshData<Real>> &);
 
 template <BoundaryType bound_type>
 TaskStatus ReceiveBoundBufs(std::shared_ptr<MeshData<Real>> &md) {
@@ -246,6 +250,8 @@ template TaskStatus
 ReceiveBoundBufs<BoundaryType::gmg_prolongate_recv>(std::shared_ptr<MeshData<Real>> &);
 template TaskStatus
 ReceiveBoundBufs<BoundaryType::flxcor_recv>(std::shared_ptr<MeshData<Real>> &);
+template TaskStatus
+ReceiveBoundBufs<BoundaryType::gmg_same>(std::shared_ptr<MeshData<Real>> &);
 
 template <BoundaryType bound_type>
 TaskStatus SetBounds(std::shared_ptr<MeshData<Real>> &md) {
@@ -356,6 +362,7 @@ template TaskStatus
 SetBounds<BoundaryType::gmg_prolongate_recv>(std::shared_ptr<MeshData<Real>> &);
 template TaskStatus
 SetBounds<BoundaryType::flxcor_recv>(std::shared_ptr<MeshData<Real>> &);
+template TaskStatus SetBounds<BoundaryType::gmg_same>(std::shared_ptr<MeshData<Real>> &);
 
 template <BoundaryType bound_type>
 TaskStatus ProlongateBounds(std::shared_ptr<MeshData<Real>> &md) {
@@ -400,56 +407,55 @@ template TaskStatus
 ProlongateBounds<BoundaryType::nonlocal>(std::shared_ptr<MeshData<Real>> &);
 template TaskStatus
 ProlongateBounds<BoundaryType::gmg_prolongate_recv>(std::shared_ptr<MeshData<Real>> &);
+template TaskStatus
+ProlongateBounds<BoundaryType::gmg_same>(std::shared_ptr<MeshData<Real>> &);
 
-// Adds all relevant boundary communication to a single task list
-template <BoundaryType bounds>
-TaskID AddBoundaryExchangeTasks(TaskID dependency, TaskList &tl,
-                                std::shared_ptr<MeshData<Real>> &md, bool multilevel) {
-  // TODO(LFR): Splitting up the boundary tasks while doing prolongation can cause some
-  //            possible issues for sparse fields. In particular, the order in which
-  //            fields are allocated and then set could potentially result in different
-  //            results if the default sparse value is non-zero.
-  // const auto any = BoundaryType::any;
-  static_assert(bounds == BoundaryType::any || bounds == BoundaryType::gmg_same);
-  // const auto local = BoundaryType::local;
-  // const auto nonlocal = BoundaryType::nonlocal;
-
-  // auto send = tl.AddTask(dependency, SendBoundBufs<nonlocal>, md);
-  // auto send_local = tl.AddTask(dependency, SendBoundBufs<local>, md);
-
-  // auto recv_local = tl.AddTask(dependency, ReceiveBoundBufs<local>, md);
-  // auto set_local = tl.AddTask(recv_local, SetBounds<local>, md);
-
-  // auto recv = tl.AddTask(dependency, ReceiveBoundBufs<nonlocal>, md);
-  // auto set = tl.AddTask(recv, SetBounds<nonlocal>, md);
+template <BoundaryType bound_type>
+TaskStatus ProlongateInternalBounds(std::shared_ptr<MeshData<Real>> &md) {
+  PARTHENON_INSTRUMENT
 
-  // auto cbound = tl.AddTask(set, ApplyCoarseBoundaryConditions, md);
+  Mesh *pmesh = md->GetMeshPointer();
+  auto &cache = md->GetBvarsCache().GetSubCache(bound_type, false);
 
-  // auto pro_local = tl.AddTask(cbound | set_local | set, ProlongateBounds<local>, md);
-  // auto pro = tl.AddTask(cbound | set_local | set, ProlongateBounds<nonlocal>, md);
+  auto [rebuild, nbound] = CheckReceiveBufferCacheForRebuild<bound_type, false>(md);
 
-  // auto out = (pro_local | pro);
+  if (rebuild) {
+    if constexpr (bound_type == BoundaryType::gmg_prolongate_recv) {
+      RebuildBufferCache<bound_type, false>(md, nbound, BndInfo::GetSetBndInfo,
+                                            ProResInfo::GetInteriorProlongate);
+    } else if constexpr (bound_type == BoundaryType::gmg_restrict_recv) {
+      RebuildBufferCache<bound_type, false>(md, nbound, BndInfo::GetSetBndInfo,
+                                            ProResInfo::GetNull);
+    } else {
+      RebuildBufferCache<bound_type, false>(md, nbound, BndInfo::GetSetBndInfo,
+                                            ProResInfo::GetSet);
+    }
+  }
 
-  auto send = tl.AddTask(dependency, TF(SendBoundBufs<bounds>), md);
-  auto recv = tl.AddTask(dependency, TF(ReceiveBoundBufs<bounds>), md);
-  auto set = tl.AddTask(recv, TF(SetBounds<bounds>), md);
+  if (nbound > 0 && pmesh->multilevel && md->NumBlocks() > 0) {
+    auto pmb = md->GetBlockData(0)->GetBlockPointer();
+    StateDescriptor *resolved_packages = pmb->resolved_packages.get();
 
-  auto pro = set;
-  if (md->GetMeshPointer()->multilevel) {
-    auto cbound = tl.AddTask(set, TF(ApplyBoundaryConditionsOnCoarseOrFineMD), md, true);
-    pro = tl.AddTask(cbound, TF(ProlongateBounds<bounds>), md);
+    // Prolongate from coarse buffer
+    refinement::ProlongateInternal(resolved_packages, cache.prores_cache, pmb->cellbounds,
+                                   pmb->c_cellbounds);
   }
-  auto fbound = tl.AddTask(pro, TF(ApplyBoundaryConditionsOnCoarseOrFineMD), md, false);
-
-  return fbound;
+  return TaskStatus::complete;
 }
-template TaskID
-AddBoundaryExchangeTasks<BoundaryType::any>(TaskID, TaskList &,
-                                            std::shared_ptr<MeshData<Real>> &, bool);
+template TaskStatus
+ProlongateInternalBounds<BoundaryType::any>(std::shared_ptr<MeshData<Real>> &);
+template TaskStatus
+ProlongateInternalBounds<BoundaryType::local>(std::shared_ptr<MeshData<Real>> &);
+template TaskStatus
+ProlongateInternalBounds<BoundaryType::nonlocal>(std::shared_ptr<MeshData<Real>> &);
+template TaskStatus ProlongateInternalBounds<BoundaryType::gmg_prolongate_recv>(
+    std::shared_ptr<MeshData<Real>> &);
+template TaskStatus
+ProlongateInternalBounds<BoundaryType::gmg_same>(std::shared_ptr<MeshData<Real>> &);
 
-template TaskID
-AddBoundaryExchangeTasks<BoundaryType::gmg_same>(TaskID, TaskList &,
-                                                 std::shared_ptr<MeshData<Real>> &, bool);
+bool IsMeshMultilevel(std::shared_ptr<MeshData<Real>> &md) {
+  return md->GetMeshPointer()->multilevel;
+}
 
 TaskID AddFluxCorrectionTasks(TaskID dependency, TaskList &tl,
                               std::shared_ptr<MeshData<Real>> &md, bool multilevel) {

src/bvals/comms/bvals_in_one.hpp

@@ -23,6 +23,7 @@
 #include <vector>
 
 #include "basic_types.hpp"
+#include "bvals/boundary_conditions.hpp"
 #include "bvals/neighbor_block.hpp"
 #include "coordinates/coordinates.hpp"
 
@@ -62,6 +63,8 @@ inline TaskStatus SetBoundaries(std::shared_ptr<MeshData<Real>> &md) {
 
 template <BoundaryType bound_type>
 TaskStatus ProlongateBounds(std::shared_ptr<MeshData<Real>> &md);
+template <BoundaryType bound_type>
+TaskStatus ProlongateInternalBounds(std::shared_ptr<MeshData<Real>> &md);
 inline TaskStatus ProlongateBoundaries(std::shared_ptr<MeshData<Real>> &md) {
   return ProlongateBounds<BoundaryType::any>(md);
 }
@@ -79,14 +82,52 @@ static TaskStatus SetFluxCorrections(std::shared_ptr<MeshData<Real>> &md) {
   return SetBounds<BoundaryType::flxcor_recv>(md);
 }
 
+// Adds all relevant flux correction tasks to a single task list
+TaskID AddFluxCorrectionTasks(TaskID dependency, TaskList &tl,
+                              std::shared_ptr<MeshData<Real>> &md, bool multilevel);
+
+using BValOnMDFunc_t = std::function<TaskStatus(std::shared_ptr<MeshData<Real>> &, bool)>;
+using BValOnMDTasks_t =
+    std::function<TaskID(TaskID, TaskList *, std::shared_ptr<MeshData<Real>>, bool)>;
+bool IsMeshMultilevel(std::shared_ptr<MeshData<Real>> &md);
+
 // Adds all relevant boundary communication to a single task list
 template <BoundaryType bounds = BoundaryType::any>
 TaskID AddBoundaryExchangeTasks(TaskID dependency, TaskList &tl,
-                                std::shared_ptr<MeshData<Real>> &md, bool multilevel);
+                                std::shared_ptr<MeshData<Real>> &md, bool multilevel,
+                                BValOnMDTasks_t ApplyBCs) {
+  static_assert(bounds == BoundaryType::any || bounds == BoundaryType::gmg_same);
 
-// Adds all relevant flux correction tasks to a single task list
-TaskID AddFluxCorrectionTasks(TaskID dependency, TaskList &tl,
-                              std::shared_ptr<MeshData<Real>> &md, bool multilevel);
+  auto send = tl.AddTask(dependency, TF(SendBoundBufs<bounds>), md);
+  auto recv = tl.AddTask(dependency, TF(ReceiveBoundBufs<bounds>), md);
+  auto set = tl.AddTask(recv, TF(SetBounds<bounds>), md);
+
+  auto pro = set;
+  if (IsMeshMultilevel(md)) {
+    auto cbound = ApplyBCs(set, &tl, md, true);
+    pro = tl.AddTask(cbound, TF(ProlongateBounds<bounds>), md);
+  }
+  auto fbound = ApplyBCs(pro, &tl, md, false);
+  if (IsMeshMultilevel(md)) {
+    // Need to prolongate internal bounds after setting physical boundary
+    // conditions, since the internal prolongation uses the values of the
+    // normal buffer on shared elements (rather than values in the coarse)
+    // buffer for prolongation.
+    fbound = tl.AddTask(fbound, TF(ProlongateInternalBounds<bounds>), md);
+  }
+  return fbound;
+}
+
+template <BoundaryType bounds = BoundaryType::any>
+TaskID AddBoundaryExchangeTasks(
+    TaskID dependency, TaskList &tl, std::shared_ptr<MeshData<Real>> &md, bool multilevel,
+    BValOnMDFunc_t ApplyBCs = ApplyBoundaryConditionsOnCoarseOrFineMD) {
+  return AddBoundaryExchangeTasks<bounds>(
+      dependency, tl, md, multilevel,
+      [&](TaskID id, TaskList *tl, std::shared_ptr<MeshData<Real>> md, bool coarse) {
+        return tl->AddTask(id, TF(ApplyBCs), md, coarse);
+      });
+}
 
 // These tasks should not be called in down stream code
 TaskStatus BuildBoundaryBuffers(std::shared_ptr<MeshData<Real>> &md);