Jmm/indexable types

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[PR470]](https://github.com/lanl/singularity-eos/pull/470) Add the ability to access lambda elements by named types
 - [[PR459]](https://github.com/lanl/singularity-eos/pull/459) Add electron and ion tables to EOSPAC and SpinerEOS backends
 - [[PR453]](https://github.com/lanl/singularity-eos/pull/453) A PT space PTE solver
 - [[PR444]](https://github.com/lanl/singularity-eos/pull/444) Add Z split modifier and electron ideal gas EOS

doc/sphinx/src/using-closures.rst

@@ -449,22 +449,32 @@ used by the solver. There are helper routines for providing the needed
 scratch space, wich will tell you how many bytes per mixed cell are
 required. For example:
 
-.. cpp:function:: int PTESolverRhoTRequiredScratch(const int nmat);
+.. cpp:function:: int PTESolverRhoTRequiredScratch(const int nmat, bool with_cache);
 
 and
 
-.. cpp:function:: int PTESolverRhoURequiredScratch(const int nmat);
+.. cpp:function:: int PTESolverRhoURequiredScratch(const int nmat, bool with_cache);
 
 provide the number of real numbers (i.e., either ``float`` or
 ``double``) required for a single cell given a number of materials in
-equilibriun for either the ``RhoT`` or ``RhoU`` solver. The equivalent
-functions
+equilibriun for either the ``RhoT`` or ``RhoU`` solver. The
+``with_cache`` flag specifies whether or not you'd like the PTE solver
+to manage memory for lambdas for you. This argument is optional and
+the default is true.
 
-.. cpp:function:: size_t PTESolverRhoTRequiredScratchInBytes(const int nmat);
+.. note::
+
+  The pre-allocated cache is used only automatically if you pass in a
+  ``singularity::NullIndexer`` type in for the lambda indexer, which
+  returns a ``nullptr`` for all lambda access.
+
+The equivalent functions
+
+.. cpp:function:: size_t PTESolverRhoTRequiredScratchInBytes(const int nmat, bool with_cache);
 
 and
 
-.. cpp:function:: int PTESolverRhoURequiredScratchInBytes(const int nmat);
+.. cpp:function:: int PTESolverRhoURequiredScratchInBytes(const int nmat, bool with_cache);
 
 give the size in bytes needed to be allocated per cell given a number
 of materials ``nmat``.

doc/sphinx/src/using-eos.rst

@@ -659,6 +659,107 @@ minimize latency due to ``malloc`` and ``free``. Several models, such
 as ``SpinerEOS`` also use the persistency of these arrays to cache
 useful quantities for a performance boost.
 
+Named Arguments to Lambda Indexers
+-----------------------------------
+
+Lambdas support a more free-form kind of indexer. In particular, you
+may define indexers that take **names** via a type system. For
+example, you can write code such as:
+
+.. code-block:: cpp
+
+  Lambda_t lambda;
+  lambda[MeanIonizationState()] = zbar;
+  Real P = eos.PressureFromDensityTemperature(rho, T, lambda);
+
+The available types currently supported by default are:
+
+.. code-block:: cpp
+
+  struct singularity::IndexableTypes::MeanIonizationState;
+  struct singularity::IndexableTypes::LogDensity;
+  struct singularity::IndexableTypes::LogTemperature;
+  struct singularity::IndexableTypes::MeanAtomicMass;
+  struct singularity::IndexableTypes::MeanAtomicNumber;
+  struct singularity::IndexableTypes::ElectronFraction;
+
+However if you are implementing your own equation of state class, you
+can define new types with the macro
+``SINGULARITY_DECLARE_INDEXABLE_TYPE``. For example:
+
+.. code-block::
+
+  namespace IndexableTypes {
+  SINGULARITY_DECLARE_INDEXABLE_TYPE(MyLambdaParameter);
+  }
+
+To use an indexable type, you must define an indexer with an overload
+of the ``[]`` operator that takes your type. For example:
+
+.. code-block:: cpp
+
+  class MyLambda_t {
+  public:
+    MyLambda_t() = default;
+    PORTABLE_FORCEINLINE_FUNCTION
+    Real &operator[](const std::size_t idx) {
+      return data_[idx];
+    }
+    PORTABLE_FORCEINLINE_FUNCTION
+    Real &operator[](const MeanIonizationState &zbar) {
+      return data_[2];
+    }
+  private:
+    std::array<Real, 3> data_;
+  };
+
+which might be used as
+
+.. code-blocK:: cpp
+
+  MyLambda_t lambda;
+  lambda[0] = lRho;
+  lambda[1] = lT;
+  lambda[MeanIonizationState()] = Z;
+
+where ``MeanIonizationState`` is shorthand for index 2, since you
+defined that overload. To more easily enable mixing and matching
+integer-based indexing with type-based indexing, the function
+
+.. code-block:: cpp
+
+  template<typename Name_t, typename Indexer_t>
+  PORTABLE_FORCEINLINE_FUNCTION
+  Real &Get(Indexer_t &&lambda, std::size_t idx = 0);
+
+will return a reference to the value at named index ``Name_t()`` if
+that overload is defined in ``Indexer_t`` and otherwise return a
+reference at index ``idx``.
+
+As a convenience tool, the struct
+
+.. code-block:: cpp
+
+  template<typename... Ts>
+  singularity::IndexableTypes::VariadicIndexer;
+
+automatically defines an indexer that accepts all named indices in the
+variadic list ``Ts...`` and also integer indexing. It's a fixed-size
+array under the hood.
+
+You can check if an equation of state is compatible with a given named
+index type by calling
+
+.. code-block:: cpp
+
+   eos.NeedsLambda(NamedIndex())
+
+which returns ``true`` if the EOS is compatible with ``NamedIndex``
+and ``false`` otherwise.
+
+All of the above functionality is available in
+the header file ``singularity-eos/base/indexable_types.hpp``.
+
 EOS Modifiers
 --------------
 

singularity-eos/CMakeLists.txt

@@ -25,6 +25,7 @@ register_headers(
 
     # Normal files
     base/fast-math/logs.hpp
+    base/indexable_types.hpp
     base/robust_utils.hpp
     base/root-finding-1d/root_finding.hpp
     base/serialization_utils.hpp

singularity-eos/base/indexable_types.hpp

@@ -0,0 +1,90 @@
+//------------------------------------------------------------------------------
+// © 2025. Triad National Security, LLC. All rights reserved.  This
+// program was produced under U.S. Government contract 89233218CNA000001
+// for Los Alamos National Laboratory (LANL), which is operated by Triad
+// National Security, LLC for the U.S.  Department of Energy/National
+// Nuclear Security Administration. All rights in the program are
+// reserved by Triad National Security, LLC, and the U.S. Department of
+// Energy/National Nuclear Security Administration. The Government is
+// granted for itself and others acting on its behalf a nonexclusive,
+// paid-up, irrevocable worldwide license in this material to reproduce,
+// prepare derivative works, distribute copies to the public, perform
+// publicly and display publicly, and to permit others to do so.
+//------------------------------------------------------------------------------
+
+#ifndef SINGULARITY_EOS_BASE_INDEXABLE_TYPES_
+#define SINGULARITY_EOS_BASE_INDEXABLE_TYPES_
+
+#include <array>
+#include <type_traits>
+#include <utility>
+
+#include <ports-of-call/portability.hpp>
+#include <singularity-eos/base/variadic_utils.hpp>
+
+#define SINGULARITY_DECLARE_INDEXABLE_TYPE(TYPE_NAME)                                    \
+  struct TYPE_NAME : public singularity::IndexableTypes::IndexableTypesBase<TYPE_NAME> {}
+
+namespace singularity {
+namespace IndexableTypes {
+// Base class for indexable types that uses CRTP
+template <typename CRTP>
+struct IndexableTypesBase {
+  template <typename, typename = void>
+  struct IsOverloadedIn : std::false_type {};
+  template <typename T>
+  struct IsOverloadedIn<T, std::void_t<decltype(std::declval<T>()[std::declval<CRTP>()])>>
+      : std::true_type {};
+};
+} // namespace IndexableTypes
+
+namespace IndexerUtils {
+// Convenience function for accessing an indexer by either type or
+// natural number index depending on what is available
+template <typename T, typename Indexer_t>
+PORTABLE_FORCEINLINE_FUNCTION auto &Get(Indexer_t &&lambda, std::size_t idx = 0) {
+  if constexpr (T::template IsOverloadedIn<Indexer_t>::value) {
+    return lambda[T()];
+  } else {
+    return lambda[idx];
+  }
+}
+
+// This is a convenience struct to easily build a small indexer with
+// a set of indexable types.
+template <typename... Ts>
+class VariadicIndexer {
+ public:
+  VariadicIndexer() = default;
+  template <typename T,
+            typename = std::enable_if_t<variadic_utils::contains<T, Ts...>::value>>
+  PORTABLE_FORCEINLINE_FUNCTION Real &operator[](const T &t) {
+    // get the index of T in the variadic type list
+    std::size_t idx = 0;
+    std::size_t i = 0;
+    ((std::is_same_v<T, Ts> ? idx = i : ++i), ...);
+
+    return data_[idx];
+  }
+  PORTABLE_FORCEINLINE_FUNCTION
+  Real &operator[](const std::size_t idx) { return data_[idx]; }
+  PORTABLE_FORCEINLINE_FUNCTION
+  Real &operator[](const int idx) { return data_[idx]; }
+  static inline constexpr std::size_t size() { return sizeof...(Ts); }
+
+ private:
+  std::array<Real, sizeof...(Ts)> data_;
+};
+
+} // namespace IndexerUtils
+
+namespace IndexableTypes {
+SINGULARITY_DECLARE_INDEXABLE_TYPE(MeanIonizationState);
+SINGULARITY_DECLARE_INDEXABLE_TYPE(LogDensity);
+SINGULARITY_DECLARE_INDEXABLE_TYPE(LogTemperature);
+SINGULARITY_DECLARE_INDEXABLE_TYPE(MeanAtomicMass);
+SINGULARITY_DECLARE_INDEXABLE_TYPE(MeanAtomicNumber);
+SINGULARITY_DECLARE_INDEXABLE_TYPE(ElectronFraction);
+} // namespace IndexableTypes
+} // namespace singularity
+#endif // SINGULARITY_EOS_BASE_INDEXABLE_TYPES_