Updates to DataCollection

CHANGELOG.md

@@ -26,6 +26,7 @@
 - [[PR 1004]](https://github.com/parthenon-hpc-lab/parthenon/pull/1004) Allow parameter modification from an input file for restarts
 
 ### Fixed (not changing behavior/API/variables/...)
+- [[PR 1092]](https://github.com/parthenon-hpc-lab/parthenon/pull/1092) Updates to DataCollection and MeshData to remove requirement of predefining MeshBlockData
 - [[PR 1104]](https://github.com/parthenon-hpc-lab/parthenon/pull/1104) Fix reading restarts due to hidden ghost var
 - [[PR 1098]](https://github.com/parthenon-hpc-lab/parthenon/pull/1098) Move to symmetrized logical coordinates and fix SMR bug
 - [[PR 1095]](https://github.com/parthenon-hpc-lab/parthenon/pull/1095) Add missing include guards in hdf5 restart

example/fine_advection/advection_driver.cpp

@@ -62,17 +62,6 @@ TaskCollection AdvectionDriver::MakeTaskCollection(BlockList_t &blocks, const in
   // Build MeshBlockData containers that will be included in MeshData containers. It is
   // gross that this has to be done by hand.
   const auto &stage_name = integrator->stage_name;
-  if (stage == 1) {
-    for (int i = 0; i < blocks.size(); i++) {
-      auto &pmb = blocks[i];
-      // first make other useful containers
-      auto &base = pmb->meshblock_data.Get();
-      pmb->meshblock_data.Add("dUdt", base);
-      for (int s = 1; s < integrator->nstages; s++)
-        pmb->meshblock_data.Add(stage_name[s], base);
-    }
-  }
-
   const Real beta = integrator->beta[stage - 1];
   const Real dt = integrator->dt;
 

src/interface/data_collection.cpp

@@ -32,30 +32,26 @@ std::shared_ptr<T> &DataCollection<T>::Add(const std::string &label) {
   return containers_[label];
 }
 
+template <>
 std::shared_ptr<MeshData<Real>> &
-GetOrAdd_impl(Mesh *pmy_mesh_,
-              std::map<std::string, std::shared_ptr<MeshData<Real>>> &containers_,
-              BlockList_t &block_list, const std::string &mbd_label,
-              const int &partition_id, const std::optional<int> gmg_level) {
-  std::string label = mbd_label + "_part-" + std::to_string(partition_id);
-  if (gmg_level) label = label + "_gmg-" + std::to_string(*gmg_level);
+DataCollection<MeshData<Real>>::GetOrAdd_impl(const std::string &mbd_label,
+                                              const int &partition_id,
+                                              const std::optional<int> gmg_level) {
+  std::string label = GetKey(mbd_label, partition_id, gmg_level);
   auto it = containers_.find(label);
   if (it == containers_.end()) {
     // TODO(someone) add caching of partitions to Mesh at some point
     const int pack_size = pmy_mesh_->DefaultPackSize();
+    auto &block_list =
+        gmg_level ? pmy_mesh_->gmg_block_lists[*gmg_level] : pmy_mesh_->block_list;
     auto partitions = partition::ToSizeN(block_list, pack_size);
     // Account for possibly empty block_list
     if (partitions.size() == 0) partitions = std::vector<BlockList_t>(1);
     for (auto i = 0; i < partitions.size(); i++) {
-      std::string md_label = mbd_label + "_part-" + std::to_string(i);
-      if (gmg_level) md_label = md_label + "_gmg-" + std::to_string(*gmg_level);
+      std::string md_label = GetKey(mbd_label, partition_id, gmg_level);
       containers_[md_label] = std::make_shared<MeshData<Real>>(mbd_label);
-      containers_[md_label]->Set(partitions[i], pmy_mesh_);
-      if (gmg_level) {
-        containers_[md_label]->grid = GridIdentifier::two_level_composite(*gmg_level);
-      } else {
-        containers_[md_label]->grid = GridIdentifier::leaf();
-      }
+      containers_[md_label]->Initialize(partitions[i], pmy_mesh_, gmg_level);
+      containers_[md_label]->partition = i;
     }
   }
   return containers_[label];
@@ -65,16 +61,14 @@ template <>
 std::shared_ptr<MeshData<Real>> &
 DataCollection<MeshData<Real>>::GetOrAdd(const std::string &mbd_label,
                                          const int &partition_id) {
-  return GetOrAdd_impl(pmy_mesh_, containers_, pmy_mesh_->block_list, mbd_label,
-                       partition_id, {});
+  return GetOrAdd_impl(mbd_label, partition_id, {});
 }
 
 template <>
 std::shared_ptr<MeshData<Real>> &
 DataCollection<MeshData<Real>>::GetOrAdd(int gmg_level, const std::string &mbd_label,
                                          const int &partition_id) {
-  return GetOrAdd_impl(pmy_mesh_, containers_, pmy_mesh_->gmg_block_lists[gmg_level],
-                       mbd_label, partition_id, gmg_level);
+  return GetOrAdd_impl(mbd_label, partition_id, gmg_level);
 }
 
 template class DataCollection<MeshData<Real>>;

src/interface/data_collection.hpp

@@ -19,10 +19,17 @@
 #include <string>
 #include <vector>
 
+#include "basic_types.hpp"
+#include "utils/concepts_lite.hpp"
 #include "utils/error_checking.hpp"
 
 namespace parthenon {
 class Mesh;
+class MeshBlock;
+template <class T>
+class MeshData;
+template <class T>
+class MeshBlockData;
 /// The DataCollection class is an abstract container that contains at least a
 /// "base" container of some type (e.g., of MeshData or MeshBlockData) plus
 /// additional containers identified by string labels.
@@ -44,31 +51,39 @@ class DataCollection {
 
   void SetMeshPointer(Mesh *pmesh) { pmy_mesh_ = pmesh; }
 
-  template <typename ID_t>
-  std::shared_ptr<T> &Add(const std::string &name, const std::shared_ptr<T> &src,
+  template <class SRC_t, typename ID_t>
+  std::shared_ptr<T> &Add(const std::string &name, const std::shared_ptr<SRC_t> &src,
                           const std::vector<ID_t> &fields, const bool shallow) {
-    auto it = containers_.find(name);
+    if constexpr (!((std::is_same_v<SRC_t, MeshBlock> &&
+                     std::is_same_v<T, MeshBlockData<Real>>) ||
+                    std::is_same_v<SRC_t, T>)) {
+      // SRC_t and T are incompatible
+      static_assert(always_false<SRC_t>, "Incompatible source and container types.");
+    }
+
+    auto key = GetKey(name, src);
+    auto it = containers_.find(key);
     if (it != containers_.end()) {
       if (fields.size() && !(it->second)->Contains(fields)) {
-        PARTHENON_THROW(name + " already exists in collection but fields do not match.");
+        PARTHENON_THROW(key + " already exists in collection but fields do not match.");
       }
       return it->second;
     }
 
     auto c = std::make_shared<T>(name);
-    c->Initialize(src.get(), fields, shallow);
+    c->Initialize(src, fields, shallow);
 
-    Set(name, c);
-
-    return containers_[name];
+    containers_[key] = c;
+    return containers_[key];
   }
-  template <typename ID_t = std::string>
-  std::shared_ptr<T> &Add(const std::string &label, const std::shared_ptr<T> &src,
+  template <class SRC_t, typename ID_t = std::string>
+  std::shared_ptr<T> &Add(const std::string &label, const std::shared_ptr<SRC_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,
+  template <class SRC_t, typename ID_t = std::string>
+  std::shared_ptr<T> &AddShallow(const std::string &label,
+                                 const std::shared_ptr<SRC_t> &src,
                                  const std::vector<ID_t> &fields = {}) {
     return Add(label, src, fields, true);
   }
@@ -105,6 +120,28 @@ class DataCollection {
   }
 
  private:
+  template <class U>
+  std::string GetKey(const std::string &stage_label, const std::shared_ptr<U> &in) {
+    if constexpr (std::is_same_v<U, MeshData<Real>>) {
+      std::string key = stage_label + "_part-" + std::to_string(in->partition);
+      if (in->grid.type == GridType::two_level_composite)
+        key = key + "_gmg-" + std::to_string(in->grid.logical_level);
+      return key;
+    } else {
+      return stage_label;
+    }
+  }
+  std::string GetKey(const std::string &stage_label, std::optional<int> partition_id,
+                     std::optional<int> gmg_level) {
+    std::string key = stage_label;
+    if (partition_id) key = key + "_part-" + std::to_string(*partition_id);
+    if (gmg_level) key = key + "_gmg-" + std::to_string(*gmg_level);
+    return key;
+  }
+
+  std::shared_ptr<T> &GetOrAdd_impl(const std::string &mbd_label, const int &partition_id,
+                                    const std::optional<int> gmg_level);
+
   Mesh *pmy_mesh_;
   std::map<std::string, std::shared_ptr<T>> containers_;
 };

src/interface/mesh_data.cpp

@@ -17,23 +17,30 @@
 namespace parthenon {
 
 template <typename T>
-void MeshData<T>::Set(BlockList_t blocks, Mesh *pmesh, int ndim) {
+void MeshData<T>::Initialize(BlockList_t blocks, Mesh *pmesh, int ndim,
+                             std::optional<int> gmg_level) {
   const int nblocks = blocks.size();
   ndim_ = ndim;
   block_data_.resize(nblocks);
   SetMeshPointer(pmesh);
   for (int i = 0; i < nblocks; i++) {
-    block_data_[i] = blocks[i]->meshblock_data.Get(stage_name_);
+    block_data_[i] = blocks[i]->meshblock_data.Add(stage_name_, blocks[i]);
+  }
+  if (gmg_level) {
+    grid = GridIdentifier::two_level_composite(*gmg_level);
+  } else {
+    grid = GridIdentifier::leaf();
   }
 }
 
 template <typename T>
-void MeshData<T>::Set(BlockList_t blocks, Mesh *pmesh) {
+void MeshData<T>::Initialize(BlockList_t blocks, Mesh *pmesh,
+                             std::optional<int> gmg_level) {
   int ndim;
   if (pmesh != nullptr) {
     ndim = pmesh->ndim;
   }
-  Set(blocks, pmesh, ndim);
+  Initialize(blocks, pmesh, ndim, gmg_level);
 }
 
 template <typename T>

src/interface/mesh_data.hpp

@@ -190,10 +190,12 @@ const MeshBlockPack<P> &PackOnMesh(M &map, BlockDataList_t<Real> &block_data_,
 template <typename T>
 class MeshData {
  public:
+  using parent_t = Mesh;
   MeshData() = default;
   explicit MeshData(const std::string &name) : stage_name_(name) {}
 
   GridIdentifier grid;
+  int partition;
 
   const auto &StageName() const { return stage_name_; }
 
@@ -241,11 +243,12 @@ class MeshData {
     }
   }
 
-  void Set(BlockList_t blocks, Mesh *pmesh, int ndim);
-  void Set(BlockList_t blocks, Mesh *pmesh);
+  void Initialize(BlockList_t blocks, Mesh *pmesh, int ndim,
+                  std::optional<int> gmg_level = {});
+  void Initialize(BlockList_t blocks, Mesh *pmesh, std::optional<int> gmg_level = {});
 
   template <typename ID_t>
-  void Initialize(const MeshData<T> *src, const std::vector<ID_t> &vars,
+  void Initialize(std::shared_ptr<MeshData<T>> src, const std::vector<ID_t> &vars,
                   const bool shallow) {
     if (src == nullptr) {
       PARTHENON_THROW("src points at null");

src/interface/meshblock_data.cpp

@@ -33,37 +33,6 @@
 #include "utils/utils.hpp"
 
 namespace parthenon {
-
-template <typename T>
-void MeshBlockData<T>::Initialize(
-    const std::shared_ptr<StateDescriptor> resolved_packages,
-    const std::shared_ptr<MeshBlock> pmb) {
-  SetBlockPointer(pmb);
-  resolved_packages_ = resolved_packages;
-
-  // clear all variables, maps, and pack caches
-  varVector_.clear();
-  varMap_.clear();
-  varUidMap_.clear();
-  flagsToVars_.clear();
-  varPackMap_.clear();
-  coarseVarPackMap_.clear();
-  varFluxPackMap_.clear();
-
-  for (auto const &q : resolved_packages->AllFields()) {
-    AddField(q.first.base_name, q.second, q.first.sparse_id);
-  }
-
-  const auto &swarm_container = GetSwarmData();
-  swarm_container->Initialize(resolved_packages, pmb);
-
-  Metadata::FlagCollection flags({Metadata::Sparse, Metadata::ForceAllocOnNewBlocks});
-  auto vars = GetVariablesByFlag(flags);
-  for (auto &v : vars.vars()) {
-    AllocateSparse(v->label());
-  }
-}
-
 ///
 /// The internal routine for adding a new field.  This subroutine
 /// is topology aware and will allocate accordingly.
@@ -270,9 +239,9 @@ MeshBlockData<T>::GetVariablesByName(const std::vector<std::string> &names,
       } else {
         var_list.Add(v, sparse_ids_set);
       }
-    } else if ((resolved_packages_ != nullptr) &&
-               (resolved_packages_->SparseBaseNamePresent(name))) {
-      const auto &sparse_pool = resolved_packages_->GetSparsePool(name);
+    } else if ((resolved_packages != nullptr) &&
+               (resolved_packages->SparseBaseNamePresent(name))) {
+      const auto &sparse_pool = resolved_packages->GetSparsePool(name);
 
       // add all sparse ids of the pool
       for (const auto iter : sparse_pool.pool()) {