WIP tracers

src/CMakeLists.txt

@@ -21,7 +21,10 @@ add_executable(
         hydro/srcterms/tabular_cooling.cpp
         refinement/gradient.cpp
         refinement/other.cpp
+        tracers/tracers.cpp
+        tracers/tracers.hpp
         utils/few_modes_ft.cpp
+        utils/few_modes_ft.hpp
 )
 
 add_subdirectory(pgen)

src/hydro/hydro.cpp

@@ -24,6 +24,7 @@
 #include "../recon/weno3_simple.hpp"
 #include "../recon/wenoz_simple.hpp"
 #include "../refinement/refinement.hpp"
+#include "../tracers/tracers.hpp"
 #include "../units.hpp"
 #include "defs.hpp"
 #include "diffusion/diffusion.hpp"
@@ -52,6 +53,7 @@ using parthenon::HistoryOutputVar;
 parthenon::Packages_t ProcessPackages(std::unique_ptr<ParameterInput> &pin) {
   parthenon::Packages_t packages;
   packages.Add(Hydro::Initialize(pin.get()));
+  packages.Add(Tracers::Initialize(pin.get()));
   return packages;
 }
 

src/hydro/hydro_driver.cpp

@@ -1,6 +1,6 @@
 //========================================================================================
 // AthenaPK - a performance portable block structured AMR astrophysical MHD code.
-// Copyright (c) 2020-2023, Athena-Parthenon Collaboration. All rights reserved.
+// Copyright (c) 2020-2024, Athena-Parthenon Collaboration. All rights reserved.
 // Licensed under the BSD 3-Clause License (the "LICENSE").
 //========================================================================================
 
@@ -12,13 +12,16 @@
 
 // Parthenon headers
 #include "amr_criteria/refinement_package.hpp"
+#include "basic_types.hpp"
 #include "bvals/comms/bvals_in_one.hpp"
 #include "prolong_restrict/prolong_restrict.hpp"
+#include "utils/error_checking.hpp"
 #include <parthenon/parthenon.hpp>
 // AthenaPK headers
 #include "../eos/adiabatic_hydro.hpp"
 #include "../pgen/cluster/agn_triggering.hpp"
 #include "../pgen/cluster/magnetic_tower.hpp"
+#include "../tracers/tracers.hpp"
 #include "diffusion/diffusion.hpp"
 #include "glmmhd/glmmhd.hpp"
 #include "hydro.hpp"
@@ -694,6 +697,48 @@ TaskCollection HydroDriver::MakeTaskCollection(BlockList_t &blocks, int stage) {
     }
   }
 
+  auto tracers_pkg = pmesh->packages.Get("tracers");
+  // First order operator split tracer advection
+  if (stage == integrator->nstages && tracers_pkg->Param<bool>("enabled")) {
+    const std::string swarm_name = "tracers";
+    TaskRegion &sync_region_tr = tc.AddRegion(1);
+    {
+      for (auto &pmb : blocks) {
+        auto &tl = sync_region_tr[0];
+        auto &sd = pmb->meshblock_data.Get()->GetSwarmData();
+        auto reset_comms =
+            tl.AddTask(none, &SwarmContainer::ResetCommunication, sd.get());
+      }
+    }
+
+    TaskRegion &async_region_tr = tc.AddRegion(blocks.size());
+    for (int n = 0; n < blocks.size(); n++) {
+      auto &tl = async_region_tr[n];
+      auto &pmb = blocks[n];
+      auto &sd = pmb->meshblock_data.Get()->GetSwarmData();
+      auto &mbd0 = pmb->meshblock_data.Get("base");
+      auto tracer_advect =
+          tl.AddTask(none, Tracers::AdvectTracers, mbd0.get(), integrator->dt);
+
+      auto send = tl.AddTask(tracer_advect, &SwarmContainer::Send, sd.get(),
+                             BoundaryCommSubset::all);
+
+      auto receive =
+          tl.AddTask(send, &SwarmContainer::Receive, sd.get(), BoundaryCommSubset::all);
+    }
+    // TODO(pgrete) Fix/cleanup once we got swarm packs.
+    // We need just a single region with a single task in order to be able to use plain
+    // reductions.
+    PARTHENON_REQUIRE_THROWS(num_partitions == 1,
+                             "Only pack_size=-1 currently supported for tracers.")
+    TaskRegion &single_tasklist_per_pack_region_4 = tc.AddRegion(num_partitions);
+    for (int i = 0; i < num_partitions; i++) {
+      auto &tl = single_tasklist_per_pack_region_4[i];
+      auto &mu0 = pmesh->mesh_data.GetOrAdd("base", i);
+      auto fill = tl.AddTask(none, Tracers::FillTracers, mu0.get(), tm);
+    }
+  }
+
   if (stage == integrator->nstages && pmesh->adaptive) {
     TaskRegion &async_region_4 = tc.AddRegion(num_task_lists_executed_independently);
     for (int i = 0; i < blocks.size(); i++) {

src/pgen/turbulence.cpp

@@ -26,6 +26,7 @@
 
 // AthenaPK headers
 #include "../main.hpp"
+#include "../tracers/tracers.hpp"
 #include "../units.hpp"
 #include "../utils/few_modes_ft.hpp"
 
@@ -333,6 +334,63 @@ void ProblemGenerator(Mesh *pmesh, ParameterInput *pin, MeshData<Real> *md) {
               0.5 * (SQR(u(IB1, k, j, i)) + SQR(u(IB2, k, j, i)) + SQR(u(IB3, k, j, i)));
         }
       });
+
+  /* tracer init section */
+  auto tracer_pkg = pmb->packages.Get("tracers");
+  if (tracer_pkg->Param<bool>("enabled")) {
+
+    for (int b = 0; b < md->NumBlocks(); b++) {
+      pmb = md->GetBlockData(b)->GetBlockPointer();
+      auto &swarm = pmb->meshblock_data.Get()->GetSwarmData()->Get("tracers");
+      auto rng_pool = tracer_pkg->Param<RNGPool>("rng_pool");
+
+      const Real &x_min = pmb->coords.Xf<1>(ib.s);
+      const Real &y_min = pmb->coords.Xf<2>(jb.s);
+      const Real &z_min = pmb->coords.Xf<3>(kb.s);
+      const Real &x_max = pmb->coords.Xf<1>(ib.e + 1);
+      const Real &y_max = pmb->coords.Xf<2>(jb.e + 1);
+      const Real &z_max = pmb->coords.Xf<3>(kb.e + 1);
+
+      // as for num_tracers on each block... will get too many on multiple blocks
+      // TODO: distribute amongst blocks.
+      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);
+
+      // Create new particles and get accessor
+      auto new_particles_context = swarm->AddEmptyParticles(num_tracers_per_block);
+
+      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 &id = swarm->Get<int>("id").Get();
+
+      auto swarm_d = swarm->GetDeviceContext();
+
+      const int gid = pmb->gid;
+      pmb->par_for(
+          "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);
+          });
+    }
+  }
 }
 
 //----------------------------------------------------------------------------------------
@@ -478,5 +536,4 @@ void UserWorkBeforeOutput(MeshBlock *pmb, ParameterInput *pin,
   auto state_dist = few_modes_ft.GetDistState();
   pin->SetString("problem/turbulence", "state_dist", state_dist);
 }
-
 } // namespace turbulence