Add support for uint64 swarm vars and add default ids

CHANGELOG.md

@@ -1,10 +1,12 @@
 # Changelog
 
 ## Current develop
+
 ### Added (new features/APIs/variables/...)
 
 
 ### Changed (changing behavior/API/variables/...)
+- [[PR 1253]](https://github.com/parthenon-hpc-lab/parthenon/pull/1253) Add support for uint64 swarm variables and add default id
 
 
 ### Fixed (not changing behavior/API/variables/...)
@@ -17,6 +19,7 @@
 
 
 ### Incompatibilities (i.e. breaking changes)
+- [[PR 1253]](https://github.com/parthenon-hpc-lab/parthenon/pull/1253) Add support for uint64 swarm variables and add default id
 
 
 

doc/sphinx/src/interface/metadata.rst

@@ -218,6 +218,14 @@ requested. ``FluxRequest::Any`` does not modify search parameters. You
 will get flux or non-flux variables, and variable associations will be
 ignored.
 
+``Swarm``
+---------
+By default all partices in a ``Swarm`` contain a persistent, unique identifier
+field. If this behavior is not desired (e.g., because anonymous particles are
+constantly created and destroyed in a given numerical method),
+the ``Metadata::NoPersistentParticleIds`` flag prevent the standard allocation
+(and communication of a ``swarm.id`` field).
+
 Application Metadata Flags
 ---------------------------
 

doc/sphinx/src/particles.rst

@@ -13,16 +13,20 @@ 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
 ``std::vector<int>{N1, N2}`` as the second argument to the
 ``ParticleVariable`` ``Metadata``. All ``Swarm``\ s by default contain
-``x``, ``y``, and ``z`` ``ParticleVariable``\ s; additional fields can
-be added as:
+positional ``x``, ``y``, and ``z`` ``ParticleVariable``\ s and a ``uint64`` particle
+``id`` that is persistent and unique for a given simulation.
+The latter field can be disabled (e.g., because it is not required for a method
+that constant creates and destroys anonymous particles) by providing the
+``Metadata::NoPersistentParticleIds`` flag;
+additional fields can be added as:
 
 .. code:: cpp
 
@@ -256,17 +260,14 @@ per-swarm list will be output for that swarm.
    Some visualization tools, like Visit and Paraview, prefer to have
    access to an ``id`` field for each particle, however it's not clear
    that a unique ID is required for each particle in
-   general. Therefore, swarms do not automatically contain an ID swarm
-   variable. However, when Parthenon outputs a swarm, it automatically
-   generates an ID variable even if one is not present or
-   requested. If a variable named ``id'' is available **and** the user
-   requests it be output, Parthenon will use it. Otherwise, Parthenon
-   will generate an ``id`` variable just for output and write it to
-   file.
+   general. Therefore, if swarms do not automatically contain a unique ID swarm
+   variable (because the ``Metadata::NoPersistentParticleIds`` has been passed when
+   creating the swarm), Parthenon will generate an ``id`` variable just for output
+   (i.e., it is still not available within the simulation itself) and write it to file.
 
 .. warning::
 
-   The automatically generted ``id`` is unique for a snapshot in time,
+   The automatically generated ``id`` is unique for a snapshot in time,
    but not guaranteed to be time invariant. Indeed it is likely
    **not** the same between dumps.
 
@@ -280,11 +281,7 @@ Putting it all together, you might have an output block that looks like this:
    swarms = swarm1, swarm2
    swarm_variables = shared_var
    swarm1_variables = per_swarm_var
-   swarm2_variables = id
 
 The result would be that both ``swarm1`` and ``swarm2`` output the
-variables ``x``, ``y``, ``z``, and ``shared_var``. But only ``swarm1``
-outputs ``per_swarm_var``. Both ``swarm1`` and ``swarm2`` will output
-an ``id`` field. But the ``id`` field for ``swarm1`` will be
-automatically generated, but the ``id`` field for ``swarm2`` will use
-the user-initialized value if such a quantity is available.
+variables ``id``, ``x``, ``y``, ``z``, and ``shared_var``. But only ``swarm1``
+outputs ``per_swarm_var``.

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

@@ -30,8 +30,10 @@
 #include "bvals/comms/bvals_in_one.hpp"
 #include "config.hpp"
 #include "globals.hpp"
+#include "interface/metadata.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;
@@ -146,10 +148,12 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
 
   // Add swarm of tracer particles
   std::string swarm_name = "tracers";
-  Metadata swarm_metadata({Metadata::Provides, Metadata::None});
+  // `NoPersistentParticleIds` is just passed to test this aspect in the regression tests.
+  // For typical tracers, persistent ids are pretty important.
+  Metadata swarm_metadata(
+      {Metadata::Provides, Metadata::None, Metadata::NoPersistentParticleIds});
   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 +376,6 @@ 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 swarm_d = swarm->GetDeviceContext();
   pmb->par_for(
@@ -390,7 +393,6 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
 
         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 * gid + n;
 
         rng_pool.free_state(rng_gen);
       });

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)                                                      \
   /************************************************/                                     \
@@ -123,6 +125,8 @@
   /** Align memory of fields to cell centered memory                                     \
       (Field will be missing one layer of ghosts if it is not cell centered) **/         \
   PARTHENON_INTERNAL_FOR_FLAG(CellMemAligned)                                            \
+  /** Particles in a Swarm will not contain a persistent, unique id field **/            \
+  PARTHENON_INTERNAL_FOR_FLAG(NoPersistentParticleIds)                                   \
   /************************************************/                                     \
   /** Vars specifying coordinates for visualization purposes **/                         \
   /** You can specify a single 3D var **/                                                \
@@ -448,6 +452,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,13 +11,15 @@
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
 #include <algorithm>
+#include <cstdint>
 #include <cstdlib>
 #include <limits>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "interface/metadata.hpp"
 #include "mesh/mesh.hpp"
 #include "pack/swarm_default_names.hpp"
 #include "swarm.hpp"
@@ -84,6 +86,9 @@ Swarm::Swarm(const std::string &label, const Metadata &metadata, const int nmax_
   uid_ = get_uid_(label_);
 
   // Add default swarm fields
+  if (!metadata.IsSet(Metadata::NoPersistentParticleIds)) {
+    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 +123,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 +134,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 +153,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 +167,7 @@ void Swarm::Remove(const std::string &label) {
     }
     idx++;
   }
-  if (found == true) {
+  if (found) {
     // first delete the variable
     int_vector[idx].reset();
 
@@ -167,28 +177,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 +244,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 +253,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 +432,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 +462,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;
         }
       });