Merge pull request #1143 from parthenon-hpc-lab/lroberts36/add-small-…

…riot

Small: Quality of life changes for downstream codes

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[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.
 - [[PR 1130]](https://github.com/parthenon-hpc-lab/parthenon/pull/1130) Enable `parthenon::par_reduce` for MD loops with Kokkos 1D Range

doc/sphinx/src/utilities.rst

@@ -0,0 +1,48 @@
+``TypeList``s
+=============
+
+Provides a wrapper class around a variadic pack of types to simplify
+performing compile time operations on the pack. There are templated 
+types defined giving the type at a particular index in the pack, types 
+for extracting sub-``TypeList``s of the original type list, and ``constexpr``
+functions for getting the index of the first instance of a type in the 
+pack. Additionally it provides a capability for iterating an ``auto`` lambda
+over the type list, which can be useful for calling a ``static`` function 
+defined for each of the types on each of the types. *In the future, it
+would be nice to have the capability to make a unique type list from
+another type list (i.e. the unique one only a single instance of each type
+in the original type list)*
+
+``TypeList``s have many applications and are commonly found in many 
+codebases, but in Parthenon one of the main use cases is for storing 
+lists of types associated with field variables that are used in type 
+based ``SparsePack``s.  
+
+Robust
+======
+
+Provides a number of functions for doing operations on floating point
+numbers that are bounded, avoid division by zero, etc. 
+
+C++11 Style Concepts Implementation
+===================================
+
+*This documentation needs to be written (see issue #695), but there are
+extensive comments in src/utlils/concepts_lite.hpp and examples of
+useage in tst/unit/test_concepts_lite.hpp*
+
+``Indexer``
+===========
+
+Provides functionality for iterating over an arbitrary dimensional
+hyper-rectangular index space using a flattened loop. Specific 
+instantiations, e.g. ``Indexer5D``, are provided up to eight 
+dimensions. Useage:
+
+.. code:: cpp
+
+  Indexer4D idxer({0, 3}, {1, 2}, {0, 5}, {10, 16});
+  for (int flat_idx = 0; flat_idx < idxer.size(); ++flat_idx) {
+    auto [i, j, k, l] = idxer(flat_idx);
+    // Do stuff in the 4D index space...
+  }

src/CMakeLists.txt

@@ -270,6 +270,7 @@ add_library(parthenon
   utils/sort.hpp
   utils/string_utils.cpp
   utils/string_utils.hpp
+  utils/type_list.hpp
   utils/unique_id.cpp
   utils/unique_id.hpp
   utils/utils.hpp

src/interface/make_pack_descriptor.hpp

@@ -33,6 +33,7 @@
 #include "interface/state_descriptor.hpp"
 #include "interface/variable.hpp"
 #include "mesh/mesh.hpp"
+#include "utils/type_list.hpp"
 
 namespace parthenon {
 
@@ -136,6 +137,15 @@ inline auto MakePackDescriptor(StateDescriptor *psd, const std::vector<Uid_t> &v
   return typename SparsePack<>::Descriptor(base_desc);
 }
 
+template <template <class...> class TL, class... Types, class... Args>
+inline auto MakePackDescriptorFromTypeList(TL<Types...>, Args &&...args) {
+  return MakePackDescriptor<Types...>(std::forward<Args>(args)...);
+}
+
+template <class TL, class... Args>
+inline auto MakePackDescriptorFromTypeList(Args &&...args) {
+  return MakePackDescriptorFromTypeList(TL(), std::forward<Args>(args)...);
+}
 } // namespace parthenon
 
 #endif // INTERFACE_MAKE_PACK_DESCRIPTOR_HPP_

src/interface/sparse_pack.hpp

@@ -33,6 +33,7 @@
 #include "interface/sparse_pack_base.hpp"
 #include "interface/variable.hpp"
 #include "utils/concepts_lite.hpp"
+#include "utils/type_list.hpp"
 #include "utils/utils.hpp"
 
 namespace parthenon {

src/interface/sparse_pack_base.cpp

@@ -250,6 +250,13 @@ SparsePackBase SparsePackBase::Build(T *pmd, const PackDescriptor &desc,
                       pack.pack_h_(0, b, idx).size() > 0,
                       "Seems like this variable might not actually be allocated.");
 
+                  // Save the tensor indices for easy access
+                  for (int l = 0; l < leading_dim; ++l) {
+                    pack.pack_h_(l, b, idx).t = t;
+                    pack.pack_h_(l, b, idx).u = u;
+                    pack.pack_h_(l, b, idx).v = v;
+                  }
+
                   if (desc.flat) {
                     coords_h(idx) = pmbd->GetBlockPointer()->coords_device;
                   }

src/interface/update.cpp

@@ -156,11 +156,6 @@ TaskStatus SparseDealloc(MeshData<Real> *md) {
   auto packIdx = desc.GetMap();
 
   ParArray2D<bool> is_zero("IsZero", pack.GetNBlocks(), pack.GetMaxNumberOfVars());
-  const int Ni = ib.e + 1 - ib.s;
-  const int Nj = jb.e + 1 - jb.s;
-  const int Nk = kb.e + 1 - kb.s;
-  const int NjNi = Nj * Ni;
-  const int NkNjNi = Nk * NjNi;
   Kokkos::parallel_for(
       PARTHENON_AUTO_LABEL,
       Kokkos::TeamPolicy<>(parthenon::DevExecSpace(), pack.GetNBlocks(), Kokkos::AUTO),

src/interface/variable_state.hpp

@@ -62,6 +62,9 @@ struct VariableState : public empty_state_t {
   int vector_component = NODIR;
   bool initialized = true;
 
+  int t{-1};
+  int u{-1};
+  int v{-1};
   TopologicalType topological_type = TopologicalType::Cell;
   TopologicalElement topological_element = TopologicalElement::CC;
   std::size_t tensor_components;

src/mesh/mesh_refinement.cpp

@@ -98,21 +98,7 @@ void MeshRefinement::SetRefinement(AmrTag flag) {
       deref_count_ = 0;
     } else {
       deref_count_++;
-      int ec = 0, js, je, ks, ke;
-      if (!pmb->block_size.symmetry(X2DIR)) {
-        js = -1;
-        je = 1;
-      } else {
-        js = 0;
-        je = 0;
-      }
-      if (!pmb->block_size.symmetry(X3DIR)) {
-        ks = -1;
-        ke = 1;
-      } else {
-        ks = 0;
-        ke = 0;
-      }
+      int ec = 0;
       for (const auto &nb : pmb->neighbors) {
         if (nb.loc.level() > pmb->loc.level()) ec++;
       }

src/solvers/bicgstab_solver.hpp

@@ -24,8 +24,8 @@
 #include "kokkos_abstraction.hpp"
 #include "solvers/mg_solver.hpp"
 #include "solvers/solver_utils.hpp"
-
 #include "tasks/tasks.hpp"
+#include "utils/type_list.hpp"
 
 namespace parthenon {
 
@@ -70,31 +70,40 @@ class BiCGSTABSolver {
   PARTHENON_INTERNALSOLVERVARIABLE(u, p);
   PARTHENON_INTERNALSOLVERVARIABLE(u, x);
 
+  using internal_types_tl = TypeList<rhat0, v, h, s, t, r, p, x>;
+  using preconditioner_t = MGSolver<u, rhs, equations>;
+  using all_internal_types_tl =
+      concatenate_type_lists_t<internal_types_tl,
+                               typename preconditioner_t::internal_types_tl>;
+
   std::vector<std::string> GetInternalVariableNames() const {
-    std::vector<std::string> names{rhat0::name(), v::name(), h::name(), s::name(),
-                                   t::name(),     r::name(), p::name(), x::name()};
+    std::vector<std::string> names;
     if (params_.precondition) {
-      auto pre_names = preconditioner.GetInternalVariableNames();
-      names.insert(names.end(), pre_names.begin(), pre_names.end());
+      all_internal_types_tl::IterateTypes(
+          [&names](auto t) { names.push_back(decltype(t)::name()); });
+    } else {
+      internal_types_tl::IterateTypes(
+          [&names](auto t) { names.push_back(decltype(t)::name()); });
     }
     return names;
   }
 
   BiCGSTABSolver(StateDescriptor *pkg, BiCGSTABParams params_in,
-                 equations eq_in = equations(), std::vector<int> shape = {})
-      : preconditioner(pkg, params_in.mg_params, eq_in, shape), params_(params_in),
-        iter_counter(0), eqs_(eq_in) {
+                 equations eq_in = equations(), std::vector<int> shape = {},
+                 const std::string &container = "base")
+      : preconditioner(pkg, params_in.mg_params, eq_in, shape, container),
+        params_(params_in), iter_counter(0), eqs_(eq_in), container_(container) {
     using namespace refinement_ops;
     auto m_no_ghost =
         Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, shape);
-    pkg->AddField(x::name(), m_no_ghost);
     pkg->AddField(rhat0::name(), m_no_ghost);
     pkg->AddField(v::name(), m_no_ghost);
     pkg->AddField(h::name(), m_no_ghost);
     pkg->AddField(s::name(), m_no_ghost);
     pkg->AddField(t::name(), m_no_ghost);
     pkg->AddField(r::name(), m_no_ghost);
     pkg->AddField(p::name(), m_no_ghost);
+    pkg->AddField(x::name(), m_no_ghost);
   }
 
   template <class TL_t>
@@ -105,8 +114,8 @@ class BiCGSTABSolver {
   TaskID AddTasks(TaskList &tl, TaskID dependence, Mesh *pmesh, const int partition) {
     using namespace utils;
     TaskID none;
-    auto &md = pmesh->mesh_data.GetOrAdd("base", partition);
-    std::string label = "bicg_comm_" + std::to_string(partition);
+    auto &md = pmesh->mesh_data.GetOrAdd(container_, partition);
+    std::string label = container_ + "bicg_comm_" + std::to_string(partition);
     auto &md_comm =
         pmesh->mesh_data.AddShallow(label, md, std::vector<std::string>{u::name()});
     iter_counter = 0;
@@ -313,14 +322,15 @@ class BiCGSTABSolver {
   BiCGSTABParams &GetParams() { return params_; }
 
  protected:
-  MGSolver<u, rhs, equations> preconditioner;
+  preconditioner_t preconditioner;
   BiCGSTABParams params_;
   int iter_counter;
   AllReduce<Real> rtr, pAp, rhat0v, rhat0r, ts, tt, residual, rhs2;
   Real rhat0r_old;
   equations eqs_;
   Real final_residual;
   int final_iteration;
+  std::string container_;
 };
 
 } // namespace solvers

src/solvers/mg_solver.hpp

@@ -27,6 +27,7 @@
 #include "solvers/solver_utils.hpp"
 #include "tasks/tasks.hpp"
 #include "utils/robust.hpp"
+#include "utils/type_list.hpp"
 
 namespace parthenon {
 
@@ -76,13 +77,18 @@ class MGSolver {
   PARTHENON_INTERNALSOLVERVARIABLE(u, temp); // Temporary storage
   PARTHENON_INTERNALSOLVERVARIABLE(u, u0);   // Storage for initial solution during FAS
   PARTHENON_INTERNALSOLVERVARIABLE(u, D);    // Storage for (approximate) diagonal
+
+  using internal_types_tl = TypeList<res_err, temp, u0, D>;
   std::vector<std::string> GetInternalVariableNames() const {
-    return {res_err::name(), temp::name(), u0::name(), D::name()};
+    std::vector<std::string> names;
+    internal_types_tl::IterateTypes(
+        [&names](auto t) { names.push_back(decltype(t)::name()); });
+    return names;
   }
 
   MGSolver(StateDescriptor *pkg, MGParams params_in, equations eq_in = equations(),
-           std::vector<int> shape = {})
-      : params_(params_in), iter_counter(0), eqs_(eq_in) {
+           std::vector<int> shape = {}, const std::string &container = "base")
+      : params_(params_in), iter_counter(0), eqs_(eq_in), container_(container) {
     using namespace parthenon::refinement_ops;
     // The ghost cells of res_err need to be filled, but this is accomplished by
     // copying res_err into u, communicating, then copying u back into res_err
@@ -128,7 +134,7 @@ class MGSolver {
     auto partitions = pmesh->GetDefaultBlockPartitions(GridIdentifier::leaf());
     if (partition >= partitions.size())
       PARTHENON_FAIL("Does not work with non-default partitioning.");
-    auto &md = pmesh->mesh_data.Add("base", partitions[partition]);
+    auto &md = pmesh->mesh_data.Add(container_, partitions[partition]);
     auto comm = AddBoundaryExchangeTasks<BoundaryType::any>(mg_finest, itl, md,
                                                             pmesh->multilevel);
     auto calc_pointwise_res = eqs_.template Ax<u, res_err>(itl, comm, md);
@@ -207,6 +213,8 @@ class MGSolver {
   equations eqs_;
   Real final_residual;
   int final_iteration;
+  std::string container_;
+
   // These functions apparently have to be public to compile with cuda since
   // they contain device side lambdas
  public:
@@ -344,7 +352,7 @@ class MGSolver {
     auto partitions =
         pmesh->GetDefaultBlockPartitions(GridIdentifier::two_level_composite(level));
     if (partition >= partitions.size()) return dependence;
-    auto &md = pmesh->mesh_data.Add("base", partitions[partition]);
+    auto &md = pmesh->mesh_data.Add(container_, partitions[partition]);
 
     auto task_out = dependence;
     if (level < max_level) {
@@ -386,20 +394,20 @@ class MGSolver {
       PARTHENON_FAIL("Unknown solver type.");
     }
 
-    auto decorate_task_name = [partition, level](const std::string &in, auto b) {
-      return std::make_tuple(in + "(p:" + std::to_string(partition) +
-                                 ", l:" + std::to_string(level) + ")",
-                             1, b);
-    };
+//    auto decorate_task_name = [partition, level](const std::string &in, auto b) {
+//      return std::make_tuple(in + "(p:" + std::to_string(partition) +
+//                                 ", l:" + std::to_string(level) + ")",
+//                             1, b);
+//    };
 
-//#define BTF(...) decorate_task_name(TF(__VA_ARGS__))
+// #define BTF(...) decorate_task_name(TF(__VA_ARGS__))
 #define BTF(...) TF(__VA_ARGS__)
     bool multilevel = (level != min_level);
 
     auto partitions =
         pmesh->GetDefaultBlockPartitions(GridIdentifier::two_level_composite(level));
     if (partition >= partitions.size()) return dependence;
-    auto &md = pmesh->mesh_data.Add("base", partitions[partition]);
+    auto &md = pmesh->mesh_data.Add(container_, partitions[partition]);
     auto &md_comm = pmesh->mesh_data.AddShallow(
         "mg_comm", md, std::vector<std::string>{u::name(), res_err::name()});
 

src/tasks/tasks.hpp

@@ -296,7 +296,10 @@ class TaskList {
                          std::forward<Args>(args)...);
     } else {
       static_assert(always_false<Arg1>, "Bad signature for AddTask.");
+      return TaskID();
     }
+    // Stops some compilers from complaining
+    return TaskID();
   }
 
   template <class... Args>

src/utils/constants.hpp

@@ -118,6 +118,10 @@ class PhysicalConstants {
                                              kb * kb * kb / (15.0 * h * h * h * c * c);
   static constexpr double sb = stefan_boltzmann;
 
+  // Radiation constant
+  static constexpr double radiation_constant = 4.0 * stefan_boltzmann / speed_of_light;
+  static constexpr double ar = radiation_constant;
+
   // Faraday constant
   static constexpr double faraday_constant = 96485.33645957 * capacitance;
   static constexpr double faraday = faraday_constant;