Fix swarm interface

parthenon-hpc-lab · Mar 6, 2024 · 1465cd4 · 1465cd4
1 parent b993570
commit 1465cd4
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 25 deletions.
diff --git a/external/parthenon b/external/parthenon
diff --git a/src/pgen/turbulence.cpp b/src/pgen/turbulence.cpp
@@ -354,11 +354,12 @@ void ProblemGenerator(Mesh *pmesh, ParameterInput *pin, MeshData<Real> *md) {
 
       // as for num_tracers on each block... will get too many on multiple blocks
       // TODO: distribute amongst blocks.
-      const auto num_tracers_total = tracer_pkg->Param<int>("num_tracers");
-      const int number_block = num_tracers_total;
+      const auto num_tracers_per_cell = tracer_pkg->Param<Real>("num_tracers_per_cell");
+      const auto num_tracers_per_block =
+          static_cast<int>(pmb->GetNumberOfMeshBlockCells() * num_tracers_per_cell);
 
-      ParArrayND<int> new_indices;
-      swarm->AddEmptyParticles(number_block, new_indices);
+      // Create new particles and get accessor
+      auto new_particles_context = swarm->AddEmptyParticles(num_tracers_per_block);
 
       auto &x = swarm->Get<Real>("x").Get();
       auto &y = swarm->Get<Real>("y").Get();
@@ -368,22 +369,26 @@ void ProblemGenerator(Mesh *pmesh, ParameterInput *pin, MeshData<Real> *md) {
       auto swarm_d = swarm->GetDeviceContext();
 
       const int gid = pmb->gid;
-      const int max_active_index = swarm->GetMaxActiveIndex();
       pmb->par_for(
-          "ProblemGenerator::Turbulence::DistributeTracers", 0, max_active_index,
-          KOKKOS_LAMBDA(const int n) {
-            if (swarm_d.IsActive(n)) {
-              auto rng_gen = rng_pool.get_state();
-
-              x(n) = x_min + rng_gen.drand() * (x_max - x_min);
-              y(n) = y_min + rng_gen.drand() * (y_max - y_min);
-              z(n) = z_min + rng_gen.drand() * (z_max - z_min);
-              id(n) = num_tracers_total * gid + n;
-
-              bool on_current_mesh_block = true;
-              swarm_d.GetNeighborBlockIndex(n, x(n), y(n), z(n), on_current_mesh_block);
-              rng_pool.free_state(rng_gen);
-            }
+          "ProblemGenerator::Turbulence::DistributeTracers", 0,
+          new_particles_context.GetNewParticlesMaxIndex(),
+          KOKKOS_LAMBDA(const int new_n) {
+            auto rng_gen = rng_pool.get_state();
+            const int n = new_particles_context.GetNewParticleIndex(new_n);
+
+            x(n) = x_min + rng_gen.drand() * (x_max - x_min);
+            y(n) = y_min + rng_gen.drand() * (y_max - y_min);
+            z(n) = z_min + rng_gen.drand() * (z_max - z_min);
+            // Note that his works only during one time init.
+            // If (somehwere else) we eventually add dynamic particles, then we need to
+            // manage ids (not indices) more globally.
+            id(n) = num_tracers_per_block * gid + n;
+
+            rng_pool.free_state(rng_gen);
+
+            // TODO(pgrete) check if this actually required
+            bool on_current_mesh_block = true;
+            swarm_d.GetNeighborBlockIndex(n, x(n), y(n), z(n), on_current_mesh_block);
           });
     }
   }

diff --git a/src/tracers/tracers.cpp b/src/tracers/tracers.cpp
@@ -33,8 +33,11 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
 
   Params &params = tracer_pkg->AllParams();
 
-  const int num_tracers = pin->GetOrAddInteger("tracers", "num_tracers", 0);
-  params.Add("num_tracers", num_tracers);
+  const auto num_tracers_per_cell =
+      pin->GetOrAddReal("tracers", "num_tracers_per_cell", 0.0);
+  PARTHENON_REQUIRE_THROWS(num_tracers_per_cell >= 0.0,
+                           "What's a negative number of particles per cell?!");
+  params.Add("num_tracers_per_cell", num_tracers_per_cell);
 
   // Initialize random number generator pool
   int rng_seed = pin->GetOrAddInteger("tracers", "rng_seed", time(NULL));
@@ -57,12 +60,15 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   tracer_pkg->AddSwarmValue("vel_y", swarm_name, real_swarmvalue_metadata);
   // TODO(pgrete) check proper handling of <3D sims
   tracer_pkg->AddSwarmValue("vel_z", swarm_name, real_swarmvalue_metadata);
+  PARTHENON_REQUIRE_THROWS(pin->GetInteger("parthenon/mesh", "nx3") > 1,
+                           "Tracers/swarms currently only supported/tested in 3D.");
 
   // TODO(pgrete) this should be safe because we call this package init after the hydro
   // one, but we should check if there's direct way to access Params of other packages.
   const bool mhd = pin->GetString("hydro", "fluid") == "glmmhd";
 
-  PARTHENON_REQUIRE_THROWS(pin->GetString("parthenon/mesh", "refinement") == "none",
+  PARTHENON_REQUIRE_THROWS(pin->GetOrAddString("parthenon/mesh", "refinement", "none") ==
+                               "none",
                            "Tracers/swarms currently only supported on uniform meshes.");
 
   if (mhd) {
@@ -79,8 +85,6 @@ TaskStatus AdvectTracers(MeshBlockData<Real> *mbd, const Real dt) {
   auto &sd = pmb->swarm_data.Get();
   auto &swarm = sd->Get("tracers");
 
-  const auto ndim = pmb->pmy_mesh->ndim;
-
   auto &x = swarm->Get<Real>("x").Get();
   auto &y = swarm->Get<Real>("y").Get();
   auto &z = swarm->Get<Real>("z").Get();