Pack on virtual fields

doc/sphinx/src/sparse_packs.rst

@@ -85,6 +85,83 @@ snippet.
    care should be taken not to assume that scratch variables can persist between tasks, even when they are directly
    dependent on each other. Issues resolved by setting ``-DPARTHENON_DEBUG_SCRATCH=ON`` can be an indication of this issue.
 
+Slicing into ``SparsePack``\ s with ``SubPack``\ s
+--------------------------------------------------
+
+A `SubPack` provdies a view into a slice of a `SparsePack` along a given dimension(s).
+`SubPack`\ s are constructed with a block + `kji` index to give a slice into the 
+fields at the meshblock + cell index. When a `SubPack` is constructed with the
+`Axis` template parameters then the `SubPack` also providies a view into slices
+of the `SparsePack` along the provided axes offset from the provided `ijk` indices.
+
+.. code:: c++
+
+  const int ni = ib.e - ib.s + 1;
+  const int ic = ib.s + ni / 2;
+  par_for(
+      PARTHENON_AUTO_LABEL, 0, sparse_pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e,
+      KOKKOS_LAMBDA(int b, int k, int j) {
+        int ltot = 0;
+        int lo = sparse_pack.GetLowerBound(b, v3());
+        int hi = sparse_pack.GetUpperBound(b, v3());
+        auto sub_pack = parthenon::SubPack<Axis::I>(sparse_pack, b, k, j, ic);
+
+        for (int i = ib.s - ni / 2; i <= ib.e - ni / 2; i++) {
+          for (int c = 0; c <= hi - lo; ++c) {
+            Real n = i + ic + 1e1 * j + 1e2 * k + 1e4 * c + 1e5 * v + 1e3 * b;
+            // indexes into sparse_pack(b, v3(c), k, j, ic + i)
+            if (n != sub_pack(v3(c), i)) ltot += 1;
+          }
+        }
+      });
+
+Using Virtual Fields
+--------------------
+
+Type-based ``SparsePack``\ s, when indexed with special ``virtual_variable_t`` types will return a value that is
+calculated from other variables in the pack from the ``evaluate`` function registered to the type. Virtual variable types
+need to export the other fields that their ``evaluate`` method depends on, allowing the virtual variable type, when included
+in the ``PackDescriptor``, to pack the real variables. 
+
+At the simplest a virtual variable's ``evaluate`` method has its signature take in a ``SparsePack<Ts...>`` as well as the ``i,j,k``
+indices and returns a ``Real`` as the value for the virtual field. 
+
+Additionally a virtual variable can declare a ``pack_type`` to refer to a 0D, 1D, 2D, or 3D ``SubPack``, which case the ``evaluate``
+signature only takes in the desired ``SubPack`` centered at teh ``i,j,k`` index.
+
+.. code:: c++
+
+   struct d1 : parthenon::variable_names::virtual_variable_t<v1, v3> {
+     KOKKOS_INLINE_FUNCTION d1() : x(0.0) {}
+
+     KOKKOS_INLINE_FUNCTION d1(const Real &xx) : x(xx) {}
+
+     template <typename Pack_t>
+     KOKKOS_INLINE_FUNCTION Real evaluate(const Pack_t &pack, const int b, const int k,
+                                          const int j, const int i) const {
+       return pack(b, v1(), k, j, i) * pack(b, v3(1), k, j, i) * pack(b, v3(2), k, j, i) + x;
+     }
+
+     Real x;
+   };
+
+   // use a subpack to index into the sparse pack
+   struct d1_subpack : public parthenon::variable_names::virtual_variable_t<v1, v3> {
+     // declare the type of subpack we want to use for our evaluate method
+     using pack_type = parthenon::SubPack0D;
+
+     KOKKOS_INLINE_FUNCTION d1_subpack() : x(0.0) {}
+
+     KOKKOS_INLINE_FUNCTION d1_subpack(const Real &xx) : x(xx) {}
+
+     template <typename Pack_t>
+     KOKKOS_INLINE_FUNCTION Real evaluate(const Pack_t &pack) const {
+       return pack(v1()) * pack(v3(1)) * pack(v3(2)) + x;
+     }
+
+     Real x;
+   };
+
 
 Building and Using a ``SparsePack``
 -----------------------------------

src/basic_types.hpp

@@ -177,6 +177,9 @@ inline bool operator<(const GridIdentifier &lhs, const GridIdentifier &rhs) {
   return lhs.logical_level() < rhs.logical_level();
 }
 
+// Enumeration for accessing spatial axes of a meshblock
+enum class Axis { I = 0, J = 1, K = 2 };
+
 // Enumeration for accessing a field on different locations of the grid:
 // CC = cell center of (i, j, k)
 // F1 = x-normal face at (i - 1/2, j, k)

src/pack/pack_utils.hpp

@@ -22,7 +22,12 @@
 #include <utility>
 #include <vector>
 
+#include <Kokkos_Core.hpp>
+
 #include "basic_types.hpp"
+#include "utils/concepts_lite.hpp"
+#include "utils/error_checking.hpp"
+#include "utils/type_list.hpp"
 
 // SFINAE for block iter so that Sparse/SwarmPacks can work for MeshBlockData and MeshData
 namespace {
@@ -161,6 +166,65 @@ struct any_nonautoflux : public base_t<true> {
   }
 };
 using any = any_nonautoflux;
+
+// Concept to state that a typed-field is dependent on other
+// fields, and is not itself an actual indexable field.
+template <typename T>
+concept DependentVariable =
+    requires { typename T::independent_vars; }; // NOLINT(readability/braces)
+
+template <typename... Ts>
+struct virtual_variable_t {
+  using type = virtual_variable_t<Ts...>;
+  using independent_vars = TypeList<Ts...>;
+};
+
+// Concept to check that a type shares an ancestor with virtual_variable_t
+template <typename T>
+concept VirtualVariable = requires {
+  typename T::type;
+  requires is_specialization_of<typename T::type, virtual_variable_t>::value;
+}; // NOLINT(readability/braces)
+
+// Concept to check that a type wants a subpack to evaluate the virtual field
+template <typename T>
+concept VirtualSubPack = requires {
+  typename T::pack_type;
+  requires std::is_same_v<decltype(T::pack_type::Naxes), const int>;
+}; // NOLINT(readability/braces)
+
+namespace impl {
+
+struct AllIndependentVariables {
+  template <typename T>
+    requires(!DependentVariable<T>)
+  static auto get(T) {
+    return TypeList<T>();
+  }
+
+  template <template <typename...> typename TL, typename... Ts>
+  static auto get(TL<Ts...>) {
+    return AllIndependentVariables::get(Ts()...);
+  }
+
+  template <typename T>
+    requires(DependentVariable<T>)
+  static auto get(T) {
+    return AllIndependentVariables::get(typename T::independent_vars());
+  }
+
+  template <typename T, typename... Ts>
+  static auto get(T, Ts...) {
+    return union_type_lists_t<decltype(AllIndependentVariables::get(T())),
+                              decltype(AllIndependentVariables::get(Ts()...))>();
+  }
+};
+} // namespace impl
+
+// Get a TypeList of all the independent variables that make up the requested types.
+template <typename... Ts>
+using all_independent_variables_t = decltype(impl::AllIndependentVariables::get(Ts()...));
+
 } // namespace variable_names
 
 // Namespace in which to put swarm variable name types that are used for indexing into

src/pack/sparse_pack/make_pack_descriptor.hpp

@@ -24,8 +24,10 @@
 #include "interface/metadata.hpp"
 #include "interface/state_descriptor.hpp"
 #include "mesh/mesh.hpp"
+#include "pack/pack_utils.hpp"
 #include "pack/sparse_pack/pack_descriptor.hpp"
 #include "pack/sparse_pack/sparse_pack.hpp"
+#include "utils/concepts_lite.hpp"
 #include "utils/type_list.hpp"
 
 namespace parthenon {
@@ -63,6 +65,7 @@ inline auto MakePackDescriptor(MT *pmd, const std::vector<std::string> &vars,
 }
 
 template <class... Ts, class MT>
+  requires(!variable_names::DependentVariable<Ts> && ...)
 inline auto MakePackDescriptor(MT *pmd, const std::vector<MetadataFlag> &flags = {},
                                const std::set<PDOpt> &options = {}) {
   const std::vector<std::string> vars{Ts::name()...};
@@ -72,6 +75,7 @@ inline auto MakePackDescriptor(MT *pmd, const std::vector<MetadataFlag> &flags =
 }
 
 template <class... Ts, class MT>
+  requires(!variable_names::DependentVariable<Ts> && ...)
 inline auto MakePackDescriptor(SparsePack<Ts...> pack, MT *pmd,
                                const std::vector<MetadataFlag> &flags = {},
                                const std::set<PDOpt> &options = {}) {
@@ -110,6 +114,14 @@ inline auto MakePackDescriptorFromTypeList(Args &&...args) {
   return MakePackDescriptorFromTypeList(TL(), std::forward<Args>(args)...);
 }
 
+template <class... Ts, class MT>
+  requires(variable_names::DependentVariable<Ts> || ...)
+inline auto MakePackDescriptor(MT *pmd, const std::vector<MetadataFlag> &flags = {},
+                               const std::set<PDOpt> &options = {}) {
+  return MakePackDescriptorFromTypeList(
+      variable_names::all_independent_variables_t<Ts...>(), pmd, flags, options);
+}
+
 struct PackDescriptorCacheBase {
   virtual ~PackDescriptorCacheBase() = default;
 };

src/pack/sparse_pack/sparse_pack.hpp

@@ -21,6 +21,7 @@
 #include "pack/block_selector.hpp"
 #include "pack/pack_utils.hpp"
 #include "pack/sparse_pack/sparse_pack_base.hpp"
+#include "pack/subpack.hpp"
 #include "utils/concepts_lite.hpp"
 #include "utils/type_list.hpp"
 
@@ -340,6 +341,51 @@ class SparsePack : public SparsePackBase {
     return pack_(static_cast<int>(el) % 3, b, vidx)(k, j, i);
   }
 
+  template <typename Tin, typename... Args>
+    requires(variable_names::VirtualVariable<Tin> && !variable_names::VirtualSubPack<Tin>)
+  KOKKOS_INLINE_FUNCTION Real operator()(const int b, const Tin &t, const int k,
+                                         const int j, const int i, Args &&...args) const {
+    static_assert(
+        TypeList<Ts...>::IsIn(typename Tin::independent_vars()),
+        "SparsePack must pack all independent_vars needed for virtual variable");
+    return t.evaluate(*this, b, k, j, i, std::forward<Args>(args)...);
+  }
+
+  template <typename Tin, typename... Args>
+    requires(variable_names::VirtualVariable<Tin> && variable_names::VirtualSubPack<Tin>)
+  KOKKOS_INLINE_FUNCTION Real operator()(const int b, const Tin &t, const int k,
+                                         const int j, const int i, Args &&...args) const {
+    static_assert(
+        TypeList<Ts...>::IsIn(typename Tin::independent_vars()),
+        "SparsePack must pack all independent_vars needed for virtual variable");
+    using sub_pack_type = decltype(SubPack<typename Tin::pack_type>(*this, b, k, j, i));
+    return t.evaluate(SubPack<typename Tin::pack_type>(*this, b, k, j, i),
+                      std::forward<Args>(args)...);
+  }
+  template <typename Tin, typename... Args>
+    requires(variable_names::VirtualVariable<Tin> && variable_names::VirtualSubPack<Tin>)
+  KOKKOS_INLINE_FUNCTION Real operator()(const int b, const TE el, const Tin &t,
+                                         const int k, const int j, const int i,
+                                         Args &&...args) const {
+    static_assert(
+        TypeList<Ts...>::IsIn(typename Tin::independent_vars()),
+        "SparsePack must pack all independent_vars needed for virtual variable");
+    using sub_pack_type = decltype(SubPack<typename Tin::pack_type>(*this, b, k, j, i));
+    return t.evaluate(SubPack<typename Tin::pack_type>(*this, b, k, j, i), el,
+                      std::forward<Args>(args)...);
+  }
+
+  template <typename Tin, typename... Args>
+    requires(variable_names::VirtualVariable<Tin> && !variable_names::VirtualSubPack<Tin>)
+  KOKKOS_INLINE_FUNCTION Real operator()(const int b, const TE el, const Tin &t,
+                                         const int k, const int j, const int i,
+                                         Args &&...args) const {
+    static_assert(
+        TypeList<Ts...>::IsIn(typename Tin::independent_vars()),
+        "SparsePack must pack all independent_vars needed for virtual variable");
+    return t.evaluate(*this, b, el, k, j, i, std::forward<Args>(args)...);
+  }
+
   // flux() overloads
   template <class... Args>
   KOKKOS_INLINE_FUNCTION auto &flux(const int b, const TopologicalElement te,