Thread UIDs through data_collection

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[PR 1060]](https://github.com/parthenon-hpc-lab/parthenon/pull/1060) Add the ability to request new MeshData/MeshBlockData objects by selecting variables by UID.
 - [[PR1039]](https://github.com/parthenon-hpc-lab/parthenon/pull/1039) Add ability to output custom coordinate positions for Visit/Paraview
 - [[PR1019](https://github.com/parthenon-hpc-lab/parthenon/pull/1019) Enable output for non-cell-centered variables
 

src/interface/data_collection.cpp

@@ -22,36 +22,14 @@
 namespace parthenon {
 
 template <typename T>
-std::shared_ptr<T> &
-DataCollection<T>::Add(const std::string &name, const std::shared_ptr<T> &src,
-                       const std::vector<std::string> &field_names, const bool shallow) {
-  auto it = containers_.find(name);
+std::shared_ptr<T> &DataCollection<T>::Add(const std::string &label) {
+  // error check for duplicate names
+  auto it = containers_.find(label);
   if (it != containers_.end()) {
-    if (!(it->second)->Contains(field_names)) {
-      PARTHENON_THROW(name +
-                      "already exists in collection but does not contain field names");
-    }
     return it->second;
   }
-
-  auto c = std::make_shared<T>(name);
-  c->Initialize(src.get(), field_names, shallow);
-
-  Set(name, c);
-
-  return containers_[name];
-}
-template <typename T>
-std::shared_ptr<T> &DataCollection<T>::Add(const std::string &label,
-                                           const std::shared_ptr<T> &src,
-                                           const std::vector<std::string> &flags) {
-  return Add(label, src, flags, false);
-}
-template <typename T>
-std::shared_ptr<T> &DataCollection<T>::AddShallow(const std::string &label,
-                                                  const std::shared_ptr<T> &src,
-                                                  const std::vector<std::string> &flags) {
-  return Add(label, src, flags, true);
+  containers_[label] = std::make_shared<T>();
+  return containers_[label];
 }
 
 std::shared_ptr<MeshData<Real>> &

src/interface/data_collection.hpp

@@ -19,6 +19,8 @@
 #include <string>
 #include <vector>
 
+#include "utils/error_checking.hpp"
+
 namespace parthenon {
 class Mesh;
 /// The DataCollection class is an abstract container that contains at least a
@@ -42,21 +44,36 @@ class DataCollection {
 
   void SetMeshPointer(Mesh *pmesh) { pmy_mesh_ = pmesh; }
 
-  std::shared_ptr<T> &Add(const std::string &label, const std::shared_ptr<T> &src,
-                          const std::vector<std::string> &flags, const bool shallow);
-  std::shared_ptr<T> &Add(const std::string &label, const std::shared_ptr<T> &src,
-                          const std::vector<std::string> &flags = {});
-  std::shared_ptr<T> &AddShallow(const std::string &label, const std::shared_ptr<T> &src,
-                                 const std::vector<std::string> &flags = {});
-  std::shared_ptr<T> &Add(const std::string &label) {
-    // error check for duplicate names
-    auto it = containers_.find(label);
+  template <typename ID_t>
+  std::shared_ptr<T> &Add(const std::string &name, const std::shared_ptr<T> &src,
+                          const std::vector<ID_t> &fields, const bool shallow) {
+    auto it = containers_.find(name);
     if (it != containers_.end()) {
+      if (!(it->second)->Contains(fields)) {
+        PARTHENON_THROW(name +
+                        " already exists in collection but does not contain fields");
+      }
       return it->second;
     }
-    containers_[label] = std::make_shared<T>();
-    return containers_[label];
+
+    auto c = std::make_shared<T>(name);
+    c->Initialize(src.get(), fields, shallow);
+
+    Set(name, c);
+
+    return containers_[name];
+  }
+  template <typename ID_t = std::string>
+  std::shared_ptr<T> &Add(const std::string &label, const std::shared_ptr<T> &src,
+                          const std::vector<ID_t> &fields = {}) {
+    return Add(label, src, fields, false);
+  }
+  template <typename ID_t = std::string>
+  std::shared_ptr<T> &AddShallow(const std::string &label, const std::shared_ptr<T> &src,
+                                 const std::vector<ID_t> &fields = {}) {
+    return Add(label, src, fields, true);
   }
+  std::shared_ptr<T> &Add(const std::string &label);
 
   auto &Stages() { return containers_; }
   const auto &Stages() const { return containers_; }

src/interface/mesh_data.cpp

@@ -16,31 +16,6 @@
 
 namespace parthenon {
 
-template <typename T>
-void MeshData<T>::Initialize(const MeshData<T> *src,
-                             const std::vector<std::string> &names, const bool shallow) {
-  if (src == nullptr) {
-    PARTHENON_THROW("src points at null");
-  }
-  pmy_mesh_ = src->GetParentPointer();
-  const int nblocks = src->NumBlocks();
-  block_data_.resize(nblocks);
-
-  grid = src->grid;
-  if (grid.type == GridType::two_level_composite) {
-    for (int i = 0; i < nblocks; i++) {
-      block_data_[i] =
-          pmy_mesh_->gmg_block_lists[src->grid.logical_level][i]->meshblock_data.Add(
-              stage_name_, src->GetBlockData(i), names, shallow);
-    }
-  } else {
-    for (int i = 0; i < nblocks; i++) {
-      block_data_[i] = pmy_mesh_->block_list[i]->meshblock_data.Add(
-          stage_name_, src->GetBlockData(i), names, shallow);
-    }
-  }
-}
-
 template <typename T>
 void MeshData<T>::Set(BlockList_t blocks, Mesh *pmesh, int ndim) {
   const int nblocks = blocks.size();
@@ -61,6 +36,17 @@ void MeshData<T>::Set(BlockList_t blocks, Mesh *pmesh) {
   Set(blocks, pmesh, ndim);
 }
 
+template <typename T>
+std::shared_ptr<MeshBlock> MeshData<T>::GetBlock(const int i, const MeshData<T> *src,
+                                                 const Mesh *pmesh) const {
+  if (src->grid.type == GridType::two_level_composite) {
+    // JMM: need .at() here, not operator[] because of const correctness.
+    return pmesh->gmg_block_lists.at(src->grid.logical_level)[i];
+  } else {
+    return pmesh->block_list[i];
+  }
+}
+
 template class MeshData<Real>;
 
 } // namespace parthenon

src/interface/mesh_data.hpp

@@ -242,8 +242,24 @@ class MeshData {
 
   void Set(BlockList_t blocks, Mesh *pmesh, int ndim);
   void Set(BlockList_t blocks, Mesh *pmesh);
-  void Initialize(const MeshData<T> *src, const std::vector<std::string> &names,
-                  const bool shallow);
+
+  template <typename ID_t>
+  void Initialize(const MeshData<T> *src, const std::vector<ID_t> &vars,
+                  const bool shallow) {
+    if (src == nullptr) {
+      PARTHENON_THROW("src points at null");
+    }
+    pmy_mesh_ = src->GetParentPointer();
+    const int nblocks = src->NumBlocks();
+    block_data_.resize(nblocks);
+
+    grid = src->grid;
+    for (int i = 0; i < nblocks; ++i) {
+      block_data_[i] =
+          GetBlock(i, src, pmy_mesh_)
+              ->meshblock_data.Add(stage_name_, src->GetBlockData(i), vars, shallow);
+    }
+  }
 
   const std::shared_ptr<MeshBlockData<T>> &GetBlockData(int n) const {
     assert(n >= 0 && n < block_data_.size());
@@ -420,16 +436,18 @@ class MeshData {
     return true;
   }
 
-  bool Contains(const std::vector<std::string> &names) const {
-    for (const auto &b : block_data_) {
-      if (!b->Contains(names)) return false;
-    }
-    return true;
+  template <typename Vars_t>
+  bool Contains(const Vars_t &vars) const noexcept {
+    return std::all_of(block_data_.begin(), block_data_.end(),
+                       [this, vars](const auto &b) { return b->Contains(vars); });
   }
 
   SparsePackCache &GetSparsePackCache() { return sparse_pack_cache_; }
 
  private:
+  std::shared_ptr<MeshBlock> GetBlock(const int i, const MeshData<T> *src,
+                                      const Mesh *pmesh) const;
+
   int ndim_;
   Mesh *pmy_mesh_;
   BlockDataList_t<T> block_data_;

src/interface/meshblock_data.cpp

@@ -79,45 +79,6 @@ void MeshBlockData<T>::AddField(const std::string &base_name, const Metadata &me
   }
 }
 
-// TODO(JMM): Move to unique IDs at some point
-template <typename T>
-void MeshBlockData<T>::Initialize(const MeshBlockData<T> *src,
-                                  const std::vector<std::string> &names,
-                                  const bool shallow_copy) {
-  assert(src != nullptr);
-  SetBlockPointer(src);
-  resolved_packages_ = src->resolved_packages_;
-  is_shallow_ = shallow_copy;
-
-  auto add_var = [=](auto var) {
-    if (shallow_copy || var->IsSet(Metadata::OneCopy)) {
-      Add(var);
-    } else {
-      Add(var->AllocateCopy(pmy_block));
-    }
-  };
-
-  // special case when the list of names is empty, copy everything
-  if (names.empty()) {
-    for (auto v : src->GetVariableVector()) {
-      add_var(v);
-    }
-  } else {
-    auto var_map = src->GetVariableMap();
-
-    for (const auto &name : names) {
-      bool found = false;
-      auto v = var_map.find(name);
-      if (v != var_map.end()) {
-        found = true;
-        add_var(v->second);
-      }
-      PARTHENON_REQUIRE_THROWS(found, "MeshBlockData::CopyFrom: Variable '" + name +
-                                          "' not found");
-    }
-  }
-}
-
 /// Queries related to variable packs
 /// This is a helper function that queries the cache for the given pack.
 /// The strings are the keys and the lists are the values.

src/interface/meshblock_data.hpp

@@ -13,6 +13,7 @@
 #ifndef INTERFACE_MESHBLOCK_DATA_HPP_
 #define INTERFACE_MESHBLOCK_DATA_HPP_
 
+#include <algorithm>
 #include <limits>
 #include <map>
 #include <memory>
@@ -109,8 +110,33 @@ class MeshBlockData {
   /// Create copy of MeshBlockData, possibly with a subset of named fields,
   /// and possibly shallow.  Note when shallow=false, new storage is allocated
   /// for non-OneCopy vars, but the data from src is not actually deep copied
-  void Initialize(const MeshBlockData<T> *src, const std::vector<std::string> &names,
-                  const bool shallow);
+  template <typename ID_t>
+  void Initialize(const MeshBlockData<T> *src, const std::vector<ID_t> &vars,
+                  const bool shallow_copy) {
+    PARTHENON_DEBUG_REQUIRE(src != nullptr, "Source data must be non-null.");
+    SetBlockPointer(src);
+    resolved_packages_ = src->resolved_packages_;
+    is_shallow_ = shallow_copy;
+
+    auto add_var = [=](auto var) {
+      if (shallow_copy || var->IsSet(Metadata::OneCopy)) {
+        Add(var);
+      } else {
+        Add(var->AllocateCopy(pmy_block));
+      }
+    };
+
+    // special case when the list of vars is empty, copy everything
+    if (vars.empty()) {
+      for (auto v : src->GetVariableVector()) {
+        add_var(v);
+      }
+    } else {
+      for (const auto &v : vars) {
+        add_var(src->GetVarPtr(v));
+      }
+    }
+  }
 
   //
   // Queries related to Variable objects
@@ -124,14 +150,13 @@ class MeshBlockData {
   const MapToVars<T> &GetVariableMap() const noexcept { return varMap_; }
 
   std::shared_ptr<Variable<T>> GetVarPtr(const std::string &label) const {
-    auto it = varMap_.find(label);
-    PARTHENON_REQUIRE_THROWS(it != varMap_.end(),
+    PARTHENON_REQUIRE_THROWS(varMap_.count(label),
                              "Couldn't find variable '" + label + "'");
-    return it->second;
+    return varMap_.at(label);
   }
   std::shared_ptr<Variable<T>> GetVarPtr(const Uid_t &uid) const {
     PARTHENON_REQUIRE_THROWS(varUidMap_.count(uid),
-                             "Variable ID " + std::to_string(uid) + "not found!");
+                             "Variable ID " + std::to_string(uid) + " not found!");
     return varUidMap_.at(uid);
   }
 
@@ -388,15 +413,12 @@ class MeshBlockData {
     return (my_keys == cmp_keys);
   }
 
-  bool Contains(const std::string &name) const noexcept {
-    if (varMap_.find(name) != varMap_.end()) return true;
-    return false;
-  }
-  bool Contains(const std::vector<std::string> &names) const noexcept {
-    for (const auto &name : names) {
-      if (!Contains(name)) return false;
-    }
-    return true;
+  bool Contains(const std::string &name) const noexcept { return varMap_.count(name); }
+  bool Contains(const Uid_t &uid) const noexcept { return varUidMap_.count(uid); }
+  template <typename ID_t>
+  bool Contains(const std::vector<ID_t> &vars) const noexcept {
+    return std::all_of(vars.begin(), vars.end(),
+                       [this](const auto &v) { return this->Contains(v); });
   }
 
   void SetAllVariablesToInitialized() {

src/mesh/mesh.hpp

@@ -123,8 +123,8 @@ class Mesh {
 
   std::map<int, BlockList_t> gmg_block_lists;
   std::map<int, DataCollection<MeshData<Real>>> gmg_mesh_data;
-  int GetGMGMaxLevel() { return current_level; }
-  int GetGMGMinLevel() { return gmg_min_logical_level_; }
+  int GetGMGMaxLevel() const { return current_level; }
+  int GetGMGMinLevel() const { return gmg_min_logical_level_; }
 
   // functions
   void Initialize(bool init_problem, ParameterInput *pin, ApplicationInput *app_in);

tst/unit/test_data_collection.cpp

@@ -116,5 +116,34 @@ TEST_CASE("Adding MeshBlockData objects to a DataCollection", "[DataCollection]"
         REQUIRE(hxv2(0) == hv2(0));
       }
     }
+    AND_WHEN("We want only a subset of variables in a new MeshBlockData by UID") {
+      // reset vars
+      par_for(
+          loop_pattern_flatrange_tag, "init vars", DevExecSpace(), 0, 0,
+          KOKKOS_LAMBDA(const int i) { v2(0) = 222; });
+      std::vector<parthenon::Uid_t> uids;
+      uids.push_back(mbd->UniqueID("var2"));
+      uids.push_back(mbd->UniqueID("var3"));
+      auto x = d.Add("part", mbd, uids);
+      THEN("Requesting the missing variables should throw") {
+        REQUIRE_THROWS(x->Get("var1"));
+      }
+      AND_THEN("Requesting the specified variables should work as expected") {
+        auto &xv2 = x->Get("var2").data;
+        auto &xv3 = x->Get("var3").data;
+        par_for(
+            loop_pattern_flatrange_tag, "init vars", DevExecSpace(), 0, 0,
+            KOKKOS_LAMBDA(const int i) {
+              xv2(0) = 22;
+              xv3(0) = 33;
+            });
+        auto hv2 = v2.GetHostMirrorAndCopy();
+        auto hv3 = v3.GetHostMirrorAndCopy();
+        auto hxv2 = xv2.GetHostMirrorAndCopy();
+        auto hxv3 = xv3.GetHostMirrorAndCopy();
+        REQUIRE(hxv3(0) != hv3(0));
+        REQUIRE(hxv2(0) == hv2(0));
+      }
+    }
   }
 }