Add GetPackDimension to StateDescriptor

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[PR 1148]](https://github.com/parthenon-hpc-lab/parthenon/pull/1148) Add `GetPackDimension` to `StateDescriptor` for calculating pack sizes before `Mesh` initialization
 - [[PR 1143]](https://github.com/parthenon-hpc-lab/parthenon/pull/1143) Add tensor indices to VariableState, add radiation constant to constants, add TypeLists, allow for arbitrary containers for solvers
 - [[PR 1140]](https://github.com/parthenon-hpc-lab/parthenon/pull/1140) Allow for relative convergence tolerance in BiCGSTAB solver.
 - [[PR 1047]](https://github.com/parthenon-hpc-lab/parthenon/pull/1047) General three- and four-valent 2D forests w/ arbitrary orientations.

src/interface/state_descriptor.cpp

@@ -566,4 +566,37 @@ StateDescriptor::GetVariableNames(const Metadata::FlagCollection &flags) {
   return GetVariableNames({}, flags, {});
 }
 
+// Get the total length of this StateDescriptor's variables when packed
+int StateDescriptor::GetPackDimension(const std::vector<std::string> &req_names,
+                                      const Metadata::FlagCollection &flags,
+                                      const std::vector<int> &sparse_ids) {
+  std::vector<std::string> names = GetVariableNames(req_names, flags, sparse_ids);
+  int nvar = 0;
+  for (auto name : names) {
+    const auto &meta = metadataMap_[VarID(name)];
+    int var_len = 0;
+    if (meta.Shape().size() < 1) {
+      var_len = 1;
+    } else {
+      var_len = meta.Shape()[0];
+    }
+    nvar += var_len;
+  }
+  return nvar;
+}
+int StateDescriptor::GetPackDimension(const std::vector<std::string> &req_names,
+                                      const std::vector<int> &sparse_ids) {
+  return GetPackDimension(req_names, Metadata::FlagCollection(), sparse_ids);
+}
+int StateDescriptor::GetPackDimension(const Metadata::FlagCollection &flags,
+                                      const std::vector<int> &sparse_ids) {
+  return GetPackDimension({}, flags, sparse_ids);
+}
+int StateDescriptor::GetPackDimension(const std::vector<std::string> &req_names) {
+  return GetPackDimension(req_names, Metadata::FlagCollection(), {});
+}
+int StateDescriptor::GetPackDimension(const Metadata::FlagCollection &flags) {
+  return GetPackDimension({}, flags, {});
+}
+
 } // namespace parthenon

src/interface/state_descriptor.hpp

@@ -276,6 +276,16 @@ class StateDescriptor {
   std::vector<std::string> GetVariableNames(const std::vector<std::string> &req_names);
   std::vector<std::string> GetVariableNames(const Metadata::FlagCollection &flags);
 
+  int GetPackDimension(const std::vector<std::string> &req_names,
+                       const Metadata::FlagCollection &flags,
+                       const std::vector<int> &sparse_ids);
+  int GetPackDimension(const std::vector<std::string> &req_names,
+                       const std::vector<int> &sparse_ids);
+  int GetPackDimension(const Metadata::FlagCollection &flags,
+                       const std::vector<int> &sparse_ids);
+  int GetPackDimension(const std::vector<std::string> &req_names);
+  int GetPackDimension(const Metadata::FlagCollection &flags);
+
   std::size_t
   RefinementFuncID(const refinement::RefinementFunctions_t &funcs) const noexcept {
     return refinementFuncMaps_.funcs_to_ids.at(funcs);