fix some issues with vector sparse fields, make restart work with tru…

…e sparse, need to update gold files for sparse_advection test

example/sparse_advection/parthenon_app_inputs.cpp

@@ -85,6 +85,43 @@ void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
           });
     }
   }
+
+  if (pkg->Param<bool>("restart_test")) {
+    bool any_nonzero = false;
+
+    for (int k = kb.s; k <= kb.e; k++) {
+      for (int j = jb.s; j <= jb.e; j++) {
+        for (int i = ib.s; i <= ib.e; i++) {
+          auto x = coords.x1v(i);
+          auto y = coords.x2v(j);
+          auto z = coords.x3v(k);
+          auto r2 = x * x + y * y + z * z;
+          if (r2 < 0.5 * size) {
+            any_nonzero = true;
+          }
+        }
+      }
+    }
+
+    if (any_nonzero) {
+      data->AllocSparseID("z_shape_shift", 1);
+      data->AllocSparseID("z_shape_shift", 3);
+      data->AllocSparseID("z_shape_shift", 4);
+
+      auto v = data->PackVariables(
+          std::vector<std::string>{"z_dense_A", "z_dense_B", "z_shape_shift"});
+
+      pmb->par_for(
+          "SparseAdvection::ProblemGenerator", 0, v.GetDim(4) - 1, kb.s, kb.e, jb.s, jb.e,
+          ib.s, ib.e, KOKKOS_LAMBDA(const int n, const int k, const int j, const int i) {
+            auto x = coords.x1v(i);
+            auto y = coords.x2v(j);
+            auto z = coords.x3v(k);
+            auto r2 = x * x + y * y + z * z;
+            v(n, k, j, i) = (r2 < 0.5 * size ? 1.0 : 0.0);
+          });
+    }
+  }
 }
 
 //========================================================================================

example/sparse_advection/parthinput.sparse_advection

@@ -58,6 +58,8 @@ integrator = rk2
 ncycle_out_mesh = -10000
 
 <SparseAdvection>
+restart_test = true
+
 cfl = 0.45
 speed = 1.5
 
@@ -71,7 +73,7 @@ dt = 0.5
 <parthenon/output0>
 file_type = hdf5
 dt = 0.1
-variables = sparse
+variables = sparse, z_dense_A, z_dense_B, z_shape_shift
 
 <parthenon/output3>
 file_type = hst

example/sparse_advection/sparse_advection_package.cpp

@@ -22,6 +22,7 @@
 #include <parthenon/package.hpp>
 
 #include "defs.hpp"
+#include "interface/metadata.hpp"
 #include "interface/sparse_pool.hpp"
 #include "kokkos_abstraction.hpp"
 #include "reconstruct/dc_inline.hpp"
@@ -42,6 +43,9 @@ using parthenon::UserHistoryOperation;
 std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   auto pkg = std::make_shared<StateDescriptor>("sparse_advection_package");
 
+  bool restart_test = pin->GetOrAddBoolean("SparseAdvection", "restart_test", false);
+  pkg->AddParam("restart_test", restart_test);
+
   Real cfl = pin->GetOrAddReal("SparseAdvection", "cfl", 0.45);
   pkg->AddParam("cfl", cfl);
   Real refine_tol = pin->GetOrAddReal("SparseAdvection", "refine_tol", 0.3);
@@ -64,15 +68,37 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   pkg->AddParam("vz", RealArr_t{0.0, 0.0, 0.0, 0.0});
 
   // add sparse field
-  Metadata m({Metadata::Cell, Metadata::Independent, Metadata::WithFluxes,
-              Metadata::FillGhost, Metadata::Sparse},
-             std::vector<int>({1}));
-  parthenon::SparsePool pool("sparse", m);
+  {
+    Metadata m({Metadata::Cell, Metadata::Independent, Metadata::WithFluxes,
+                Metadata::FillGhost, Metadata::Sparse},
+               std::vector<int>({1}));
+    SparsePool pool("sparse", m);
+
+    for (int sid = 0; sid < NUM_FIELDS; ++sid) {
+      pool.Add(sid);
+    }
+    pkg->AddSparsePool(pool);
+  }
+
+  // add fields for restart test ("Z" prefix so they are after sparse in alphabetical
+  // list, helps with reusing velocity vectors)
+  if (restart_test) {
+    Metadata m_dense({Metadata::Cell, Metadata::Independent, Metadata::WithFluxes,
+                      Metadata::FillGhost});
+    pkg->AddField("z_dense_A", m_dense);
+    pkg->AddField("z_dense_B", m_dense);
+
+    Metadata m_sparse({Metadata::Cell, Metadata::Independent, Metadata::WithFluxes,
+                       Metadata::FillGhost, Metadata::Sparse});
 
-  for (int sid = 0; sid < NUM_FIELDS; ++sid) {
-    pool.Add(sid);
+    SparsePool pool("z_shape_shift", m_sparse);
+    pool.Add(1, std::vector<int>{1}, std::vector<std::string>{"scalar"});
+    pool.Add(3, std::vector<int>{3}, Metadata::Vector,
+             std::vector<std::string>{"vec_x", "vec_y", "vec_z"});
+    pool.Add(4, std::vector<int>{4}, Metadata::Vector);
+
+    pkg->AddSparsePool(pool);
   }
-  pkg->AddSparsePool(pool);
 
   pkg->CheckRefinementBlock = CheckRefinement;
   pkg->EstimateTimestepBlock = EstimateTimestepBlock;
@@ -182,22 +208,19 @@ TaskStatus CalculateFluxes(std::shared_ptr<MeshBlockData<Real>> &rc) {
       KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int k, const int j) {
         parthenon::ScratchPad2D<Real> ql(member.team_scratch(scratch_level), nvar, nx1);
         parthenon::ScratchPad2D<Real> qr(member.team_scratch(scratch_level), nvar, nx1);
+
         // get reconstructed state on faces
         parthenon::DonorCellX1(member, k, j, ib.s - 1, ib.e + 1, v, ql, qr);
+
         // Sync all threads in the team so that scratch memory is consistent
         member.team_barrier();
 
         for (int n = 0; n < nvar; n++) {
           if (!v.IsAllocated(n)) continue;
-          if (vx[n] > 0.0) {
-            par_for_inner(member, ib.s, ib.e + 1, [&](const int i) {
-              v.flux(X1DIR, n, k, j, i) = ql(n, i) * vx[n];
-            });
-          } else {
-            par_for_inner(member, ib.s, ib.e + 1, [&](const int i) {
-              v.flux(X1DIR, n, k, j, i) = qr(n, i) * vx[n];
-            });
-          }
+          const auto this_v = vx[n % NUM_FIELDS];
+          par_for_inner(member, ib.s, ib.e + 1, [&](const int i) {
+            v.flux(X1DIR, n, k, j, i) = (this_v > 0.0 ? ql(n, i) : qr(n, i)) * this_v;
+          });
         }
       });
 
@@ -206,31 +229,28 @@ TaskStatus CalculateFluxes(std::shared_ptr<MeshBlockData<Real>> &rc) {
     pmb->par_for_outer(
         "x2 flux", 3 * scratch_size_in_bytes, scratch_level, kb.s, kb.e, jb.s, jb.e + 1,
         KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int k, const int j) {
-          // the overall algorithm/use of scratch pad here is clear inefficient and kept
+          // the overall algorithm/use of scratch pad here is clearly inefficient and kept
           // just for demonstrating purposes. The key point is that we cannot reuse
           // reconstructed arrays for different `j` with `j` being part of the outer
           // loop given that this loop can be handled by multiple threads simultaneously.
-
           parthenon::ScratchPad2D<Real> ql(member.team_scratch(scratch_level), nvar, nx1);
           parthenon::ScratchPad2D<Real> qr(member.team_scratch(scratch_level), nvar, nx1);
           parthenon::ScratchPad2D<Real> q_unused(member.team_scratch(scratch_level), nvar,
                                                  nx1);
+
           // get reconstructed state on faces
           parthenon::DonorCellX2(member, k, j - 1, ib.s, ib.e, v, ql, q_unused);
           parthenon::DonorCellX2(member, k, j, ib.s, ib.e, v, q_unused, qr);
+
           // Sync all threads in the team so that scratch memory is consistent
           member.team_barrier();
+
           for (int n = 0; n < nvar; n++) {
             if (!v.IsAllocated(n)) continue;
-            if (vy[n] > 0.0) {
-              par_for_inner(member, ib.s, ib.e, [&](const int i) {
-                v.flux(X2DIR, n, k, j, i) = ql(n, i) * vy[n];
-              });
-            } else {
-              par_for_inner(member, ib.s, ib.e, [&](const int i) {
-                v.flux(X2DIR, n, k, j, i) = qr(n, i) * vy[n];
-              });
-            }
+            const auto this_v = vy[n % NUM_FIELDS];
+            par_for_inner(member, ib.s, ib.e, [&](const int i) {
+              v.flux(X2DIR, n, k, j, i) = (this_v > 0.0 ? ql(n, i) : qr(n, i)) * this_v;
+            });
           }
         });
   }

scripts/python/packages/parthenon_tools/parthenon_tools/phdf_diff.py

@@ -172,11 +172,11 @@ def compare_attribute_group(f0, f1, name):
 
             print("\nFirst file:")
             for k in diffs:
-                print("%20s: " % k, group0[k])
+                print("%20s: " % k, group0[k] if k in group0 else "<does not exist>")
 
             print("\nSecond file:")
             for k in diffs:
-                print("%20s: " % k, group1[k])
+                print("%20s: " % k, group1[k] if k in group1 else "<does not exist>")
         else:
             print("  %20s: no diffs" % name)
 

src/interface/meshblock_data.hpp

@@ -142,10 +142,8 @@ class MeshBlockData {
 
   std::shared_ptr<CellVariable<T>> GetCellVarPtr(const std::string &label) const {
     auto it = varMap_.find(label);
-    if (it == varMap_.end()) {
-      PARTHENON_THROW(std::string("\n") + std::string(label) +
-                      std::string(" array not found in Get()\n"));
-    }
+    PARTHENON_REQUIRE_THROWS(it != varMap_.end(),
+                             "Couldn't find variable '" + label + "'");
     return it->second;
   }
 

src/interface/metadata.cpp

@@ -120,7 +120,6 @@ std::vector<MetadataFlag> Metadata::Flags() const {
 std::array<int, 6> Metadata::GetArrayDims(std::weak_ptr<MeshBlock> wpmb,
                                           bool coarse) const {
   std::array<int, 6> arrDims;
-
   const auto &shape = shape_;
   const int N = shape.size();
 

src/interface/metadata.hpp

@@ -225,6 +225,16 @@ class Metadata {
       PARTHENON_REQUIRE_THROWS(
           shape_.size() <= 3,
           "Variables tied to mesh entities can only have a shape of rank <= 3");
+
+      int num_comp = 1;
+      for (auto s : shape) {
+        num_comp *= s;
+      }
+
+      PARTHENON_REQUIRE_THROWS(component_labels.size() == 0 ||
+                                   (component_labels.size() == num_comp),
+                               "Must provide either 0 component labels or the same "
+                               "number as the number of components");
     }
   }
 

src/interface/update.cpp

@@ -20,6 +20,7 @@
 #include "globals.hpp"
 #include "interface/meshblock_data.hpp"
 #include "interface/metadata.hpp"
+#include "interface/variable_pack.hpp"
 #include "mesh/mesh.hpp"
 #include "mesh/meshblock.hpp"
 
@@ -147,12 +148,9 @@ TaskStatus SparseDeallocCheck(MeshData<Real> *md) {
   const IndexRange jb = md->GetBoundsJ(IndexDomain::entire);
   const IndexRange kb = md->GetBoundsK(IndexDomain::entire);
 
+  PackIndexMap map;
   const std::vector<MetadataFlag> flags({Metadata::Sparse});
-  const auto &pack = md->PackVariables(flags);
-
-  // get list of variables in same order as they appear in the pack, since variable
-  // metadata is the same on all blocks, we can just use the first block
-  const auto &var_list = md->GetBlockData(0)->GetVariablesByFlag(flags, true);
+  const auto &pack = md->PackVariables(flags, map);
 
   const int num_blocks = pack.GetDim(5);
   const int num_vars = pack.GetDim(4);
@@ -205,10 +203,19 @@ TaskStatus SparseDeallocCheck(MeshData<Real> *md) {
   auto is_zero_h = Kokkos::create_mirror_view_and_copy(HostMemSpace(), is_zero);
 
   for (int b = 0; b < num_blocks; ++b) {
-    for (int v = 0; v < num_vars; ++v) {
-      auto &counter = md->GetBlockData(b)->Get(var_list.labels()[v]).dealloc_count;
-      if (is_zero_h(b, v)) {
-        // this var is zero, increment dealloc counter
+    for (auto var_itr : map.Map()) {
+      const auto label = var_itr.first;
+      auto &counter = md->GetBlockData(b)->Get(label).dealloc_count;
+      bool all_zero = true;
+      for (int v = var_itr.second.first; v <= var_itr.second.second; ++v) {
+        if (!is_zero_h(b, v)) {
+          all_zero = false;
+          break;
+        }
+      }
+
+      if (all_zero) {
+        // all components of this var are zero, increment dealloc counter
         counter += 1;
       } else {
         counter = 0;
@@ -217,9 +224,9 @@ TaskStatus SparseDeallocCheck(MeshData<Real> *md) {
       if (counter > Globals::sparse_config.deallocation_count) {
         // this variable has been flagged for deallocation deallocation_count times in a
         // row, now deallocate it
-        // printf("Deallocating var %s on block %i\n", var_list.labels()[v].c_str(),
+        // printf("Deallocating var %s on block %i\n", label.c_str(),
         //        md->GetBlockData(b)->GetBlockPointer()->gid);
-        md->GetBlockData(b)->GetBlockPointer()->DeallocateSparse(var_list.labels()[v]);
+        md->GetBlockData(b)->GetBlockPointer()->DeallocateSparse(label);
       }
     }
   }

src/interface/variable.hpp

@@ -72,6 +72,9 @@ class CellVariable {
     }
 
     if (IsSet(Metadata::FillGhost)) {
+      PARTHENON_REQUIRE_THROWS(
+          GetDim(4) == NumComponents(),
+          "CellCenteredBoundaryVariable currently only supports rank-1 variables");
       vbvar = std::make_shared<CellCenteredBoundaryVariable>(wpmb.lock(), IsSparse(),
                                                              label(), GetDim(4));
     }

src/interface/variable_pack.hpp

@@ -93,6 +93,8 @@ class PackIndexMap {
  public:
   PackIndexMap() = default;
 
+  const auto &Map() const { return map_; }
+
   const auto &get(const std::string &base_name, int sparse_id = InvalidSparseID) const {
     const auto &key = MakeVarLabel(base_name, sparse_id);
     auto itr = map_.find(key);

src/outputs/parthenon_hdf5.cpp

@@ -107,7 +107,7 @@ struct VarInfo {
   }
 
   explicit VarInfo(const std::shared_ptr<CellVariable<Real>> &var)
-      : VarInfo(var->label(), var->metadata().getComponentLabels(), var->GetDim(4),
+      : VarInfo(var->label(), var->metadata().getComponentLabels(), var->NumComponents(),
                 var->IsSparse(), var->IsSet(Metadata::Vector)) {}
 };
 
@@ -137,23 +137,25 @@ static void writeXdmfArrayRef(std::ofstream &fid, const std::string &prefix,
 }
 
 static void writeXdmfSlabVariableRef(std::ofstream &fid, const std::string &name,
+                                     const std::vector<std::string> &component_labels,
                                      std::string &hdfFile, int iblock, const int &vlen,
                                      int &ndims, hsize_t *dims,
                                      const std::string &dims321, bool isVector) {
   // writes a slab reference to file
-
   std::vector<std::string> names;
   int nentries = 1;
-  int vector_size = 1;
   if (vlen == 1 || isVector) {
+    // we only make one entry, because either vlen == 1, or we write this as a vector
     names.push_back(name);
   } else {
     nentries = vlen;
     for (int i = 0; i < vlen; i++) {
-      names.push_back(name + "_" + std::to_string(i));
+      names.push_back(
+          name + "_" +
+          (component_labels.empty() ? std::to_string(i) : component_labels[i]));
     }
   }
-  if (isVector) vector_size = vlen;
+  const int vector_size = isVector ? vlen : 1;
 
   const std::string prefix = "      ";
   for (int i = 0; i < nentries; i++) {
@@ -261,8 +263,8 @@ void genXDMF(std::string hdfFile, Mesh *pm, SimTime *tm, int nx1, int nx2, int n
     for (const auto &vinfo : var_list) {
       const int vlen = vinfo.vlen;
       dims[4] = vlen;
-      writeXdmfSlabVariableRef(xdmf, vinfo.label, hdfFile, ib, vlen, ndims, dims, dims321,
-                               vinfo.is_vector);
+      writeXdmfSlabVariableRef(xdmf, vinfo.label, vinfo.component_labels, hdfFile, ib,
+                               vlen, ndims, dims, dims321, vinfo.is_vector);
     }
     xdmf << "      </Grid>" << std::endl;
   }
@@ -325,7 +327,7 @@ void PHDF5Output::WriteOutputFileImpl(Mesh *pm, ParameterInput *pin, SimTime *tm
   filename.append(restart_ ? ".rhdf" : ".phdf");
 
   // After file has been opened with the current number, already advance output
-  // parameters so that for restarts the file is not immediatly overwritten again.
+  // parameters so that for restarts the file is not immediately overwritten again.
   output_params.file_number++;
   output_params.next_time += output_params.dt;
   pin->SetInteger(output_params.block_name, "file_number", output_params.file_number);