WIP: Towards Curvilinear Swarms

CHANGELOG.md

@@ -5,6 +5,8 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[PR 1391]](https://github.com/parthenon-hpc-lab/parthenon/pull/1391) Towards Curvilinear Swarms
+- [[PR 1382]](https://github.com/parthenon-hpc-lab/parthenon/pull/1382) Support Particle AMR
 - [[PR 1378]](https://github.com/parthenon-hpc-lab/parthenon/pull/1378) MeshData Swarm Tasks
 - [[PR 1377]](https://github.com/parthenon-hpc-lab/parthenon/pull/1377) Extend Initialization Hierarchy
 - [[PR 1332]](https://github.com/parthenon-hpc-lab/parthenon/pull/1332) Add global WatchDog
@@ -36,6 +38,7 @@
 
 
 ### Incompatibilities (i.e. breaking changes)
+- [[PR 1391]](https://github.com/parthenon-hpc-lab/parthenon/pull/1391) Towards Curvilinear Swarms
 - [[PR 1377]](https://github.com/parthenon-hpc-lab/parthenon/pull/1377) Extend Initialization Hierarchy
 - [[PR 1376]](https://github.com/parthenon-hpc-lab/parthenon/pull/1376) Refactor SwarmPacks
 

dev/plan_histories/1382.md

@@ -0,0 +1,42 @@
+<!-- This file was made in part with generative AI. -->
+
+## Goal
+
+Add correct AMR remeshing support for swarms so particle records are preserved across
+refinement, derefinement, and same-level redistribution, while keeping the data path
+compatible with device-resident execution and GPU-aware MPI.
+
+## Functional Plan
+
+- Add a dedicated swarm remesh phase that runs after the new leaf mesh and field AMR
+  transfers are already complete.
+- Build the remesh around block ownership changes:
+  - same-level remap sends every particle in an old block to one new block
+  - derefinement sends every particle in a child block to its parent block
+  - refinement uses particle positions only to choose among that old block's daughters
+- Represent each transported particle as one complete mixed-type record containing the
+  destination block identifier plus all swarm variables needed to reconstruct the
+  particle.
+- Pack and unpack those particle records in a precision-independent wire format.
+- Keep the swarm payload buffers device-resident during communication so off-rank
+  transfers can use direct device-buffer MPI, and keep same-rank transfers on the
+  device as well.
+- Recreate the received particles in the destination swarms, restore all particle
+  variables, and clear any cached swarm pack state invalidated by the ownership change.
+- Add the minimal generic swarm-schema helpers needed to enumerate variable names and
+  compute packed record sizes without hard-coding a particular swarm layout.
+- Exercise the new capability in the `particle_tracers` example by enabling AMR tagging
+  and retaining an analytically checkable particle motion.
+- Add focused validation for the new functionality:
+  - a refinement unit test that checks particles land on the correct daughter blocks
+  - a derefinement unit test that checks particles reunite on the correct parent block
+  - regression coverage that checks AMR tracer motion against the analytic translation
+- Correct any supporting same-rank swarm communication issue needed to make these AMR
+  particle transfers reliable.
+
+## Scope Boundaries
+
+- Do not redesign the field AMR transfer path; only integrate the swarm remesh phase at
+  the correct point in the existing AMR sequence.
+- Keep the implementation generic over swarm variable layout rather than tailoring it to
+  one example problem.

doc/sphinx/src/particles.rst

@@ -1,3 +1,5 @@
+.. This file was made in part with generative AI.
+
 Particles
 =========
 
@@ -11,8 +13,8 @@ Swarms
 
 A ``Swarm`` contains all the particle data for all particles of a given
 species. It owns a set of ``ParticleVariable``\ s, one for each value of
-each particle. For example, the spatial positions ``x``, ``y``, and
-``z`` of the particles in a swarm are three separate
+each particle. For example, the spatial positions ``x1``, ``x2``, and
+``x3`` of the particles in a swarm are three separate
 ``ParticleVariable``\ s. ``ParticleVariable``\ s can be either ``Real``-,
 ``uint64_t``-, or ``int``-valued, which is specified by the metadata values
 ``Metadata::Real``, ``MetaData::UInt64``, and ``Metadata::Integer``. ``ParticleVariable``\ s
@@ -21,7 +23,7 @@ should also contain the ``Metadata::Particle`` flag. By default,
 to 2D data per particle is currently supported, by passing
 ``std::vector<int>{N1, N2}`` as the second argument to the
 ``ParticleVariable`` ``Metadata``. All ``Swarm``\ s by default contain
-positional ``x``, ``y``, and ``z`` ``ParticleVariable``\ s and a ``uint64`` particle
+positional ``x1``, ``x2``, and ``x3`` ``ParticleVariable``\ s and a ``uint64`` particle
 ``id`` that is persistent and unique for a given simulation.
 The latter field can be disabled (e.g., because it is not required for a method
 that constant creates and destroys anonymous particles) by providing the
@@ -47,6 +49,13 @@ permit ``Swarm``s to be retrieved from ``"base"`` ``MeshBlockData`` and ``MeshDa
    particles that the reservation holds for a given swarm by calling
    ``metadata.SetInitialSwarmPoolReservation(nparticles);``.
 
+.. note::
+
+   Particle infrastructure now operates in the native ``x1``, ``x2``, and ``x3``
+   coordinate directions rather than assuming Cartesian ``x``, ``y``, and ``z``.
+   That said, particles and swarms in curvilinear coordinate systems remain mostly
+   untested.
+
 The ``Swarm`` is a host-side object, but some of its data members are
 required for device- side compution. To access this data, a
 ``SwarmDeviceContext`` object is created via
@@ -92,11 +101,11 @@ computations on currently active particles:
 
 .. code:: cpp
 
-   auto &x = swarm.Get("x").Get();
+   auto &x1 = swarm.Get(swarm_position::x1::name()).Get();
    swarm.pmy_block->par_for("Simple loop", 0, swarm.GetMaxActiveIndex(),
      KOKKOS_LAMBDA(const int n) {
        if (swarm_d.IsActive(n)) {
-         x(n) += 1.0;
+         x1(n) += 1.0;
        }
      });
 
@@ -152,7 +161,11 @@ must be repeated in the driver’s evolution function until all particles
 are completed. See the ``particles`` example for further details. Note
 that this pattern is blocking, and may be replaced in the future.
 
-AMR is currently not supported, but support will be added in the future.
+AMR support for swarms is available. When the mesh changes, swarm
+particles are redistributed to the correct new leaf block after
+same-level redistribution, refinement, or derefinement. The
+``particle_tracers`` example demonstrates this path for advected tracer
+particles on an adaptively refined mesh.
 
 Variable Packing
 ----------------
@@ -173,14 +186,14 @@ Similar to grid variables, swarms can be packed over ``MeshBlock``\ s via ``Swar
 a ``SwarmPack`` via a ``std::vector<std::string>`` or the type-based variable prescription
 previously used by ``SparsePack``\ s (see :ref:`sparse_packs`).
 
-For packing via string (wherein below, ``swarm_position::x::name()`` returns a string),
+For packing via string (wherein below, ``swarm_position::x1::name()`` returns a string),
 one must specify the data type by template argument:
 
 .. code:: cpp
 
-   std::vector<std::string> vars{swarm_position::x::name(),
-                                 swarm_position::y::name(),
-                                 swarm_position::z::name()};
+   std::vector<std::string> vars{swarm_position::x1::name(),
+                                 swarm_position::x2::name(),
+                                 swarm_position::x3::name()};
    static auto desc = MakeSwarmPackDescriptor<Real>(swarm_name, vars);
    auto pack = desc.GetPack(md);
 
@@ -190,9 +203,9 @@ example), the type can be inferred automatically:
 
 .. code:: cpp
 
-   static auto desc = MakeSwarmPackDescriptor<swarm_position::x,
-                                              swarm_position::y,
-                                              swarm_position::z>(swarm_name);
+   static auto desc = MakeSwarmPackDescriptor<swarm_position::x1,
+                                              swarm_position::x2,
+                                              swarm_position::x3>(swarm_name);
    auto pack = desc.GetPack(md);
 
 
@@ -292,3 +305,6 @@ Putting it all together, you might have an output block that looks like this:
 The result would be that both ``swarm1`` and ``swarm2`` output the
 variables ``id``, ``x``, ``y``, ``z``, and ``shared_var``. But only ``swarm1``
 outputs ``per_swarm_var``.
+
+.. note::
+    This file was made in part with generative AI.

example/particle_leapfrog/particle_leapfrog.cpp

@@ -183,9 +183,9 @@ void PostInitialization(MeshBlock *pmb, ParameterInput *pin) {
   auto new_particles_context = swarm->AddEmptyParticles(num_particles_this_block);
 
   auto &id = swarm->Get<std::uint64_t>(swarm_position::id::name()).Get();
-  auto &x = swarm->Get<Real>(swarm_position::x::name()).Get();
-  auto &y = swarm->Get<Real>(swarm_position::y::name()).Get();
-  auto &z = swarm->Get<Real>(swarm_position::z::name()).Get();
+  auto &x = swarm->Get<Real>(swarm_position::x1::name()).Get();
+  auto &y = swarm->Get<Real>(swarm_position::x2::name()).Get();
+  auto &z = swarm->Get<Real>(swarm_position::x3::name()).Get();
   auto &v = swarm->Get<Real>("v").Get();
   auto &vv = swarm->Get<Real>("vv").Get();
 
@@ -225,7 +225,7 @@ TaskStatus TransportParticles(MeshData<Real> *md, const StagedIntegrator *integr
   // NOTE(@pdmullen): the data type for Positions (Real) are automatically deduced from
   // the variable typing
   static auto desc_pos =
-      MakeSwarmPackDescriptor<swarm_position::x, swarm_position::y, swarm_position::z>(
+      MakeSwarmPackDescriptor<swarm_position::x1, swarm_position::x2, swarm_position::x3>(
           swarm_name);
   auto pack_pos = desc_pos.GetPack(md);
 
@@ -250,19 +250,19 @@ TaskStatus TransportParticles(MeshData<Real> *md, const StagedIntegrator *integr
         const auto swarm_d = pack_pos.GetContext(b);
         if (swarm_d.IsActive(n)) {
           // drift
-          pack_pos(b, swarm_position::x(), n) += pack_v(b, iv + 0, n) * 0.5 * dt;
-          pack_pos(b, swarm_position::y(), n) += pack_v(b, iv + 1, n) * 0.5 * dt;
-          pack_pos(b, swarm_position::z(), n) += pack_v(b, iv + 2, n) * 0.5 * dt;
+          pack_pos(b, swarm_position::x1(), n) += pack_v(b, iv + 0, n) * 0.5 * dt;
+          pack_pos(b, swarm_position::x2(), n) += pack_v(b, iv + 1, n) * 0.5 * dt;
+          pack_pos(b, swarm_position::x3(), n) += pack_v(b, iv + 2, n) * 0.5 * dt;
 
           // kick
           pack_v(b, iv + 0, n) += ax * dt;
           pack_v(b, iv + 1, n) += ay * dt;
           pack_v(b, iv + 2, n) += az * dt;
 
           // drift
-          pack_pos(b, swarm_position::x(), n) += pack_v(b, iv + 0, n) * 0.5 * dt;
-          pack_pos(b, swarm_position::y(), n) += pack_v(b, iv + 1, n) * 0.5 * dt;
-          pack_pos(b, swarm_position::z(), n) += pack_v(b, iv + 2, n) * 0.5 * dt;
+          pack_pos(b, swarm_position::x1(), n) += pack_v(b, iv + 0, n) * 0.5 * dt;
+          pack_pos(b, swarm_position::x2(), n) += pack_v(b, iv + 1, n) * 0.5 * dt;
+          pack_pos(b, swarm_position::x3(), n) += pack_v(b, iv + 2, n) * 0.5 * dt;
 
           // id
           PARTHENON_REQUIRE(pack_id(b, swarm_position::id(), n) >= 0,

example/particle_tracers/particle_tracers.cpp

@@ -28,6 +28,7 @@
 #include <vector>
 
 // Parthenon includes
+#include "amr_criteria/refinement_package.hpp"
 #include "basic_types.hpp"
 #include "bvals/comms/bvals_in_one.hpp"
 #include "config.hpp"
@@ -115,6 +116,9 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   const Real advected_amp = 0.5;
   adv_pkg->AddParam("advected_mean", advected_mean);
   adv_pkg->AddParam("advected_amp", advected_amp);
+  adv_pkg->AddParam("refine_tol", pin->GetOrAddReal("Background", "refine_tol", 1.25));
+  adv_pkg->AddParam("derefine_tol",
+                    pin->GetOrAddReal("Background", "derefine_tol", 1.05));
   PARTHENON_REQUIRE(advected_mean > advected_amp,
                     "Advected field must be everywere positive!");
 
@@ -130,9 +134,41 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   // Assign package timestep hook
   adv_pkg->EstimateTimestepMesh = EstimateTimestepMesh;
 
+  adv_pkg->CheckRefinementBlock = CheckRefinementBlock;
+
   return adv_pkg;
 }
 
+AmrTag CheckRefinementBlock(MeshBlockData<Real> *rc) {
+  auto pmb = rc->GetBlockPointer();
+  auto pkg = pmb->packages.Get("advection_package");
+  static auto desc =
+      parthenon::MakePackDescriptor<field::advected>(pmb->resolved_packages.get());
+  auto pack = desc.GetPack(rc);
+
+  const IndexRange ib = pmb->cellbounds.GetBoundsI(IndexDomain::entire);
+  const IndexRange jb = pmb->cellbounds.GetBoundsJ(IndexDomain::entire);
+  const IndexRange kb = pmb->cellbounds.GetBoundsK(IndexDomain::entire);
+
+  typename Kokkos::MinMax<Real>::value_type minmax;
+  pmb->par_reduce(
+      PARTHENON_AUTO_LABEL, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+      KOKKOS_LAMBDA(const int k, const int j, const int i,
+                    typename Kokkos::MinMax<Real>::value_type &lminmax) {
+        const Real val = pack(0, field::advected(), k, j, i);
+        lminmax.min_val = val < lminmax.min_val ? val : lminmax.min_val;
+        lminmax.max_val = val > lminmax.max_val ? val : lminmax.max_val;
+      },
+      Kokkos::MinMax<Real>(minmax));
+
+  const auto &refine_tol = pkg->Param<Real>("refine_tol");
+  const auto &derefine_tol = pkg->Param<Real>("derefine_tol");
+
+  if (minmax.max_val > refine_tol) return AmrTag::refine;
+  if (minmax.max_val < derefine_tol) return AmrTag::derefine;
+  return AmrTag::same;
+}
+
 } // namespace advection_package
 
 namespace particles_package {
@@ -248,7 +284,7 @@ void SourceTracers(MeshBlock *pmb, ParameterInput *pin) {
 
   // Create pack
   static auto desc =
-      MakeSwarmPackDescriptor<swarm_position::x, swarm_position::y, swarm_position::z>(
+      MakeSwarmPackDescriptor<swarm_position::x1, swarm_position::x2, swarm_position::x3>(
           swarm_name);
   auto pack = desc.GetPack(mbd.get());
 
@@ -260,9 +296,9 @@ void SourceTracers(MeshBlock *pmb, ParameterInput *pin) {
         auto rng_gen = rng_pool.get_state();
 
         // Extract particle position
-        Real &xx = pack(0, swarm_position::x(), n);
-        Real &yy = pack(0, swarm_position::y(), n);
-        Real &zz = pack(0, swarm_position::z(), n);
+        Real &xx = pack(0, swarm_position::x1(), n);
+        Real &yy = pack(0, swarm_position::x2(), n);
+        Real &zz = pack(0, swarm_position::x3(), n);
 
         // Rejection sample the x position
         Real val;
@@ -298,7 +334,7 @@ TaskStatus AdvectTracers(MeshData<Real> *md, const Real dt) {
 
   // Create pack
   static auto desc =
-      MakeSwarmPackDescriptor<swarm_position::x, swarm_position::y, swarm_position::z>(
+      MakeSwarmPackDescriptor<swarm_position::x1, swarm_position::x2, swarm_position::x3>(
           swarm_name);
   auto pack = desc.GetPack(md);
 
@@ -309,9 +345,9 @@ TaskStatus AdvectTracers(MeshData<Real> *md, const Real dt) {
         auto [b, n] = pack.GetBlockParticleIndices(idx);
         const auto &swarm_d = pack.GetContext(b);
         if (swarm_d.IsActive(n)) {
-          pack(b, swarm_position::x(), n) += vx * dt;
-          pack(b, swarm_position::y(), n) += vy * dt;
-          pack(b, swarm_position::z(), n) += vz * dt;
+          pack(b, swarm_position::x1(), n) += vx * dt;
+          pack(b, swarm_position::x2(), n) += vy * dt;
+          pack(b, swarm_position::x3(), n) += vz * dt;
         }
       });
 
@@ -330,7 +366,7 @@ TaskStatus DepositTracers(MeshData<Real> *md) {
   // Create packs
   static auto desc = MakePackDescriptor<field::deposition>(resolved_pkgs.get());
   static auto pdesc =
-      MakeSwarmPackDescriptor<swarm_position::x, swarm_position::y, swarm_position::z>(
+      MakeSwarmPackDescriptor<swarm_position::x1, swarm_position::x2, swarm_position::x3>(
           swarm_name);
   auto vmesh = desc.GetPack(md);
   auto vpart = pdesc.GetPack(md);
@@ -351,9 +387,9 @@ TaskStatus DepositTracers(MeshData<Real> *md) {
         const auto &swarm_d = vpart.GetContext(b);
         if (swarm_d.IsActive(n)) {
           int ip, jp, kp;
-          const Real &xx = vpart(b, swarm_position::x(), n);
-          const Real &yy = vpart(b, swarm_position::y(), n);
-          const Real &zz = vpart(b, swarm_position::z(), n);
+          const Real &xx = vpart(b, swarm_position::x1(), n);
+          const Real &yy = vpart(b, swarm_position::x2(), n);
+          const Real &zz = vpart(b, swarm_position::x3(), n);
           swarm_d.Xtoijk(xx, yy, zz, ip, jp, kp);
 
           // For testing in this example we make sure the indices are correct; these
@@ -505,6 +541,15 @@ TaskCollection ParticleDriver::StepTasks() {
     auto defrag = tl.AddTask(deposit, parthenon::DefragSwarmsMesh, base, 0.9);
   }
 
+  if (pmesh->adaptive) {
+    auto &amr_region = tc.AddRegion(num_partitions);
+    for (int i = 0; i < num_partitions; ++i) {
+      auto &tl = amr_region[i];
+      auto &base = pmesh->mesh_data.GetOrAdd("base", i);
+      tl.AddTask(none, parthenon::Refinement::Tag<MeshData<Real>>, base.get());
+    }
+  }
+
   return tc;
 }
 
@@ -533,8 +578,9 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
 
   // Mesh physical size
   const auto mesh_size = pmb->pmy_mesh->mesh_size;
-  const Real mesh_lx1 = mesh_size.xmax(X1DIR) - mesh_size.xmin(X1DIR);
-  const Real kwave = 2.0 * M_PI / mesh_lx1;
+  const Real x_min_mesh = mesh_size.xmin(X1DIR);
+  const Real x_max_mesh = mesh_size.xmax(X1DIR);
+  const Real kwave = 2.0 * M_PI / (x_max_mesh - x_min_mesh);
 
   // Create pack
   static auto desc = parthenon::MakePackDescriptor<field::advected>(resolved_pkgs.get());

example/particle_tracers/particle_tracers.hpp

@@ -14,9 +14,12 @@
 // license in this material to reproduce, prepare derivative works, distribute copies to
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
+
 #ifndef EXAMPLE_PARTICLE_TRACERS_PARTICLE_TRACERS_HPP_
 #define EXAMPLE_PARTICLE_TRACERS_PARTICLE_TRACERS_HPP_
 
+// This file was made in part with generative AI.
+
 #include <memory>
 
 #include "Kokkos_Random.hpp"
@@ -59,6 +62,7 @@ namespace advection_package {
 
 std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin);
 Real EstimateTimestepMesh(MeshData<Real> *rc);
+AmrTag CheckRefinementBlock(MeshBlockData<Real> *rc);
 
 } // namespace advection_package