WIP tracers

src/CMakeLists.txt

@@ -18,6 +18,8 @@ add_executable(
         hydro/srcterms/tabular_cooling.cpp
         refinement/gradient.cpp
         refinement/other.cpp
+        tracers/tracers.cpp
+        tracers/tracers.hpp
         utils/few_modes_ft.cpp
 )
 

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

@@ -19,6 +19,7 @@
 #include "../eos/adiabatic_hydro.hpp"
 #include "../pgen/cluster/agn_triggering.hpp"
 #include "../pgen/cluster/magnetic_tower.hpp"
+#include "../tracers/tracers.hpp"
 #include "glmmhd/glmmhd.hpp"
 #include "hydro.hpp"
 #include "hydro_driver.hpp"
@@ -382,6 +383,36 @@ TaskCollection HydroDriver::MakeTaskCollection(BlockList_t &blocks, int stage) {
           fill_derived, parthenon::Update::EstimateTimestep<MeshData<Real>>, mu0.get());
     }
   }
+  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->swarm_data.Get();
+        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->swarm_data.Get();
+      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);
+    }
+  }
 
   if (stage == integrator->nstages && pmesh->adaptive) {
     TaskRegion &async_region_4 = tc.AddRegion(num_task_lists_executed_independently);

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,59 @@ 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(0)->GetBlockPointer();
+      auto &sd = pmb->swarm_data.Get();
+      auto &swarm = pmb->swarm_data.Get()->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_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 &x = swarm->Get<Real>("x").Get();
+      auto &y = swarm->Get<Real>("y").Get();
+      auto &z = swarm->Get<Real>("z").Get();
+      auto &id = swarm->Get<int>("id").Get();
+
+      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);
+            }
+          });
+    }
+  }
 }
 
 //----------------------------------------------------------------------------------------
@@ -476,6 +530,13 @@ void UserWorkBeforeOutput(MeshBlock *pmb, ParameterInput *pin) {
   // store state of distribution
   auto state_dist = few_modes_ft.GetDistState();
   pin->SetString("problem/turbulence", "state_dist", state_dist);
+
+  // TODO(pgrete) move this to a more central location in the driver
+  auto tracer_pkg = pmb->packages.Get("tracers");
+  if (tracer_pkg->Param<bool>("enabled")) {
+    auto &mbd = pmb->meshblock_data.Get();
+    Tracers::FillTracers(mbd.get());
+  }
 }
 
 } // namespace turbulence

src/tracers/tracers.cpp

@@ -0,0 +1,186 @@
+//========================================================================================
+// AthenaPK - a performance portable block structured AMR astrophysical MHD code.
+// Copyright (c) 2024, Athena-Parthenon Collaboration. All rights reserved.
+// Licensed under the BSD 3-Clause License (the "LICENSE").
+//========================================================================================
+// Tracer implementation refacored from https://github.com/lanl/phoebus
+//========================================================================================
+// © 2021-2023. 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 for the U.S.
+// Department of Energy/National Nuclear Security Administration. All
+// rights in the program are reserved by Triad National Security, LLC,
+// and the U.S. Department of Energy/National Nuclear Security
+// Administration. The Government is granted for itself and others
+// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide
+// 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.
+
+#include "tracers.hpp"
+#include "../main.hpp"
+#include "utils/error_checking.hpp"
+
+namespace Tracers {
+using namespace parthenon::package::prelude;
+
+std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
+  auto tracer_pkg = std::make_shared<StateDescriptor>("tracers");
+  const bool enabled = pin->GetOrAddBoolean("tracers", "enabled", false);
+  tracer_pkg->AddParam<bool>("enabled", enabled);
+  if (!enabled) return tracer_pkg;
+
+  Params &params = tracer_pkg->AllParams();
+
+  const int num_tracers = pin->GetOrAddInteger("tracers", "num_tracers", 0);
+  params.Add("num_tracers", num_tracers);
+
+  // Initialize random number generator pool
+  int rng_seed = pin->GetOrAddInteger("tracers", "rng_seed", time(NULL));
+  tracer_pkg->AddParam<>("rng_seed", rng_seed);
+  RNGPool rng_pool(rng_seed);
+  tracer_pkg->AddParam<>("rng_pool", rng_pool);
+
+  // Add swarm of tracers
+  std::string swarm_name = "tracers";
+  Metadata swarm_metadata({Metadata::Provides, Metadata::None});
+  tracer_pkg->AddSwarm(swarm_name, swarm_metadata);
+  Metadata real_swarmvalue_metadata({Metadata::Real});
+  tracer_pkg->AddSwarmValue("id", swarm_name, Metadata({Metadata::Integer}));
+
+  // TODO(pgrete) Add CheckDesired/required for vars
+  // thermo variables
+  tracer_pkg->AddSwarmValue("rho", swarm_name, real_swarmvalue_metadata);
+  tracer_pkg->AddSwarmValue("pressure", swarm_name, real_swarmvalue_metadata);
+  tracer_pkg->AddSwarmValue("vel_x", swarm_name, real_swarmvalue_metadata);
+  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);
+
+  // 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",
+                           "Tracers/swarms currently only supported on uniform meshes.");
+
+  if (mhd) {
+    tracer_pkg->AddSwarmValue("B_x", swarm_name, real_swarmvalue_metadata);
+    tracer_pkg->AddSwarmValue("B_y", swarm_name, real_swarmvalue_metadata);
+    tracer_pkg->AddSwarmValue("B_z", swarm_name, real_swarmvalue_metadata);
+  }
+
+  return tracer_pkg;
+} // Initialize
+
+TaskStatus AdvectTracers(MeshBlockData<Real> *mbd, const Real dt) {
+  auto *pmb = mbd->GetParentPointer();
+  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();
+
+  auto swarm_d = swarm->GetDeviceContext();
+
+  const auto &prim_pack = mbd->PackVariables(std::vector<std::string>{"prim"});
+
+  // update loop. RK2
+  const int max_active_index = swarm->GetMaxActiveIndex();
+  pmb->par_for(
+      "Advect Tracers", 0, max_active_index, KOKKOS_LAMBDA(const int n) {
+        if (swarm_d.IsActive(n)) {
+          int k, j, i;
+          swarm_d.Xtoijk(x(n), y(n), z(n), i, j, k);
+
+          // Predictor/corrector will first make sense with non constant interpolation
+          // TODO(pgrete) add non-cnonst interpolation
+          // predictor
+          // const Real kx = x(n) + 0.5 * dt * rhs1;
+          // const Real ky = y(n) + 0.5 * dt * rhs2;
+          // const Real kz = z(n) + 0.5 * dt * rhs3;
+
+          // corrector
+          x(n) += prim_pack(IV1, k, j, i) * dt;
+          y(n) += prim_pack(IV2, k, j, i) * dt;
+          z(n) += prim_pack(IV3, k, j, i) * dt;
+
+          bool unused_temp = true;
+          swarm_d.GetNeighborBlockIndex(n, x(n), y(n), z(n), unused_temp);
+        }
+      });
+
+  return TaskStatus::complete;
+} // AdvectTracers
+
+/**
+ * FillDerived function for tracers.
+ * Registered Quantities (in addition to t, x, y, z):
+ * rho, vel, B
+ **/
+void FillTracers(MeshBlockData<Real> *mbd) {
+
+  auto *pmb = mbd->GetParentPointer();
+  auto hydro_pkg = pmb->packages.Get("Hydro");
+  auto &sd = pmb->swarm_data.Get();
+  auto &swarm = sd->Get("tracers");
+
+  const auto mhd = hydro_pkg->Param<Fluid>("fluid") == Fluid::glmmhd;
+
+  // TODO(pgrete) cleanup once get swarm packs (currently in development upstream)
+  // pull swarm vars
+  auto &x = swarm->Get<Real>("x").Get();
+  auto &y = swarm->Get<Real>("y").Get();
+  auto &z = swarm->Get<Real>("z").Get();
+  auto &vel_x = swarm->Get<Real>("vel_x").Get();
+  auto &vel_y = swarm->Get<Real>("vel_y").Get();
+  auto &vel_z = swarm->Get<Real>("vel_z").Get();
+  // Assign some (definitely existing) default var
+  auto B_x = vel_x.Get();
+  auto B_y = vel_x.Get();
+  auto B_z = vel_x.Get();
+  if (mhd) {
+    B_x = swarm->Get<Real>("B_x").Get();
+    B_y = swarm->Get<Real>("B_y").Get();
+    B_z = swarm->Get<Real>("B_z").Get();
+  }
+  auto &rho = swarm->Get<Real>("rho").Get();
+  auto &pressure = swarm->Get<Real>("pressure").Get();
+
+  // Get hydro/mhd fluid vars
+  const auto &prim_pack = mbd->PackVariables(std::vector<std::string>{"prim"});
+
+  auto swarm_d = swarm->GetDeviceContext();
+
+  // update loop.
+  const int max_active_index = swarm->GetMaxActiveIndex();
+  pmb->par_for(
+      "Fill Tracers", 0, max_active_index, KOKKOS_LAMBDA(const int n) {
+        if (swarm_d.IsActive(n)) {
+          int k, j, i;
+          swarm_d.Xtoijk(x(n), y(n), z(n), i, j, k);
+
+          // TODO(pgrete) Interpolate
+          rho(n) = prim_pack(IDN, k, j, i);
+          vel_x(n) = prim_pack(IV1, k, j, i);
+          vel_y(n) = prim_pack(IV2, k, j, i);
+          vel_z(n) = prim_pack(IV3, k, j, i);
+          pressure(n) = prim_pack(IPR, k, j, i);
+          if (mhd) {
+            B_x(n) = prim_pack(IB1, k, j, i);
+            B_y(n) = prim_pack(IB2, k, j, i);
+            B_z(n) = prim_pack(IB3, k, j, i);
+          }
+
+          bool unsed_tmp = true;
+          swarm_d.GetNeighborBlockIndex(n, x(n), y(n), z(n), unsed_tmp);
+        }
+      });
+
+} // FillTracers
+
+} // namespace Tracers

src/tracers/tracers.hpp

@@ -0,0 +1,46 @@
+//========================================================================================
+// AthenaPK - a performance portable block structured AMR astrophysical MHD code.
+// Copyright (c) 2024, Athena-Parthenon Collaboration. All rights reserved.
+// Licensed under the BSD 3-Clause License (the "LICENSE").
+//========================================================================================
+// Tracer implementation refacored from https://github.com/lanl/phoebus
+//========================================================================================
+// © 2021-2023. 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 for the U.S.
+// Department of Energy/National Nuclear Security Administration. All
+// rights in the program are reserved by Triad National Security, LLC,
+// and the U.S. Department of Energy/National Nuclear Security
+// Administration. The Government is granted for itself and others
+// acting on its behalf a nonexclusive, paid-up, irrevocable worldwide
+// 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.
+
+#ifndef TRACERS_HPP_
+#define TRACERS_HPP_
+
+#include <memory>
+
+#include "Kokkos_Random.hpp"
+
+#include <parthenon/driver.hpp>
+#include <parthenon/package.hpp>
+
+using namespace parthenon::driver::prelude;
+using namespace parthenon::package::prelude;
+
+using RNGPool = Kokkos::Random_XorShift64_Pool<>;
+
+namespace Tracers {
+
+std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin);
+
+TaskStatus AdvectTracers(MeshBlockData<Real> *mbd, const Real dt);
+
+void FillTracers(MeshBlockData<Real> *rc);
+
+} // namespace Tracers
+
+#endif // TRACERS_HPP_