Provide more flexible ways to add subset MeshBlockData objects to DataCollections

src/amr_criteria/amr_criteria.hpp

@@ -17,12 +17,13 @@
 #include <string>
 
 #include "defs.hpp"
-#include "interface/meshblock_data.hpp"
 #include "mesh/domain.hpp"
 
 namespace parthenon {
 
 class ParameterInput;
+template <class>
+class MeshBlockData;
 
 struct AMRBounds {
   AMRBounds(const IndexRange &ib, const IndexRange &jb, const IndexRange &kb)

src/interface/data_collection.cpp

@@ -22,43 +22,54 @@
 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> &flags) {
+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);
   if (it != containers_.end()) {
-    if (!(it->second)->Contains(flags)) {
-      PARTHENON_THROW(name + "already exists in collection but does not contain flags");
+    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>();
-  c->Copy(src, flags);
+  c->Copy(src.get(), field_names, shallow);
 
-  containers_[name] = c;
-  return containers_[name];
-}
+  Set(name, c);
 
-template <typename T>
-std::shared_ptr<T> DataCollection<T>::Add(const std::string &name,
-                                          const std::shared_ptr<T> &src) {
-  // error check for duplicate names
-  auto it = containers_.find(name);
-  if (it != containers_.end()) {
-    // check to make sure they are the same
-    if (!(*src == *(it->second))) {
-      PARTHENON_THROW("Error attempting to add a Container to a Collection");
+  if constexpr (std::is_same<T, MeshData<Real>>::value) {
+    for (int b = 0; b < pmy_mesh_->block_list.size(); b++) {
+      auto &mbd = pmy_mesh_->block_list[b]->meshblock_data;
+      mbd.Set(name, c->GetBlockData(b));
     }
-    return it->second;
   }
 
-  auto c = std::make_shared<T>();
-  c->Copy(src);
-
-  containers_[name] = c;
   return containers_[name];
 }
+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) {
+  return Add(name, src, field_names, false);
+}
+template <typename T>
+std::shared_ptr<T> &
+DataCollection<T>::AddShallow(const std::string &name, const std::shared_ptr<T> &src,
+                              const std::vector<std::string> &field_names) {
+  return Add(name, src, field_names, true);
+}
+template <typename T>
+std::shared_ptr<T> &DataCollection<T>::Add(const std::string &name,
+                                           const std::shared_ptr<T> &src) {
+  return Add(name, src, {}, false);
+}
+template <typename T>
+std::shared_ptr<T> &DataCollection<T>::AddShallow(const std::string &name,
+                                                  const std::shared_ptr<T> &src) {
+  return Add(name, src, {}, true);
+}
 
 template <>
 std::shared_ptr<MeshData<Real>> &

src/interface/data_collection.hpp

@@ -38,10 +38,15 @@ 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);
-  std::shared_ptr<T> Add(const std::string &label, const std::shared_ptr<T> &src);
-  std::shared_ptr<T> Add(const std::string &label) {
+  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, const std::shared_ptr<T> &src);
+  std::shared_ptr<T> &AddShallow(const std::string &label, const std::shared_ptr<T> &src);
+  std::shared_ptr<T> &Add(const std::string &label) {
     // error check for duplicate names
     auto it = containers_.find(label);
     if (it != containers_.end()) {
@@ -64,6 +69,8 @@ class DataCollection {
     return it->second;
   }
 
+  void Set(const std::string &name, std::shared_ptr<T> &d) { containers_[name] = d; }
+
   std::shared_ptr<T> &GetOrAdd(const std::string &mbd_label, const int &partition_id);
 
   void PurgeNonBase() {

src/interface/make_pack_descriptor.hpp

@@ -31,10 +31,13 @@
 #include "interface/metadata.hpp"
 #include "interface/sparse_pack.hpp"
 #include "interface/state_descriptor.hpp"
+#include "interface/variable.hpp"
 #include "mesh/mesh.hpp"
 
 namespace parthenon {
 
+inline auto MakeDefaultPackDescriptor() { return typename SparsePack<>::Descriptor(); }
+
 inline auto MakePackDescriptor(StateDescriptor *psd, const std::vector<std::string> &vars,
                                const std::vector<bool> &use_regex,
                                const std::vector<MetadataFlag> &flags = {},
@@ -116,6 +119,23 @@ inline auto MakePackDescriptor(
   return MakePackDescriptor(psd, vars, use_regex, flags, options);
 }
 
+inline auto MakePackDescriptor(StateDescriptor *psd, const std::vector<Uid_t> &var_ids,
+                               const std::vector<MetadataFlag> &flags = {},
+                               const std::set<PDOpt> &options = {}) {
+  auto selector = [&](int vidx, const VarID &id, const Metadata &md) {
+    if (flags.size() > 0) {
+      for (const auto &flag : flags) {
+        if (!md.IsSet(flag)) return false;
+      }
+    }
+    if (Variable<Real>::GetUniqueID(id.label()) == var_ids[vidx]) return true;
+    return false;
+  };
+
+  impl::PackDescriptor base_desc(psd, var_ids, selector, options);
+  return typename SparsePack<>::Descriptor(base_desc);
+}
+
 } // namespace parthenon
 
 #endif // INTERFACE_MAKE_PACK_DESCRIPTOR_HPP_

src/interface/mesh_data.hpp

@@ -31,6 +31,7 @@
 #include "utils/communication_buffer.hpp"
 #include "utils/error_checking.hpp"
 #include "utils/object_pool.hpp"
+#include "utils/unique_id.hpp"
 #include "utils/utils.hpp"
 
 namespace parthenon {
@@ -236,14 +237,15 @@ class MeshData {
   }
 
   template <typename... Args>
-  void Copy(const std::shared_ptr<MeshData<T>> src, Args &&...args) {
-    if (src.get() == nullptr) {
+  void Copy(const MeshData<T> *src, Args &&...args) {
+    if (src == nullptr) {
       PARTHENON_THROW("src points at null");
     }
     const int nblocks = src->NumBlocks();
     block_data_.resize(nblocks);
     for (int i = 0; i < nblocks; i++) {
-      block_data_[i]->Copy(src->GetBlockData(i), std::forward<Args>(args)...);
+      block_data_[i] = std::make_shared<MeshBlockData<T>>();
+      block_data_[i]->Copy(src->GetBlockData(i).get(), std::forward<Args>(args)...);
     }
   }
 
@@ -435,6 +437,13 @@ class MeshData {
   BvarsCache_t bvars_cache_;
 };
 
+template <typename T, typename... Args>
+std::vector<Uid_t> UidIntersection(MeshData<T> *md1, MeshData<T> *md2, Args &&...args) {
+  auto u1 = md1->GetBlockData(0)->GetVariableUIDs(std::forward<Args>(args)...);
+  auto u2 = md2->GetBlockData(0)->GetVariableUIDs(std::forward<Args>(args)...);
+  return UidIntersection(u1, u2);
+}
+
 } // namespace parthenon
 
 #endif // INTERFACE_MESH_DATA_HPP_

src/interface/meshblock_data.cpp

@@ -73,39 +73,31 @@ void MeshBlockData<T>::AddField(const std::string &base_name, const Metadata &me
   }
 }
 
-// TODO(JMM): Add CopyFrom that takes unique IDs
+// TODO(JMM): Move to unique IDs at some point
 template <typename T>
-void MeshBlockData<T>::CopyFrom(const MeshBlockData<T> &src, bool shallow_copy,
-                                const std::vector<std::string> &names,
-                                const std::vector<MetadataFlag> &flags,
-                                const std::vector<int> &sparse_ids) {
+void MeshBlockData<T>::Copy(const MeshBlockData<T> *src,
+                            const std::vector<std::string> &names,
+                            const bool shallow_copy) {
+  assert(src != nullptr);
   SetBlockPointer(src);
-  resolved_packages_ = src.resolved_packages_;
-  std::unordered_set<int> sparse_ids_set(sparse_ids.begin(), sparse_ids.end());
-
-  auto add_var = [=, &flags, &sparse_ids](auto var) {
-    if (!flags.empty() && !var->metadata().AnyFlagsSet(flags)) {
-      return;
-    }
-
-    if (!sparse_ids.empty() && var->IsSparse() &&
-        (sparse_ids_set.count(var->GetSparseID()) == 0)) {
-      return;
-    }
+  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()) {
+    for (auto v : src->GetVariableVector()) {
       add_var(v);
     }
   } else {
-    auto var_map = src.GetVariableMap();
+    auto var_map = src->GetVariableMap();
 
     for (const auto &name : names) {
       bool found = false;
@@ -114,54 +106,12 @@ void MeshBlockData<T>::CopyFrom(const MeshBlockData<T> &src, bool shallow_copy,
         found = true;
         add_var(v->second);
       }
-
-      if (!found && (resolved_packages_ != nullptr)) {
-        // check if this is a sparse base name, if so we get its pool of sparse_ids,
-        // otherwise we get an empty pool
-        const auto &sparse_pool = resolved_packages_->GetSparsePool(name);
-
-        // add all sparse ids of the pool
-        for (const auto iter : sparse_pool.pool()) {
-          // this variable must exist, if it doesn't something is very wrong
-          const auto &v = varMap_.at(MakeVarLabel(name, iter.first));
-          add_var(v);
-          found = true;
-        }
-      }
-
       PARTHENON_REQUIRE_THROWS(found, "MeshBlockData::CopyFrom: Variable '" + name +
                                           "' not found");
     }
   }
 }
 
-// Constructor for getting sub-containers
-// the variables returned are all shallow copies of the src container.
-// Optionally extract only some of the sparse ids of src variable.
-template <typename T>
-MeshBlockData<T>::MeshBlockData(const MeshBlockData<T> &src,
-                                const std::vector<std::string> &names,
-                                const std::vector<int> &sparse_ids) {
-  CopyFrom(src, true, names, {}, sparse_ids);
-}
-
-// TODO(JMM): Add constructor that takes unique IDs
-template <typename T>
-MeshBlockData<T>::MeshBlockData(const MeshBlockData<T> &src,
-                                const std::vector<MetadataFlag> &flags,
-                                const std::vector<int> &sparse_ids) {
-  CopyFrom(src, true, {}, flags, sparse_ids);
-}
-
-// provides a container that has a single sparse slice
-template <typename T>
-std::shared_ptr<MeshBlockData<T>>
-MeshBlockData<T>::SparseSlice(const std::vector<int> &sparse_ids) const {
-  auto c = std::make_shared<MeshBlockData<T>>();
-  c->CopyFrom(*this, true, {}, {}, sparse_ids);
-  return c;
-}
-
 /// 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.