Merge pull request #655 from lanl/brryan/extend_init_user_mesh_data

Enable user boundary conditions for particles

CHANGELOG.md

@@ -11,6 +11,7 @@
 - [[PR 586]](https://github.com/lanl/parthenon/pull/586) Implement true sparse capability with automatic allocation and deallocation of sparse
 
 ### Changed (changing behavior/API/variables/...)
+- [[PR 655]](https://github.com/lanl/parthenon/pull/655) Enable user boundary conditions for particles
 - [[PR 623]](https://github.com/lanl/parthenon/pull/623) Enable Params to optionally return non-const pointers
 - [[PR 604]](https://github.com/lanl/parthenon/pull/604) Allow modification of SimTime in PreStepUserWorkInLoop
 - [[PR 617]](https://github.com/lanl/parthenon/pull/617) Unify the coordinates API for MeshBlockPack and VariablePack

docs/particles.md

@@ -111,5 +111,6 @@ classes containing boundary condition functions for updating particles or removi
 when they are in boundary regions are allocated depending on the boundary flags specified
 in the input file. Currently, outflow and periodic boundaries are supported natively.
 User-specified boundary conditions must be set by specifying the "user" flag in the input
-parameter file and then updating the appropriate Swarm::bounds array entries to separately
-allocated boundary condition objects.
+parameter file and then updating the appropriate Swarm::bounds array entries to factory
+functions that allocate device-side boundary condition objects. An example is given in the
+`particles` example when ix1 and ox1 are set to `user` in the input parameter file.

example/particles/main.cpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2020. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2022. 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
@@ -20,12 +20,8 @@ int main(int argc, char *argv[]) {
   using parthenon::ParthenonStatus;
   ParthenonManager pman;
 
-  // Redefine parthenon defaults
-  pman.app_input->ProcessPackages = particles_example::ProcessPackages;
-  pman.app_input->ProblemGenerator = particles_example::ProblemGenerator;
-
   // call ParthenonInit to initialize MPI and Kokkos, parse the input deck, and set up
-  auto manager_status = pman.ParthenonInit(argc, argv);
+  auto manager_status = pman.ParthenonInitEnv(argc, argv);
   if (manager_status == ParthenonStatus::complete) {
     pman.ParthenonFinalize();
     return 0;
@@ -34,8 +30,31 @@ int main(int argc, char *argv[]) {
     pman.ParthenonFinalize();
     return 1;
   }
-  // Now that ParthenonInit has been called and setup succeeded, the code can now
-  // make use of MPI and Kokkos
+
+  // Redefine parthenon defaults
+  pman.app_input->ProcessPackages = particles_example::ProcessPackages;
+  pman.app_input->ProblemGenerator = particles_example::ProblemGenerator;
+  if (pman.pinput->GetString("parthenon/mesh", "ix1_bc") == "user") {
+    // In this case, we are setting a custom swarm boundary condition while still using
+    // a default parthenon boundary condition for cell variables. In general, one can
+    // provide both custom cell variable and swarm boundary conditions. However, to use
+    // custom boundary conditions for either cell variables or swarms, the parthenon
+    // boundary must be set to "user" and both cell variable and swarm boundaries provided
+    // as here.
+    pman.app_input->boundary_conditions[parthenon::BoundaryFace::inner_x1] =
+        parthenon::BoundaryFunction::OutflowInnerX1;
+    pman.app_input->swarm_boundary_conditions[parthenon::BoundaryFace::inner_x1] =
+        particles_example::SetSwarmIX1UserBC;
+  }
+  if (pman.pinput->GetString("parthenon/mesh", "ox1_bc") == "user") {
+    // Again, we use a default parthenon boundary condition for cell variables but a
+    // custom swarm boundary condition.
+    pman.app_input->boundary_conditions[parthenon::BoundaryFace::outer_x1] =
+        parthenon::BoundaryFunction::OutflowOuterX1;
+    pman.app_input->swarm_boundary_conditions[parthenon::BoundaryFace::outer_x1] =
+        particles_example::SetSwarmOX1UserBC;
+  }
+  pman.ParthenonInitPackagesAndMesh();
 
   // Initialize the driver
   particles_example::ParticleDriver driver(pman.pinput.get(), pman.app_input.get(),

example/particles/particles.cpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2020-2021. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2022. 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
@@ -25,6 +25,16 @@
 // *************************************************//
 namespace particles_example {
 
+std::unique_ptr<ParticleBound, DeviceDeleter<parthenon::DevMemSpace>>
+SetSwarmIX1UserBC() {
+  return DeviceAllocate<ParticleBoundIX1User>();
+}
+
+std::unique_ptr<ParticleBound, DeviceDeleter<parthenon::DevMemSpace>>
+SetSwarmOX1UserBC() {
+  return DeviceAllocate<ParticleBoundOX1User>();
+}
+
 Packages_t ProcessPackages(std::unique_ptr<ParameterInput> &pin) {
   Packages_t packages;
   packages.Add(particles_example::Particles::Initialize(pin.get()));

example/particles/particles.hpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2020-2021. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2022. 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
@@ -26,6 +26,28 @@ using namespace parthenon;
 
 namespace particles_example {
 
+// Duplicates of ParticleBoundIX1Outflow and ParticleBoundOX1Outflow to illustrate custom
+// particle boundary conditions
+class ParticleBoundIX1User : public ParticleBound {
+ public:
+  KOKKOS_INLINE_FUNCTION void Apply(const int n, double &x, double &y, double &z,
+                                    const SwarmDeviceContext &swarm_d) const override {
+    if (x < swarm_d.x_min_global_) {
+      swarm_d.MarkParticleForRemoval(n);
+    }
+  }
+};
+
+class ParticleBoundOX1User : public ParticleBound {
+ public:
+  KOKKOS_INLINE_FUNCTION void Apply(const int n, double &x, double &y, double &z,
+                                    const SwarmDeviceContext &swarm_d) const override {
+    if (x > swarm_d.x_max_global_) {
+      swarm_d.MarkParticleForRemoval(n);
+    }
+  }
+};
+
 typedef Kokkos::Random_XorShift64_Pool<> RNGPool;
 
 class ParticleDriver : public EvolutionDriver {
@@ -44,6 +66,10 @@ class ParticleDriver : public EvolutionDriver {
 void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin);
 Packages_t ProcessPackages(std::unique_ptr<ParameterInput> &pin);
 
+std::unique_ptr<ParticleBound, DeviceDeleter<parthenon::DevMemSpace>> SetSwarmIX1UserBC();
+
+std::unique_ptr<ParticleBound, DeviceDeleter<parthenon::DevMemSpace>> SetSwarmOX1UserBC();
+
 namespace Particles {
 
 std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin);

src/application_input.hpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2020-2021. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2022. 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
@@ -33,7 +33,7 @@ struct ApplicationInput {
   std::function<Packages_t(std::unique_ptr<ParameterInput> &)> ProcessPackages = nullptr;
 
   // Mesh functions
-  std::function<void(ParameterInput *)> InitUserMeshData = nullptr;
+  std::function<void(Mesh *, ParameterInput *)> InitUserMeshData = nullptr;
 
   std::function<void(Mesh *, ParameterInput *, SimTime &)> PreStepMeshUserWorkInLoop =
       nullptr;
@@ -46,8 +46,8 @@ struct ApplicationInput {
       PostStepDiagnosticsInLoop = nullptr;
 
   std::function<void(Mesh *, ParameterInput *, SimTime &)> UserWorkAfterLoop = nullptr;
-  BValFunc boundary_conditions[BOUNDARY_NFACES] = {nullptr, nullptr, nullptr,
-                                                   nullptr, nullptr, nullptr};
+  BValFunc boundary_conditions[BOUNDARY_NFACES] = {nullptr};
+  SBValFunc swarm_boundary_conditions[BOUNDARY_NFACES] = {nullptr};
 
   // MeshBlock functions
   std::function<std::unique_ptr<MeshBlockApplicationData>(MeshBlock *, ParameterInput *)>

src/bvals/boundary_conditions.hpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2020. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2022. 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
@@ -20,13 +20,16 @@
 
 #include "basic_types.hpp"
 #include "interface/meshblock_data.hpp"
+#include "interface/swarm_boundaries.hpp"
 #include "mesh/domain.hpp"
 
 namespace parthenon {
 
 // Physical boundary conditions
 
 using BValFunc = std::function<void(std::shared_ptr<MeshBlockData<Real>> &, bool)>;
+using SBValFunc = std::function<
+    std::unique_ptr<ParticleBound, DeviceDeleter<parthenon::DevMemSpace>>()>;
 
 TaskStatus ProlongateBoundaries(std::shared_ptr<MeshBlockData<Real>> &rc);
 

src/interface/swarm.cpp

@@ -1,5 +1,5 @@
 //========================================================================================
-// (C) (or copyright) 2020-2021. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2022. 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
@@ -92,65 +92,47 @@ Swarm::Swarm(const std::string &label, const Metadata &metadata, const int nmax_
   marked_for_removal_.data.DeepCopy(marked_for_removal_h);
 }
 
-void Swarm::AllocateBoundaries() {
-  auto pmb = GetBlockPointer();
+template <class BOutflow, class BPeriodic, int iFace>
+void Swarm::AllocateBoundariesImpl_(MeshBlock *pmb) {
   std::stringstream msg;
-
   auto &bcs = pmb->pmy_mesh->mesh_bcs;
-
-  if (bcs[0] == BoundaryFlag::outflow) {
-    bounds_uptrs[0] = DeviceAllocate<ParticleBoundIX1Outflow>();
-  } else if (bcs[0] == BoundaryFlag::periodic) {
-    bounds_uptrs[0] = DeviceAllocate<ParticleBoundIX1Periodic>();
-  } else if (bcs[0] != BoundaryFlag::user) {
-    msg << "ix1 boundary flag " << static_cast<int>(bcs[0]) << " not supported!";
-    PARTHENON_THROW(msg);
-  }
-
-  if (bcs[1] == BoundaryFlag::outflow) {
-    bounds_uptrs[1] = DeviceAllocate<ParticleBoundOX1Outflow>();
-  } else if (bcs[1] == BoundaryFlag::periodic) {
-    bounds_uptrs[1] = DeviceAllocate<ParticleBoundOX1Periodic>();
-  } else if (bcs[1] != BoundaryFlag::user) {
-    msg << "ox1 boundary flag " << static_cast<int>(bcs[1]) << " not supported!";
-    PARTHENON_THROW(msg);
-  }
-
-  if (bcs[2] == BoundaryFlag::outflow) {
-    bounds_uptrs[2] = DeviceAllocate<ParticleBoundIX2Outflow>();
-  } else if (bcs[2] == BoundaryFlag::periodic) {
-    bounds_uptrs[2] = DeviceAllocate<ParticleBoundIX2Periodic>();
-  } else if (bcs[2] != BoundaryFlag::user) {
-    msg << "ix2 boundary flag " << static_cast<int>(bcs[2]) << " not supported!";
+  if (bcs[iFace] == BoundaryFlag::outflow) {
+    bounds_uptrs[iFace] = DeviceAllocate<BOutflow>();
+  } else if (bcs[iFace] == BoundaryFlag::periodic) {
+    bounds_uptrs[iFace] = DeviceAllocate<BPeriodic>();
+  } else if (bcs[iFace] == BoundaryFlag::user) {
+    if (pmb->pmy_mesh->SwarmBndryFnctn[iFace] != nullptr) {
+      bounds_uptrs[iFace] = pmb->pmy_mesh->SwarmBndryFnctn[iFace]();
+    } else {
+      msg << "ix" << iFace + 1
+          << " user boundary requested but provided function is null!";
+      PARTHENON_THROW(msg);
+    }
+  } else {
+    msg << "ix" << iFace + 1 << " boundary flag " << static_cast<int>(bcs[iFace])
+        << " not supported!";
     PARTHENON_THROW(msg);
   }
+}
 
-  if (bcs[3] == BoundaryFlag::outflow) {
-    bounds_uptrs[3] = DeviceAllocate<ParticleBoundOX2Outflow>();
-  } else if (bcs[3] == BoundaryFlag::periodic) {
-    bounds_uptrs[3] = DeviceAllocate<ParticleBoundOX2Periodic>();
-  } else if (bcs[3] != BoundaryFlag::user) {
-    msg << "ox2 boundary flag " << static_cast<int>(bcs[3]) << " not supported!";
-    PARTHENON_THROW(msg);
-  }
+void Swarm::AllocateBoundaries() {
+  auto pmb = GetBlockPointer();
+  std::stringstream msg;
 
-  if (bcs[4] == BoundaryFlag::outflow) {
-    bounds_uptrs[4] = DeviceAllocate<ParticleBoundIX3Outflow>();
-  } else if (bcs[4] == BoundaryFlag::periodic) {
-    bounds_uptrs[4] = DeviceAllocate<ParticleBoundIX3Periodic>();
-  } else if (bcs[4] != BoundaryFlag::user) {
-    msg << "ix3 boundary flag " << static_cast<int>(bcs[4]) << " not supported!";
-    PARTHENON_THROW(msg);
-  }
+  auto &bcs = pmb->pmy_mesh->mesh_bcs;
 
-  if (bcs[5] == BoundaryFlag::outflow) {
-    bounds_uptrs[5] = DeviceAllocate<ParticleBoundOX3Outflow>();
-  } else if (bcs[5] == BoundaryFlag::periodic) {
-    bounds_uptrs[5] = DeviceAllocate<ParticleBoundOX3Periodic>();
-  } else if (bcs[5] != BoundaryFlag::user) {
-    msg << "ox3 boundary flag " << static_cast<int>(bcs[5]) << " not supported!";
-    PARTHENON_THROW(msg);
-  }
+  AllocateBoundariesImpl_<ParticleBoundIX1Outflow, ParticleBoundIX1Periodic, 0>(
+      pmb.get());
+  AllocateBoundariesImpl_<ParticleBoundOX1Outflow, ParticleBoundOX1Periodic, 1>(
+      pmb.get());
+  AllocateBoundariesImpl_<ParticleBoundIX2Outflow, ParticleBoundIX2Periodic, 2>(
+      pmb.get());
+  AllocateBoundariesImpl_<ParticleBoundOX2Outflow, ParticleBoundOX2Periodic, 3>(
+      pmb.get());
+  AllocateBoundariesImpl_<ParticleBoundIX3Outflow, ParticleBoundIX3Periodic, 4>(
+      pmb.get());
+  AllocateBoundariesImpl_<ParticleBoundOX3Outflow, ParticleBoundOX3Periodic, 5>(
+      pmb.get());
 
   for (int n = 0; n < 6; n++) {
     bounds_d.bounds[n] = bounds_uptrs[n].get();
@@ -556,6 +538,10 @@ void Swarm::SortParticlesByCell() {
 
           if (cellSorted(start_index).cell_idx_1d_ >= cell_idx_1d) {
             start_index--;
+            if (start_index < 0) {
+              start_index = -1;
+              break;
+            }
             if (cellSorted(start_index).cell_idx_1d_ < cell_idx_1d) {
               start_index = -1;
               break;
@@ -564,6 +550,10 @@ void Swarm::SortParticlesByCell() {
           }
           if (cellSorted(start_index).cell_idx_1d_ < cell_idx_1d) {
             start_index++;
+            if (start_index > max_active_index) {
+              start_index = -1;
+              break;
+            }
             if (cellSorted(start_index).cell_idx_1d_ > cell_idx_1d) {
               start_index = -1;
               break;

src/interface/swarm.hpp

@@ -210,6 +210,9 @@ class Swarm {
   template <class T>
   vpack_types::SwarmVarList<T> MakeVarListAll_();
 
+  template <class BOutflow, class BPeriodic, int iFace>
+  void AllocateBoundariesImpl_(MeshBlock *pmb);
+
   void SetNeighborIndices1D_();
   void SetNeighborIndices2D_();
   void SetNeighborIndices3D_();

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-2021. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2022. 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
@@ -290,7 +290,7 @@ Mesh::Mesh(ParameterInput *pin, ApplicationInput *app_in, Packages_t &packages,
     max_level = 63;
   }
 
-  InitUserMeshData(pin);
+  InitUserMeshData(this, pin);
 
   if (multilevel) {
     if (block_size.nx1 % 2 == 1 || (block_size.nx2 % 2 == 1 && (ndim >= 2)) ||
@@ -675,7 +675,7 @@ Mesh::Mesh(ParameterInput *pin, ApplicationInput *app_in, RestartReader &rr,
     max_level = 63;
   }
 
-  InitUserMeshData(pin);
+  InitUserMeshData(this, pin);
 
   // Populate logical locations
   auto lx123 = rr.ReadDataset<int64_t>("/Blocks/loc.lx123");
@@ -960,6 +960,21 @@ void Mesh::EnrollBndryFncts_(ApplicationInput *app_in) {
     default: // periodic/block BCs handled elsewhere.
       break;
     }
+
+    switch (mesh_bcs[f]) {
+    case BoundaryFlag::user:
+      if (app_in->swarm_boundary_conditions[f] != nullptr) {
+        SwarmBndryFnctn[f] = app_in->swarm_boundary_conditions[f];
+      } else {
+        std::stringstream msg;
+        msg << "A user boundary condition for face " << f
+            << " was requested, but no swarm condition was enrolled." << std::endl;
+        PARTHENON_THROW(msg);
+      }
+      break;
+    default: // Default BCs handled elsewhere
+      break;
+    }
   }
 }