From 98b322eea5565abdc6351ce20b7c3a8c33dd6804 Mon Sep 17 00:00:00 2001
From: Jonah Miller <jonahm@lanl.gov>
Date: Thu, 18 Apr 2024 11:43:44 -0600
Subject: [PATCH 1/9] template Add and AddShallow on either varname or UID

---
 src/interface/data_collection.cpp | 32 ++++------------------
 src/interface/data_collection.hpp | 39 +++++++++++++++++++--------
 src/interface/mesh_data.cpp       | 36 ++++++++-----------------
 src/interface/mesh_data.hpp       | 32 +++++++++++++++++-----
 src/interface/meshblock_data.cpp  | 39 ---------------------------
 src/interface/meshblock_data.hpp  | 45 +++++++++++++++++++++++--------
 src/mesh/mesh.hpp                 |  4 +--
 7 files changed, 105 insertions(+), 122 deletions(-)

diff --git a/src/interface/data_collection.cpp b/src/interface/data_collection.cpp
index 3ac429fd6df4..38e14d2f536e 100644
--- a/src/interface/data_collection.cpp
+++ b/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>> &
diff --git a/src/interface/data_collection.hpp b/src/interface/data_collection.hpp
index 6119bf0a8ae2..ae3788364bd4 100644
--- a/src/interface/data_collection.hpp
+++ b/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_; }
diff --git a/src/interface/mesh_data.cpp b/src/interface/mesh_data.cpp
index 27a4dc520180..f42d697622ee 100644
--- a/src/interface/mesh_data.cpp
+++ b/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) {
-    int gmg_level = src->grid.logical_level - pmy_mesh_->GetGMGMinLogicalLevel();
-    for (int i = 0; i < nblocks; i++) {
-      block_data_[i] = pmy_mesh_->gmg_block_lists[gmg_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) {
+    int gmg_level = src->grid.logical_level - pmesh->GetGMGMinLogicalLevel();
+    return pmesh->gmg_block_lists[gmg_level][i];
+  } else {
+    return pmesh->block_list[i];
+  }
+}
+
 template class MeshData<Real>;
 
 } // namespace parthenon
diff --git a/src/interface/mesh_data.hpp b/src/interface/mesh_data.hpp
index 07726e864f56..74e21f8f5d2c 100644
--- a/src/interface/mesh_data.hpp
+++ b/src/interface/mesh_data.hpp
@@ -248,8 +248,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());
@@ -426,16 +442,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_;
diff --git a/src/interface/meshblock_data.cpp b/src/interface/meshblock_data.cpp
index 5f685c0e119d..1b8185b4abce 100644
--- a/src/interface/meshblock_data.cpp
+++ b/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.
diff --git a/src/interface/meshblock_data.hpp b/src/interface/meshblock_data.hpp
index caf3be9c9d3c..ff5132a1b4e3 100644
--- a/src/interface/meshblock_data.hpp
+++ b/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) {
+    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 vars is empty, copy everything
+    if (vars.empty()) {
+      for (auto v : src->GetVariableVector()) {
+        add_var(v);
+      }
+    } else {
+      for (const auto &v : vars) {
+        add_var(GetVarPtr(v));
+      }
+    }
+  }
 
   //
   // Queries related to Variable objects
@@ -388,15 +414,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() {
diff --git a/src/mesh/mesh.hpp b/src/mesh/mesh.hpp
index 81a2e9673be6..94189edb8c3f 100644
--- a/src/mesh/mesh.hpp
+++ b/src/mesh/mesh.hpp
@@ -123,8 +123,8 @@ class Mesh {
   std::vector<LogicalLocMap_t> gmg_grid_locs;
   std::vector<BlockList_t> gmg_block_lists;
   std::vector<DataCollection<MeshData<Real>>> gmg_mesh_data;
-  int GetGMGMaxLevel() { return gmg_grid_locs.size() - 1; }
-  int GetGMGMinLogicalLevel() { return gmg_min_logical_level_; }
+  int GetGMGMaxLevel() const { return gmg_grid_locs.size() - 1; }
+  int GetGMGMinLogicalLevel() const { return gmg_min_logical_level_; }
 
   // functions
   void Initialize(bool init_problem, ParameterInput *pin, ApplicationInput *app_in);

From 392b51e0741d7f3a9747fbe17dfcef2e42a05679 Mon Sep 17 00:00:00 2001
From: Jonah Miller <jonahm@lanl.gov>
Date: Thu, 18 Apr 2024 16:18:44 -0600
Subject: [PATCH 2/9] fix Intialize bug... missing a src->

---
 src/interface/meshblock_data.hpp  | 11 +++++------
 tst/unit/test_data_collection.cpp | 29 +++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/interface/meshblock_data.hpp b/src/interface/meshblock_data.hpp
index ff5132a1b4e3..90134020a789 100644
--- a/src/interface/meshblock_data.hpp
+++ b/src/interface/meshblock_data.hpp
@@ -113,7 +113,7 @@ class MeshBlockData {
   template <typename ID_t>
   void Initialize(const MeshBlockData<T> *src, const std::vector<ID_t> &vars,
                   const bool shallow_copy) {
-    assert(src != nullptr);
+    PARTHENON_DEBUG_REQUIRE(src != nullptr, "Source data must be non-null.");
     SetBlockPointer(src);
     resolved_packages_ = src->resolved_packages_;
     is_shallow_ = shallow_copy;
@@ -133,7 +133,7 @@ class MeshBlockData {
       }
     } else {
       for (const auto &v : vars) {
-        add_var(GetVarPtr(v));
+        add_var(src->GetVarPtr(v));
       }
     }
   }
@@ -150,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);
   }
 
diff --git a/tst/unit/test_data_collection.cpp b/tst/unit/test_data_collection.cpp
index 2e6c3d108133..8b5c87ca22a4 100644
--- a/tst/unit/test_data_collection.cpp
+++ b/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));
+      }
+    }
   }
 }

From b1c4c48a565e68fc7e1851b39e68ddfad5f419fc Mon Sep 17 00:00:00 2001
From: Jonah Miller <jonahm@lanl.gov>
Date: Thu, 18 Apr 2024 16:58:03 -0600
Subject: [PATCH 3/9] fix bug

---
 src/interface/mesh_data.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/interface/mesh_data.cpp b/src/interface/mesh_data.cpp
index f42d697622ee..a17779da4432 100644
--- a/src/interface/mesh_data.cpp
+++ b/src/interface/mesh_data.cpp
@@ -40,8 +40,8 @@ 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) {
-    int gmg_level = src->grid.logical_level - pmesh->GetGMGMinLogicalLevel();
-    return pmesh->gmg_block_lists[gmg_level][i];
+    // 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];
   }

From fdfb405ad5e4e48be516b97768032bda76a312a7 Mon Sep 17 00:00:00 2001
From: Jonah Miller <jonahm@lanl.gov>
Date: Thu, 18 Apr 2024 17:01:29 -0600
Subject: [PATCH 4/9] CHANGELOG

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e58da8591d5f..6f162b90e0ad 100644
--- a/CHANGELOG.md
+++ b/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
 

From 3ee78d152bf1dd5e4b242e8c92e72f7d22e12d9a Mon Sep 17 00:00:00 2001
From: Jonah Maxwell Miller <jonahm@lanl.gov>
Date: Tue, 23 Apr 2024 09:14:51 -0600
Subject: [PATCH 5/9] pgrete comments

---
 src/interface/mesh_data.cpp | 5 +++--
 src/interface/mesh_data.hpp | 6 +++---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/interface/mesh_data.cpp b/src/interface/mesh_data.cpp
index a17779da4432..212d762a2461 100644
--- a/src/interface/mesh_data.cpp
+++ b/src/interface/mesh_data.cpp
@@ -37,8 +37,9 @@ void MeshData<T>::Set(BlockList_t blocks, Mesh *pmesh) {
 }
 
 template <typename T>
-std::shared_ptr<MeshBlock> MeshData<T>::GetBlock(const int i, const MeshData<T> *src,
-                                                 const Mesh *pmesh) const {
+std::shared_ptr<MeshBlock> MeshData<T>::GetSourceBlockPointer_(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];
diff --git a/src/interface/mesh_data.hpp b/src/interface/mesh_data.hpp
index 1779769ac8ef..8e3b81278d6d 100644
--- a/src/interface/mesh_data.hpp
+++ b/src/interface/mesh_data.hpp
@@ -256,7 +256,7 @@ class MeshData {
     grid = src->grid;
     for (int i = 0; i < nblocks; ++i) {
       block_data_[i] =
-          GetBlock(i, src, pmy_mesh_)
+          GetSourceBlockPointer_(i, src, pmy_mesh_)
               ->meshblock_data.Add(stage_name_, src->GetBlockData(i), vars, shallow);
     }
   }
@@ -445,8 +445,8 @@ class MeshData {
   SparsePackCache &GetSparsePackCache() { return sparse_pack_cache_; }
 
  private:
-  std::shared_ptr<MeshBlock> GetBlock(const int i, const MeshData<T> *src,
-                                      const Mesh *pmesh) const;
+  std::shared_ptr<MeshBlock> GetSourceBlockPointer_(const int i, const MeshData<T> *src,
+                                                    const Mesh *pmesh) const;
 
   int ndim_;
   Mesh *pmy_mesh_;

From 9c7f92019932f5a84f2070fdfe6b05e70e3cb818 Mon Sep 17 00:00:00 2001
From: Jonah Maxwell Miller <jonahm@lanl.gov>
Date: Tue, 23 Apr 2024 09:15:56 -0600
Subject: [PATCH 6/9] add comment in test

---
 tst/unit/test_data_collection.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tst/unit/test_data_collection.cpp b/tst/unit/test_data_collection.cpp
index 8b5c87ca22a4..3e84480c1e7e 100644
--- a/tst/unit/test_data_collection.cpp
+++ b/tst/unit/test_data_collection.cpp
@@ -91,7 +91,8 @@ TEST_CASE("Adding MeshBlockData objects to a DataCollection", "[DataCollection]"
       }
     }
     AND_WHEN("We want only a subset of variables in a new MeshBlockData") {
-      // reset vars
+      // reset vars so that we can check this is overwritten/or is a
+      // new stage
       par_for(
           loop_pattern_flatrange_tag, "init vars", DevExecSpace(), 0, 0,
           KOKKOS_LAMBDA(const int i) { v2(0) = 222; });
@@ -117,7 +118,8 @@ TEST_CASE("Adding MeshBlockData objects to a DataCollection", "[DataCollection]"
       }
     }
     AND_WHEN("We want only a subset of variables in a new MeshBlockData by UID") {
-      // reset vars
+      // reset vars so that we can check this is overwritten/or is a
+      // new stage
       par_for(
           loop_pattern_flatrange_tag, "init vars", DevExecSpace(), 0, 0,
           KOKKOS_LAMBDA(const int i) { v2(0) = 222; });

From 88466b44fc1d34c298b2a68935f6e883a8c8a990 Mon Sep 17 00:00:00 2001
From: Jonah Maxwell Miller <jonahm@lanl.gov>
Date: Tue, 23 Apr 2024 17:13:46 -0600
Subject: [PATCH 7/9] lroberts comments

---
 src/interface/data_collection.hpp |  5 ++---
 src/interface/mesh_data.cpp       | 12 ------------
 src/interface/mesh_data.hpp       | 24 ++++++++++++++++++------
 src/interface/meshblock_data.hpp  | 10 ++++++++++
 4 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/src/interface/data_collection.hpp b/src/interface/data_collection.hpp
index ae3788364bd4..362650fc7ff3 100644
--- a/src/interface/data_collection.hpp
+++ b/src/interface/data_collection.hpp
@@ -49,9 +49,8 @@ class DataCollection {
                           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");
+      if (!(it->second)->ContainsExactly(fields)) {
+        PARTHENON_THROW(name + " already exists in collection but fields do not match.");
       }
       return it->second;
     }
diff --git a/src/interface/mesh_data.cpp b/src/interface/mesh_data.cpp
index 212d762a2461..8e9172024b8b 100644
--- a/src/interface/mesh_data.cpp
+++ b/src/interface/mesh_data.cpp
@@ -36,18 +36,6 @@ void MeshData<T>::Set(BlockList_t blocks, Mesh *pmesh) {
   Set(blocks, pmesh, ndim);
 }
 
-template <typename T>
-std::shared_ptr<MeshBlock> MeshData<T>::GetSourceBlockPointer_(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
diff --git a/src/interface/mesh_data.hpp b/src/interface/mesh_data.hpp
index 8e3b81278d6d..e7f8c16e85b4 100644
--- a/src/interface/mesh_data.hpp
+++ b/src/interface/mesh_data.hpp
@@ -253,11 +253,19 @@ class MeshData {
     const int nblocks = src->NumBlocks();
     block_data_.resize(nblocks);
 
+    // TODO(JMM/LFR): There is an edge case where if you call
+    // Initialize() on a set of meshblocks where some blocks contain
+    // the desired MeshBlockData object and some don't, this call will
+    // fail. (It will raise a runtime error due to a dictionary not
+    // being found.) This was present in the previous iteration of
+    // this code, as well as this iteration. Fixing this requires
+    // modifying DataCollection::GetOrAdd. In the future we should
+    // make that "just work (tm)."
     grid = src->grid;
     for (int i = 0; i < nblocks; ++i) {
-      block_data_[i] =
-          GetSourceBlockPointer_(i, src, pmy_mesh_)
-              ->meshblock_data.Add(stage_name_, src->GetBlockData(i), vars, shallow);
+      auto pmbd = src->GetBlockData(i);
+      block_data_[i] = pmbd->GetBlockSharedPointer()->meshblock_data.Add(
+          stage_name_, pmbd, vars, shallow);
     }
   }
 
@@ -436,18 +444,22 @@ class MeshData {
     return true;
   }
 
+  // vars may be a subset of the MeshData object
   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); });
   }
+  // MeshData object must contain these vars and only these vars
+  template <typename Vars_t>
+  bool ContainsExactly(const Vars_t &vars) const noexcept {
+    return std::all_of(block_data_.begin(), block_data_.end(),
+                       [this, vars](const auto &b) { return b->ContainsExactly(vars); });
+  }
 
   SparsePackCache &GetSparsePackCache() { return sparse_pack_cache_; }
 
  private:
-  std::shared_ptr<MeshBlock> GetSourceBlockPointer_(const int i, const MeshData<T> *src,
-                                                    const Mesh *pmesh) const;
-
   int ndim_;
   Mesh *pmy_mesh_;
   BlockDataList_t<T> block_data_;
diff --git a/src/interface/meshblock_data.hpp b/src/interface/meshblock_data.hpp
index 90134020a789..f3e9583133b8 100644
--- a/src/interface/meshblock_data.hpp
+++ b/src/interface/meshblock_data.hpp
@@ -17,6 +17,7 @@
 #include <limits>
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
 #include <utility>
 #include <vector>
@@ -420,6 +421,12 @@ class MeshBlockData {
     return std::all_of(vars.begin(), vars.end(),
                        [this](const auto &v) { return this->Contains(v); });
   }
+  template <typename ID_t>
+  bool ContainsExactly(const std::vector<ID_t> &vars) const noexcept {
+    // JMM: Assumes vars contains no duplicates. But that would have
+    // been caught elsewhere because `MeshBlockData::Add` would have failed.
+    return Contains(vars) && (vars.size() == varVector_.size());
+  }
 
   void SetAllVariablesToInitialized() {
     std::for_each(varVector_.begin(), varVector_.end(),
@@ -441,6 +448,9 @@ class MeshBlockData {
                 int sparse_id = InvalidSparseID);
 
   void Add(std::shared_ptr<Variable<T>> var) noexcept {
+    if (varUidMap_.count(var->GetUniqueID())) {
+      PARTHENON_THROW("Tried to add variable " + var->label() + " twice!");
+    }
     varVector_.push_back(var);
     varMap_[var->label()] = var;
     varUidMap_[var->GetUniqueID()] = var;

From f3df0b5b89c08cb4c5d70a90ae9d7806393ab9dd Mon Sep 17 00:00:00 2001
From: Jonah Maxwell Miller <jonahm@lanl.gov>
Date: Tue, 23 Apr 2024 17:29:23 -0600
Subject: [PATCH 8/9] special case fields empty

---
 src/interface/data_collection.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/interface/data_collection.hpp b/src/interface/data_collection.hpp
index 362650fc7ff3..d9067c768d08 100644
--- a/src/interface/data_collection.hpp
+++ b/src/interface/data_collection.hpp
@@ -49,7 +49,7 @@ class DataCollection {
                           const std::vector<ID_t> &fields, const bool shallow) {
     auto it = containers_.find(name);
     if (it != containers_.end()) {
-      if (!(it->second)->ContainsExactly(fields)) {
+      if ((fields.size() > 0) && !(it->second)->ContainsExactly(fields)) {
         PARTHENON_THROW(name + " already exists in collection but fields do not match.");
       }
       return it->second;

From 3b24f96ef9bb0d82bdd7eab00354fb7bd7150e31 Mon Sep 17 00:00:00 2001
From: Jonah Maxwell Miller <jonahm@lanl.gov>
Date: Tue, 23 Apr 2024 17:31:06 -0600
Subject: [PATCH 9/9] no need to > 0 a size_t in a bool

---
 src/interface/data_collection.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/interface/data_collection.hpp b/src/interface/data_collection.hpp
index d9067c768d08..8dea579728c4 100644
--- a/src/interface/data_collection.hpp
+++ b/src/interface/data_collection.hpp
@@ -49,7 +49,7 @@ class DataCollection {
                           const std::vector<ID_t> &fields, const bool shallow) {
     auto it = containers_.find(name);
     if (it != containers_.end()) {
-      if ((fields.size() > 0) && !(it->second)->ContainsExactly(fields)) {
+      if (fields.size() && !(it->second)->ContainsExactly(fields)) {
         PARTHENON_THROW(name + " already exists in collection but fields do not match.");
       }
       return it->second;