clear particles in pgen

CHANGELOG.md

@@ -29,6 +29,7 @@
 
 
 ### Changed (changing behavior/API/variables/...)
+- [[PR 1441]](https://github.com/parthenon-hpc-lab/parthenon/pull/1441) Make it so you can safely create particles in a problem generator.
 - [[PR 1438]](https://github.com/parthenon-hpc-lab/parthenon/pull/1438) Performance tuning for the loop abstraction machinery and add loop abstraction OpenMP support
 - [[PR 1416][(https://github.com/parthenon-hpc-lab/parthenon/pull/1416) Remove virtual tag from destructors in sparse and swarm pack base classes
 - [[PR 1401]](https://github.com/parthenon-hpc-lab/parthenon/pull/1401) Sparse Field Component Names
@@ -68,6 +69,7 @@
 
 
 ### Incompatibilities (i.e. breaking changes)
+- [[PR 1441]](https://github.com/parthenon-hpc-lab/parthenon/pull/1441) Clear particles in problem generator, meaning that particles can no longer be seeded in the problem generator at the root level and be appropriately refined at pgen.
 - [[PR 1385]](https://github.com/parthenon-hpc-lab/parthenon/pull/1385) ParameterInput internal storage refactor removes direct access to linked list (`pfirst_block`). Use `GetBlocksWithPrefix()` or `GetBlockNames()` instead.
 - [[PR 1351]](https://github.com/parthenon-hpc-lab/parthenon/pull/1351) Bump Kokkos 5 & C++20
 - [[PR 1377]](https://github.com/parthenon-hpc-lab/parthenon/pull/1377) Extend Initialization Hierarchy

example/particle_tracers/parthinput.particle_tracers

@@ -56,6 +56,9 @@ variables = advected, tracer_deposition
 
 <Tracers>
 num_tracers = 10000
+# Source particles in each problem-generator pass. Set false to source only after
+# initialization AMR has resolved.
+source_in_problem_generator = true
 
 <parthenon/output0>
 file_type = hdf5

example/particle_tracers/particle_tracers.cpp

@@ -212,6 +212,9 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   auto tr_pkg = std::make_shared<StateDescriptor>("particles_package");
 
   tr_pkg->AddParam("num_tracers", pin->GetOrAddInteger("Tracers", "num_tracers", 100));
+  const bool source_in_problem_generator =
+      pin->GetOrAddBoolean("Tracers", "source_in_problem_generator", true);
+  tr_pkg->AddParam("source_in_problem_generator", source_in_problem_generator);
 
   // `NoPersistentParticleIds` is just passed to test this aspect in the regression tests.
   // For typical tracers, persistent ids might be important.
@@ -221,8 +224,11 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   // Assign package timestep hook
   tr_pkg->EstimateTimestepMesh = EstimateTimestepMesh;
 
-  // Assign package final initialization hook
-  tr_pkg->PostInitializationBlock = SourceTracers;
+  // Source particles from the ProblemGenerator by default. Alternatively, source them
+  // after initialization AMR has resolved.
+  if (!source_in_problem_generator) {
+    tr_pkg->PostInitializationBlock = SourceTracers;
+  }
 
   return tr_pkg;
 }
@@ -554,10 +560,9 @@ TaskCollection ParticleDriver::StepTasks() {
 }
 
 // *************************************************//
-// Define the ProblemGenerator. Initializing the,  *//
-// advected field.  Recall that initial particle   *//
-// sourcing is handled in FinalInitialization      *//
-// owned by the particles package.                 */
+// Define the ProblemGenerator, initializing the   *//
+// advected field and, when selected, the initial  *//
+// tracer particles.                               *//
 // *************************************************//
 
 void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
@@ -568,6 +573,7 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
 
   // Advection package params
   auto &adv_pkg = pmb->packages.Get("advection_package");
+  auto &tr_pkg = pmb->packages.Get("particles_package");
   const Real &advected_mean = adv_pkg->Param<Real>("advected_mean");
   const Real &advected_amp = adv_pkg->Param<Real>("advected_amp");
 
@@ -594,6 +600,10 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
         pack(0, field::advected(), k, j, i) =
             advected_mean + advected_amp * std::sin(kwave * x1v);
       });
+
+  if (tr_pkg->Param<bool>("source_in_problem_generator")) {
+    particles_package::SourceTracers(pmb, pin);
+  }
 }
 
 } // namespace tracers_example

src/interface/swarm.cpp

@@ -3,7 +3,7 @@
 // Copyright(C) 2020-2026 The Parthenon collaboration
 // Licensed under the 3-clause BSD License, see LICENSE file for details
 //========================================================================================
-// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved.
+// (C) (or copyright) 2020-2026. 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
@@ -14,6 +14,8 @@
 // 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.
 //========================================================================================
+// This file was modified with the assistance of generative AI.
+
 #include <algorithm>
 #include <cstdint>
 #include <cstdlib>
@@ -386,6 +388,15 @@ void Swarm::RemoveMarkedParticles() {
   UpdateEmptyIndices();
 }
 
+void Swarm::ClearParticles() {
+  Kokkos::deep_copy(mask_, false);
+  Kokkos::deep_copy(marked_for_removal_, false);
+  num_active_ = 0;
+  max_active_index_ = inactive_max_active_index;
+  new_indices_max_idx_ = -1;
+  UpdateEmptyIndices();
+}
+
 void Swarm::Defrag() {
   if (GetNumActive() == 0) {
     return;

src/interface/swarm.hpp

@@ -198,6 +198,9 @@ class Swarm {
   /// Remove particles marked for removal and update internal indexing
   void RemoveMarkedParticles();
 
+  /// Mark all particle slots inactive while retaining the allocated pool.
+  void ClearParticles();
+
   /// Open up memory for new empty particles, return a mask to these particles
   NewParticlesContext AddEmptyParticles(const int num_to_add);
 

src/interface/swarm_container.cpp

@@ -10,6 +10,8 @@
 // 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.
 //========================================================================================
+// This file was modified with the assistance of generative AI.
+
 #include <cstdlib>
 #include <iostream>
 #include <memory>
@@ -95,6 +97,12 @@ TaskStatus SwarmContainer::DefragAll() {
   return TaskStatus::complete;
 }
 
+void SwarmContainer::ClearParticles() {
+  for (auto &s : swarmVector_) {
+    s->ClearParticles();
+  }
+}
+
 TaskStatus SwarmContainer::SortParticlesByCell() {
   PARTHENON_INSTRUMENT
 

src/interface/swarm_container.hpp

@@ -10,6 +10,8 @@
 // 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.
 //========================================================================================
+// This file was modified with the assistance of generative AI.
+
 #ifndef INTERFACE_SWARM_CONTAINER_HPP_
 #define INTERFACE_SWARM_CONTAINER_HPP_
 
@@ -140,6 +142,9 @@ class SwarmContainer {
   TaskStatus Defrag(double min_occupancy);
   TaskStatus DefragAll();
 
+  // Remove all particles while retaining each swarm's allocated pool.
+  void ClearParticles();
+
   // Sort-by-cell task
   TaskStatus SortParticlesByCell();
 

src/mesh/mesh.cpp

@@ -770,6 +770,13 @@ void Mesh::Initialize(bool init_problem, ParameterInput *pin, ApplicationInput *
             (nmb != 0 && block_list[0]->PostProblemGenerator != nullptr)),
           "Mesh and MeshBlock PostProblemGenerators are defined. Please use only one.");
 
+      // Problem generation is repeated while initialization AMR resolves. Clear the
+      // particles remeshed from the preceding pass so this pass regenerates them rather
+      // than appending to them. Keep the allocated pools for reuse.
+      for (int i = 0; i < nmb; ++i) {
+        block_list[i]->meshblock_data.Get()->GetSwarmData()->ClearParticles();
+      }
+
       // Call Mesh ProblemGenerator
       if (ProblemGenerator != nullptr) {
         for (auto &partition : GetDefaultBlockPartitions()) {

tst/regression/CMakeLists.txt

@@ -172,7 +172,7 @@ if (ENABLE_HDF5)
   list(APPEND TEST_PROCS ${NUM_MPI_PROC_TESTING})
   list(APPEND TEST_ARGS "--driver ${PROJECT_BINARY_DIR}/example/particle_tracers/particle-tracers \
     --driver_input ${CMAKE_CURRENT_SOURCE_DIR}/test_suites/particle_tracers_amr/parthinput.particle_tracers_amr \
-    --num_steps 2")
+    --num_steps 4")
   list(APPEND EXTRA_TEST_LABELS "")
 
   list(APPEND TEST_DIRS diffusion)

tst/regression/test_suites/particle_tracers_amr/parthinput.particle_tracers_amr

@@ -57,6 +57,8 @@ derefine_tol = 1.05
 
 <Tracers>
 num_tracers = 4096
+# The regression overrides this to false for the PostInitialization coverage runs.
+source_in_problem_generator = true
 
 <parthenon/output0>
 file_type = hdf5

tst/regression/test_suites/particle_tracers_amr/particle_tracers_amr.py

@@ -28,18 +28,28 @@ def sorted_positions(positions):
 
 class TestCase(utils.test_case.TestCaseAbs):
     def Prepare(self, parameters, step):
-        if step == 1:
+        if step in (1, 3):
+            source_location = (
+                "problem_generator" if step == 1 else "post_initialization"
+            )
+            source_in_problem_generator = "true" if step == 1 else "false"
             parameters.driver_cmd_line_args = [
-                "parthenon/job/problem_id=particle_tracers_amr_init",
+                f"parthenon/job/problem_id=particle_tracers_amr_{source_location}_init",
                 "parthenon/mesh/refinement=adaptive",
                 "parthenon/mesh/numlevel=2",
                 "parthenon/time/tlim=0.0",
+                f"Tracers/source_in_problem_generator={source_in_problem_generator}",
             ]
-        elif step == 2:
+        elif step in (2, 4):
+            source_location = (
+                "problem_generator" if step == 2 else "post_initialization"
+            )
+            source_in_problem_generator = "true" if step == 2 else "false"
             parameters.driver_cmd_line_args = [
-                "parthenon/job/problem_id=particle_tracers_amr",
+                f"parthenon/job/problem_id=particle_tracers_amr_{source_location}",
                 "parthenon/mesh/refinement=adaptive",
                 "parthenon/mesh/numlevel=2",
+                f"Tracers/source_in_problem_generator={source_in_problem_generator}",
             ]
         return parameters
 
@@ -51,40 +61,75 @@ def Analyse(self, parameters):
         )
         from phdf import phdf
 
-        initial = phdf("particle_tracers_amr_init.out0.final.phdf")
-        amr = phdf("particle_tracers_amr.out0.final.phdf")
-
-        amr_swarm = amr.GetSwarm("tracers")
-        initial_swarm = initial.GetSwarm("tracers")
-
-        initial_pos = np.vstack(
-            (initial_swarm.x, initial_swarm.y, initial_swarm.z)
-        ).transpose()
-        amr_pos = np.vstack((amr_swarm.x, amr_swarm.y, amr_swarm.z)).transpose()
-
-        initial_pos[:, 0] = ((initial_pos[:, 0] + 0.5 + 0.35) % 1.0) - 0.5
-        initial_pos = sorted_positions(initial_pos)
-        amr_pos = sorted_positions(amr_pos)
-
-        if initial_pos.shape != amr_pos.shape:
-            print("Particle count changed during AMR tracer evolution.")
-            print("initial:", initial_pos.shape, "final:", amr_pos.shape)
-            return False
-
-        if not np.allclose(initial_pos, amr_pos, atol=1.0e-10, rtol=0.0):
-            diff = np.max(np.abs(initial_pos - amr_pos))
-            print("AMR tracer positions differ from the analytic translation.")
-            print("max difference:", diff)
-            return False
-
-        initial_bounds = np.array(initial.BlockBounds)
-        final_bounds = np.array(amr.BlockBounds)
-        mesh_changed = initial.NumBlocks != amr.NumBlocks
-        mesh_changed = mesh_changed or initial_bounds.shape != final_bounds.shape
-        if not mesh_changed:
-            mesh_changed = not np.allclose(initial_bounds, final_bounds)
-        if not mesh_changed:
-            print("AMR mesh did not change between initialization and final output.")
+        initial_positions = {}
+        for source_location in ("problem_generator", "post_initialization"):
+            initial = phdf(
+                f"particle_tracers_amr_{source_location}_init.out0.final.phdf"
+            )
+            amr = phdf(f"particle_tracers_amr_{source_location}.out0.final.phdf")
+
+            initial_swarm = initial.GetSwarm("tracers")
+            amr_swarm = amr.GetSwarm("tracers")
+            initial_pos = np.vstack(
+                (initial_swarm.x, initial_swarm.y, initial_swarm.z)
+            ).transpose()
+            amr_pos = np.vstack((amr_swarm.x, amr_swarm.y, amr_swarm.z)).transpose()
+
+            # SourceTracers rounds each block's share independently, so the sum of
+            # rounded allocations need not equal the requested global count. The 40-block
+            # initialization mesh therefore contains 4104 particles for the requested 4096.
+            expected_num_tracers = 4104
+            if initial_pos.shape[0] != expected_num_tracers:
+                print(
+                    f"Incorrect tracer count after {source_location} initialization AMR."
+                )
+                print(
+                    "expected:", expected_num_tracers, "actual:", initial_pos.shape[0]
+                )
+                return False
+
+            initial_positions[source_location] = sorted_positions(initial_pos.copy())
+            translated_initial_pos = initial_pos.copy()
+            translated_initial_pos[:, 0] = (
+                (translated_initial_pos[:, 0] + 0.5 + 0.35) % 1.0
+            ) - 0.5
+            translated_initial_pos = sorted_positions(translated_initial_pos)
+            amr_pos = sorted_positions(amr_pos)
+
+            if translated_initial_pos.shape != amr_pos.shape:
+                print(
+                    f"Particle count changed during {source_location} AMR tracer evolution."
+                )
+                print("initial:", translated_initial_pos.shape, "final:", amr_pos.shape)
+                return False
+
+            if not np.allclose(translated_initial_pos, amr_pos, atol=1.0e-10, rtol=0.0):
+                diff = np.max(np.abs(translated_initial_pos - amr_pos))
+                print(
+                    f"{source_location} AMR tracer positions differ from the analytic translation."
+                )
+                print("max difference:", diff)
+                return False
+
+            initial_bounds = np.array(initial.BlockBounds)
+            final_bounds = np.array(amr.BlockBounds)
+            mesh_changed = initial.NumBlocks != amr.NumBlocks
+            mesh_changed = mesh_changed or initial_bounds.shape != final_bounds.shape
+            if not mesh_changed:
+                mesh_changed = not np.allclose(initial_bounds, final_bounds)
+            if not mesh_changed:
+                print(f"AMR mesh did not change during the {source_location} run.")
+                return False
+
+        if not np.allclose(
+            initial_positions["problem_generator"],
+            initial_positions["post_initialization"],
+            atol=1.0e-10,
+            rtol=0.0,
+        ):
+            print(
+                "Particle setup differs between ProblemGenerator and PostInitialization."
+            )
             return False
 
         return True