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(),
        axom::fmt::format("{} parameter spaces given in the template argument but {} input mfem spaces were supplied.",
                          sizeof...(InputSpaces), input_mfem_spaces.size()));
 
    if constexpr (sizeof...(InputSpaces) > 0) {
      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 t, auto X, auto V, auto... inputs) {
          auto orig_tuple = integrand(t, 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 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(double time, double dt, ConstFieldPtr shape_disp, const std::vector<ConstFieldPtr>& fields,
                        [[maybe_unused]] const std::vector<ConstQuadratureFieldPtr>& quad_fields = {},
                        int block_row = 0) const override
  {
    SLIC_ERROR_IF(block_row != 0, "Invalid block row and column requested in fieldJacobian for FunctionalResidual");
    dt_ = dt;
    auto ret = (*weak_form_)(time, *shape_disp, *fields[input_indices]...);
    return ret;
  }
 
  std::unique_ptr<mfem::HypreParMatrix> jacobian(
      double time, double dt, ConstFieldPtr shape_disp, const std::vector<ConstFieldPtr>& fields,
      const std::vector<double>& jacobian_weights,
      [[maybe_unused]] const std::vector<ConstQuadratureFieldPtr>& quad_fields = {}, int block_row = 0) const override
  {
    SLIC_ERROR_IF(block_row != 0, "Invalid block row and column requested in fieldJacobian for FunctionalResidual");
 
    dt_ = dt;
 
    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, 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, shape_disp, fields));
        addToJ(jacobian_weights[input_col], assemble(K));
      }
    }
 
    return J;
  }
 
  void jvp(double time, double dt, 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,
           const std::vector<DualFieldPtr>& jvp_reactions) const override
  {
    SLIC_ERROR_IF(v_fields.size() != fields.size(),
                  "Invalid number of field sensitivities relative to the number of fields");
    SLIC_ERROR_IF(jvp_reactions.size() != 1, "FunctionalResidual nonlinear systems only supports 1 output residual");
 
    dt_ = dt;
    auto jacs =
        jacobianFunctions(std::make_integer_sequence<int, sizeof...(input_indices)>{}, time, shape_disp, fields);
 
    *jvp_reactions[0] = 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, shape_disp, fields));
        K.AddMult(*v_fields[input_col], *jvp_reactions[0]);
      }
    }
  }
 
  void vjp(double time, double dt, ConstFieldPtr shape_disp, const std::vector<ConstFieldPtr>& fields,
           [[maybe_unused]] const std::vector<ConstQuadratureFieldPtr>& quad_fields,
           const std::vector<ConstFieldPtr>& v_fields, DualFieldPtr vjp_shape_disp_sensitivity,
           const std::vector<DualFieldPtr>& vjp_sensitivities,
           [[maybe_unused]] const std::vector<QuadratureFieldPtr>& vjp_quad_field_sensitivities) const override
  {
    SLIC_ERROR_IF(vjp_sensitivities.size() != fields.size(),
                  "Invalid number of field sensitivities relative to the number of fields");
    SLIC_ERROR_IF(v_fields.size() != 1, "FunctionalResidual nonlinear systems only supports 1 output residual");
 
    dt_ = dt;
    auto vecJacs = vectorJacobianFunctions(std::make_integer_sequence<int, sizeof...(input_indices)>{}, time,
                                           shape_disp, v_fields[0], fields);
    {
      auto shape_vjp = smith::get<DERIVATIVE>((*v_dot_weak_form_residual_)(DifferentiateWRT<0>{}, time, *shape_disp,
                                                                           *v_fields[0], *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, shape_disp, v_fields[0], 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 <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)>{
        [=](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)>{
        [=](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();
 
  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