Sort particles by cell

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[PR 619]](https://github.com/lanl/parthenon/pull/619) Sort particles by cell
 - [[PR 605]](https://github.com/lanl/parthenon/pull/605) Add output triggering by signaling.
 - [[PR 602]](https://github.com/lanl/parthenon/pull/602) Added tuning functionality for HDF5 output
 

docs/particles.md

@@ -60,6 +60,16 @@ swarm.pmy_block->par_for("Simple loop", 0, swarm.GetMaxActiveIndex(),
   });
 ```
 
+## Sorting
+
+By default, particles are stored in per-meshblock pools of memory. However, one frequently wants
+convenient access to all the particles in each computational cell separately. To facilitate this,
+the Swarm provides the method `SortParticlesByCell` (and the `SwarmContainer` provides the matching
+task `SortParticlesByCell`). Calling this function populates internal data structures that map from
+per-cell indices to the per-meshblock data array. These are accessed by the `SwarmDeviceContext`
+member functions `GetParticleCountPerCell` and `GetFullIndex`. See `examples/particles` for example
+usage.
+
 ## Defragmenting
 
 Because one typically loops over particles from 0 to `max_active_index`, if only a small

example/particles/parthinput.particles

@@ -57,4 +57,5 @@ num_particles = 10
 particle_speed = 1.0
 rng_seed = 23487
 const_dt = 0.5
+deposition_method = per_cell
 destroy_particles_frac = 0.1

example/particles/particles.cpp

@@ -35,6 +35,8 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
   // Don't do anything for now
 }
 
+enum class DepositionMethod { per_particle, per_cell };
+
 // *************************************************//
 // define the "physics" package particles_package, *//
 // which includes defining various functions that  *//
@@ -61,6 +63,16 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
       destroy_particles_frac >= 0. && destroy_particles_frac <= 1.,
       "Fraction of particles to destroy each timestep must be between 0 and 1");
 
+  std::string deposition_method =
+      pin->GetOrAddString("Particles", "deposition_method", "per_particle");
+  if (deposition_method == "per_particle") {
+    pkg->AddParam<>("deposition_method", DepositionMethod::per_particle);
+  } else if (deposition_method == "per_cell") {
+    pkg->AddParam<>("deposition_method", DepositionMethod::per_cell);
+  } else {
+    PARTHENON_THROW("deposition method not recognized");
+  }
+
   bool orbiting_particles =
       pin->GetOrAddBoolean("Particles", "orbiting_particles", false);
   pkg->AddParam<>("orbiting_particles", orbiting_particles);
@@ -139,11 +151,25 @@ TaskStatus DestroySomeParticles(MeshBlock *pmb) {
   return TaskStatus::complete;
 }
 
+TaskStatus SortParticlesIfUsingPerCellDeposition(MeshBlock *pmb) {
+  auto pkg = pmb->packages.Get("particles_package");
+  const auto deposition_method = pkg->Param<DepositionMethod>("deposition_method");
+  if (deposition_method == DepositionMethod::per_cell) {
+    auto swarm = pmb->swarm_data.Get()->Get("my particles");
+    swarm->SortParticlesByCell();
+  }
+
+  return TaskStatus::complete;
+}
+
 TaskStatus DepositParticles(MeshBlock *pmb) {
   Kokkos::Profiling::pushRegion("Task_Particles_DepositParticles");
 
   auto swarm = pmb->swarm_data.Get()->Get("my particles");
 
+  auto pkg = pmb->packages.Get("particles_package");
+  const auto deposition_method = pkg->Param<DepositionMethod>("deposition_method");
+
   // Meshblock geometry
   const IndexRange &ib = pmb->cellbounds.GetBoundsI(IndexDomain::interior);
   const IndexRange &jb = pmb->cellbounds.GetBoundsJ(IndexDomain::interior);
@@ -162,34 +188,46 @@ TaskStatus DepositParticles(MeshBlock *pmb) {
   auto swarm_d = swarm->GetDeviceContext();
 
   auto &particle_dep = pmb->meshblock_data.Get()->Get("particle_deposition").data;
-  // Reset particle count
-  pmb->par_for(
-      "ZeroParticleDep", kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
-      KOKKOS_LAMBDA(const int k, const int j, const int i) {
-        particle_dep(k, j, i) = 0.;
-      });
-
   const int ndim = pmb->pmy_mesh->ndim;
 
-  pmb->par_for(
-      "DepositParticles", 0, swarm->GetMaxActiveIndex(), KOKKOS_LAMBDA(const int n) {
-        if (swarm_d.IsActive(n)) {
-          int i = static_cast<int>(std::floor((x(n) - minx_i) / dx_i) + ib.s);
-          int j = 0;
-          if (ndim > 1) {
-            j = static_cast<int>(std::floor((y(n) - minx_j) / dx_j) + jb.s);
-          }
-          int k = 0;
-          if (ndim > 2) {
-            k = static_cast<int>(std::floor((z(n) - minx_k) / dx_k) + kb.s);
-          }
+  if (deposition_method == DepositionMethod::per_particle) {
+    // Reset particle count
+    pmb->par_for(
+        "ZeroParticleDep", kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+        KOKKOS_LAMBDA(const int k, const int j, const int i) {
+          particle_dep(k, j, i) = 0.;
+        });
 
-          if (i >= ib.s && i <= ib.e && j >= jb.s && j <= jb.e && k >= kb.s &&
-              k <= kb.e) {
-            Kokkos::atomic_add(&particle_dep(k, j, i), weight(n));
+    pmb->par_for(
+        "DepositParticles", 0, swarm->GetMaxActiveIndex(), KOKKOS_LAMBDA(const int n) {
+          if (swarm_d.IsActive(n)) {
+            int i = static_cast<int>(std::floor((x(n) - minx_i) / dx_i) + ib.s);
+            int j = 0;
+            if (ndim > 1) {
+              j = static_cast<int>(std::floor((y(n) - minx_j) / dx_j) + jb.s);
+            }
+            int k = 0;
+            if (ndim > 2) {
+              k = static_cast<int>(std::floor((z(n) - minx_k) / dx_k) + kb.s);
+            }
+
+            if (i >= ib.s && i <= ib.e && j >= jb.s && j <= jb.e && k >= kb.s &&
+                k <= kb.e) {
+              Kokkos::atomic_add(&particle_dep(k, j, i), weight(n));
+            }
           }
-        }
-      });
+        });
+  } else if (deposition_method == DepositionMethod::per_cell) {
+    pmb->par_for(
+        "DepositParticlesByCell", kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+        KOKKOS_LAMBDA(const int k, const int j, const int i) {
+          particle_dep(k, j, i) = 0.;
+          for (int n = 0; n < swarm_d.GetParticleCountPerCell(k, j, i); n++) {
+            const int idx = swarm_d.GetFullIndex(k, j, i, n);
+            particle_dep(k, j, i) += weight(idx);
+          }
+        });
+  }
 
   Kokkos::Profiling::popRegion(); // Task_Particles_DepositParticles
   return TaskStatus::complete;
@@ -595,8 +633,10 @@ TaskCollection ParticleDriver::MakeFinalizationTaskCollection() const {
 
     auto destroy_some_particles = tl.AddTask(none, DestroySomeParticles, pmb.get());
 
-    auto deposit_particles =
-        tl.AddTask(destroy_some_particles, DepositParticles, pmb.get());
+    auto sort_particles = tl.AddTask(destroy_some_particles,
+                                     SortParticlesIfUsingPerCellDeposition, pmb.get());
+
+    auto deposit_particles = tl.AddTask(sort_particles, DepositParticles, pmb.get());
 
     // Defragment if swarm memory pool occupancy is 90%
     auto defrag = tl.AddTask(deposit_particles, &SwarmContainer::Defrag, sc.get(), 0.9);

src/CMakeLists.txt

@@ -206,6 +206,7 @@ add_library(parthenon
   utils/reductions.hpp
   utils/show_config.cpp
   utils/signal_handler.cpp
+  utils/sort.hpp
   utils/string_utils.cpp
   utils/string_utils.hpp
   utils/utils.hpp

src/interface/swarm.cpp

@@ -12,12 +12,15 @@
 //========================================================================================
 #include <algorithm>
 #include <cstdlib>
+#include <limits>
 #include <memory>
 #include <utility>
 #include <vector>
 
 #include "mesh/mesh.hpp"
 #include "swarm.hpp"
+#include "utils/error_checking.hpp"
+#include "utils/sort.hpp"
 
 namespace parthenon {
 
@@ -27,6 +30,9 @@ SwarmDeviceContext Swarm::GetDeviceContext() const {
   context.mask_ = mask_.data;
   context.blockIndex_ = blockIndex_;
   context.neighborIndices_ = neighborIndices_;
+  context.cellSorted_ = cellSorted_;
+  context.cellSortedBegin_ = cellSortedBegin_;
+  context.cellSortedNumber_ = cellSortedNumber_;
 
   auto pmb = GetBlockPointer();
   auto pmesh = pmb->pmy_mesh;
@@ -62,7 +68,8 @@ Swarm::Swarm(const std::string &label, const Metadata &metadata, const int nmax_
       marked_for_removal_("mfr", nmax_pool_, Metadata({Metadata::Boolean})),
       neighbor_send_index_("nsi", nmax_pool_, Metadata({Metadata::Integer})),
       blockIndex_("blockIndex_", nmax_pool_),
-      neighborIndices_("neighborIndices_", 4, 4, 4), mpiStatus(true) {
+      neighborIndices_("neighborIndices_", 4, 4, 4),
+      cellSorted_("cellSorted_", nmax_pool_), mpiStatus(true) {
   PARTHENON_REQUIRE_THROWS(typeid(Coordinates_t) == typeid(UniformCartesian),
                            "SwarmDeviceContext only supports a uniform Cartesian mesh!");
 
@@ -275,6 +282,8 @@ void Swarm::setPoolMax(const int nmax_pool) {
       "setPoolMax_marked_for_removal", nmax_pool_, nmax_pool - 1,
       KOKKOS_LAMBDA(const int n) { marked_for_removal_data(n) = false; });
 
+  Kokkos::resize(cellSorted_, nmax_pool);
+
   neighbor_send_index_.Get().Resize(nmax_pool);
 
   blockIndex_.Resize(nmax_pool);
@@ -472,6 +481,113 @@ void Swarm::Defrag() {
   max_active_index_ = num_active_ - 1;
 }
 
+///
+/// Routine to sort particles by cell. Updates internal swarm variables:
+///  cellSorted_: 1D Per-cell sorted array of swarm memory indices
+///  (SwarmKey::swarm_index_) cellSortedBegin_: Per-cell array of starting indices in
+///  cellSorted_ cellSortedNumber_: Per-cell array of number of particles in each cell
+///
+void Swarm::SortParticlesByCell() {
+  auto pmb = GetBlockPointer();
+
+  auto &x = Get<Real>("x").Get();
+  auto &y = Get<Real>("y").Get();
+  auto &z = Get<Real>("z").Get();
+
+  const int nx1 = pmb->cellbounds.ncellsi(IndexDomain::entire);
+  const int nx2 = pmb->cellbounds.ncellsj(IndexDomain::entire);
+  const int nx3 = pmb->cellbounds.ncellsk(IndexDomain::entire);
+  PARTHENON_REQUIRE(nx1 * nx2 * nx3 < std::numeric_limits<int>::max(),
+                    "Too many cells for an int32 to store cell_idx_1d below!");
+
+  auto cellSorted = cellSorted_;
+  int ncells = pmb->cellbounds.GetTotal(IndexDomain::entire);
+  int num_active = num_active_;
+  int max_active_index = max_active_index_;
+
+  // Allocate data if necessary
+  if (cellSortedBegin_.GetDim(1) == 0) {
+    cellSortedBegin_ = ParArrayND<int>("cellSortedBegin_", nx3, nx2, nx1);
+    cellSortedNumber_ = ParArrayND<int>("cellSortedNumber_", nx3, nx2, nx1);
+  }
+  auto cellSortedBegin = cellSortedBegin_;
+  auto cellSortedNumber = cellSortedNumber_;
+  auto swarm_d = GetDeviceContext();
+
+  // Write an unsorted list
+  pmb->par_for(
+      "Write unsorted list", 0, max_active_index_, KOKKOS_LAMBDA(const int n) {
+        int i, j, k;
+        swarm_d.Xtoijk(x(n), y(n), z(n), i, j, k);
+        const int64_t cell_idx_1d = i + nx1 * (j + nx2 * k);
+        cellSorted(n) = SwarmKey(static_cast<int>(cell_idx_1d), n);
+      });
+
+  sort(cellSorted, SwarmKeyComparator(), 0, max_active_index);
+
+  // Update per-cell arrays for easier accessing later
+  const IndexRange &ib = pmb->cellbounds.GetBoundsI(IndexDomain::entire);
+  const IndexRange &jb = pmb->cellbounds.GetBoundsJ(IndexDomain::entire);
+  const IndexRange &kb = pmb->cellbounds.GetBoundsK(IndexDomain::entire);
+  pmb->par_for(
+      "Update per-cell arrays", kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+      KOKKOS_LAMBDA(const int k, const int j, const int i) {
+        int cell_idx_1d = i + nx1 * (j + nx2 * k);
+        // Find starting index, first by guessing
+        int start_index = static_cast<int>((cell_idx_1d * num_active / ncells));
+        int n = 0;
+        while (true) {
+          n++;
+          // Check if we left the list
+          if (start_index < 0 || start_index > max_active_index) {
+            start_index = -1;
+            break;
+          }
+
+          if (cellSorted(start_index).cell_idx_1d_ == cell_idx_1d) {
+            if (start_index == 0) {
+              break;
+            } else if (cellSorted(start_index - 1).cell_idx_1d_ != cell_idx_1d) {
+              break;
+            } else {
+              start_index--;
+              continue;
+            }
+          }
+
+          if (cellSorted(start_index).cell_idx_1d_ >= cell_idx_1d) {
+            start_index--;
+            if (cellSorted(start_index).cell_idx_1d_ < cell_idx_1d) {
+              start_index = -1;
+              break;
+            }
+            continue;
+          }
+          if (cellSorted(start_index).cell_idx_1d_ < cell_idx_1d) {
+            start_index++;
+            if (cellSorted(start_index).cell_idx_1d_ > cell_idx_1d) {
+              start_index = -1;
+              break;
+            }
+            continue;
+          }
+        }
+        cellSortedBegin(k, j, i) = start_index;
+        if (start_index == -1) {
+          cellSortedNumber(k, j, i) = 0;
+        } else {
+          int number = 0;
+          int current_index = start_index;
+          while (current_index <= max_active_index &&
+                 cellSorted(current_index).cell_idx_1d_ == cell_idx_1d) {
+            current_index++;
+            number++;
+            cellSortedNumber(k, j, i) = number;
+          }
+        }
+      });
+}
+
 ///
 /// Routine for precomputing neighbor indices to efficiently compute particle
 /// position in terms of neighbor blocks based on spatial position. See

src/interface/swarm.hpp

@@ -161,6 +161,10 @@ class Swarm {
   /// memory
   void Defrag();
 
+  /// Sort particle list by cell each particle belongs to, according to 1D cell
+  /// index (i + nx*(j + ny*k))
+  void SortParticlesByCell();
+
   // used in case of swarm boundary communication
   void SetupPersistentMPI();
   std::shared_ptr<BoundarySwarm> vbswarm;
@@ -248,6 +252,13 @@ class Swarm {
   int total_received_particles_;
 
   ParArrayND<int> neighbor_buffer_index_; // Map from neighbor index to neighbor bufid
+
+  ParArray1D<SwarmKey>
+      cellSorted_; // 1D per-cell sorted array of key-value swarm memory indices
+
+  ParArrayND<int> cellSortedBegin_; // Per-cell array of starting indices in cell_sorted_
+
+  ParArrayND<int> cellSortedNumber_; // Per-cell array of number of particles in each cell
 };
 
 template <class T>

src/interface/swarm_container.cpp

@@ -80,7 +80,6 @@ TaskStatus SwarmContainer::Defrag(double min_occupancy) {
                            "Max fractional occupancy of swarm must be >= 0 and <= 1");
 
   for (auto &s : swarmVector_) {
-    s->SetupPersistentMPI();
     if (s->GetNumActive() > 0 &&
         s->GetNumActive() / (s->GetMaxActiveIndex() + 1.0) < min_occupancy) {
       s->Defrag();
@@ -92,6 +91,18 @@ TaskStatus SwarmContainer::Defrag(double min_occupancy) {
   return TaskStatus::complete;
 }
 
+TaskStatus SwarmContainer::SortParticlesByCell() {
+  Kokkos::Profiling::pushRegion("Task_SwarmContainer_SortParticlesByCell");
+
+  for (auto &s : swarmVector_) {
+    s->SortParticlesByCell();
+  }
+
+  Kokkos::Profiling::popRegion();
+
+  return TaskStatus::complete;
+}
+
 void SwarmContainer::SendBoundaryBuffers() {}
 
 void SwarmContainer::SetupPersistentMPI() {

src/interface/swarm_container.hpp

@@ -136,6 +136,9 @@ class SwarmContainer {
   // Defragmentation task
   TaskStatus Defrag(double min_occupancy);
 
+  // Sort-by-cell task
+  TaskStatus SortParticlesByCell();
+
   // Communication routines
   void SetupPersistentMPI();
   [[deprecated("Not yet implemented")]] void SetBoundaries();