heat_transfer.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 "mfem.hpp"
 
#include "smith/physics/common.hpp"
#include "smith/physics/heat_transfer_input.hpp"
#include "smith/physics/base_physics.hpp"
#include "smith/numerics/odes.hpp"
#include "smith/numerics/stdfunction_operator.hpp"
#include "smith/numerics/functional/shape_aware_functional.hpp"
#include "smith/physics/state/state_manager.hpp"
#include "smith/physics/mesh.hpp"
#include "smith/physics/materials/thermal_material.hpp"
 
namespace smith {
 
namespace heat_transfer {
 
const LinearSolverOptions default_linear_options = {.linear_solver = LinearSolver::GMRES,
                                                    .preconditioner = Preconditioner::HypreL1Jacobi,
                                                    .relative_tol = 1.0e-6,
                                                    .absolute_tol = 1.0e-12,
                                                    .max_iterations = 200};
 
#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 = 500,
                                                          .print_level = 1};
 
const TimesteppingOptions default_timestepping_options = {TimestepMethod::BackwardEuler,
                                                          DirichletEnforcementMethod::RateControl};
 
const TimesteppingOptions default_static_options = {TimestepMethod::QuasiStatic};
 
}  // namespace heat_transfer
 
template <int order, int dim, typename parameters = Parameters<>,
          typename parameter_indices = std::make_integer_sequence<int, parameters::n>>
class HeatTransfer;
 
template <int order, int dim, typename... parameter_space, int... parameter_indices>
class HeatTransfer<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;
 
  HeatTransfer(const NonlinearSolverOptions nonlinear_opts, const LinearSolverOptions lin_opts,
               const smith::TimesteppingOptions timestepping_opts, const std::string& physics_name,
               std::shared_ptr<smith::Mesh> smith_mesh, std::vector<std::string> parameter_names = {}, int cycle = 0,
               double time = 0.0, bool checkpoint_to_disk = false)
      : HeatTransfer(std::make_unique<EquationSolver>(nonlinear_opts, lin_opts, smith_mesh->getComm()),
                     timestepping_opts, physics_name, smith_mesh, parameter_names, cycle, time, checkpoint_to_disk)
  {
  }
 
  HeatTransfer(std::unique_ptr<smith::EquationSolver> solver, const smith::TimesteppingOptions timestepping_opts,
               const std::string& physics_name, std::shared_ptr<smith::Mesh> smith_mesh,
               std::vector<std::string> parameter_names = {}, int cycle = 0, double time = 0.0,
               bool checkpoint_to_disk = false)
      : BasePhysics(physics_name, smith_mesh, cycle, time, checkpoint_to_disk),
        temperature_(StateManager::newState(H1<order>{}, detail::addPrefix(physics_name, "temperature"), mesh_->tag())),
        temperature_rate_(
            StateManager::newState(H1<order>{}, detail::addPrefix(physics_name, "temperature_rate"), mesh_->tag())),
        adjoint_temperature_(
            StateManager::newState(H1<order>{}, detail::addPrefix(physics_name, "adjoint_temperature"), mesh_->tag())),
        implicit_sensitivity_temperature_start_of_step_(adjoint_temperature_.space(),
                                                        detail::addPrefix(physics_name, "total_deriv_wrt_temperature")),
        temperature_adjoint_load_(temperature_.space(), detail::addPrefix(physics_name, "temperature_adjoint_load")),
        temperature_rate_adjoint_load_(temperature_.space(),
                                       detail::addPrefix(physics_name, "temperature_rate_adjoint_load")),
        residual_with_bcs_(temperature_.space().TrueVSize()),
        nonlin_solver_(std::move(solver)),
        ode_(temperature_.space().TrueVSize(),
             {.time = time_, .u = u_, .dt = dt_, .du_dt = temperature_rate_, .previous_dt = previous_dt_},
             *nonlin_solver_, bcs_)
  {
    SLIC_ERROR_ROOT_IF(
        mfemParMesh().Dimension() != dim,
        axom::fmt::format("Compile time class dimension template parameter and runtime mesh dimension do not match"));
 
    SLIC_ERROR_ROOT_IF(
        !nonlin_solver_,
        "EquationSolver argument is nullptr in HeatTransfer constructor. It is possible that it was previously moved.");
 
    // Check for dynamic mode
    if (timestepping_opts.timestepper != TimestepMethod::QuasiStatic) {
      ode_.SetTimestepper(timestepping_opts.timestepper);
      ode_.SetEnforcementMethod(timestepping_opts.enforcement_method);
      is_quasistatic_ = false;
    } else {
      is_quasistatic_ = true;
    }
 
    states_.push_back(&temperature_);
    if (!is_quasistatic_) {
      states_.push_back(&temperature_rate_);
    }
 
    adjoints_.push_back(&adjoint_temperature_);
 
    // Create a pack of the primal field and parameter finite element spaces
    const mfem::ParFiniteElementSpace* test_space = &temperature_.space();
    const mfem::ParFiniteElementSpace* shape_space = &mesh_->shapeDisplacementSpace();
 
    std::array<const mfem::ParFiniteElementSpace*, sizeof...(parameter_space) + NUM_STATE_VARS> trial_spaces;
    trial_spaces[0] = &temperature_.space();
    trial_spaces[1] = &temperature_.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(mfemParMesh(), 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(scalar_trial, scalar_trial, parameter_space...)>>(
            shape_space, test_space, trial_spaces);
 
    nonlin_solver_->setOperator(residual_with_bcs_);
 
    int true_size = temperature_.space().TrueVSize();
    u_.SetSize(true_size);
    u_predicted_.SetSize(true_size);
 
    initializeThermalStates();
  }
 
  HeatTransfer(const HeatTransferInputOptions& input_options, const std::string& physics_name,
               std::shared_ptr<smith::Mesh> smith_mesh, int cycle = 0, double time = 0.0)
      : HeatTransfer(input_options.nonlin_solver_options, input_options.lin_solver_options,
                     input_options.timestepping_options, physics_name, smith_mesh, {}, cycle, time)
  {
    for (const auto& mat : input_options.materials) {
      if (std::holds_alternative<smith::heat_transfer::LinearIsotropicConductor>(mat)) {
        setMaterial(std::get<smith::heat_transfer::LinearIsotropicConductor>(mat));
      } else if (std::holds_alternative<smith::heat_transfer::LinearConductor<dim>>(mat)) {
        setMaterial(std::get<smith::heat_transfer::LinearConductor<dim>>(mat));
      }
    }
 
    if (input_options.initial_temperature) {
      auto temp = input_options.initial_temperature->constructScalar();
      temperature_.project(*temp);
    }
 
    if (input_options.source_coef) {
      // TODO: Not implemented yet in input files
      // NOTE: cannot use std::functions that use mfem::vector
      SLIC_ERROR("'source' is not implemented yet in input files.");
    }
 
    // Process the BCs in sorted order for correct behavior with repeated attributes
    std::map<std::string, input::BoundaryConditionInputOptions> sorted_bcs(input_options.boundary_conditions.begin(),
                                                                           input_options.boundary_conditions.end());
    for (const auto& [bc_name, bc] : sorted_bcs) {
      // FIXME: Better naming for boundary conditions?
      if (bc_name.find("temperature") != std::string::npos) {
        std::shared_ptr<mfem::Coefficient> temp_coef(bc.coef_opts.constructScalar());
        bcs_.addEssential(bc.attrs, temp_coef, temperature_.space(), *bc.coef_opts.component);
      } else if (bc_name.find("flux") != std::string::npos) {
        // TODO: Not implemented yet in input files
        // NOTE: cannot use std::functions that use mfem::vector
        SLIC_ERROR("'flux' is not implemented yet in input files.");
      } else {
        SLIC_WARNING_ROOT("Ignoring boundary condition with unknown name: " << physics_name);
      }
    }
  }
 
  void initializeThermalStates()
  {
    dt_ = 0.0;
    previous_dt_ = -1.0;
    time_end_step_ = 0.0;
 
    u_ = 0.0;
    temperature_ = 0.0;
    temperature_rate_ = 0.0;
    adjoint_temperature_ = 0.0;
    implicit_sensitivity_temperature_start_of_step_ = 0.0;
    temperature_adjoint_load_ = 0.0;
    temperature_rate_adjoint_load_ = 0.0;
  }
 
  void resetStates(int cycle = 0, double time = 0.0) override
  {
    BasePhysics::initializeBasePhysicsStates(cycle, time);
    initializeThermalStates();
 
    if (!checkpoint_to_disk_) {
      checkpoint_states_.clear();
      auto state_names = stateNames();
      for (const auto& state_name : state_names) {
        checkpoint_states_[state_name].push_back(state(state_name));
      }
    }
  }
 
  void setTemperatureBCs(const std::set<int>& temp_bdr, std::function<double(const mfem::Vector& x, double t)> temp)
  {
    // Project the coefficient onto the grid function
    temp_bdr_coef_ = std::make_shared<mfem::FunctionCoefficient>(temp);
 
    bcs_.addEssential(temp_bdr, temp_bdr_coef_, temperature_.space());
  }
 
  void advanceTimestep(double dt) override
  {
    if (is_quasistatic_) {
      time_ += dt;
 
      // Project the essential boundary coefficients
      for (auto& bc : bcs_.essentials()) {
        bc.setDofs(temperature_, time_);
      }
      nonlin_solver_->solve(temperature_);
    } else {
      // Step the time integrator
      // Note that the ODE solver handles the essential boundary condition application itself
 
      // The current ode interface tracks 2 times, one internally which we have a handle to via time_,
      // and one here via the step interface.
      // We are ignoring this one, and just using the internal version for now.
      // This may need to be revisited when more complex time integrators are required,
      // but at the moment, the double times creates a lot of confusion, so
      // we short circuit the extra time here by passing a dummy time and ignoring it.
      double time_tmp = time_;
      ode_.Step(temperature_, time_tmp, dt);
    }
 
    cycle_ += 1;
 
    if (checkpoint_to_disk_) {
      outputStateToDisk();
    } else {
      auto state_names = stateNames();
      for (const auto& state_name : state_names) {
        checkpoint_states_[state_name].push_back(state(state_name));
      }
    }
 
    if (cycle_ > max_cycle_) {
      timesteps_.push_back(dt);
      max_cycle_ = cycle_;
      max_time_ = time_;
    }
  }
 
  template <typename MaterialType>
  struct ThermalMaterialIntegrand {
    ThermalMaterialIntegrand(MaterialType material) : material_(material) {}
 
    template <typename X, typename T, typename dT_dt, typename... Params>
    auto operator()(double /*time*/, X x, T temperature, dT_dt dtemp_dt, Params... params) const
    {
      // Get the value and the gradient from the input tuple
      auto [u, du_dX] = temperature;
      auto du_dt = get<VALUE>(dtemp_dt);
 
      auto [heat_capacity, heat_flux] = material_(x, u, du_dX, params...);
 
      return smith::tuple{heat_capacity * du_dt, -1.0 * heat_flux};
    }
 
   private:
    MaterialType material_;
  };
 
  template <int... active_parameters, typename MaterialType>
  void setMaterial(DependsOn<active_parameters...>, const MaterialType& material, Domain& domain)
  {
    residual_->AddDomainIntegral(Dimension<dim>{}, DependsOn<0, 1, NUM_STATE_VARS + active_parameters...>{},
                                 ThermalMaterialIntegrand<MaterialType>(material), domain);
  }
 
  template <typename MaterialType>
  void setMaterial(const MaterialType& material, Domain& domain)
  {
    setMaterial(DependsOn<>{}, material, domain);
  }
 
  void setTemperature(std::function<double(const mfem::Vector& x, double t)> temp)
  {
    // Project the coefficient onto the grid function
    mfem::FunctionCoefficient temp_coef(temp);
 
    temp_coef.SetTime(time_);
    temperature_.project(temp_coef);
  }
 
  void setTemperature(const FiniteElementState temp) { temperature_ = temp; }
 
  template <int... active_parameters, typename SourceType>
  void setSource(DependsOn<active_parameters...>, SourceType source_function, Domain& domain)
  {
    residual_->AddDomainIntegral(
        Dimension<dim>{}, DependsOn<0, 1, active_parameters + NUM_STATE_VARS...>{},
        [source_function](double t, auto x, auto temperature, auto /* dtemp_dt */, auto... params) {
          // Get the value and the gradient from the input tuple
          auto [u, du_dX] = temperature;
 
          auto source = source_function(x, t, u, du_dX, params...);
 
          // Return the source and the flux as a tuple
          return smith::tuple{-1.0 * source, smith::zero{}};
        },
        domain);
  }
 
  template <typename SourceType>
  void setSource(SourceType source_function, Domain& domain)
  {
    setSource(DependsOn<>{}, source_function, domain);
  }
 
  template <int... active_parameters, typename FluxType>
  void setFluxBCs(DependsOn<active_parameters...>, FluxType flux_function, Domain& domain)
  {
    residual_->AddBoundaryIntegral(
        Dimension<dim - 1>{}, DependsOn<0, 1, active_parameters + NUM_STATE_VARS...>{},
        [flux_function](double t, auto X, auto u, auto /* dtemp_dt */, auto... params) {
          auto temp = get<VALUE>(u);
          auto n = cross(get<DERIVATIVE>(X));
 
          return flux_function(X, normalize(n), t, temp, params...);
        },
        domain);
  }
 
  template <typename FluxType>
  void setFluxBCs(FluxType flux_function, Domain& domain)
  {
    setFluxBCs(DependsOn<>{}, flux_function, domain);
  }
 
  const smith::FiniteElementState& temperature() const { return temperature_; };
 
  const smith::FiniteElementState& temperatureRate() const { return temperature_rate_; };
 
  const FiniteElementState& state(const std::string& state_name) const override
  {
    if (state_name == "temperature") {
      return temperature_;
    } else if (state_name == "temperature_rate") {
      return temperature_rate_;
    }
 
    SLIC_ERROR_ROOT(axom::fmt::format("State '{}' requested from solid mechanics module '{}', but it doesn't exist",
                                      state_name, name_));
    return temperature_;
  }
 
  void setState(const std::string& state_name, const FiniteElementState& state) override
  {
    if (state_name == "temperature") {
      temperature_ = state;
      if (!checkpoint_to_disk_) {
        checkpoint_states_["temperature"][static_cast<size_t>(cycle_)] = temperature_;
      }
      return;
    }
 
    SLIC_ERROR_ROOT(axom::fmt::format(
        "setState for state named '{}' requested from heat transfer module '{}', but it doesn't exist", state_name,
        name_));
  }
 
  virtual std::vector<std::string> stateNames() const override
  {
    if (is_quasistatic_) {
      return std::vector<std::string>{"temperature"};
    } else {
      return std::vector<std::string>{"temperature", "temperature_rate"};
    }
  }
 
  template <int... active_parameters, typename callable>
  void addCustomBoundaryIntegral(DependsOn<active_parameters...>, callable qfunction, Domain& domain)
  {
    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 == "temperature") {
      return adjoint_temperature_;
    }
 
    SLIC_ERROR_ROOT(axom::fmt::format("Adjoint '{}' requested from solid mechanics module '{}', but it doesn't exist",
                                      state_name, name_));
    return adjoint_temperature_;
  }
 
  void completeSetup() override
  {
    if (is_quasistatic_) {
      residual_with_bcs_ = mfem_ext::StdFunctionOperator(
          temperature_.space().TrueVSize(),
 
          [this](const mfem::Vector& u, mfem::Vector& r) {
            const mfem::Vector res = (*residual_)(time_, shapeDisplacement(), u, temperature_rate_,
                                                  *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& u) -> mfem::Operator& {
            auto [r, drdu] = (*residual_)(time_, shapeDisplacement(), differentiate_wrt(u), temperature_rate_,
                                          *parameters_[parameter_indices].state...);
            J_ = assemble(drdu);
            J_e_ = bcs_.eliminateAllEssentialDofsFromMatrix(*J_);
            return *J_;
          });
    } else {
      residual_with_bcs_ = mfem_ext::StdFunctionOperator(
          temperature_.space().TrueVSize(),
 
          [this](const mfem::Vector& du_dt, mfem::Vector& r) {
            add(1.0, u_, dt_, du_dt, u_predicted_);
            const mfem::Vector res =
                (*residual_)(time_, shapeDisplacement(), u_predicted_, du_dt, *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& du_dt) -> mfem::Operator& {
            add(1.0, u_, dt_, du_dt, u_predicted_);
 
            // K := dR/du
            auto K = smith::get<DERIVATIVE>((*residual_)(time_, shapeDisplacement(), differentiate_wrt(u_predicted_),
                                                         du_dt, *parameters_[parameter_indices].state...));
            std::unique_ptr<mfem::HypreParMatrix> k_mat(assemble(K));
 
            // M := dR/du_dot
            auto M =
                smith::get<DERIVATIVE>((*residual_)(time_, shapeDisplacement(), u_predicted_, differentiate_wrt(du_dt),
                                                    *parameters_[parameter_indices].state...));
            std::unique_ptr<mfem::HypreParMatrix> m_mat(assemble(M));
 
            // J := M + dt K
            J_.reset(mfem::Add(1.0, *m_mat, dt_, *k_mat));
            J_e_ = bcs_.eliminateAllEssentialDofsFromMatrix(*J_);
 
            return *J_;
          });
    }
 
    if (checkpoint_to_disk_) {
      outputStateToDisk();
    } else {
      checkpoint_states_.clear();
      auto state_names = stateNames();
      for (const auto& state_name : state_names) {
        checkpoint_states_[state_name].push_back(state(state_name));
      }
    }
  }
 
  virtual void setAdjointLoad(std::unordered_map<std::string, const smith::FiniteElementDual&> loads) override
  {
    SLIC_ERROR_ROOT_IF(loads.size() == 0,
                       "Adjoint load container size must be greater than 0 in the heat transfer module.");
 
    auto temp_adjoint_load = loads.find("temperature");
    auto temp_rate_adjoint_load = loads.find("temperature_rate");  // does not need to be specified
 
    SLIC_ERROR_ROOT_IF(temp_adjoint_load == loads.end(), "Adjoint load for \"temperature\" not found.");
 
    temperature_adjoint_load_ = temp_adjoint_load->second;
    // Add the sign correction to move the term to the RHS
    temperature_adjoint_load_ *= -1.0;
 
    if (temp_rate_adjoint_load != loads.end()) {
      temperature_rate_adjoint_load_ = temp_rate_adjoint_load->second;
      temperature_rate_adjoint_load_ *= -1.0;
    }
  }
 
  void reverseAdjointTimestep() override
  {
    auto& lin_solver = nonlin_solver_->linearSolver();
 
    mfem::HypreParVector adjoint_essential(temperature_adjoint_load_);
    adjoint_essential = 0.0;
 
    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");
 
    cycle_--;  // cycle is now at n \in [0,N-1]
 
    double dt = getCheckpointedTimestep(cycle_);
    auto end_step_solution = getCheckpointedStates(cycle_ + 1);
 
    temperature_ = end_step_solution.at("temperature");
 
    if (is_quasistatic_) {
      // We store the previous timestep's temperature as the current temperature for use in the lambdas computing the
      // sensitivities.
 
      auto [_, drdu] = (*residual_)(time_, shapeDisplacement(), differentiate_wrt(temperature_), temperature_rate_,
                                    *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, temperature_adjoint_load_, adjoint_essential);
      }
 
      lin_solver.SetOperator(*J_T);
      lin_solver.Mult(temperature_adjoint_load_, adjoint_temperature_);
    } else {
      SLIC_ERROR_ROOT_IF(ode_.GetTimestepper() != TimestepMethod::BackwardEuler,
                         "Only backward Euler implemented for transient adjoint heat conduction.");
 
      temperature_rate_ = end_step_solution.at("temperature_rate");
 
      // K := dR/du
      auto K = smith::get<DERIVATIVE>((*residual_)(time_, shapeDisplacement(), differentiate_wrt(temperature_),
                                                   temperature_rate_, *parameters_[parameter_indices].state...));
      std::unique_ptr<mfem::HypreParMatrix> k_mat(assemble(K));
 
      // M := dR/du_dot
      auto M = smith::get<DERIVATIVE>((*residual_)(time_, shapeDisplacement(), temperature_,
                                                   differentiate_wrt(temperature_rate_),
                                                   *parameters_[parameter_indices].state...));
      std::unique_ptr<mfem::HypreParMatrix> m_mat(assemble(M));
 
      J_.reset(mfem::Add(1.0, *m_mat, dt, *k_mat));
      auto J_T = std::unique_ptr<mfem::HypreParMatrix>(J_->Transpose());
 
      // recall that temperature_adjoint_load_vector and d_temperature_dt_adjoint_load_vector were already multiplied by
      // -1 above
      mfem::HypreParVector modified_RHS(temperature_adjoint_load_);
      modified_RHS *= dt;
      modified_RHS.Add(1.0, temperature_rate_adjoint_load_);
      modified_RHS.Add(-dt, implicit_sensitivity_temperature_start_of_step_);
 
      for (const auto& bc : bcs_.essentials()) {
        bc.apply(*J_T, modified_RHS, adjoint_essential);
      }
 
      lin_solver.SetOperator(*J_T);
      lin_solver.Mult(modified_RHS, adjoint_temperature_);
 
      // This multiply is technically on M transposed.  However, this matrix should be symmetric unless
      // the thermal capacity is a function of the temperature rate of change, which is thermodynamically
      // impossible, and fortunately not possible with our material interface.
      // Not doing the transpose here to avoid doing unnecessary work.
      m_mat->Mult(adjoint_temperature_, implicit_sensitivity_temperature_start_of_step_);
      implicit_sensitivity_temperature_start_of_step_ *= -1.0 / dt;
      implicit_sensitivity_temperature_start_of_step_.Add(1.0 / dt,
                                                          temperature_rate_adjoint_load_);  // already multiplied by -1
    }
 
    time_end_step_ = time_;
    time_ -= dt;
  }
 
  FiniteElementDual computeTimestepSensitivity(size_t parameter_field) override
  {
    // TODO: the time is likely not being handled correctly on the reverse pass, but we don't
    //       have tests to confirm.
    auto drdparam = smith::get<DERIVATIVE>(d_residual_d_[parameter_field](time_end_step_));
    auto drdparam_mat = assemble(drdparam);
 
    drdparam_mat->MultTranspose(adjoint_temperature_, *parameters_[parameter_field].sensitivity);
 
    return *parameters_[parameter_field].sensitivity;
  }
 
  const FiniteElementDual& computeTimestepShapeSensitivity() override
  {
    auto drdshape =
        smith::get<DERIVATIVE>((*residual_)(time_end_step_, differentiate_wrt(shapeDisplacement()), temperature_,
                                            temperature_rate_, *parameters_[parameter_indices].state...));
 
    auto drdshape_mat = assemble(drdshape);
 
    drdshape_mat->MultTranspose(adjoint_temperature_, shape_displacement_dual_);
 
    return shapeDisplacementSensitivity();
  }
 
  const std::unordered_map<std::string, const smith::FiniteElementDual&> computeInitialConditionSensitivity()
      const override
  {
    return {{"temperature", implicit_sensitivity_temperature_start_of_step_}};
  }
 
  virtual ~HeatTransfer() = default;
 
 protected:
  using scalar_trial = H1<order>;
 
  using shape_trial = H1<SHAPE_ORDER, dim>;
 
  using test = H1<order>;
 
  smith::FiniteElementState temperature_;
 
  FiniteElementState temperature_rate_;
 
  smith::FiniteElementState adjoint_temperature_;
 
  smith::FiniteElementDual implicit_sensitivity_temperature_start_of_step_;
 
  smith::FiniteElementDual temperature_adjoint_load_;
 
  smith::FiniteElementDual temperature_rate_adjoint_load_;
 
  std::unique_ptr<ShapeAwareFunctional<shape_trial, test(scalar_trial, scalar_trial, parameter_space...)>> residual_;
 
  std::unique_ptr<mfem::HypreParMatrix> M_;
 
  std::shared_ptr<mfem::Coefficient> temp_bdr_coef_;
 
  mfem_ext::StdFunctionOperator residual_with_bcs_;
 
  std::unique_ptr<EquationSolver> nonlin_solver_;
 
  mfem_ext::FirstOrderODE ode_;
 
  std::unique_ptr<mfem::HypreParMatrix> J_;
 
  std::unique_ptr<mfem::HypreParMatrix> J_e_;
 
  double dt_;
 
  double previous_dt_;
 
  double time_end_step_;
 
  mfem::Vector u_;
 
  mfem::Vector u_predicted_;
 
  std::array<std::function<decltype((*residual_)(DifferentiateWRT<1>{}, 0.0, shape_displacement_, temperature_,
                                                 temperature_rate_, *parameters_[parameter_indices].state...))(double)>,
             sizeof...(parameter_indices)>
      d_residual_d_ = {[&](double TIME) {
        return (*residual_)(DifferentiateWRT<NUM_STATE_VARS + 1 + parameter_indices>{}, TIME, shape_displacement_,
                            temperature_, temperature_rate_, *parameters_[parameter_indices].state...);
      }...};
};
 
}  // namespace smith