Add support for uint64 swarm vars and add default ids

CHANGELOG.md

@@ -19,6 +19,7 @@
 - [[PR 1161]](https://github.com/parthenon-hpc-lab/parthenon/pull/1161) Make flux field Metadata accessible, add Metadata::CellMemAligned flag, small perfomance upgrades
 
 ### Changed (changing behavior/API/variables/...)
+- [[PR 1253]](https://github.com/parthenon-hpc-lab/parthenon/pull/1253) Add support for uint64 swarm variables and enforce id
 - [[PR 1243]](https://github.com/parthenon-hpc-lab/parthenon/pull/1243) Increase MPI timeout counter by 10x
 - [[PR 1223]](https://github.com/parthenon-hpc-lab/parthenon/pull/1223) Update to v4 artifacts
 - [[PR 1216]](https://github.com/parthenon-hpc-lab/parthenon/pull/1216) Move to Kokkos 4.5

doc/sphinx/src/particles.rst

@@ -13,9 +13,9 @@ A ``Swarm`` contains all the particle data for all particles of a given
 species. It owns a set of ``ParticleVariable``\ s, one for each value of
 each particle. For example, the spatial positions ``x``, ``y``, and
 ``z`` of the particles in a swarm are three separate
-``ParticleVariable``\ s. ``ParticleVariable``\ s can be either ``Real``-
-or ``int``-valued, which is specified by the metadata values
-``Metadata::Real`` and ``Metadata::Integer``. ``ParticleVariable``\ s
+``ParticleVariable``\ s. ``ParticleVariable``\ s can be either ``Real``-,
+ ``uint64_t``-, or ``int``-valued, which is specified by the metadata values
+``Metadata::Real``, ``MetaData::UInt64``, and ``Metadata::Integer``. ``ParticleVariable``\ s
 should also contain the ``Metadata::Particle`` flag. By default,
 ``ParticleVariable``\ s provide one scalar quantity per particle, but up
 to 2D data per particle is currently supported, by passing

example/particle_leapfrog/particle_leapfrog.cpp

@@ -65,7 +65,6 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   std::string swarm_name = "my_particles";
   Metadata swarm_metadata({Metadata::Provides, Metadata::None, Metadata::Independent});
   pkg->AddSwarm(swarm_name, swarm_metadata);
-  pkg->AddSwarmValue("id", swarm_name, Metadata({Metadata::Integer}));
   Metadata vreal_swarmvalue_metadata({Metadata::Real, Metadata::Vector},
                                      std::vector<int>{3});
   pkg->AddSwarmValue("v", swarm_name, vreal_swarmvalue_metadata);
@@ -177,7 +176,7 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
 
   auto new_particles_context = swarm->AddEmptyParticles(num_particles_this_block);
 
-  auto &id = swarm->Get<int>("id").Get();
+  auto &id = swarm->Get<std::uint64_t>(swarm_position::id::name()).Get();
   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();
@@ -224,16 +223,9 @@ TaskStatus TransportParticles(MeshData<Real> *md, const StagedIntegrator *integr
           swarm_name);
   auto pack_pos = desc_pos.GetPack(md);
 
-  // Make a SwarmPack via strings to get ids
-  // NOTE(@pdmullen): since we are constructing the pack via strings, we must specify
-  // the datatype associated with ids (i.e., int).  We also extract an indexing map.
-  std::vector<std::string> vars_id{"id"};
-  static auto desc_id = MakeSwarmPackDescriptor<int>(swarm_name, vars_id);
-  auto pack_id = desc_id.GetPack(md);
-  auto pack_id_map = desc_id.GetMap();
-  parthenon::PackIdx spi_id(pack_id_map["id"]);
-
   // Make a SwarmPack via strings to get v (note that v is a vector!)
+  // NOTE(@pdmullen): since we are constructing the pack via strings, we must specify
+  // the datatype associated with ids (i.e., Real).  We also extract an indexing map.
   std::vector<std::string> vars_v{"v"};
   static auto desc_v = MakeSwarmPackDescriptor<Real>(swarm_name, vars_v);
   auto pack_v = desc_v.GetPack(md);
@@ -244,7 +236,6 @@ TaskStatus TransportParticles(MeshData<Real> *md, const StagedIntegrator *integr
       DEFAULT_LOOP_PATTERN, "TestSwarmPack", DevExecSpace(), 0,
       pack_pos.GetMaxFlatIndex(), KOKKOS_LAMBDA(const int idx) {
         auto [b, n] = pack_pos.GetBlockParticleIndices(idx);
-        const int iid = pack_id.GetLowerBound(b, spi_id);
         const int iv = pack_v.GetLowerBound(b, spi_v);
         const auto swarm_d = pack_pos.GetContext(b);
         if (swarm_d.IsActive(n)) {

example/particle_tracers/particle_tracers.cpp

@@ -32,6 +32,7 @@
 #include "globals.hpp"
 #include "interface/update.hpp"
 #include "kokkos_abstraction.hpp"
+#include "pack/swarm_default_names.hpp"
 #include "prolong_restrict/prolong_restrict.hpp"
 
 using namespace parthenon::driver::prelude;
@@ -149,7 +150,6 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   Metadata swarm_metadata({Metadata::Provides, Metadata::None});
   pkg->AddSwarm(swarm_name, swarm_metadata);
   Metadata real_swarmvalue_metadata({Metadata::Real});
-  pkg->AddSwarmValue("id", swarm_name, Metadata({Metadata::Integer}));
 
   pkg->EstimateTimestepBlock = EstimateTimestepBlock;
 
@@ -372,7 +372,7 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
   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 &id = swarm->Get<std::uint64_t>(swarm_position::id::name()).Get();
 
   auto swarm_d = swarm->GetDeviceContext();
   pmb->par_for(

scripts/python/packages/parthenon_tools/parthenon_tools/phdf.py

@@ -57,6 +57,10 @@ def Get(self, variable):
             print(f"ERROR: {variable} not found in {self.name}")
             return None
 
+    @property
+    def id(self):
+        return self.Get("swarm.id")
+
     @property
     def x(self):
         return self.Get("swarm.x")

src/interface/metadata.hpp

@@ -80,6 +80,8 @@
   PARTHENON_INTERNAL_FOR_FLAG(Boolean)                                                   \
   /** Integer-valued quantity */                                                         \
   PARTHENON_INTERNAL_FOR_FLAG(Integer)                                                   \
+  /** uint64-t-valued quantity */                                                        \
+  PARTHENON_INTERNAL_FOR_FLAG(UInt64)                                                    \
   /** Real-valued quantity */                                                            \
   PARTHENON_INTERNAL_FOR_FLAG(Real)                                                      \
   /************************************************/                                     \
@@ -448,6 +450,8 @@ class Metadata {
       return Boolean;
     } else if (IsSet(Integer)) {
       return Integer;
+    } else if (IsSet(UInt64)) {
+      return UInt64;
     } else if (IsSet(Real)) {
       return Real;
     }

src/interface/swarm.cpp

@@ -11,6 +11,7 @@
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
 #include <algorithm>
+#include <cstdint>
 #include <cstdlib>
 #include <limits>
 #include <memory>
@@ -84,6 +85,7 @@ Swarm::Swarm(const std::string &label, const Metadata &metadata, const int nmax_
   uid_ = get_uid_(label_);
 
   // Add default swarm fields
+  Add(swarm_position::id::name(), Metadata({Metadata::UInt64}));
   Add(swarm_position::x::name(), Metadata({Metadata::Real}));
   Add(swarm_position::y::name(), Metadata({Metadata::Real}));
   Add(swarm_position::z::name(), Metadata({Metadata::Real}));
@@ -118,7 +120,8 @@ void Swarm::Add(const std::string &label, const Metadata &metadata) {
   // labels must be unique, even between different types of data
   //  if (intMap_.count(label) > 0 || realMap_.count(label) > 0) {
   if (std::get<getType<int>()>(maps_).count(label) > 0 ||
-      std::get<getType<Real>()>(maps_).count(label) > 0) {
+      std::get<getType<Real>()>(maps_).count(label) > 0 ||
+      std::get<getType<std::uint64_t>()>(maps_).count(label) > 0) {
     throw std::invalid_argument("swarm variable " + label +
                                 " already enrolled during Add()!");
   }
@@ -128,6 +131,8 @@ void Swarm::Add(const std::string &label, const Metadata &metadata) {
 
   if (newm.Type() == Metadata::Integer) {
     Add_<int>(label, newm);
+  } else if (newm.Type() == Metadata::UInt64) {
+    Add_<std::uint64_t>(label, newm);
   } else if (newm.Type() == Metadata::Real) {
     Add_<Real>(label, newm);
   } else {
@@ -145,6 +150,8 @@ void Swarm::Remove(const std::string &label) {
 
   auto &int_map = std::get<getType<int>()>(maps_);
   auto &int_vector = std::get<getType<int>()>(vectors_);
+  auto &uint64_map = std::get<getType<std::uint64_t>()>(maps_);
+  auto &uint64_vector = std::get<getType<std::uint64_t>()>(vectors_);
   auto &real_map = std::get<getType<Real>()>(maps_);
   auto &real_vector = std::get<getType<Real>()>(vectors_);
 
@@ -157,7 +164,7 @@ void Swarm::Remove(const std::string &label) {
     }
     idx++;
   }
-  if (found == true) {
+  if (found) {
     // first delete the variable
     int_vector[idx].reset();
 
@@ -167,28 +174,44 @@ void Swarm::Remove(const std::string &label) {
 
     // Also remove variable from map
     int_map.erase(label);
+    return;
   }
 
-  if (found == false) {
-    idx = 0;
-    for (const auto &v : real_vector) {
-      if (label == v->label()) {
-        found = true;
-        break;
-      }
-      idx++;
+  // search next variable type (real)
+  idx = 0;
+  for (const auto &v : real_vector) {
+    if (label == v->label()) {
+      found = true;
+      break;
     }
+    idx++;
   }
-  if (found == true) {
+  if (found) {
     real_vector[idx].reset();
     if (real_vector.size() > 1) real_vector[idx] = std::move(real_vector.back());
     real_vector.pop_back();
     real_map.erase(label);
+    return;
   }
 
-  if (found == false) {
-    throw std::invalid_argument("swarm variable not found in Remove()");
+  // search next variable type (uint64_t)
+  idx = 0;
+  for (const auto &v : uint64_vector) {
+    if (label == v->label()) {
+      found = true;
+      break;
+    }
+    idx++;
+  }
+  if (found) {
+    uint64_vector[idx].reset();
+    if (uint64_vector.size() > 1) uint64_vector[idx] = std::move(uint64_vector.back());
+    uint64_vector.pop_back();
+    uint64_map.erase(label);
+    return;
   }
+
+  throw std::invalid_argument("swarm variable not found in Remove()");
 }
 
 void Swarm::SetPoolMax(const std::int64_t nmax_pool) {
@@ -218,6 +241,7 @@ void Swarm::SetPoolMax(const std::int64_t nmax_pool) {
   pmb->LogMemUsage(n_new * sizeof(int));
 
   auto &int_vector = std::get<getType<int>()>(vectors_);
+  auto &uint64_vector = std::get<getType<std::uint64_t>()>(vectors_);
   auto &real_vector = std::get<getType<Real>()>(vectors_);
 
   for (auto &d : int_vector) {
@@ -226,6 +250,12 @@ void Swarm::SetPoolMax(const std::int64_t nmax_pool) {
     pmb->LogMemUsage(n_new * sizeof(int));
   }
 
+  for (auto &d : uint64_vector) {
+    d->data.Resize(d->data.GetDim(6), d->data.GetDim(5), d->data.GetDim(4),
+                   d->data.GetDim(3), d->data.GetDim(2), nmax_pool);
+    pmb->LogMemUsage(n_new * sizeof(std::uint64_t));
+  }
+
   for (auto &d : real_vector) {
     d->data.Resize(d->data.GetDim(6), d->data.GetDim(5), d->data.GetDim(4),
                    d->data.GetDim(3), d->data.GetDim(2), nmax_pool);
@@ -399,17 +429,23 @@ void Swarm::Defrag() {
 
   // Get all dynamical variables in swarm
   auto &int_vector = std::get<getType<int>()>(vectors_);
+  auto &uint64_vector = std::get<getType<std::uint64_t>()>(vectors_);
   auto &real_vector = std::get<getType<Real>()>(vectors_);
   PackIndexMap real_imap;
   PackIndexMap int_imap;
+  PackIndexMap uint64_imap;
   auto vreal = PackAllVariables_<Real>(real_imap);
   auto vint = PackAllVariables_<int>(int_imap);
+  auto vuint64 = PackAllVariables_<std::uint64_t>(uint64_imap);
   int real_vars_size = real_vector.size();
   int int_vars_size = int_vector.size();
+  int uint64_vars_size = uint64_vector.size();
   auto real_map = real_imap.Map();
   auto int_map = int_imap.Map();
+  auto uint64_map = uint64_imap.Map();
   const int realPackDim = vreal.GetDim(2);
   const int intPackDim = vint.GetDim(2);
+  const int uint64PackDim = vuint64.GetDim(2);
 
   // Loop over only the active number of particles, and if mask is empty, copy in particle
   // using address from prefix sum
@@ -423,6 +459,9 @@ void Swarm::Defrag() {
           for (int vidx = 0; vidx < intPackDim; vidx++) {
             vint(vidx, n) = vint(vidx, nread);
           }
+          for (int vidx = 0; vidx < uint64PackDim; vidx++) {
+            vuint64(vidx, n) = vuint64(vidx, nread);
+          }
           mask(n) = true;
         }
       });

src/interface/swarm.hpp

@@ -79,11 +79,14 @@ class Swarm {
  private:
   static const int IntVec = 0;
   static const int RealVec = 1;
+  static const int UInt64Vec = 2;
 
   template <class T>
   static constexpr int getType() {
     if (std::is_same<T, int>::value) {
       return IntVec;
+    } else if (std::is_same<T, std::uint64_t>::value) {
+      return UInt64Vec;
     }
     return RealVec;
   }
@@ -214,6 +217,9 @@ class Swarm {
     for (auto &v : std::get<1>(vectors_)) {
       size += v->NumComponents();
     }
+    for (auto &v : std::get<2>(vectors_)) {
+      size += v->NumComponents();
+    }
 
     return size;
   }
@@ -268,9 +274,11 @@ class Swarm {
   Metadata m_;
   int nmax_pool_;
   std::string info_;
-  std::tuple<ParticleVariableVector<int>, ParticleVariableVector<Real>> vectors_;
+  std::tuple<ParticleVariableVector<int>, ParticleVariableVector<Real>,
+             ParticleVariableVector<std::uint64_t>>
+      vectors_;
 
-  std::tuple<MapToParticle<int>, MapToParticle<Real>> maps_;
+  std::tuple<MapToParticle<int>, MapToParticle<Real>, MapToParticle<uint64_t>> maps_;
 
   ParArray1D<bool> mask_;
   ParArray1D<bool> marked_for_removal_;

src/interface/swarm_comms.cpp

@@ -166,13 +166,17 @@ void Swarm::LoadBuffers_() {
   pmb->exec_space.fence();
 
   auto &int_vector = std::get<getType<int>()>(vectors_);
+  auto &uint64_vector = std::get<getType<std::uint64_t>()>(vectors_);
   auto &real_vector = std::get<getType<Real>()>(vectors_);
   PackIndexMap real_imap;
   PackIndexMap int_imap;
+  PackIndexMap uint64_imap;
   auto vreal = PackAllVariables_<Real>(real_imap);
   auto vint = PackAllVariables_<int>(int_imap);
+  auto vuint64 = PackAllVariables_<std::uint64_t>(uint64_imap);
   const int realPackDim = vreal.GetDim(2);
   const int intPackDim = vint.GetDim(2);
+  const int uint64PackDim = vuint64.GetDim(2);
 
   auto &x = Get<Real>(swarm_position::x::name()).Get();
   auto &y = Get<Real>(swarm_position::y::name()).Get();
@@ -263,6 +267,10 @@ void Swarm::LoadBuffers_() {
                 bdvar.send[bufid](buffer_index) = static_cast<Real>(vint(i, p_index));
                 buffer_index++;
               }
+              for (int i = 0; i < uint64PackDim; i++) {
+                bdvar.send[bufid](buffer_index) = static_cast<Real>(vuint64(i, p_index));
+                buffer_index++;
+              }
             }
           }
         });
@@ -319,13 +327,17 @@ void Swarm::UnloadBuffers_() {
     auto neighbor_buffer_index = neighbor_buffer_index_;
 
     auto &int_vector = std::get<getType<int>()>(vectors_);
+    auto &uint64_vector = std::get<getType<std::uint64_t>()>(vectors_);
     auto &real_vector = std::get<getType<Real>()>(vectors_);
     PackIndexMap real_imap;
     PackIndexMap int_imap;
+    PackIndexMap uint64_imap;
     auto vreal = PackAllVariables_<Real>(real_imap);
     auto vint = PackAllVariables_<int>(int_imap);
-    int realPackDim = vreal.GetDim(2);
-    int intPackDim = vint.GetDim(2);
+    auto vuint64 = PackAllVariables_<std::uint64_t>(uint64_imap);
+    const int realPackDim = vreal.GetDim(2);
+    const int intPackDim = vint.GetDim(2);
+    const int uint64PackDim = vuint64.GetDim(2);
 
     const int particle_size = GetParticleDataSize();
     auto swarm_d = GetDeviceContext();
@@ -340,10 +352,6 @@ void Swarm::UnloadBuffers_() {
     }
     neighbor_received_particles.DeepCopy(neighbor_received_particles_h);
 
-    auto &x = Get<Real>(swarm_position::x::name()).Get();
-    auto &y = Get<Real>(swarm_position::y::name()).Get();
-    auto &z = Get<Real>(swarm_position::z::name()).Get();
-
     pmb->par_for(
         PARTHENON_AUTO_LABEL, 0, newParticlesContext.GetNewParticlesMaxIndex(),
         // n is both new particle index and index over buffer values
@@ -367,6 +375,10 @@ void Swarm::UnloadBuffers_() {
             vint(i, sid) = static_cast<int>(bdvar.recv[nbid](bid));
             bid++;
           }
+          for (int i = 0; i < uint64PackDim; i++) {
+            vuint64(i, sid) = static_cast<std::uint64_t>(bdvar.recv[nbid](bid));
+            bid++;
+          }
         });
   }
 }