shape_aware_functional.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 "smith/infrastructure/accelerator.hpp"
#include "smith/numerics/functional/functional.hpp"
 
namespace smith {
 
namespace detail {
 
template <int dim, typename shape_type>
struct ShapeCorrection {
 public:
  SMITH_HOST_DEVICE ShapeCorrection(Dimension<dim>, shape_type p)
      : J_(Identity<dim>() + get<DERIVATIVE>(p)), detJ_(det(J_)), inv_J_(inv(J_)), inv_JT_(transpose(inv_J_))
  {
  }
 
  template <typename trial_type, typename space_type>
  SMITH_HOST_DEVICE auto modify_trial_argument(space_type /* space */, const trial_type& u) const
  {
    if constexpr (space_type{}.family == Family::H1 || space_type{}.family == Family::L2) {
      // Our updated spatial coordinate is x' = x + p. We want to calculate
      // du/dx' = du/dx * dx/dx'
      //        = du/dx * (dx'/dx)^-1
      //        = du_dx * (J)^-1
      auto trial_derivative = dot(get<DERIVATIVE>(u), inv_J_);
 
      return smith::tuple{get<VALUE>(u), trial_derivative};
    }
 
    if constexpr (space_type{}.family == Family::HCURL) {
      auto modified_val = dot(get<VALUE>(u), inv_J_);
      if constexpr (dim == 2) {
        auto modified_derivative = get<DERIVATIVE>(u) / detJ_;
        return smith::tuple{modified_val, modified_derivative};
      }
      if constexpr (dim == 3) {
        auto modified_derivative = dot(get<DERIVATIVE>(u), transpose(J_));
        return smith::tuple{modified_val, modified_derivative};
      }
    }
  }
 
  template <typename test_space, typename q_func_type>
  SMITH_HOST_DEVICE auto modify_shape_aware_qf_return(test_space /*test*/, const q_func_type& v)
  {
    if constexpr (std::is_same_v<test_space, double>) {
      return detJ_ * v;
    } else {
      if constexpr (test_space{}.family == Family::H1 || test_space{}.family == Family::L2) {
        auto modified_flux = dot(get<FLUX>(v), inv_JT_) * detJ_;
        auto modified_source = get<SOURCE>(v) * detJ_;
 
        return smith::tuple{modified_source, modified_flux};
      }
      if constexpr (test_space{}.family == Family::HCURL) {
        auto modified_source = dot(get<SOURCE>(v), inv_JT_) * detJ_;
        if constexpr (dim == 3) {
          auto modified_flux = dot(get<FLUX>(v), J_);
          return smith::tuple{modified_source, modified_flux};
        } else {
          return smith::tuple{modified_source, get<FLUX>(v)};
        }
      }
    }
  }
 
 private:
  using jacobian_type = std::remove_reference_t<decltype(get<DERIVATIVE>(std::declval<shape_type>()))>;
  using detJ_type = decltype(det(std::declval<jacobian_type>()));
  using inv_J_type = decltype(inv(std::declval<jacobian_type>()));
  using inv_JT_type = decltype(inv(transpose(std::declval<jacobian_type>())));
 
  jacobian_type J_;
 
  detJ_type detJ_;
 
  inv_J_type inv_J_;
 
  inv_JT_type inv_JT_;
};
 
template <typename position_type, typename shape_type>
SMITH_HOST_DEVICE auto compute_boundary_area_correction(const position_type& X, const shape_type& shape)
{
  auto x_prime = X + shape;
 
  auto n = cross(get<DERIVATIVE>(x_prime));
 
  // smith::Functional's boundary integrals multiply the q-function output by
  // norm(cross(dX_dxi)) at that quadrature point, but if we impose a shape displacement
  // then that weight needs to be corrected. The new weight should be
  // norm(cross(dX_dxi + dp_dxi)), so we multiply by the ratio w_new / w_old
  // to get
  //   q * area_correction * w_old
  // = q * (w_new / w_old) * w_old
  // = q * w_new
 
  auto area_correction = norm(n) / norm(cross(get<DERIVATIVE>(X)));
 
  return area_correction;
}
 
template <typename lambda, typename coord_type, typename shape_type, typename space_types, typename trial_types,
          typename correction_type, int... i>
SMITH_HOST_DEVICE auto apply_shape_aware_qf_helper(const lambda& qf, double t, const coord_type& position,
                                                   const shape_type& shape, const space_types& space_tuple,
                                                   const trial_types& arg_tuple, const correction_type& correction,
                                                   std::integer_sequence<int, i...>)
{
  static_assert(tuple_size<trial_types>::value == tuple_size<space_types>::value,
                "Argument and finite element space tuples are not the same size.");
 
  auto x = smith::tuple{get<VALUE>(position) + get<VALUE>(shape),
 
                        // x := X + u,
                        // so, dx/dxi = dX/dxi + du/dxi
                        //            = dX/dxi + du/dX * dX/dxi
                        get<DERIVATIVE>(position) + get<DERIVATIVE>(shape) * get<DERIVATIVE>(position)};
 
  return qf(t, x, correction.modify_trial_argument(smith::get<i>(space_tuple), smith::get<i>(arg_tuple))...);
}
 
template <typename lambda, typename coord_type, typename state_type, typename shape_type, typename space_types,
          typename trial_types, typename correction_type, int... i>
SMITH_HOST_DEVICE auto apply_shape_aware_qf_helper_with_state(const lambda& qf, double t, const coord_type& position,
                                                              state_type& state, const shape_type& shape,
                                                              const space_types& space_tuple,
                                                              const trial_types& arg_tuple,
                                                              const correction_type& correction,
                                                              std::integer_sequence<int, i...>)
{
  static_assert(tuple_size<trial_types>::value == tuple_size<space_types>::value,
                "Argument and finite element space tuples are not the same size.");
 
  auto x = smith::tuple{get<VALUE>(position) + get<VALUE>(shape),
 
                        // x := X + u,
                        // so, dx/dxi = dX/dxi + du/dxi
                        //            = dX/dxi + du/dX * dX/dxi
                        get<DERIVATIVE>(position) + get<DERIVATIVE>(shape) * get<DERIVATIVE>(position)};
 
  return qf(t, x, state, correction.modify_trial_argument(smith::get<i>(space_tuple), smith::get<i>(arg_tuple))...);
}
 
}  // namespace detail
 
template <typename T1, typename T2, ExecutionSpace exec = smith::ExecutionSpace::CPU>
class ShapeAwareFunctional;
 
template <typename test, typename shape, typename... trials, ExecutionSpace exec>
class ShapeAwareFunctional<shape, test(trials...), exec> {
  static constexpr test test_space{};
 
  static constexpr shape shape_space{};
 
  static constexpr uint32_t num_trial_spaces = sizeof...(trials);
 
 public:
  template <typename test_type = test, typename = std::enable_if_t<!std::is_same_v<double, test_type>>>
  ShapeAwareFunctional(const mfem::ParFiniteElementSpace* shape_fes, const mfem::ParFiniteElementSpace* test_fes,
                       std::array<const mfem::ParFiniteElementSpace*, num_trial_spaces> trial_fes)
  {
    static_assert(shape_space.family == Family::H1, "Only H1 spaces allowed for shape displacements");
 
    std::array<const mfem::ParFiniteElementSpace*, num_trial_spaces + 1> prepended_spaces;
 
    prepended_spaces[0] = shape_fes;
 
    for (uint32_t i = 0; i < num_trial_spaces; ++i) {
      prepended_spaces[1 + i] = trial_fes[i];
    }
 
    functional_ = std::make_unique<Functional<test(shape, trials...), exec>>(test_fes, prepended_spaces);
  }
 
  template <typename test_type = test, typename = std::enable_if_t<std::is_same_v<double, test_type>>>
  ShapeAwareFunctional(const mfem::ParFiniteElementSpace* shape_fes,
                       std::array<const mfem::ParFiniteElementSpace*, num_trial_spaces> trial_fes)
  {
    static_assert(shape_space.family == Family::H1, "Only H1 spaces allowed for shape displacements");
 
    std::array<const mfem::ParFiniteElementSpace*, num_trial_spaces + 1> prepended_spaces;
 
    prepended_spaces[0] = shape_fes;
 
    for (uint32_t i = 0; i < num_trial_spaces; ++i) {
      prepended_spaces[1 + i] = trial_fes[i];
    }
 
    functional_ = std::make_unique<Functional<double(shape, trials...), exec>>(prepended_spaces);
  }
 
  template <typename Integrand, int dim, int... args>
  struct ShapeAwareIntegrandWrapper {
    ShapeAwareIntegrandWrapper(Integrand integrand) : integrand_(integrand) {}
    Integrand integrand_;
    template <typename PositionType, typename ShapeValueType, typename... QFuncArgs>
    SMITH_HOST_DEVICE auto operator()(double time, PositionType x, ShapeValueType shape_val,
                                      QFuncArgs... qfunc_args) const
    {
      auto qfunc_tuple = make_tuple(qfunc_args...);
      [[maybe_unused]] smith::tuple<trials...> local_trial_spaces{};
      auto reduced_trial_space_tuple = make_tuple(get<args>(local_trial_spaces)...);
 
      detail::ShapeCorrection shape_correction(Dimension<dim>{}, shape_val);
      // TODO(CUDA): When this is compiled to device code, the below make_integer_sequence will
      // need to change to a camp integer sequence.
      auto unmodified_qf_return = detail::apply_shape_aware_qf_helper(
          integrand_, time, x, shape_val, reduced_trial_space_tuple, qfunc_tuple, shape_correction,
          std::make_integer_sequence<int, sizeof...(qfunc_args)>{});
      return shape_correction.modify_shape_aware_qf_return(test_space, unmodified_qf_return);
    }
  };
 
  template <typename Integrand, int dim, int... args>
  struct ShapeAwareIntegrandWrapperWithState {
    ShapeAwareIntegrandWrapperWithState(Integrand integrand) : integrand_(integrand) {}
    Integrand integrand_;
    template <typename PositionType, typename StateType, typename ShapeValueType, typename... QFuncArgs>
    SMITH_HOST_DEVICE auto operator()(double time, PositionType x, StateType& state, ShapeValueType shape_val,
                                      QFuncArgs... qfunc_args) const
    {
      auto qfunc_tuple = make_tuple(qfunc_args...);
      [[maybe_unused]] smith::tuple<trials...> local_trial_spaces{};
      auto reduced_trial_space_tuple = make_tuple(get<args>(local_trial_spaces)...);
 
      detail::ShapeCorrection shape_correction(Dimension<dim>{}, shape_val);
      // TODO(CUDA): When this is compiled to device code, the below make_integer_sequence will
      // need to change to a camp integer sequence.
      auto unmodified_qf_return = detail::apply_shape_aware_qf_helper_with_state(
          integrand_, time, x, state, shape_val, reduced_trial_space_tuple, qfunc_tuple, shape_correction,
          std::make_integer_sequence<int, sizeof...(qfunc_args)>{});
      return shape_correction.modify_shape_aware_qf_return(test_space, unmodified_qf_return);
    }
  };
 
  template <int dim, int... args, typename lambda, typename qpt_data_type = Nothing>
  void AddDomainIntegral(Dimension<dim>, DependsOn<args...>, const lambda& integrand, Domain& domain,
                         std::shared_ptr<QuadratureData<qpt_data_type>> qdata = NoQData)
  {
    if constexpr (std::is_same_v<qpt_data_type, Nothing>) {
      functional_->AddDomainIntegral(Dimension<dim>{}, DependsOn<0, (args + 1)...>{},
                                     ShapeAwareIntegrandWrapper<lambda, dim, args...>(integrand), domain, qdata);
    } else {
      functional_->AddDomainIntegral(Dimension<dim>{}, DependsOn<0, (args + 1)...>{},
                                     ShapeAwareIntegrandWrapperWithState<lambda, dim, args...>(integrand), domain,
                                     qdata);
    }
  }
 
  template <typename test_type, typename Integrand, int dim, int... args>
  struct ShapeAwareInteriorBoundaryIntegrandWrapper {
    ShapeAwareInteriorBoundaryIntegrandWrapper(Integrand integrand) : integrand_(integrand) {}
    Integrand integrand_;
    template <typename PositionType, typename ShapeValueType, typename... QFuncArgs>
    SMITH_HOST_DEVICE auto operator()(double time, PositionType x, ShapeValueType shape_val,
                                      QFuncArgs... qfunc_args) const
    {
      auto unmodified_qf_returns = integrand_(time, x + shape_val, qfunc_args...);
      auto correction = detail::compute_boundary_area_correction(x, shape_val);
      return smith::tuple{correction * smith::get<0>(unmodified_qf_returns),
                          correction * smith::get<1>(unmodified_qf_returns)};
    }
  };
 
  template <typename Integrand, int dim, int... args>
  struct ShapeAwareInteriorBoundaryIntegrandWrapper<double, Integrand, dim, args...> {
    ShapeAwareInteriorBoundaryIntegrandWrapper(Integrand integrand) : integrand_(integrand) {}
    Integrand integrand_;
    template <typename PositionType, typename ShapeValueType, typename... QFuncArgs>
    SMITH_HOST_DEVICE auto operator()(double time, PositionType x, ShapeValueType shape_val,
                                      QFuncArgs... qfunc_args) const
    {
      auto unmodified_qf_returns = integrand_(time, x + shape_val, qfunc_args...);
      auto correction = detail::compute_boundary_area_correction(x, shape_val);
      return correction * unmodified_qf_returns;
    }
  };
 
  template <int dim, int... args, typename lambda>
  void AddInteriorFaceIntegral(Dimension<dim>, DependsOn<args...>, const lambda& integrand, Domain& domain)
  {
    functional_->AddInteriorFaceIntegral(
        Dimension<dim>{}, DependsOn<0, (args + 1)...>{},
        ShapeAwareInteriorBoundaryIntegrandWrapper<test, lambda, dim, args...>(integrand), domain);
  }
 
  template <typename Integrand, int dim, int... args>
  struct ShapeAwareBoundaryIntegrandWrapper {
    ShapeAwareBoundaryIntegrandWrapper(Integrand integrand) : integrand_(integrand) {}
    Integrand integrand_;
    template <typename PositionType, typename ShapeValueType, typename... QFuncArgs>
    SMITH_HOST_DEVICE auto operator()(double time, PositionType x, ShapeValueType shape_val,
                                      QFuncArgs... qfunc_args) const
    {
      auto unmodified_qf_return = integrand_(time, x + shape_val, qfunc_args...);
      return unmodified_qf_return * detail::compute_boundary_area_correction(x, shape_val);
    }
  };
 
  template <int dim, int... args, typename lambda>
  void AddBoundaryIntegral(Dimension<dim>, DependsOn<args...>, const lambda& integrand, Domain& domain)
  {
    functional_->AddBoundaryIntegral(Dimension<dim>{}, DependsOn<0, (args + 1)...>{},
                                     ShapeAwareBoundaryIntegrandWrapper<lambda, dim, args...>(integrand), domain);
  }
 
  template <uint32_t wrt, typename... T>
  auto operator()(DifferentiateWRT<wrt>, double t, const T&... args)
  {
    return (*functional_)(DifferentiateWRT<wrt>{}, t, args...);
  }
 
  template <typename... T>
  auto operator()(double t, const T&... args)
  {
    return (*functional_)(t, args...);
  }
 
  void updateQdata(bool update_flag) { functional_->updateQdata(update_flag); }
 
 private:
  std::unique_ptr<Functional<test(shape, trials...), exec>> functional_;
};
 
}  // namespace smith