solid_mechanics.hpp

Go to the documentation of this file.

// Copyright (c) 2019-2023, Lawrence Livermore National Security, LLC and
// other Serac Project Developers. See the top-level LICENSE file for
// details.
//
// SPDX-License-Identifier: (BSD-3-Clause)
 
#pragma once
 
#include "mfem.hpp"
 
#include "serac/infrastructure/initialize.hpp"
#include "serac/physics/common.hpp"
#include "serac/physics/solid_mechanics_input.hpp"
#include "serac/physics/base_physics.hpp"
#include "serac/numerics/odes.hpp"
#include "serac/numerics/stdfunction_operator.hpp"
#include "serac/numerics/functional/shape_aware_functional.hpp"
#include "serac/physics/state/state_manager.hpp"
#include "serac/physics/materials/solid_material.hpp"
 
namespace serac {
 
namespace solid_mechanics {
 
namespace detail {
void adjoint_integrate(double dt_n, double dt_np1, mfem::HypreParMatrix* m_mat, mfem::HypreParMatrix* k_mat,
                       mfem::HypreParVector& disp_adjoint_load_vector, mfem::HypreParVector& velo_adjoint_load_vector,
                       mfem::HypreParVector& accel_adjoint_load_vector, mfem::HypreParVector& adjoint_displacement_,
                       mfem::HypreParVector& implicit_sensitivity_displacement_start_of_step_,
                       mfem::HypreParVector& implicit_sensitivity_velocity_start_of_step_,
                       mfem::HypreParVector& adjoint_essential, BoundaryConditionManager& bcs_,
                       mfem::Solver& lin_solver);
}  // namespace detail
 
const LinearSolverOptions default_linear_options = {.linear_solver  = LinearSolver::GMRES,
                                                    .preconditioner = Preconditioner::HypreAMG,
                                                    .relative_tol   = 1.0e-6,
                                                    .absolute_tol   = 1.0e-16,
                                                    .max_iterations = 500,
                                                    .print_level    = 0};
 
#ifdef MFEM_USE_STRUMPACK
const LinearSolverOptions direct_linear_options = {.linear_solver = LinearSolver::Strumpack, .print_level = 0};
#else
const LinearSolverOptions direct_linear_options = {.linear_solver = LinearSolver::SuperLU, .print_level = 0};
#endif
 
const NonlinearSolverOptions default_nonlinear_options = {.nonlin_solver  = NonlinearSolver::Newton,
                                                          .relative_tol   = 1.0e-4,
                                                          .absolute_tol   = 1.0e-8,
                                                          .max_iterations = 10,
                                                          .print_level    = 1};
 
const TimesteppingOptions default_quasistatic_options = {TimestepMethod::QuasiStatic};
 
const TimesteppingOptions default_timestepping_options = {TimestepMethod::Newmark,
                                                          DirichletEnforcementMethod::RateControl};
 
}  // namespace solid_mechanics
 
template <int order, int dim, typename parameters = Parameters<>,
          typename parameter_indices = std::make_integer_sequence<int, parameters::n>>
class SolidMechanics;
 
template <int order, int dim, typename... parameter_space, int... parameter_indices>
class SolidMechanics<order, dim, Parameters<parameter_space...>, std::integer_sequence<int, parameter_indices...>>
    : public BasePhysics {
public:
  static constexpr int  VALUE = 0, DERIVATIVE = 1;
  static constexpr int  SHAPE = 0;
  static constexpr auto I     = Identity<dim>();
 
  static constexpr auto NUM_STATE_VARS = 2;
 
  template <typename T>
  using qdata_type = std::shared_ptr<QuadratureData<T>>;
 
  SolidMechanics(const NonlinearSolverOptions nonlinear_opts, const LinearSolverOptions lin_opts,
                 const serac::TimesteppingOptions timestepping_opts, const GeometricNonlinearities geom_nonlin,
                 const std::string& physics_name, std::string mesh_tag, std::vector<std::string> parameter_names = {},
                 int cycle = 0, double time = 0.0, bool checkpoint_to_disk = false)
      : SolidMechanics(
            std::make_unique<EquationSolver>(nonlinear_opts, lin_opts, StateManager::mesh(mesh_tag).GetComm()),
            timestepping_opts, geom_nonlin, physics_name, mesh_tag, parameter_names, cycle, time, checkpoint_to_disk)
  {
  }
 
  SolidMechanics(std::unique_ptr<serac::EquationSolver> solver, const serac::TimesteppingOptions timestepping_opts,
                 const GeometricNonlinearities geom_nonlin, const std::string& physics_name, std::string mesh_tag,
                 std::vector<std::string> parameter_names = {}, int cycle = 0, double time = 0.0,
                 bool checkpoint_to_disk = false)
      : BasePhysics(physics_name, mesh_tag, cycle, time, checkpoint_to_disk),
        displacement_(
            StateManager::newState(H1<order, dim>{}, detail::addPrefix(physics_name, "displacement"), mesh_tag_)),
        velocity_(StateManager::newState(H1<order, dim>{}, detail::addPrefix(physics_name, "velocity"), mesh_tag_)),
        acceleration_(
            StateManager::newState(H1<order, dim>{}, detail::addPrefix(physics_name, "acceleration"), mesh_tag_)),
        adjoint_displacement_(StateManager::newState(
            H1<order, dim>{}, detail::addPrefix(physics_name, "adjoint_displacement"), mesh_tag_)),
        displacement_adjoint_load_(displacement_.space(), detail::addPrefix(physics_name, "displacement_adjoint_load")),
        velocity_adjoint_load_(displacement_.space(), detail::addPrefix(physics_name, "velocity_adjoint_load")),
        acceleration_adjoint_load_(displacement_.space(), detail::addPrefix(physics_name, "acceleration_adjoint_load")),
        implicit_sensitivity_displacement_start_of_step_(displacement_.space(), "total_deriv_wrt_displacement."),
        implicit_sensitivity_velocity_start_of_step_(displacement_.space(), "total_deriv_wrt_velocity."),
        reactions_(StateManager::newDual(H1<order, dim>{}, detail::addPrefix(physics_name, "reactions"), mesh_tag_)),
        nonlin_solver_(std::move(solver)),
        ode2_(displacement_.space().TrueVSize(),
              {.time = ode_time_point_, .c0 = c0_, .c1 = c1_, .u = u_, .du_dt = v_, .d2u_dt2 = acceleration_},
              *nonlin_solver_, bcs_),
        geom_nonlin_(geom_nonlin)
  {
    SLIC_ERROR_ROOT_IF(mesh_.Dimension() != dim,
                       axom::fmt::format("Compile time dimension, {0}, and runtime mesh dimension, {1}, mismatch", dim,
                                         mesh_.Dimension()));
 
    SLIC_ERROR_ROOT_IF(!nonlin_solver_,
                       "EquationSolver argument is nullptr in SolidMechanics constructor. It is possible that it was "
                       "previously moved.");
 
    // Check for dynamic mode
    if (timestepping_opts.timestepper != TimestepMethod::QuasiStatic) {
      ode2_.SetTimestepper(timestepping_opts.timestepper);
      ode2_.SetEnforcementMethod(timestepping_opts.enforcement_method);
      is_quasistatic_ = false;
    } else {
      is_quasistatic_ = true;
    }
 
    states_.push_back(&displacement_);
    if (!is_quasistatic_) {
      states_.push_back(&velocity_);
      states_.push_back(&acceleration_);
    }
 
    adjoints_.push_back(&adjoint_displacement_);
 
    duals_.push_back(&reactions_);
 
    // Create a pack of the primal field and parameter finite element spaces
    mfem::ParFiniteElementSpace* test_space  = &displacement_.space();
    mfem::ParFiniteElementSpace* shape_space = &shape_displacement_.space();
 
    std::array<const mfem::ParFiniteElementSpace*, NUM_STATE_VARS + sizeof...(parameter_space)> trial_spaces;
    trial_spaces[0] = &displacement_.space();
    trial_spaces[1] = &displacement_.space();
 
    SLIC_ERROR_ROOT_IF(
        sizeof...(parameter_space) != parameter_names.size(),
        axom::fmt::format("{} parameter spaces given in the template argument but {} parameter names were supplied.",
                          sizeof...(parameter_space), parameter_names.size()));
 
    if constexpr (sizeof...(parameter_space) > 0) {
      tuple<parameter_space...> types{};
      for_constexpr<sizeof...(parameter_space)>([&](auto i) {
        parameters_.emplace_back(mesh_, get<i>(types), detail::addPrefix(name_, parameter_names[i]));
 
        trial_spaces[i + NUM_STATE_VARS] = &(parameters_[i].state->space());
      });
    }
 
    residual_ = std::make_unique<ShapeAwareFunctional<shape_trial, test(trial, trial, parameter_space...)>>(
        shape_space, test_space, trial_spaces);
 
    // If the user wants the AMG preconditioner with a linear solver, set the pfes
    // to be the displacement
    auto* amg_prec = dynamic_cast<mfem::HypreBoomerAMG*>(nonlin_solver_->preconditioner());
    if (amg_prec) {
      // amg_prec->SetElasticityOptions(&displacement_.space());
 
      // TODO: The above call was seg faulting in the HYPRE_BoomerAMGSetInterpRefine(amg_precond, interp_refine)
      // method as of Hypre version v2.26.0. Instead, we just set the system size for Hypre. This is a temporary work
      // around as it will decrease the effectiveness of the preconditioner.
      amg_prec->SetSystemsOptions(dim, true);
    }
 
    int true_size = velocity_.space().TrueVSize();
 
    u_.SetSize(true_size);
    v_.SetSize(true_size);
 
    du_.SetSize(true_size);
    dr_.SetSize(true_size);
    predicted_displacement_.SetSize(true_size);
 
    shape_displacement_ = 0.0;
    initializeSolidMechanicsStates();
  }
 
  SolidMechanics(const SolidMechanicsInputOptions& input_options, const std::string& physics_name, std::string mesh_tag,
                 int cycle = 0, double time = 0.0)
      : SolidMechanics(input_options.nonlin_solver_options, input_options.lin_solver_options,
                       input_options.timestepping_options, input_options.geom_nonlin, physics_name, mesh_tag, {}, cycle,
                       time)
  {
    for (auto& mat : input_options.materials) {
      if (std::holds_alternative<serac::solid_mechanics::NeoHookean>(mat)) {
        setMaterial(std::get<serac::solid_mechanics::NeoHookean>(mat));
      } else if (std::holds_alternative<serac::solid_mechanics::LinearIsotropic>(mat)) {
        setMaterial(std::get<serac::solid_mechanics::LinearIsotropic>(mat));
      } else if (std::holds_alternative<serac::solid_mechanics::J2>(mat)) {
        if constexpr (dim == 3) {
          solid_mechanics::J2::State initial_state{};
          setMaterial(std::get<serac::solid_mechanics::J2>(mat), createQuadratureDataBuffer(initial_state));
        } else {
          SLIC_ERROR_ROOT("J2 materials only work for 3D simulations");
        }
      } else if (std::holds_alternative<serac::solid_mechanics::J2Nonlinear<serac::solid_mechanics::PowerLawHardening>>(
                     mat)) {
        if constexpr (dim == 3) {
          serac::solid_mechanics::J2Nonlinear<serac::solid_mechanics::PowerLawHardening>::State initial_state{};
          setMaterial(std::get<serac::solid_mechanics::J2Nonlinear<serac::solid_mechanics::PowerLawHardening>>(mat),
                      createQuadratureDataBuffer(initial_state));
        } else {
          SLIC_ERROR_ROOT("J2Nonlinear materials only work for 3D simulations");
        }
      } else if (std::holds_alternative<serac::solid_mechanics::J2Nonlinear<serac::solid_mechanics::VoceHardening>>(
                     mat)) {
        if constexpr (dim == 3) {
          serac::solid_mechanics::J2Nonlinear<serac::solid_mechanics::VoceHardening>::State initial_state{};
          setMaterial(std::get<serac::solid_mechanics::J2Nonlinear<serac::solid_mechanics::VoceHardening>>(mat),
                      createQuadratureDataBuffer(initial_state));
        } else {
          SLIC_ERROR_ROOT("J2Nonlinear materials only work for 3D simulations");
        }
      } else {
        SLIC_ERROR("Invalid material type.");
      }
    }
 
    if (input_options.initial_displacement) {
      displacement_.project(input_options.initial_displacement->constructVector(dim));
    }
 
    if (input_options.initial_velocity) {
      velocity_.project(input_options.initial_velocity->constructVector(dim));
    }
 
    for (const auto& [bc_name, bc] : input_options.boundary_conditions) {
      // FIXME: Better naming for boundary conditions?
      if (bc_name.find("displacement") != std::string::npos) {
        if (bc.coef_opts.isVector()) {
          std::shared_ptr<mfem::VectorCoefficient> disp_coef(bc.coef_opts.constructVector(dim));
          bcs_.addEssential(bc.attrs, disp_coef, displacement_.space());
        } else {
          SLIC_ERROR_ROOT_IF(
              !bc.coef_opts.component,
              "Component not specified with scalar coefficient when setting the displacement condition.");
          std::shared_ptr<mfem::Coefficient> disp_coef(bc.coef_opts.constructScalar());
          bcs_.addEssential(bc.attrs, disp_coef, displacement_.space(), *bc.coef_opts.component);
        }
      } else if (bc_name.find("traction") != std::string::npos) {
        // TODO: Not implemented yet in input files
        SLIC_ERROR("'traction' is not implemented yet in input files.");
      } else if (bc_name.find("traction_ref") != std::string::npos) {
        // TODO: Not implemented yet in input files
        SLIC_ERROR("'traction_ref' is not implemented yet in input files.");
      } else if (bc_name.find("pressure") != std::string::npos) {
        // TODO: Not implemented yet in input files
        SLIC_ERROR("'pressure' is not implemented yet in input files.");
      } else if (bc_name.find("pressure_ref") != std::string::npos) {
        // TODO: Not implemented yet in input files
        SLIC_ERROR("'pressure_ref' is not implemented yet in input files.");
      } else {
        SLIC_WARNING_ROOT("Ignoring boundary condition with unknown name: " << bc_name);
      }
    }
  }
 
  virtual ~SolidMechanics() {}
 
  void initializeSolidMechanicsStates()
  {
    c0_ = 0.0;
    c1_ = 0.0;
 
    displacement_ = 0.0;
    velocity_     = 0.0;
    acceleration_ = 0.0;
 
    adjoint_displacement_      = 0.0;
    displacement_adjoint_load_ = 0.0;
    velocity_adjoint_load_     = 0.0;
    acceleration_adjoint_load_ = 0.0;
 
    implicit_sensitivity_displacement_start_of_step_ = 0.0;
    implicit_sensitivity_velocity_start_of_step_     = 0.0;
 
    reactions_ = 0.0;
 
    u_                      = 0.0;
    v_                      = 0.0;
    du_                     = 0.0;
    dr_                     = 0.0;
    predicted_displacement_ = 0.0;
 
    if (!checkpoint_to_disk_) {
      checkpoint_states_.clear();
 
      checkpoint_states_["displacement"].push_back(displacement_);
      checkpoint_states_["velocity"].push_back(velocity_);
      checkpoint_states_["acceleration"].push_back(acceleration_);
    }
  }
 
  void resetStates(int cycle = 0, double time = 0.0) override
  {
    BasePhysics::initializeBasePhysicsStates(cycle, time);
    initializeSolidMechanicsStates();
  }
 
  template <typename T>
  qdata_type<T> createQuadratureDataBuffer(T initial_state)
  {
    constexpr auto Q = order + 1;
 
    std::array<uint32_t, mfem::Geometry::NUM_GEOMETRIES> elems = geometry_counts(mesh_);
    std::array<uint32_t, mfem::Geometry::NUM_GEOMETRIES> qpts_per_elem{};
 
    std::vector<mfem::Geometry::Type> geometries;
    if (dim == 2) {
      geometries = {mfem::Geometry::TRIANGLE, mfem::Geometry::SQUARE};
    } else {
      geometries = {mfem::Geometry::TETRAHEDRON, mfem::Geometry::CUBE};
    }
 
    for (auto geom : geometries) {
      qpts_per_elem[size_t(geom)] = uint32_t(num_quadrature_points(geom, Q));
    }
 
    return std::make_shared<QuadratureData<T>>(elems, qpts_per_elem, initial_state);
  }
 
  void setDisplacementBCs(const std::set<int>& disp_bdr, std::function<void(const mfem::Vector&, mfem::Vector&)> disp)
  {
    // Project the coefficient onto the grid function
    disp_bdr_coef_ = std::make_shared<mfem::VectorFunctionCoefficient>(dim, disp);
 
    bcs_.addEssential(disp_bdr, disp_bdr_coef_, displacement_.space());
  }
 
  void setDisplacementBCs(const std::set<int>&                                            disp_bdr,
                          std::function<void(const mfem::Vector&, double, mfem::Vector&)> disp)
  {
    // Project the coefficient onto the grid function
    disp_bdr_coef_ = std::make_shared<mfem::VectorFunctionCoefficient>(dim, disp);
 
    bcs_.addEssential(disp_bdr, disp_bdr_coef_, displacement_.space());
  }
 
  void setDisplacementBCs(const std::set<int>& disp_bdr, std::function<double(const mfem::Vector& x)> disp,
                          int component)
  {
    // Project the coefficient onto the grid function
    component_disp_bdr_coef_ = std::make_shared<mfem::FunctionCoefficient>(disp);
 
    bcs_.addEssential(disp_bdr, component_disp_bdr_coef_, displacement_.space(), component);
  }
 
  void setDisplacementBCsByDofList(const mfem::Array<int>                                          true_dofs,
                                   std::function<void(const mfem::Vector&, double, mfem::Vector&)> disp)
  {
    disp_bdr_coef_ = std::make_shared<mfem::VectorFunctionCoefficient>(dim, disp);
 
    bcs_.addEssential(true_dofs, disp_bdr_coef_, displacement_.space());
  }
 
  void setDisplacementBCsByDofList(const mfem::Array<int>                                  true_dofs,
                                   std::function<void(const mfem::Vector&, mfem::Vector&)> disp)
  {
    disp_bdr_coef_ = std::make_shared<mfem::VectorFunctionCoefficient>(dim, disp);
 
    bcs_.addEssential(true_dofs, disp_bdr_coef_, displacement_.space());
  }
 
  void setDisplacementBCs(std::function<bool(const mfem::Vector&)>                        is_node_constrained,
                          std::function<void(const mfem::Vector&, double, mfem::Vector&)> disp)
  {
    auto constrained_dofs = calculateConstrainedDofs(is_node_constrained);
 
    setDisplacementBCsByDofList(constrained_dofs, disp);
  }
 
  void setDisplacementBCs(std::function<bool(const mfem::Vector&)>                is_node_constrained,
                          std::function<void(const mfem::Vector&, mfem::Vector&)> disp)
  {
    auto constrained_dofs = calculateConstrainedDofs(is_node_constrained);
 
    setDisplacementBCsByDofList(constrained_dofs, disp);
  }
 
  void setDisplacementBCs(std::function<bool(const mfem::Vector&)>           is_node_constrained,
                          std::function<double(const mfem::Vector&, double)> disp, int component)
  {
    auto constrained_dofs = calculateConstrainedDofs(is_node_constrained, component);
 
    auto vector_function = [disp, component](const mfem::Vector& x, double time, mfem::Vector& displacement) {
      displacement            = 0.0;
      displacement(component) = disp(x, time);
    };
 
    setDisplacementBCsByDofList(constrained_dofs, vector_function);
  }
 
  void setDisplacementBCs(std::function<bool(const mfem::Vector& x)>   is_node_constrained,
                          std::function<double(const mfem::Vector& x)> disp, int component)
  {
    auto constrained_dofs = calculateConstrainedDofs(is_node_constrained, component);
 
    auto vector_function = [disp, component](const mfem::Vector& x, mfem::Vector& displacement) {
      displacement            = 0.0;
      displacement(component) = disp(x);
    };
 
    setDisplacementBCsByDofList(constrained_dofs, vector_function);
  }
 
  const FiniteElementState& state(const std::string& state_name) const override
  {
    if (state_name == "displacement") {
      return displacement_;
    } else if (state_name == "velocity") {
      return velocity_;
    } else if (state_name == "acceleration") {
      return acceleration_;
    }
 
    SLIC_ERROR_ROOT(axom::fmt::format("State '{}' requested from solid mechanics module '{}', but it doesn't exist",
                                      state_name, name_));
    return displacement_;
  }
 
  void setState(const std::string& state_name, const FiniteElementState& state) override
  {
    if (state_name == "displacement") {
      displacement_ = state;
      if (!checkpoint_to_disk_) {
        checkpoint_states_["displacement"][static_cast<size_t>(cycle_)] = displacement_;
      }
      return;
    } else if (state_name == "velocity") {
      velocity_ = state;
      if (!checkpoint_to_disk_) {
        checkpoint_states_["velocity"][static_cast<size_t>(cycle_)] = velocity_;
      }
      return;
    }
 
    SLIC_ERROR_ROOT(axom::fmt::format(
        "setState for state named '{}' requested from solid mechanics module '{}', but it doesn't exist", state_name,
        name_));
  }
 
  std::vector<std::string> stateNames() const override
  {
    if (is_quasistatic_) {
      return std::vector<std::string>{"displacement"};
    } else {
      return std::vector<std::string>{"displacement", "velocity", "acceleration"};
    }
  }
 
  template <int... active_parameters, typename callable>
  void addCustomBoundaryIntegral(DependsOn<active_parameters...>, callable qfunction,
                                 const std::optional<Domain>& optional_domain = std::nullopt)
  {
    Domain domain = (optional_domain) ? *optional_domain : EntireBoundary(mesh_);
 
    residual_->AddBoundaryIntegral(Dimension<dim - 1>{}, DependsOn<0, 1, active_parameters + NUM_STATE_VARS...>{},
                                   qfunction, domain);
  }
 
  const FiniteElementState& adjoint(const std::string& state_name) const override
  {
    if (state_name == "displacement") {
      return adjoint_displacement_;
    }
 
    SLIC_ERROR_ROOT(axom::fmt::format("Adjoint '{}' requested from solid mechanics module '{}', but it doesn't exist",
                                      state_name, name_));
    return adjoint_displacement_;
  }
 
  std::vector<std::string> adjointNames() const override { return std::vector<std::string>{{"displacement"}}; }
 
  template <int... active_parameters, typename callable, typename StateType = Nothing>
  void addCustomDomainIntegral(DependsOn<active_parameters...>, callable qfunction,
                               qdata_type<StateType> qdata = NoQData)
  {
    residual_->AddDomainIntegral(Dimension<dim>{}, DependsOn<0, 1, active_parameters + NUM_STATE_VARS...>{}, qfunction,
                                 mesh_, qdata);
  }
 
  template <int... active_parameters, typename MaterialType, typename StateType = Empty>
  void setMaterial(DependsOn<active_parameters...>, MaterialType material, qdata_type<StateType> qdata = EmptyQData)
  {
    residual_->AddDomainIntegral(
        Dimension<dim>{},
        DependsOn<0, 1,
                  active_parameters + NUM_STATE_VARS...>{},  // the magic number "+ NUM_STATE_VARS" accounts for the
                                                             // fact that the displacement, acceleration, and shape
                                                             // fields are always-on and come first, so the `n`th
                                                             // parameter will actually be argument `n + NUM_STATE_VARS`
        [material, geom_nonlin = geom_nonlin_](double /*t*/, auto /*x*/, auto& state, auto displacement,
                                               auto acceleration, auto... params) {
          auto du_dX   = get<DERIVATIVE>(displacement);
          auto d2u_dt2 = get<VALUE>(acceleration);
 
          auto stress = material(state, du_dX, params...);
 
          auto dx_dX = 0.0 * du_dX + I;
 
          if (geom_nonlin == GeometricNonlinearities::On) {
            dx_dX += du_dX;
          }
 
          auto flux = dot(stress, transpose(inv(dx_dX))) * det(dx_dX);
 
          return serac::tuple{material.density * d2u_dt2, flux};
        },
        mesh_, qdata);
  }
 
  template <typename MaterialType, typename StateType = Empty>
  void setMaterial(MaterialType material, std::shared_ptr<QuadratureData<StateType>> qdata = EmptyQData)
  {
    setMaterial(DependsOn<>{}, material, qdata);
  }
 
  void setDisplacement(std::function<void(const mfem::Vector& x, mfem::Vector& disp)> disp)
  {
    // Project the coefficient onto the grid function
    mfem::VectorFunctionCoefficient disp_coef(dim, disp);
    displacement_.project(disp_coef);
  }
 
  void setDisplacement(const FiniteElementState& temp) { displacement_ = temp; }
 
  void setVelocity(std::function<void(const mfem::Vector& x, mfem::Vector& vel)> vel)
  {
    // Project the coefficient onto the grid function
    mfem::VectorFunctionCoefficient vel_coef(dim, vel);
    velocity_.project(vel_coef);
  }
 
  void setVelocity(const FiniteElementState& temp) { velocity_ = temp; }
 
  template <int... active_parameters, typename BodyForceType>
  void addBodyForce(DependsOn<active_parameters...>, BodyForceType body_force,
                    const std::optional<Domain>& optional_domain = std::nullopt)
  {
    Domain domain = (optional_domain) ? *optional_domain : EntireDomain(mesh_);
 
    residual_->AddDomainIntegral(
        Dimension<dim>{}, DependsOn<0, 1, active_parameters + NUM_STATE_VARS...>{},
        [body_force](double t, auto x, auto /* displacement */, auto /* acceleration */, auto... params) {
          return serac::tuple{-1.0 * body_force(x, t, params...), zero{}};
        },
        domain);
  }
 
  template <typename BodyForceType>
  void addBodyForce(BodyForceType body_force, const std::optional<Domain>& optional_domain = std::nullopt)
  {
    addBodyForce(DependsOn<>{}, body_force, optional_domain);
  }
 
  template <int... active_parameters, typename TractionType>
  void setTraction(DependsOn<active_parameters...>, TractionType traction_function,
                   const std::optional<Domain>& optional_domain = std::nullopt)
  {
    Domain domain = (optional_domain) ? *optional_domain : EntireBoundary(mesh_);
 
    residual_->AddBoundaryIntegral(
        Dimension<dim - 1>{}, DependsOn<0, 1, active_parameters + NUM_STATE_VARS...>{},
        [traction_function](double t, auto X, auto /* displacement */, auto /* acceleration */, auto... params) {
          auto n = cross(get<DERIVATIVE>(X));
 
          return -1.0 * traction_function(get<VALUE>(X), normalize(n), t, params...);
        },
        domain);
  }
 
  template <typename TractionType>
  void setTraction(TractionType traction_function, const std::optional<Domain>& optional_domain = std::nullopt)
  {
    setTraction(DependsOn<>{}, traction_function, optional_domain);
  }
 
  template <int... active_parameters, typename PressureType>
  void setPressure(DependsOn<active_parameters...>, PressureType pressure_function,
                   const std::optional<Domain>& optional_domain = std::nullopt)
  {
    Domain domain = (optional_domain) ? *optional_domain : EntireBoundary(mesh_);
 
    residual_->AddBoundaryIntegral(
        Dimension<dim - 1>{}, DependsOn<0, 1, active_parameters + NUM_STATE_VARS...>{},
        [pressure_function, geom_nonlin = geom_nonlin_](double t, auto X, auto displacement, auto /* acceleration */,
                                                        auto... params) {
          // Calculate the position and normal in the shape perturbed deformed configuration
          auto x = X + 0.0 * displacement;
 
          if (geom_nonlin == GeometricNonlinearities::On) {
            x = x + displacement;
          }
 
          auto n = cross(get<DERIVATIVE>(x));
 
          // serac::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 + du_dxi + dp_dxi)) where u is displacement and p is shape displacement. This implies:
          //
          //   pressure * normalize(normal_new) * w_new
          // = pressure * normalize(normal_new) * (w_new / w_old) * w_old
          // = pressure * normalize(normal_new) * (norm(normal_new) / norm(normal_old)) * w_old
          // = pressure * (normal_new / norm(normal_new)) * (norm(normal_new) / norm(normal_old)) * w_old
          // = pressure * (normal_new / norm(normal_old)) * w_old
 
          // We always query the pressure function in the undeformed configuration
          return pressure_function(get<VALUE>(X), t, params...) * (n / norm(cross(get<DERIVATIVE>(X))));
        },
        domain);
  }
 
  template <typename PressureType>
  void setPressure(PressureType pressure_function, const std::optional<Domain>& optional_domain = std::nullopt)
  {
    setPressure(DependsOn<>{}, pressure_function, optional_domain);
  }
 
  virtual std::unique_ptr<mfem_ext::StdFunctionOperator> buildQuasistaticOperator()
  {
    // the quasistatic case is entirely described by the residual,
    // there is no ordinary differential equation
    return std::make_unique<mfem_ext::StdFunctionOperator>(
        displacement_.space().TrueVSize(),
 
        // residual function
        [this](const mfem::Vector& u, mfem::Vector& r) {
          const mfem::Vector res = (*residual_)(ode_time_point_, shape_displacement_, u, acceleration_,
                                                *parameters_[parameter_indices].state...);
 
          // TODO this copy is required as the sundials solvers do not allow move assignments because of their memory
          // tracking strategy
          // See https://github.com/mfem/mfem/issues/3531
          r = res;
          r.SetSubVector(bcs_.allEssentialTrueDofs(), 0.0);
        },
 
        // gradient of residual function
        [this](const mfem::Vector& u) -> mfem::Operator& {
          auto [r, drdu] = (*residual_)(ode_time_point_, shape_displacement_, differentiate_wrt(u), acceleration_,
                                        *parameters_[parameter_indices].state...);
          J_             = assemble(drdu);
          J_e_           = bcs_.eliminateAllEssentialDofsFromMatrix(*J_);
          return *J_;
        });
  }
 
  std::pair<const mfem::HypreParMatrix&, const mfem::HypreParMatrix&> stiffnessMatrix() const
  {
    SLIC_ERROR_ROOT_IF(!J_ || !J_e_, "Stiffness matrix has not yet been assembled.");
 
    return {*J_, *J_e_};
  }
 
  void completeSetup() override
  {
    // Build the dof array lookup tables
    displacement_.space().BuildDofToArrays();
 
    if (is_quasistatic_) {
      residual_with_bcs_ = buildQuasistaticOperator();
    } else {
      // the dynamic case is described by a residual function and a second order
      // ordinary differential equation. Here, we define the residual function in
      // terms of an acceleration.
      residual_with_bcs_ = std::make_unique<mfem_ext::StdFunctionOperator>(
          displacement_.space().TrueVSize(),
 
          [this](const mfem::Vector& d2u_dt2, mfem::Vector& r) {
            add(1.0, u_, c0_, d2u_dt2, predicted_displacement_);
            const mfem::Vector res = (*residual_)(ode_time_point_, shape_displacement_, predicted_displacement_,
                                                  d2u_dt2, *parameters_[parameter_indices].state...);
 
            // TODO this copy is required as the sundials solvers do not allow move assignments because of their memory
            // tracking strategy
            // See https://github.com/mfem/mfem/issues/3531
            r = res;
            r.SetSubVector(bcs_.allEssentialTrueDofs(), 0.0);
          },
 
          [this](const mfem::Vector& d2u_dt2) -> mfem::Operator& {
            add(1.0, u_, c0_, d2u_dt2, predicted_displacement_);
 
            // K := dR/du
            auto K = serac::get<DERIVATIVE>((*residual_)(ode_time_point_, shape_displacement_,
                                                         differentiate_wrt(predicted_displacement_), d2u_dt2,
                                                         *parameters_[parameter_indices].state...));
            std::unique_ptr<mfem::HypreParMatrix> k_mat(assemble(K));
 
            // M := dR/da
            auto M = serac::get<DERIVATIVE>((*residual_)(ode_time_point_, shape_displacement_, predicted_displacement_,
                                                         differentiate_wrt(d2u_dt2),
                                                         *parameters_[parameter_indices].state...));
            std::unique_ptr<mfem::HypreParMatrix> m_mat(assemble(M));
 
            // J = M + c0 * K
            J_.reset(mfem::Add(1.0, *m_mat, c0_, *k_mat));
            J_e_ = bcs_.eliminateAllEssentialDofsFromMatrix(*J_);
 
            return *J_;
          });
    }
 
    nonlin_solver_->setOperator(*residual_with_bcs_);
 
    if (checkpoint_to_disk_) {
      outputStateToDisk();
    } else {
      checkpoint_states_.clear();
 
      checkpoint_states_["displacement"].push_back(displacement_);
      checkpoint_states_["velocity"].push_back(velocity_);
      checkpoint_states_["acceleration"].push_back(acceleration_);
    }
  }
 
  void zeroEssentials(FiniteElementVector& field) const
  {
    for (const auto& essential : bcs_.essentials()) {
      field.SetSubVector(essential.getTrueDofList(), 0.0);
    }
  }
 
  virtual void quasiStaticSolve(double dt)
  {
    time_ += dt;
 
    // Set the ODE time point for the time-varying loads in quasi-static problems
    ode_time_point_ = time_;
 
    // this method is essentially equivalent to the 1-liner
    // u += dot(inv(J), dot(J_elim[:, dofs], (U(t + dt) - u)[dofs]));
    warmStartDisplacement();
 
    nonlin_solver_->solve(displacement_);
 
    cycle_ += 1;
  }
 
  void advanceTimestep(double dt) override
  {
    SLIC_ERROR_ROOT_IF(!residual_, "completeSetup() must be called prior to advanceTimestep(dt) in SolidMechanics.");
 
    // If this is the first call, initialize the previous parameter values as the initial values
    if (cycle_ == 0) {
      for (auto& parameter : parameters_) {
        *parameter.previous_state = *parameter.state;
      }
    }
 
    if (is_quasistatic_) {
      quasiStaticSolve(dt);
    } else {
      ode2_.Step(displacement_, velocity_, time_, dt);
 
      cycle_ += 1;
 
      if (checkpoint_to_disk_) {
        outputStateToDisk();
      } else {
        checkpoint_states_["displacement"].push_back(displacement_);
        checkpoint_states_["velocity"].push_back(velocity_);
        checkpoint_states_["acceleration"].push_back(acceleration_);
      }
    }
 
    {
      // after finding displacements that satisfy equilibrium,
      // compute the residual one more time, this time enabling
      // the material state buffers to be updated
      residual_->updateQdata(true);
 
      reactions_ = (*residual_)(ode_time_point_, shape_displacement_, displacement_, acceleration_,
                                *parameters_[parameter_indices].state...);
 
      residual_->updateQdata(false);
    }
 
    if (cycle_ > max_cycle_) {
      timesteps_.push_back(dt);
      max_cycle_ = cycle_;
      max_time_  = time_;
    }
  }
 
  virtual void setAdjointLoad(std::unordered_map<std::string, const serac::FiniteElementDual&> loads) override
  {
    SLIC_ERROR_ROOT_IF(loads.size() == 0, "Adjoint load container size must be greater than 0 in the solid mechanics.");
 
    auto disp_adjoint_load = loads.find("displacement");
 
    SLIC_ERROR_ROOT_IF(disp_adjoint_load == loads.end(), "Adjoint load for \"displacement\" not found.");
 
    displacement_adjoint_load_ = disp_adjoint_load->second;
    // Add the sign correction to move the term to the RHS
    displacement_adjoint_load_ *= -1.0;
 
    auto velo_adjoint_load = loads.find("velocity");
 
    if (velo_adjoint_load != loads.end()) {
      velocity_adjoint_load_ = velo_adjoint_load->second;
      // Add the sign correction to move the term to the RHS
      velocity_adjoint_load_ *= -1.0;
    }
 
    auto accel_adjoint_load = loads.find("acceleration");
 
    if (accel_adjoint_load != loads.end()) {
      acceleration_adjoint_load_ = accel_adjoint_load->second;
      // Add the sign correction to move the term to the RHS
      acceleration_adjoint_load_ *= -1.0;
    }
  }
 
  void reverseAdjointTimestep() override
  {
    auto& lin_solver = nonlin_solver_->linearSolver();
 
    // By default, use a homogeneous essential boundary condition
    mfem::HypreParVector adjoint_essential(displacement_adjoint_load_);
    adjoint_essential = 0.0;
 
    if (is_quasistatic_) {
      auto [_, drdu] = (*residual_)(ode_time_point_, shape_displacement_, differentiate_wrt(displacement_),
                                    acceleration_, *parameters_[parameter_indices].state...);
      auto jacobian  = assemble(drdu);
      auto J_T       = std::unique_ptr<mfem::HypreParMatrix>(jacobian->Transpose());
 
      for (const auto& bc : bcs_.essentials()) {
        bc.apply(*J_T, displacement_adjoint_load_, adjoint_essential);
      }
 
      lin_solver.SetOperator(*J_T);
      lin_solver.Mult(displacement_adjoint_load_, adjoint_displacement_);
 
      // Reset the equation solver to use the full nonlinear residual operator.  MRT, is this needed?
      nonlin_solver_->setOperator(*residual_with_bcs_);
 
      return;
    }
 
    SLIC_ERROR_ROOT_IF(ode2_.GetTimestepper() != TimestepMethod::Newmark,
                       "Only Newmark implemented for transient adjoint solid mechanics.");
 
    SLIC_ERROR_ROOT_IF(cycle_ <= min_cycle_,
                       "Maximum number of adjoint timesteps exceeded! The number of adjoint timesteps must equal the "
                       "number of forward timesteps");
 
    // Load the end of step disp, velo, accel from the previous cycle
    {
      auto previous_states_n = getCheckpointedStates(cycle_);
 
      displacement_ = previous_states_n.at("displacement");
      velocity_     = previous_states_n.at("velocity");
      acceleration_ = previous_states_n.at("acceleration");
    }
 
    double dt_np1 = getCheckpointedTimestep(cycle_);
    double dt_n   = getCheckpointedTimestep(cycle_ - 1);
 
    // K := dR/du
    auto K = serac::get<DERIVATIVE>((*residual_)(ode_time_point_, shape_displacement_, differentiate_wrt(displacement_),
                                                 acceleration_, *parameters_[parameter_indices].state...));
    std::unique_ptr<mfem::HypreParMatrix> k_mat(assemble(K));
 
    // M := dR/da
    auto M = serac::get<DERIVATIVE>((*residual_)(ode_time_point_, shape_displacement_, displacement_,
                                                 differentiate_wrt(acceleration_),
                                                 *parameters_[parameter_indices].state...));
    std::unique_ptr<mfem::HypreParMatrix> m_mat(assemble(M));
 
    solid_mechanics::detail::adjoint_integrate(
        dt_n, dt_np1, m_mat.get(), k_mat.get(), displacement_adjoint_load_, velocity_adjoint_load_,
        acceleration_adjoint_load_, adjoint_displacement_, implicit_sensitivity_displacement_start_of_step_,
        implicit_sensitivity_velocity_start_of_step_, adjoint_essential, bcs_, lin_solver);
 
    time_ -= dt_n;
    cycle_--;
  }
 
  std::unordered_map<std::string, FiniteElementState> getCheckpointedStates(int cycle_to_load) const override
  {
    std::unordered_map<std::string, FiniteElementState> previous_states;
 
    if (checkpoint_to_disk_) {
      previous_states.emplace("displacement", displacement_);
      previous_states.emplace("velocity", velocity_);
      previous_states.emplace("acceleration", acceleration_);
      StateManager::loadCheckpointedStates(
          cycle_to_load,
          {previous_states.at("displacement"), previous_states.at("velocity"), previous_states.at("acceleration")});
      return previous_states;
    } else {
      previous_states.emplace("displacement",
                              checkpoint_states_.at("displacement")[static_cast<size_t>(cycle_to_load)]);
      previous_states.emplace("velocity", checkpoint_states_.at("velocity")[static_cast<size_t>(cycle_to_load)]);
      previous_states.emplace("acceleration",
                              checkpoint_states_.at("acceleration")[static_cast<size_t>(cycle_to_load)]);
    }
 
    return previous_states;
  }
 
  FiniteElementDual& computeTimestepSensitivity(size_t parameter_field) override
  {
    SLIC_ASSERT_MSG(parameter_field < sizeof...(parameter_indices),
                    axom::fmt::format("Invalid parameter index '{}' requested for sensitivity."));
 
    // TODO: the time is likely not being handled correctly on the reverse pass, but we don't
    //       have tests to confirm.
    auto drdparam     = serac::get<DERIVATIVE>(d_residual_d_[parameter_field](ode_time_point_));
    auto drdparam_mat = assemble(drdparam);
 
    drdparam_mat->MultTranspose(adjoint_displacement_, *parameters_[parameter_field].sensitivity);
 
    return *parameters_[parameter_field].sensitivity;
  }
 
  FiniteElementDual& computeTimestepShapeSensitivity() override
  {
    auto drdshape =
        serac::get<DERIVATIVE>((*residual_)(ode_time_point_, differentiate_wrt(shape_displacement_), displacement_,
                                            acceleration_, *parameters_[parameter_indices].state...));
 
    auto drdshape_mat = assemble(drdshape);
 
    drdshape_mat->MultTranspose(adjoint_displacement_, *shape_displacement_sensitivity_);
 
    return *shape_displacement_sensitivity_;
  }
 
  const std::unordered_map<std::string, const serac::FiniteElementDual&> computeInitialConditionSensitivity() override
  {
    return {{"displacement", implicit_sensitivity_displacement_start_of_step_},
            {"velocity", implicit_sensitivity_velocity_start_of_step_}};
  }
 
  const serac::FiniteElementState& displacement() const { return displacement_; };
 
  const serac::FiniteElementState& velocity() const { return velocity_; };
 
  const serac::FiniteElementState& acceleration() const { return acceleration_; };
 
  const serac::FiniteElementDual& reactions() { return reactions_; };
 
protected:
  using trial = H1<order, dim>;
 
  using test = H1<order, dim>;
 
  using shape_trial = H1<SHAPE_ORDER, dim>;
 
  FiniteElementState displacement_;
 
  FiniteElementState velocity_;
 
  FiniteElementState acceleration_;
 
  // In the case of transient dynamics, this is more like an adjoint_acceleration
  FiniteElementState adjoint_displacement_;
 
  FiniteElementDual displacement_adjoint_load_;
 
  FiniteElementDual velocity_adjoint_load_;
 
  FiniteElementDual acceleration_adjoint_load_;
 
  FiniteElementDual implicit_sensitivity_displacement_start_of_step_;
 
  FiniteElementDual implicit_sensitivity_velocity_start_of_step_;
 
  FiniteElementDual reactions_;
 
  std::unique_ptr<ShapeAwareFunctional<shape_trial, test(trial, trial, parameter_space...)>> residual_;
 
  std::unique_ptr<mfem_ext::StdFunctionOperator> residual_with_bcs_;
 
  std::unique_ptr<EquationSolver> nonlin_solver_;
 
  mfem_ext::SecondOrderODE ode2_;
 
  std::unique_ptr<mfem::HypreParMatrix> J_;
 
  std::unique_ptr<mfem::HypreParMatrix> J_e_;
 
  mfem::Vector predicted_displacement_;
 
  mfem::Vector du_;
 
  mfem::Vector dr_;
 
  mfem::Vector u_;
 
  mfem::Vector v_;
 
  double c0_;
 
  double c1_;
 
  GeometricNonlinearities geom_nonlin_;
 
  std::shared_ptr<mfem::VectorCoefficient> disp_bdr_coef_;
 
  std::shared_ptr<mfem::Coefficient> component_disp_bdr_coef_;
 
  std::array<std::function<decltype((*residual_)(DifferentiateWRT<1>{}, 0.0, shape_displacement_, displacement_,
                                                 acceleration_, *parameters_[parameter_indices].state...))(double)>,
             sizeof...(parameter_indices)>
      d_residual_d_ = {[&](double t) {
        return (*residual_)(DifferentiateWRT<NUM_STATE_VARS + 1 + parameter_indices>{}, t, shape_displacement_,
                            displacement_, acceleration_, *parameters_[parameter_indices].state...);
      }...};
 
  mfem::Array<int> calculateConstrainedDofs(std::function<bool(const mfem::Vector&)> is_node_constrained,
                                            std::optional<int>                       component = {})
  {
    // Get the nodal positions for the displacement vector in grid function form
    mfem::ParGridFunction nodal_positions(&displacement_.space());
    mesh_.GetNodes(nodal_positions);
 
    const int        num_nodes = nodal_positions.Size() / dim;
    mfem::Array<int> constrained_dofs;
 
    for (int i = 0; i < num_nodes; i++) {
      // Determine if this "local" node (L-vector node) is in the local true vector. I.e. ensure this node is not a
      // shared node owned by another processor
      if (nodal_positions.ParFESpace()->GetLocalTDofNumber(i) >= 0) {
        mfem::Vector     node_coords(dim);
        mfem::Array<int> node_dofs;
        for (int d = 0; d < dim; d++) {
          // Get the local dof number for the prescribed component
          int local_vector_dof = mfem::Ordering::Map<mfem::Ordering::byNODES>(
              nodal_positions.FESpace()->GetNDofs(), nodal_positions.FESpace()->GetVDim(), i, d);
 
          // Save the spatial position for this coordinate dof
          node_coords(d) = nodal_positions(local_vector_dof);
 
          // Check if this component of the displacement vector is constrained
          bool is_active_component = true;
          if (component) {
            if (*component != d) {
              is_active_component = false;
            }
          }
 
          if (is_active_component) {
            // Add the true dof for this component to the related dof list
            node_dofs.Append(nodal_positions.ParFESpace()->GetLocalTDofNumber(local_vector_dof));
          }
        }
 
        // Do the user-defined spatial query to determine if these dofs should be constrained
        if (is_node_constrained(node_coords)) {
          constrained_dofs.Append(node_dofs);
        }
 
        // Reset the temporary container for the dofs associated with a particular node
        node_dofs.DeleteAll();
      }
    }
    return constrained_dofs;
  }
 
  void warmStartDisplacement()
  {
    // Update the linearized Jacobian matrix
    auto [r, drdu] = (*residual_)(ode_time_point_, shape_displacement_, differentiate_wrt(displacement_), acceleration_,
                                  *parameters_[parameter_indices].state...);
    J_             = assemble(drdu);
    J_e_           = bcs_.eliminateAllEssentialDofsFromMatrix(*J_);
 
    du_ = 0.0;
    for (auto& bc : bcs_.essentials()) {
      bc.setDofs(du_, time_);
    }
 
    auto& constrained_dofs = bcs_.allEssentialTrueDofs();
    for (int i = 0; i < constrained_dofs.Size(); i++) {
      int j = constrained_dofs[i];
      du_[j] -= displacement_(j);
    }
 
    dr_ = 0.0;
    mfem::EliminateBC(*J_, *J_e_, constrained_dofs, du_, dr_);
 
    // Update the initial guess for changes in the parameters if this is not the first solve
    for (std::size_t parameter_index = 0; parameter_index < parameters_.size(); ++parameter_index) {
      // Compute the change in parameters parameter_diff = parameter_new - parameter_old
      serac::FiniteElementState parameter_difference = *parameters_[parameter_index].state;
      parameter_difference -= *parameters_[parameter_index].previous_state;
 
      // Compute a linearized estimate of the residual forces due to this change in parameter
      auto drdparam        = serac::get<DERIVATIVE>(d_residual_d_[parameter_index](ode_time_point_));
      auto residual_update = drdparam(parameter_difference);
 
      // Flip the sign to get the RHS of the Newton update system
      // J^-1 du = - residual
      residual_update *= -1.0;
 
      dr_ += residual_update;
 
      // Save the current parameter value for the next timestep
      *parameters_[parameter_index].previous_state = *parameters_[parameter_index].state;
    }
 
    for (int i = 0; i < constrained_dofs.Size(); i++) {
      int j  = constrained_dofs[i];
      dr_[j] = du_[j];
    }
 
    auto& lin_solver = nonlin_solver_->linearSolver();
 
    lin_solver.SetOperator(*J_);
 
    lin_solver.Mult(dr_, du_);
    displacement_ += du_;
  }
};
 
}  // namespace serac