functional_weak_form.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/physics/weak_form.hpp"
#include "smith/physics/mesh.hpp"
#include "smith/numerics/functional/shape_aware_functional.hpp"
#include "smith/physics/state/finite_element_state.hpp"
#include "smith/physics/state/finite_element_dual.hpp"
 
namespace smith {
 
template <int spatial_dim, typename OutputSpace, typename inputs = Parameters<>,
          typename input_indices = std::make_integer_sequence<int, inputs::n>>
class FunctionalWeakForm;
 
template <int spatial_dim, typename OutputSpace, typename... InputSpaces, int... input_indices>
class FunctionalWeakForm<spatial_dim, OutputSpace, Parameters<InputSpaces...>,
                         std::integer_sequence<int, input_indices...>> : public WeakForm {
 public:
  using SpacesT = std::vector<const mfem::ParFiniteElementSpace*>;  
 
  using ShapeDispSpace = H1<SHAPE_ORDER, spatial_dim>;  
 
  FunctionalWeakForm(std::string physics_name, std::shared_ptr<Mesh> mesh,
                     const mfem::ParFiniteElementSpace& output_mfem_space, const SpacesT& input_mfem_spaces)
      : WeakForm(physics_name), mesh_(mesh)
  {
    std::array<const mfem::ParFiniteElementSpace*, sizeof...(InputSpaces)> trial_spaces;
    std::array<const mfem::ParFiniteElementSpace*, sizeof...(InputSpaces) + 1> vector_residual_trial_spaces{
        &output_mfem_space};
 
    SLIC_ERROR_ROOT_IF(
        sizeof...(InputSpaces) != input_mfem_spaces.size(),
        std::format("{} parameter spaces given in the template argument but {} input mfem spaces were supplied.",
                    sizeof...(InputSpaces), input_mfem_spaces.size()));
 
    // Validate output space
    validateSpace<OutputSpace>(output_mfem_space, "output");
 
    // Validate input spaces
    if constexpr (sizeof...(InputSpaces) > 0) {
      // Validate each input space using a helper
      validateInputSpaces<0>(input_mfem_spaces);
 
      for_constexpr<sizeof...(InputSpaces)>([&](auto i) { trial_spaces[i] = input_mfem_spaces[i]; });
      for_constexpr<sizeof...(InputSpaces)>(
          [&](auto i) { vector_residual_trial_spaces[i + 1] = input_mfem_spaces[i]; });
    }
 
    const auto& shape_disp_space = mesh->shapeDisplacementSpace();
 
    weak_form_ = std::make_unique<ShapeAwareFunctional<ShapeDispSpace, OutputSpace(InputSpaces...)>>(
        &shape_disp_space, &output_mfem_space, trial_spaces);
 
    v_dot_weak_form_residual_ =
        std::make_unique<ShapeAwareFunctional<ShapeDispSpace, double(OutputSpace, InputSpaces...)>>(
            &shape_disp_space, vector_residual_trial_spaces);
  }
 
  template <int... active_parameters, typename BodyIntegralType>
  void addBodyIntegral(DependsOn<active_parameters...>, std::string body_name, BodyIntegralType integrand)
  {
    weak_form_->AddDomainIntegral(Dimension<spatial_dim>{}, DependsOn<active_parameters...>{}, integrand,
                                  mesh_->domain(body_name));
 
    v_dot_weak_form_residual_->AddDomainIntegral(
        Dimension<spatial_dim>{}, DependsOn<0, 1 + active_parameters...>{},
        [integrand](double time, auto X, auto V, auto... inputs) {
          auto orig_tuple = integrand(time, X, inputs...);
          return smith::inner(get<VALUE>(V), get<VALUE>(orig_tuple)) +
                 smith::inner(get<DERIVATIVE>(V), get<DERIVATIVE>(orig_tuple));
        },
        mesh_->domain(body_name));
  }
 
  template <typename BodyForceType>
  void addBodyIntegral(std::string body_name, BodyForceType body_integral)
  {
    addBodyIntegral(DependsOn<>{}, body_name, body_integral);
  }
 
  template <int... active_parameters, typename BodyLoadType>
  void addBodySource(DependsOn<active_parameters...> depends_on, std::string body_name, BodyLoadType load_function)
  {
    addBodyIntegral(depends_on, body_name, [load_function](double t, auto X, auto... inputs) {
      return smith::tuple{-load_function(t, get<VALUE>(X), get<VALUE>(inputs)...), smith::zero{}};
    });
  }
 
  template <int... active_parameters, typename BodyLoadType>
  void addBodySource(std::string body_name, BodyLoadType load_function)
  {
    return addBodySource(DependsOn<>{}, body_name, load_function);
  }
 
  template <int... active_parameters, typename BoundaryIntegrandType>
  void addBoundaryIntegral(DependsOn<active_parameters...>, std::string boundary_name, BoundaryIntegrandType integrand)
  {
    weak_form_->AddBoundaryIntegral(Dimension<spatial_dim - 1>{}, DependsOn<active_parameters...>{}, integrand,
                                    mesh_->domain(boundary_name));
 
    v_dot_weak_form_residual_->AddBoundaryIntegral(
        Dimension<spatial_dim - 1>{}, DependsOn<0, 1 + active_parameters...>{},
        [integrand](double t, auto X, auto V, auto... params) {
          auto orig_surface_flux = integrand(t, X, params...);
          return smith::inner(get<VALUE>(V), orig_surface_flux);
        },
        mesh_->domain(boundary_name));
  }
 
  template <typename BoundaryIntegrandType>
  void addBoundaryIntegral(std::string boundary_name, const BoundaryIntegrandType& integrand)
  {
    addBoundaryIntegral(DependsOn<>{}, boundary_name, integrand);
  }
 
  template <int... active_parameters, typename InteriorIntegrandType>
  void addInteriorBoundaryIntegral(DependsOn<active_parameters...>, std::string interior_name,
                                   InteriorIntegrandType integrand)
  {
    weak_form_->AddInteriorFaceIntegral(Dimension<spatial_dim - 1>{}, DependsOn<active_parameters...>{}, integrand,
                                        mesh_->domain(interior_name));
 
    v_dot_weak_form_residual_->AddInteriorFaceIntegral(
        Dimension<spatial_dim - 1>{}, DependsOn<0, 1 + active_parameters...>{},
        [integrand](double t, auto X, auto V, auto... params) {
          auto [V1, V2] = V;
          auto orig_surface_flux = integrand(t, X, params...);
          auto [flux_pos, flux_neg] = orig_surface_flux;
          return smith::inner(V1, flux_pos) + smith::inner(V2, flux_neg);
        },
        mesh_->domain(interior_name));
  }
 
  template <typename InteriorIntegrandType>
  void addInteriorBoundaryIntegral(std::string interior_name, const InteriorIntegrandType& integrand)
  {
    addInteriorBoundaryIntegral(DependsOn<>{}, interior_name, integrand);
  }
 
  template <int... active_parameters, typename BoundaryFluxType>
  void addBoundaryFlux(DependsOn<active_parameters...> depends_on, std::string boundary_name,
                       BoundaryFluxType flux_function)
  {
    addBoundaryIntegral(depends_on, boundary_name, [flux_function](double t, auto X, auto... inputs) {
      auto n = cross(get<DERIVATIVE>(X));
      return -flux_function(t, get<VALUE>(X), normalize(n), get<VALUE>(inputs)...);
    });
  }
 
  template <typename BoundaryFluxType>
  void addBoundaryFlux(std::string boundary_name, const BoundaryFluxType& integrand)
  {
    addBoundaryFlux(DependsOn<>{}, boundary_name, integrand);
  }
 
  mfem::Vector residual(TimeInfo time_info, ConstFieldPtr shape_disp, const std::vector<ConstFieldPtr>& fields,
                        [[maybe_unused]] const std::vector<ConstQuadratureFieldPtr>& quad_fields = {}) const override
  {
    validateFields(fields, "residual");
    dt_ = time_info.dt();
    cycle_ = time_info.cycle();
    auto ret = (*weak_form_)(time_info.time(), *shape_disp, *fields[input_indices]...);
    return ret;
  }
 
  std::unique_ptr<mfem::HypreParMatrix> jacobian(
      TimeInfo time_info, ConstFieldPtr shape_disp, const std::vector<ConstFieldPtr>& fields,
      const std::vector<double>& jacobian_weights,
      [[maybe_unused]] const std::vector<ConstQuadratureFieldPtr>& quad_fields = {}) const override
  {
    validateFields(fields, "jacobian");
    dt_ = time_info.dt();
    cycle_ = time_info.cycle();
 
    std::unique_ptr<mfem::HypreParMatrix> J;
 
    auto addToJ = [&J](double factor, std::unique_ptr<mfem::HypreParMatrix> jac_contrib) {
      if (J) {
        SLIC_ERROR_IF(J->N() != jac_contrib->N(),
                      "Multiple nonzero jacobian weights are being used on inconsistently sized input arguments.");
        SLIC_ERROR_IF(J->M() != jac_contrib->M(),
                      "Multiple nonzero jacobian weights are being used on inconsistently sized input arguments.");
        J->Add(factor, *jac_contrib);
      } else {
        J.reset(jac_contrib.release());
        if (factor != 1.0) (*J) *= factor;
      }
    };
 
    auto jacs = jacobianFunctions(std::make_integer_sequence<int, sizeof...(input_indices)>{}, time_info.time(),
                                  shape_disp, fields);
 
    for (size_t input_col = 0; input_col < jacobian_weights.size(); ++input_col) {
      if (jacobian_weights[input_col] != 0.0) {
        auto K = smith::get<DERIVATIVE>(jacs[input_col](time_info.time(), shape_disp, fields));
        addToJ(jacobian_weights[input_col], assemble(K));
      }
    }
 
    return J;
  }
 
  void jvp(TimeInfo time_info, ConstFieldPtr shape_disp, const std::vector<ConstFieldPtr>& fields,
           [[maybe_unused]] const std::vector<ConstQuadratureFieldPtr>& quad_fields,
           [[maybe_unused]] ConstFieldPtr v_shape_disp, const std::vector<ConstFieldPtr>& v_fields,
           [[maybe_unused]] const std::vector<ConstQuadratureFieldPtr>& v_quad_fields,
           DualFieldPtr jvp_reaction) const override
  {
    validateFields(fields, "jvp");
    SLIC_ERROR_IF(v_fields.size() != fields.size(),
                  "Invalid number of field sensitivities relative to the number of fields");
 
    dt_ = time_info.dt();
    cycle_ = time_info.cycle();
 
    auto jacs = jacobianFunctions(std::make_integer_sequence<int, sizeof...(input_indices)>{}, time_info.time(),
                                  shape_disp, fields);
 
    *jvp_reaction = 0.0;
 
    for (size_t input_col = 0; input_col < fields.size(); ++input_col) {
      if (v_fields[input_col] != nullptr) {
        auto K = smith::get<DERIVATIVE>(jacs[input_col](time_info.time(), shape_disp, fields));
        K.AddMult(*v_fields[input_col], *jvp_reaction);
      }
    }
  }
 
  void vjp(TimeInfo time_info, ConstFieldPtr shape_disp, const std::vector<ConstFieldPtr>& fields,
           [[maybe_unused]] const std::vector<ConstQuadratureFieldPtr>& quad_fields, ConstFieldPtr v_field,
           DualFieldPtr vjp_shape_disp_sensitivity, const std::vector<DualFieldPtr>& vjp_sensitivities,
           [[maybe_unused]] const std::vector<QuadratureFieldPtr>& vjp_quad_field_sensitivities) const override
  {
    validateFields(fields, "vjp");
    SLIC_ERROR_IF(vjp_sensitivities.size() != fields.size(),
                  "Invalid number of field sensitivities relative to the number of fields");
 
    dt_ = time_info.dt();
    cycle_ = time_info.cycle();
 
    auto vecJacs = vectorJacobianFunctions(std::make_integer_sequence<int, sizeof...(input_indices)>{},
                                           time_info.time(), shape_disp, v_field, fields);
    {
      auto shape_vjp = smith::get<DERIVATIVE>((*v_dot_weak_form_residual_)(
          DifferentiateWRT<0>{}, time_info.time(), *shape_disp, *v_field, *fields[input_indices]...));
      auto shape_vjp_vector = assemble(shape_vjp);
      *vjp_shape_disp_sensitivity += *shape_vjp_vector;
    }
 
    for (size_t input_col = 0; input_col < fields.size(); ++input_col) {
      if (vjp_sensitivities[input_col] != nullptr) {
        auto vec_jac = smith::get<DERIVATIVE>(vecJacs[input_col](time_info.time(), shape_disp, v_field, fields));
        auto vec_jac_mfem_vector = assemble(vec_jac);
        *vjp_sensitivities[input_col] += *vec_jac_mfem_vector;
      }
    }
  }
 
  ShapeAwareFunctional<ShapeDispSpace, OutputSpace(InputSpaces...)>& getShapeAwareResidual() { return *weak_form_; }
 
  ShapeAwareFunctional<ShapeDispSpace, double(OutputSpace, InputSpaces...)>& getShapeAwareVectorTimesResidual()
  {
    return *v_dot_weak_form_residual_;
  }
 
 protected:
  template <size_t I>
  void validateInputSpaces(const SpacesT& input_mfem_spaces) const
  {
    if constexpr (I < sizeof...(InputSpaces)) {
      using Space = typename std::tuple_element<I, std::tuple<InputSpaces...>>::type;
      validateSpace<Space>(*input_mfem_spaces[I], axom::fmt::format("input[{}]", I));
      validateInputSpaces<I + 1>(input_mfem_spaces);
    }
  }
 
  template <size_t I>
  void validateFieldsRecursive(const std::vector<ConstFieldPtr>& fields, const std::string& method_name) const
  {
    if constexpr (I < sizeof...(InputSpaces)) {
      using Space = typename std::tuple_element<I, std::tuple<InputSpaces...>>::type;
      validateSpace<Space>(fields[I]->space(),
                           axom::fmt::format("{}(): field[{}] ('{}')", method_name, I, fields[I]->name()));
      validateFieldsRecursive<I + 1>(fields, method_name);
    }
  }
 
  void validateFields(const std::vector<ConstFieldPtr>& fields, const std::string& method_name) const
  {
    SLIC_ERROR_ROOT_IF(fields.size() != sizeof...(InputSpaces),
                       axom::fmt::format("{}(): fields.size()={} but weak form expects {} InputSpaces", method_name,
                                         fields.size(), sizeof...(InputSpaces)));
 
    if constexpr (sizeof...(InputSpaces) > 0) {
      validateFieldsRecursive<0>(fields, method_name);
    }
  }
 
  template <typename Space>
  static void validateSpace(const mfem::ParFiniteElementSpace& mfem_space, const std::string& space_name)
  {
    const auto* fec = mfem_space.FEColl();
 
    // Note: We check using the name string rather than dynamic_cast because MFEM
    // has multiple FECollection types that represent H1 and L2 spaces
    if constexpr (Space::family == Family::H1) {
      std::string fec_name = fec->Name();
      bool is_h1 =
          (fec_name.find("H1") != std::string::npos || fec_name.find("ND_") != std::string::npos ||  // Nedelec elements
           fec_name.find("Linear") != std::string::npos);
      SLIC_ERROR_ROOT_IF(!is_h1, axom::fmt::format("Space '{}': Template specifies H1 family but mfem space uses '{}'",
                                                   space_name, fec_name));
    } else if constexpr (Space::family == Family::L2) {
      std::string fec_name = fec->Name();
      bool is_l2 = (fec_name.find("L2") != std::string::npos || fec_name.find("DG") != std::string::npos ||
                    fec_name.find("Const") != std::string::npos);
      SLIC_ERROR_ROOT_IF(
          !is_l2, axom::fmt::format("Space '{}': Template specifies L2/DG family but mfem space uses '{}'", space_name,
                                    fec_name));
    }
 
    // Validate order
    SLIC_ERROR_ROOT_IF(fec->GetOrder() != Space::order,
                       axom::fmt::format("Space '{}': Template specifies order {} but mfem space has order {}",
                                         space_name, Space::order, fec->GetOrder()));
 
    // Validate vector dimension
    SLIC_ERROR_ROOT_IF(
        mfem_space.GetVDim() != Space::components,
        axom::fmt::format("Space '{}': Template specifies {} components but mfem space has {} components (VDim)",
                          space_name, Space::components, mfem_space.GetVDim()));
  }
 
  template <int... i>
  auto jacobianFunctions(std::integer_sequence<int, i...>, double time, ConstFieldPtr shape_disp,
                         const std::vector<ConstFieldPtr>& fs) const
  {
    using JacFuncType = std::function<decltype((*weak_form_)(DifferentiateWRT<1>{}, time, *shape_disp, *fs[i]...))(
        double, ConstFieldPtr, const std::vector<ConstFieldPtr>&)>;
    return std::array<JacFuncType, sizeof...(i)>{
        [this](double _time, ConstFieldPtr _shape_disp, const std::vector<ConstFieldPtr>& _fs) {
          return (*weak_form_)(DifferentiateWRT<i + 1>{}, _time, *_shape_disp, *_fs[i]...);
        }...};
  }
 
  template <int... i>
  auto vectorJacobianFunctions(std::integer_sequence<int, i...>, double time, ConstFieldPtr shape_disp, ConstFieldPtr v,
                               const std::vector<ConstFieldPtr>& fs) const
  {
    using GradFuncType =
        std::function<decltype((*v_dot_weak_form_residual_)(DifferentiateWRT<1>{}, time, *shape_disp, *v, *fs[i]...))(
            double, ConstFieldPtr, ConstFieldPtr, const std::vector<ConstFieldPtr>&)>;
    return std::array<GradFuncType, sizeof...(i)>{
        [this](double _time, ConstFieldPtr _shape_disp, ConstFieldPtr _v, const std::vector<ConstFieldPtr>& _fs) {
          return (*v_dot_weak_form_residual_)(DifferentiateWRT<i + 2>{}, _time, *_shape_disp, *_v, *_fs[i]...);
        }...};
  }
 
  mutable double dt_ = std::numeric_limits<double>::max();
 
  mutable size_t cycle_ = 0;
 
  std::shared_ptr<Mesh> mesh_;
 
  std::unique_ptr<ShapeAwareFunctional<ShapeDispSpace, OutputSpace(InputSpaces...)>> weak_form_;
 
  std::unique_ptr<ShapeAwareFunctional<ShapeDispSpace, double(OutputSpace, InputSpaces...)>> v_dot_weak_form_residual_;
};
 
inline std::vector<const mfem::ParFiniteElementSpace*> getSpaces(const std::vector<smith::FiniteElementState>& states)
{
  std::vector<const mfem::ParFiniteElementSpace*> spaces;
  for (auto& f : states) {
    spaces.push_back(&f.space());
  }
  return spaces;
}
 
}  // namespace smith