coupling_params.hpp

Go to the documentation of this file.

// Copyright (c) Lawrence Livermore National Security, LLC and
// other Smith Project Developers. See the top-level LICENSE file for
// details.
//
// SPDX-License-Identifier: (BSD-3-Clause)
 
#pragma once
 
#include <tuple>
#include <type_traits>
#include <utility>
 
#include "smith/differentiable_numerics/field_store.hpp"
#include "smith/differentiable_numerics/system_base.hpp"
 
namespace smith {
 
template <int Dim, int Order, typename TimeRule, typename... Spaces>
struct PhysicsFields {
  using time_rule_type = TimeRule;                              
  static constexpr int dim = Dim;                               
  static constexpr int order = Order;                           
  static constexpr std::size_t num_fields = sizeof...(Spaces);  
  std::shared_ptr<FieldStore> field_store;                      
  std::tuple<FieldType<Spaces>...> fields;                      
 
  PhysicsFields(std::shared_ptr<FieldStore> fs, FieldType<Spaces>... f)
      : field_store(std::move(fs)), fields(std::move(f)...)
  {
  }
};
 
template <typename... Spaces>
struct ParamFields {
  static constexpr std::size_t num_fields = sizeof...(Spaces);  
  std::tuple<FieldType<Spaces>...> fields;                      
  ParamFields(FieldType<Spaces>... fs) : fields(std::move(fs)...) {}
};
 
template <typename... Spaces>
ParamFields(FieldType<Spaces>...) -> ParamFields<Spaces...>;
 
template <typename... PFs>
struct CouplingFields {
  std::tuple<PFs...> packs;  
};
 
namespace detail {
 
template <typename T>
struct is_physics_fields_impl : std::false_type {};
 
template <int D, int O, typename R, typename... S>
struct is_physics_fields_impl<PhysicsFields<D, O, R, S...>> : std::true_type {};
 
template <typename T>
inline constexpr bool is_physics_fields_v = is_physics_fields_impl<std::decay_t<T>>::value;
 
}  // namespace detail
 
template <typename... PFs>
auto couplingFields(const PFs&... pfs)
{
  static_assert((detail::is_physics_fields_v<PFs> && ...), "couplingFields(...) only accepts PhysicsFields packs");
  return CouplingFields<PFs...>{std::make_tuple(pfs...)};
}
 
template <typename... ParamSpaces>
auto registerParameterFields(const std::shared_ptr<FieldStore>& field_store, FieldType<ParamSpaces>... param_types)
{
  auto register_one = [&](auto param_type) {
    param_type.name = "param_" + param_type.name;
    field_store->addParameter(param_type);
    return param_type;
  };
  return ParamFields<ParamSpaces...>{register_one(std::move(param_types))...};
}
 
namespace detail {
 
template <typename T>
struct is_parameter_pack_impl : std::false_type {};
 
template <typename... S>
struct is_parameter_pack_impl<ParamFields<S...>> : std::true_type {};
 
template <typename T>
inline constexpr bool is_parameter_pack_v = is_parameter_pack_impl<std::decay_t<T>>::value;
 
template <typename T>
struct is_coupling_fields_impl : std::false_type {};
 
template <typename... PFs>
struct is_coupling_fields_impl<CouplingFields<PFs...>> : std::true_type {};
 
template <typename T>
inline constexpr bool is_coupling_fields_v = is_coupling_fields_impl<std::decay_t<T>>::value;
 
template <typename T>
struct is_coupling_packs_impl : std::false_type {};
 
template <typename... Packs>
struct is_coupling_packs_impl<std::tuple<Packs...>> : std::true_type {};
 
template <typename T>
inline constexpr bool is_coupling_packs_v = is_coupling_packs_impl<std::decay_t<T>>::value;
 
template <typename T, typename = void>
inline constexpr bool has_time_rule_v = false;
 
template <typename T>
inline constexpr bool has_time_rule_v<T, std::enable_if_t<is_physics_fields_v<T>>> = true;
 
// -------------------------------------------------------------------------
// Trailing-arg extraction
// -------------------------------------------------------------------------
 
template <typename PacksTuple>
auto flattenCouplingFields(const PacksTuple& packs)
{
  return std::apply([](const auto&... pack) { return std::tuple_cat(pack.fields...); }, packs);
}
 
inline auto collectCouplingFields() { return std::tuple<>{}; }
 
template <typename... PFs>
auto collectCouplingFields(const CouplingFields<PFs...>& coupled)
{
  return coupled.packs;
}
 
template <typename... Spaces>
auto collectCouplingFields(const ParamFields<Spaces...>& params)
{
  return std::make_tuple(params);
}
 
template <typename... PFs, typename... Spaces>
auto collectCouplingFields(const CouplingFields<PFs...>& coupled, const ParamFields<Spaces...>& params)
{
  return std::tuple_cat(coupled.packs, std::make_tuple(params));
}
 
// -------------------------------------------------------------------------
// Time-rule interpolation
// -------------------------------------------------------------------------
 
template <typename Rule, typename TimeInfoT, typename ArgsTuple, typename Callback, std::size_t... StateIs,
          std::size_t... TailIs>
decltype(auto) applyTimeRuleToPrefixImpl(const Rule& rule, const TimeInfoT& t_info, const ArgsTuple& raw_args,
                                         Callback&& callback, std::index_sequence<StateIs...>,
                                         std::index_sequence<TailIs...>)
{
  auto interpolated = rule.interpolate(t_info, std::get<StateIs>(raw_args)...);
  return std::apply(
      [&](auto&&... values) -> decltype(auto) {
        return std::forward<Callback>(callback)(std::forward<decltype(values)>(values)...,
                                                std::get<Rule::num_states + TailIs>(raw_args)...);
      },
      interpolated);
}
 
template <typename Rule, typename TimeInfoT, typename Callback, typename... RawArgs>
decltype(auto) applyTimeRuleToPrefix(const Rule& rule, const TimeInfoT& t_info, Callback&& callback,
                                     const RawArgs&... raw_args)
{
  static_assert(sizeof...(RawArgs) >= Rule::num_states, "Not enough raw arguments for time-rule interpolation");
  auto raw_tuple = std::forward_as_tuple(raw_args...);
  constexpr std::size_t tail_count = sizeof...(RawArgs) - Rule::num_states;
  return applyTimeRuleToPrefixImpl(rule, t_info, raw_tuple, std::forward<Callback>(callback),
                                   std::make_index_sequence<Rule::num_states>{},
                                   std::make_index_sequence<tail_count>{});
}
 
template <std::size_t Offset, typename Pack, typename TimeInfoT, typename RawTuple, std::size_t... Is>
auto evaluateCouplingPack(const Pack& /*pack*/, const TimeInfoT& t_info, const RawTuple& raw_args,
                          std::index_sequence<Is...>)
{
  if constexpr (is_physics_fields_v<Pack>) {
    using Rule = typename Pack::time_rule_type;
    Rule rule;
    return rule.interpolate(t_info, std::get<Offset + Is>(raw_args)...);
  } else {
    return std::forward_as_tuple(std::get<Offset + Is>(raw_args)...);
  }
}
 
template <std::size_t I, std::size_t Offset, typename PacksTuple, typename TimeInfoT, typename RawTuple>
auto evaluateCouplingPacks(const PacksTuple& packs, const TimeInfoT& t_info, const RawTuple& raw_args)
{
  if constexpr (I == std::tuple_size_v<std::decay_t<PacksTuple>>) {
    return std::tuple{};
  } else {
    const auto& pack = std::get<I>(packs);
    using Pack = std::decay_t<decltype(pack)>;
    auto head = evaluateCouplingPack<Offset>(pack, t_info, raw_args, std::make_index_sequence<Pack::num_fields>{});
    auto tail = evaluateCouplingPacks<I + 1, Offset + Pack::num_fields>(packs, t_info, raw_args);
    return std::tuple_cat(head, tail);
  }
}
 
template <typename PacksTuple, typename TimeInfoT, typename Callback, typename... RawArgs>
decltype(auto) applyCouplingTimeRules(const PacksTuple& packs, const TimeInfoT& t_info, Callback&& callback,
                                      const RawArgs&... raw_args)
{
  auto raw_tuple = std::forward_as_tuple(raw_args...);
  auto interpolated_tail = evaluateCouplingPacks<0, 0>(packs, t_info, raw_tuple);
  return std::apply(std::forward<Callback>(callback), interpolated_tail);
}
 
template <typename Rule, typename Coupling, typename TimeInfoT, typename Callback, typename... RawArgs>
decltype(auto) applyTimeRuleAndCoupling(const Rule& rule, const Coupling& coupling, const TimeInfoT& t_info,
                                        Callback&& callback, const RawArgs&... raw_args)
{
  constexpr std::size_t tail_count = sizeof...(RawArgs) - Rule::num_states;
  return applyTimeRuleToPrefix(
      rule, t_info,
      [&](auto... self_states_and_tail) {
        constexpr std::size_t n_self = sizeof...(self_states_and_tail) - tail_count;
        auto all = std::forward_as_tuple(self_states_and_tail...);
        return [&]<std::size_t... Si, std::size_t... Ti>(std::index_sequence<Si...>, std::index_sequence<Ti...>) {
          return applyCouplingTimeRules(
              coupling, t_info,
              [&](auto... interpolated_coupling) {
                return std::forward<Callback>(callback)(std::get<Si>(all)..., interpolated_coupling...);
              },
              std::get<n_self + Ti>(all)...);
        }(std::make_index_sequence<n_self>{}, std::make_index_sequence<tail_count>{});
      },
      raw_args...);
}
 
// -------------------------------------------------------------------------
// Type-level coupling-space extraction (used by weak-form parameter type computation)
// -------------------------------------------------------------------------
 
template <typename PacksTuple>
struct FlattenCoupling;
 
template <typename Tuple>
struct TupleToParameters;
 
template <typename... Spaces>
struct TupleToParameters<std::tuple<Spaces...>> {
  using type = Parameters<Spaces...>;  
};
 
template <typename Tuple>
using tuple_to_parameters_t = typename TupleToParameters<std::decay_t<Tuple>>::type;
 
template <typename Pack>
struct pack_tuple;
 
template <int D, int O, typename R, typename... Spaces>
struct pack_tuple<PhysicsFields<D, O, R, Spaces...>> {
  using type = std::tuple<Spaces...>;  
};
 
template <typename... Spaces>
struct pack_tuple<ParamFields<Spaces...>> {
  using type = std::tuple<Spaces...>;  
};
 
template <typename Pack>
using pack_tuple_t = typename pack_tuple<std::decay_t<Pack>>::type;
 
template <typename... Tuples>
using tuple_cat_t = decltype(std::tuple_cat(std::declval<Tuples>()...));
 
template <typename CouplingParams, typename FixedParams>
struct AppendParameters;
 
template <typename... Coupled, typename... Fixed>
struct AppendParameters<Parameters<Coupled...>, Parameters<Fixed...>> {
  using type = Parameters<Fixed..., Coupled...>;  
};
 
template <typename... Packs>
struct FlattenCoupling<std::tuple<Packs...>> {
 public:
  using tuple_type = tuple_cat_t<pack_tuple_t<Packs>...>;  
  using parameters = tuple_to_parameters_t<tuple_type>;    
};
 
template <typename PacksTuple>
using flatten_coupling_t = typename FlattenCoupling<std::decay_t<PacksTuple>>::parameters;
 
template <typename Rule, typename Space, typename PacksTuple>
struct TimeRuleParamsImpl;
 
template <typename Rule, typename Space, typename... CS>
struct TimeRuleParamsImpl<Rule, Space, Parameters<CS...>> {
  using type = smith::TimeRuleParams<Rule, Space, CS...>;  
};
 
template <typename Rule, typename Space, typename PacksTuple>
using TimeRuleParams = typename TimeRuleParamsImpl<Rule, Space, flatten_coupling_t<PacksTuple>>::type;
 
template <typename PacksTuple, typename FixedParams>
struct AppendCouplingToParams;
 
template <typename PacksTuple, typename... Fixed>
struct AppendCouplingToParams<PacksTuple, Parameters<Fixed...>> {
  using type = typename AppendParameters<flatten_coupling_t<PacksTuple>, Parameters<Fixed...>>::type;
};
 
}  // namespace detail
 
}  // namespace smith