Merge pull request #202 from lanl/blb/swarm_refactor

lanl · Mar 5, 2024 · 553f9e7 · 553f9e7
2 parents 5bc3b2e + e460ec4
commit 553f9e7
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 16 deletions.
diff --git a/external/parthenon b/external/parthenon
diff --git a/src/pgen/advection.cpp b/src/pgen/advection.cpp
@@ -123,8 +123,7 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
     const auto num_tracers_total = tracer_pkg->Param<int>("num_tracers");
     const int number_block = num_tracers_total;
 
-    ParArrayND<int> new_indices;
-    swarm->AddEmptyParticles(number_block, new_indices);
+    auto new_particles_context = swarm->AddEmptyParticles(number_block);
 
     auto &x = swarm->Get<Real>("x").Get();
     auto &y = swarm->Get<Real>("y").Get();
@@ -134,7 +133,7 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
     auto swarm_d = swarm->GetDeviceContext();
 
     const int gid = pmb->gid;
-    const int max_active_index = swarm->GetMaxActiveIndex();
+    const int max_active_index = new_particles_context.GetNewParticlesMaxIndex();
     pmb->par_for(
         "ProblemGenerator::Torus::DistributeTracers", 0, max_active_index,
         KOKKOS_LAMBDA(const int n) {

diff --git a/src/pgen/torus.cpp b/src/pgen/torus.cpp
@@ -665,8 +665,7 @@ void PostInitializationModifier(ParameterInput *pin, Mesh *pmesh) {
 
       const int num_tracers = std::round(num_tracers_total * number_block / number_mesh);
 
-      ParArrayND<int> new_indices;
-      swarm->AddEmptyParticles(num_tracers, new_indices);
+      auto new_particles_context = swarm->AddEmptyParticles(num_tracers);
 
       auto &x = swarm->Get<Real>("x").Get();
       auto &y = swarm->Get<Real>("y").Get();
@@ -678,7 +677,7 @@ void PostInitializationModifier(ParameterInput *pin, Mesh *pmesh) {
       auto swarm_d = swarm->GetDeviceContext();
 
       const int gid = pmb->gid;
-      const int max_active_index = swarm->GetMaxActiveIndex();
+      const int max_active_index = new_particles_context.GetNewParticlesMaxIndex();
       pmb->par_for(
           "ProblemGenerator::Torus::DistributeTracers", 0, max_active_index,
           KOKKOS_LAMBDA(const int n) {

diff --git a/src/radiation/mocmc.cpp b/src/radiation/mocmc.cpp
@@ -128,8 +128,7 @@ void MOCMCInitSamples(T *rc) {
       },
       Kokkos::Sum<int>(nsamp_tot));
 
-  ParArrayND<int> new_indices;
-  auto new_mask = swarm->AddEmptyParticles(nsamp_tot, new_indices);
+  auto new_particles_context = swarm->AddEmptyParticles(nsamp_tot);
 
   // Calculate array of starting index for each zone to compute particles
   ParArrayND<int> starting_index("Starting index", nx_k, nx_j, nx_i);
@@ -185,7 +184,7 @@ void MOCMCInitSamples(T *rc) {
         Geometry::Tetrads tetrads(ucon, trial, cov_g);
 
         for (int nsamp = 0; nsamp < static_cast<int>(v(b, dn, k, j, i)); nsamp++) {
-          const int n = new_indices(start_idx + nsamp);
+          const int n = new_particles_context.GetNewParticleIndex(start_idx + nsamp);
 
           // Create particles at zone centers
           x(n) = minx_i + (i - ib.s + rng_gen.drand()) * dx_i;

diff --git a/src/radiation/monte_carlo.cpp b/src/radiation/monte_carlo.cpp
@@ -282,8 +282,7 @@ TaskStatus MonteCarloSourceParticles(MeshBlock *pmb, MeshBlockData<Real> *rc,
   const auto num_emitted = rad->Param<Real>("num_emitted");
   rad->UpdateParam<Real>("num_emitted", num_emitted + Nstot);
 
-  ParArrayND<int> new_indices;
-  const auto new_particles_mask = swarm->AddEmptyParticles(Nstot, new_indices);
+  const auto new_particles_context = swarm->AddEmptyParticles(Nstot);
 
   auto &t = swarm->Get<Real>("t").Get();
   auto &x = swarm->Get<Real>("x").Get();
@@ -338,8 +337,8 @@ TaskStatus MonteCarloSourceParticles(MeshBlock *pmb, MeshBlockData<Real> *rc,
 
           // Loop over particles to create in this zone
           for (int n = 0; n < dNs; n++) {
-            const int m =
-                new_indices(starting_index(sidx, k - kb.s, j - jb.s, i - ib.s) + n);
+            const int m = new_particles_context.GetNewParticleIndex(
+                starting_index(sidx, k - kb.s, j - jb.s, i - ib.s) + n);
 
             // Set particle species
             swarm_species(m) = static_cast<int>(s);
@@ -403,8 +402,8 @@ TaskStatus MonteCarloSourceParticles(MeshBlock *pmb, MeshBlockData<Real> *rc,
           int dNs = v(iNs + sidx, k, j, i);
           // Loop over particles to create in this zone
           for (int n = 0; n < static_cast<int>(dNs); n++) {
-            const int m =
-                new_indices(starting_index(sidx, k - kb.s, j - jb.s, i - ib.s) + n);
+            const int m = new_particles_context.GetNewParticleIndex(
+                starting_index(sidx, k - kb.s, j - jb.s, i - ib.s) + n);
             swarm_d.MarkParticleForRemoval(m);
           }
         });