Allow for accessing flux Metadata from Metadata::WithFluxes Metadata

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[PR 1161]]((https://github.com/parthenon-hpc-lab/parthenon/pull/1161)) Add Metadata::FluxNotOneCopy
 - [[PR 1151]]((https://github.com/parthenon-hpc-lab/parthenon/pull/1151)) Add time offset `c` to LowStorageIntegrator
 - [[PR 1147]](https://github.com/parthenon-hpc-lab/parthenon/pull/1147) Add `par_reduce_inner` functions
 - [[PR 1159]](https://github.com/parthenon-hpc-lab/parthenon/pull/1159) Add additional timestep controllers in parthenon/time.

src/bvals/comms/boundary_communication.cpp

@@ -213,7 +213,7 @@ TaskStatus ReceiveBoundBufs(std::shared_ptr<MeshData<Real>> &md) {
       [&all_received](auto pbuf) { all_received = pbuf->TryReceive() && all_received; });
 
   int ibound = 0;
-  if (Globals::sparse_config.enabled) {
+  if (Globals::sparse_config.enabled && all_received) {
     ForEachBoundary<bound_type>(
         md, [&](auto pmb, sp_mbd_t rc, nb_t &nb, const sp_cv_t v) {
           const std::size_t ibuf = cache.idx_vec[ibound];

src/interface/meshblock_data.hpp

@@ -76,6 +76,9 @@ class MeshBlockData {
   MeshBlock *GetParentPointer() const { return GetBlockPointer(); }
   void SetAllowedDt(const Real dt) const { GetBlockPointer()->SetAllowedDt(dt); }
   Mesh *GetMeshPointer() const { return GetBlockPointer()->pmy_mesh; }
+
+  // This mirrors a MeshBlockData routine
+  int NumBlocks() const { return 1;}
 
   template <class... Ts>
   IndexRange GetBoundsI(Ts &&...args) const {

src/interface/metadata.cpp

@@ -251,6 +251,13 @@ bool Metadata::IsValid(bool throw_on_fail) const {
       PARTHENON_THROW("Either the Independent or Derived flag must be set");
     }
   }
+
+  if (IsSet(FillGhost) && IsSet(CellMemAligned) && (!IsSet(Cell))) {
+    valid = false;
+    if (throw_on_fail) {
+      PARTHENON_THROW("Cannot communicate ghosts of non-cell fields that have cell aligned memory.");
+    } 
+  }
 
   // Prolongation/restriction
   if (HasRefinementOps()) {
@@ -263,6 +270,13 @@ bool Metadata::IsValid(bool throw_on_fail) const {
     }
   }
 
+  // Associated fluxes
+  if (IsSet(FluxNotOneCopy)) {
+    PARTHENON_REQUIRE(
+        IsSet(WithFluxes),
+        "Asking for non-OneCopy associated fluxes without asking for associated fluxes.");
+  }
+
   return valid;
 }
 
@@ -306,19 +320,15 @@ Metadata::GetArrayDims(std::weak_ptr<MeshBlock> wpmb, bool coarse) const {
       arrDims[i + 3] = 1;
     if (IsSet(Cell)) {
       arrDims[MAX_VARIABLE_DIMENSION - 1] = 1; // Only one cell center per cell
-    } else if (IsSet(Face) && IsSet(Flux)) {
-      // 3 directions but keep the same ijk shape as cell var for performance
-      arrDims[MAX_VARIABLE_DIMENSION - 1] = 3;
     } else if (IsSet(Face) || IsSet(Edge)) {
       arrDims[MAX_VARIABLE_DIMENSION - 1] = 3; // Three faces and edges per cell
-      arrDims[0]++;
-      if (arrDims[1] > 1) arrDims[1]++;
-      if (arrDims[2] > 1) arrDims[2]++;
     } else if (IsSet(Node)) {
       arrDims[MAX_VARIABLE_DIMENSION - 1] = 1; // Only one lower left node per cell
+    }
+    if (!IsSet(CellMemAligned) && !IsSet(Cell)) { 
       arrDims[0]++;
       if (arrDims[1] > 1) arrDims[1]++;
-      if (arrDims[2] > 1) arrDims[2]++;
+      if (arrDims[2] > 1) arrDims[2]++; 
     }
   } else if (IsSet(Particle)) {
     assert(N >= 0 && N <= MAX_VARIABLE_DIMENSION - 1);

src/interface/metadata.hpp

@@ -120,6 +120,11 @@
   PARTHENON_INTERNAL_FOR_FLAG(Fine)                                                      \
   /** this variable is the flux for another variable **/                                 \
   PARTHENON_INTERNAL_FOR_FLAG(Flux)                                                      \
+  /** allocate a separate flux array for each stage if WithFluxes is specified**/        \
+  PARTHENON_INTERNAL_FOR_FLAG(FluxNotOneCopy)                                            \
+  /** Align memory of fields to cell centered memory                                     \
+      (Field will be missing one layer of ghosts if it is not cell centered) **/         \
+  PARTHENON_INTERNAL_FOR_FLAG(CellMemAligned)                                            \
   /************************************************/                                     \
   /** Vars specifying coordinates for visualization purposes **/                         \
   /** You can specify a single 3D var **/                                                \

src/interface/sparse_pack_base.cpp

@@ -67,19 +67,15 @@ SparsePackBase::GetAllocStatus(T *pmd, const PackDescriptor &desc,
                                const std::vector<bool> &include_block) {
   using mbd_t = MeshBlockData<Real>;
 
-  int nvar = desc.nvar_groups;
-
-  std::vector<int> astat;
+  const int nvar = desc.nvar_groups;
+  const int nblock = pmd->NumBlocks();
+  std::vector<int> astat(nblock * desc.nvar_tot);
+  int idx = 0;
   ForEachBlock(pmd, include_block, [&](int b, mbd_t *pmbd) {
     const auto &uid_map = pmbd->GetUidMap();
     for (int i = 0; i < nvar; ++i) {
       for (const auto &[var_name, uid] : desc.var_groups[i]) {
-        if (uid_map.count(uid) > 0) {
-          const auto pv = uid_map.at(uid);
-          astat.push_back(pv->GetAllocationStatus());
-        } else {
-          astat.push_back(-1);
-        }
+        astat[idx++] = uid_map.count(uid) > 0 ? (uid_map.at(uid))->GetAllocationStatus() : -1;
       }
     }
   });

src/interface/sparse_pack_base.hpp

@@ -138,7 +138,7 @@ struct PackDescriptor {
   // default constructor needed for certain use cases
   PackDescriptor()
       : nvar_groups(0), var_group_names({}), var_groups({}), with_fluxes(false),
-        coarse(false), flat(false), identifier("") {}
+        coarse(false), flat(false), identifier(""), nvar_tot(0) {}
 
   template <class GROUP_t, class SELECTOR_t>
   PackDescriptor(StateDescriptor *psd, const std::vector<GROUP_t> &var_groups_in,
@@ -147,7 +147,7 @@ struct PackDescriptor {
         var_groups(BuildUids(var_groups_in.size(), psd, selector)),
         with_fluxes(options.count(PDOpt::WithFluxes)),
         coarse(options.count(PDOpt::Coarse)), flat(options.count(PDOpt::Flatten)),
-        identifier(GetIdentifier()) {
+        identifier(GetIdentifier()), nvar_tot(GetNVarsTotal(var_groups)) {
     PARTHENON_REQUIRE(!(with_fluxes && coarse),
                       "Probably shouldn't be making a coarse pack with fine fluxes.");
   }
@@ -159,8 +159,18 @@ struct PackDescriptor {
   const bool coarse;
   const bool flat;
   const std::string identifier;
+  const std::size_t nvar_tot;
 
  private:
+  static int GetNVarsTotal(const std::vector<VariableGroup_t> &var_groups) { 
+    int nvar_tot = 0; 
+    for (const auto &group : var_groups) { 
+      for (const auto &[a, b] : group) { 
+        nvar_tot++;
+      }
+    }
+    return nvar_tot;
+  }
   std::string GetIdentifier() {
     std::string ident("");
     for (const auto &vgroup : var_groups) {

src/interface/state_descriptor.cpp

@@ -274,16 +274,18 @@ bool StateDescriptor::AddFieldImpl(const VarID &vid, const Metadata &m_in,
     return false; // this field has already been added
   } else {
     if (m.IsSet(Metadata::WithFluxes) && m.GetFluxName() == "") {
-      std::vector<MetadataFlag> mFlags = {Metadata::OneCopy, Metadata::Flux};
+      std::vector<MetadataFlag> mFlags = {Metadata::Flux};
+      if (!m.IsSet(Metadata::FluxNotOneCopy)) mFlags.push_back(Metadata::OneCopy);
       if (m.IsSet(Metadata::Sparse)) mFlags.push_back(Metadata::Sparse);
       if (m.IsSet(Metadata::Fine)) mFlags.push_back(Metadata::Fine);
-      if (m.IsSet(Metadata::Cell))
+      if (m.IsSet(Metadata::Cell)) {
         mFlags.push_back(Metadata::Face);
-      else if (m.IsSet(Metadata::Face))
+        mFlags.push_back(Metadata::CellMemAligned);
+      } else if (m.IsSet(Metadata::Face)) {
         mFlags.push_back(Metadata::Edge);
-      else if (m.IsSet(Metadata::Edge))
+      } else if (m.IsSet(Metadata::Edge)) {
         mFlags.push_back(Metadata::Node);
-
+      }
       Metadata mf;
       if (m.GetRefinementFunctions().label().size() > 0) {
         // Propagate custom refinement ops to flux field