differentiable_physics.cpp

// 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)
 
#include "gretl/data_store.hpp"
#include "smith/differentiable_numerics/differentiable_physics.hpp"
#include "smith/physics/weak_form.hpp"
#include "smith/physics/mesh.hpp"
#include "smith/differentiable_numerics/state_advancer.hpp"
#include "smith/differentiable_numerics/reaction.hpp"
#include "gretl/upstream_state.hpp"
 
namespace smith {
 
gretl::State<int> make_milestone(const std::vector<FieldState>& states, const std::vector<ReactionState>& reactions)
{
  std::vector<gretl::StateBase> base_states;
  for (const auto& s : states) {
    base_states.push_back(s);
  }
  for (const auto& r : reactions) {
    base_states.push_back(r);
  }
 
  auto milestone = states[0].create_state<int, int>(base_states);
 
  milestone.set_eval(
      []([[maybe_unused]] const gretl::UpstreamStates& inputs, gretl::DownstreamState& output) { output.set<int>(0); });
  milestone.set_vjp(
      []([[maybe_unused]] gretl::UpstreamStates& inputs, [[maybe_unused]] const gretl::DownstreamState& output) {});
 
  return milestone.finalize();
}
 
// mesh, equation, fields, parameters, state advancer, solver
DifferentiablePhysics::DifferentiablePhysics(std::shared_ptr<Mesh> mesh, std::shared_ptr<gretl::DataStore> graph,
                                             const FieldState& shape_disp, const std::vector<FieldState>& states,
                                             const std::vector<FieldState>& params,
                                             std::shared_ptr<StateAdvancer> advancer, std::string mech_name,
                                             const std::vector<ReactionInfo>& reaction_infos)
    : BasePhysics(mech_name, mesh, 0, 0.0, false),  // the false is checkpoint_to_disk
      checkpointer_(graph),
      advancer_(advancer),
      reaction_infos_(reaction_infos)
{
  SLIC_ERROR_IF(states.size() == 0, "Must have a least 1 state for a mechanics.");
  field_shape_displacement_ = std::make_unique<FieldState>(shape_disp);
  for (size_t i = 0; i < states.size(); ++i) {
    const auto& s = states[i];
    field_states_.push_back(s);
    initial_field_states_.push_back(s);
    state_name_to_field_index_[s.get()->name()] = i;
    state_names_.push_back(s.get()->name());
  }
 
  for (size_t i = 0; i < params.size(); ++i) {
    const auto& p = params[i];
    field_params_.push_back(p);
    param_name_to_field_index_[p.get()->name()] = i;
    param_names_.push_back(p.get()->name());
  }
 
  reaction_names_.reserve(reaction_infos_.size());
  for (size_t i = 0; i < reaction_infos_.size(); ++i) {
    SLIC_ERROR_IF(
        reaction_infos_[i].space == nullptr,
        axom::fmt::format("Dual '{}' in physics module '{}' has null FE space.", reaction_infos_[i].name, name_));
    reaction_names_.push_back(reaction_infos_[i].name);
    reaction_name_to_reaction_index_[reaction_infos_[i].name] = i;
  }
 
  completeSetup();
}
 
void DifferentiablePhysics::completeSetup()
{
  SLIC_ERROR_IF(field_states_.empty(), "Empty field state during completeSetup()");
  initializeReactionStates();
}
 
void DifferentiablePhysics::resetStates(int cycle, double time)
{
  for (size_t i = 0; i < initial_field_states_.size(); ++i) {
    field_states_[i] = initial_field_states_[i];
  }
  milestones_.clear();
  checkpointer_->reset_graph();
  initializeReactionStates();
  time_ = time;
  cycle_ = cycle;
}
 
void DifferentiablePhysics::resetAdjointStates()
{
  checkpointer_->finalize_graph();
  checkpointer_->reset_for_backprop();
  gretl_assert(checkpointer_->check_validity());
}
 
std::vector<std::string> DifferentiablePhysics::stateNames() const { return state_names_; }
 
std::vector<std::string> DifferentiablePhysics::parameterNames() const { return param_names_; }
 
std::vector<std::string> DifferentiablePhysics::dualNames() const { return reaction_names_; }
 
const FiniteElementState& DifferentiablePhysics::state([[maybe_unused]] const std::string& field_name) const
{
  SLIC_ERROR_IF(
      state_name_to_field_index_.find(field_name) == state_name_to_field_index_.end(),
      std::format("Could not find field named {0} in mesh with tag \"{1}\" to get", field_name, mesh_->tag()));
  size_t state_index = state_name_to_field_index_.at(field_name);
  return *field_states_[state_index].get();
}
 
const FiniteElementDual& DifferentiablePhysics::dual(const std::string& reaction_name) const
{
  SLIC_ERROR_IF(reaction_name_to_reaction_index_.find(reaction_name) == reaction_name_to_reaction_index_.end(),
                axom::fmt::format("Could not find reaction named {0} in mesh with tag \"{1}\" to get", reaction_name,
                                  mesh_->tag()));
  size_t reaction_index = reaction_name_to_reaction_index_.at(reaction_name);
 
  SLIC_ERROR_IF(reaction_states_.empty() && !reaction_names_.empty(),
                "Reactions were not computed during advanceState, but were requested.");
 
  SLIC_ERROR_IF(
      reaction_index >= reaction_states_.size(),
      "Reaction reactions not correctly allocated yet, cannot get reaction until after initializationStep is called.");
 
  return *reaction_states_[reaction_index].get();
}
 
FiniteElementDual DifferentiablePhysics::loadCheckpointedDual(const std::string& reaction_name, int cycle)
{
  SLIC_ERROR_IF(
      cycle != cycle_,
      axom::fmt::format("Due to checkpointing restrictions in smith::DifferentiablePhysics, cannot ask for "
                        "an arbitrary checkpointed reaction cycle, asking for cycle {}, but physics is at cycle {}",
                        cycle, cycle_));
  return dual(reaction_name);
}
 
FiniteElementState DifferentiablePhysics::loadCheckpointedState(const std::string& state_name, int cycle)
{
  SLIC_ERROR_IF(cycle != cycle_,
                std::format("Due to checkpointing restrictions in smith::Mechanics, cannot ask for an arbitrary "
                            "checkpointed cycle, asking for cycle {}, but physics is at cycle {}",
                            cycle, cycle_));
  return state(state_name);
}
 
const FiniteElementState& DifferentiablePhysics::shapeDisplacement() const { return *field_shape_displacement_->get(); }
 
const FiniteElementState& DifferentiablePhysics::parameter(std::size_t parameter_index) const
{
  SLIC_ERROR_IF(parameter_index >= field_params_.size(),
                std::format("Parameter index {} requested, but only {} parameters exist in physics module {}.",
                            parameter_index, field_params_.size(), name_));
  return *field_params_[parameter_index].get();
}
 
const FiniteElementState& DifferentiablePhysics::parameter(const std::string& parameter_name) const
{
  SLIC_ERROR_IF(
      param_name_to_field_index_.find(parameter_name) == param_name_to_field_index_.end(),
      std::format("Could not find parameter named {0} in mesh with tag \"{1}\" to get", parameter_name, mesh_->tag()));
  size_t param_index = param_name_to_field_index_.at(parameter_name);
  return parameter(param_index);
}
 
void DifferentiablePhysics::setParameter(const size_t parameter_index, const FiniteElementState& parameter_state)
{
  SLIC_ERROR_IF(parameter_index >= field_params_.size(),
                std::format("Parameter '{}' requested when only '{}' parameters exist in physics module '{}'",
                            parameter_index, field_params_.size(), name_));
  *field_params_[parameter_index].get() = parameter_state;
}
 
void DifferentiablePhysics::setShapeDisplacement(const FiniteElementState& shape_displacement)
{
  *field_shape_displacement_->get() = shape_displacement;
}
 
void DifferentiablePhysics::setState([[maybe_unused]] const std::string& field_name,
                                     [[maybe_unused]] const FiniteElementState& s)
{
  SLIC_ERROR_IF(state_name_to_field_index_.find(field_name) == state_name_to_field_index_.end(),
                std::format("Could not find field named {0} in mesh with tag {1} to set", field_name, mesh_->tag()));
  size_t state_index = state_name_to_field_index_.at(field_name);
  *field_states_[state_index].get() = s;
  *initial_field_states_[state_index].get() = s;
}
 
void DifferentiablePhysics::setAdjointLoad(
    std::unordered_map<std::string, const smith::FiniteElementDual&> string_to_reaction)
{
  for (auto string_reaction_pair : string_to_reaction) {
    std::string field_name = string_reaction_pair.first;
    const smith::FiniteElementDual& reaction = string_reaction_pair.second;
    SLIC_ERROR_IF(
        state_name_to_field_index_.find(field_name) == state_name_to_field_index_.end(),
        axom::fmt::format("Could not find reaction named {0} in mesh with tag {1}", field_name, mesh_->tag()));
    size_t state_index = state_name_to_field_index_.at(field_name);
    *field_states_[state_index].get_dual() += reaction;
  }
}
 
void DifferentiablePhysics::setDualAdjointBcs(
    std::unordered_map<std::string, const smith::FiniteElementState&> string_to_bc)
{
  for (auto string_bc_pair : string_to_bc) {
    std::string reaction_name = string_bc_pair.first;
    const smith::FiniteElementState& reaction_adjoint_load = string_bc_pair.second;
    SLIC_ERROR_IF(
        reaction_name_to_reaction_index_.find(reaction_name) == reaction_name_to_reaction_index_.end(),
        axom::fmt::format("When calling setDualAdjointBcs, could not find reaction named {0} in mesh with tag {1}",
                          reaction_name, mesh_->tag()));
    size_t reaction_index = reaction_name_to_reaction_index_.at(reaction_name);
    *reaction_states_[reaction_index].get_dual() += reaction_adjoint_load;
  }
}
 
const FiniteElementState& DifferentiablePhysics::adjoint([[maybe_unused]] const std::string& adjoint_name) const
{
  // MRT, not implemented
  SLIC_ERROR("What is the use case for asking for the adjoint solution field directly?");
  return *adjoints_[0];
}
 
void DifferentiablePhysics::advanceTimestep(double dt)
{
  if (cycle_ == 0) {
    field_states_ = initial_field_states_;
    milestones_.push_back(make_milestone(field_states_, reaction_states_).step());
  }
 
  cycle_prev_ = cycle_;
  time_prev_ = time_;
  dt_prev_ = dt;
 
  TimeInfo time_info(time_, dt, static_cast<size_t>(cycle_));
  auto [states, reactions] =
      advancer_->advanceState(time_info, *field_shape_displacement_, field_states_, field_params_);
  field_states_ = states;
  reaction_states_ = reactions;
 
  cycle_++;
  time_ += dt;
  milestones_.push_back(make_milestone(field_states_, reaction_states_).step());
}
 
void DifferentiablePhysics::reverseAdjointTimestep()
{
  --cycle_;
  const gretl::Int milestone = milestones_[static_cast<size_t>(cycle_)];
 
  field_shape_displacement_->clear_dual();
  for (auto& p : field_params_) {
    p.clear_dual();
  }
 
  gretl::Int current_step = checkpointer_->currentStep_;
  while (milestone != current_step) {
    checkpointer_->reverse_state();
    current_step = checkpointer_->currentStep_;
  }
 
  gretl::UpstreamStates upstreams(*checkpointer_, checkpointer_->upstreamSteps_[milestone]);
 
  const size_t expected_upstreams = field_states_.size() + reaction_states_.size();
  SLIC_ERROR_IF(expected_upstreams != upstreams.size(), "field/reaction states and upstream sizes do not match.");
  // recreate the upstream field states with upstream step, field, and dual values.
  for (size_t s = 0; s < field_states_.size(); ++s) {
    field_states_[s].reset_step(upstreams[s].step_);
    field_states_[s].set(upstreams[s].get<FEFieldPtr>());
    field_states_[s].set_dual(upstreams[s].get_dual<FEDualPtr, FEFieldPtr>());
  }
  for (size_t r = 0; r < reaction_states_.size(); ++r) {
    const size_t upstream_index = field_states_.size() + r;
    reaction_states_[r].reset_step(upstreams[upstream_index].step_);
    reaction_states_[r].set(upstreams[upstream_index].get<FEDualPtr>());
    reaction_states_[r].set_dual(upstreams[upstream_index].get_dual<FEFieldPtr, FEDualPtr>());
  }
}
 
FiniteElementDual DifferentiablePhysics::computeTimestepSensitivity(size_t parameter_index)
{
  return *field_params_[parameter_index].get_dual();
}
 
const FiniteElementDual& DifferentiablePhysics::computeTimestepShapeSensitivity()
{
  return *field_shape_displacement_->get_dual();
}
 
const std::unordered_map<std::string, const smith::FiniteElementDual&>
DifferentiablePhysics::computeInitialConditionSensitivity() const
{
  std::unordered_map<std::string, const smith::FiniteElementDual&> map;
  for (auto& name : stateNames()) {
    auto state_index = state_name_to_field_index_.at(name);
    map.insert({name, *initial_field_states_[state_index].get_dual()});
  }
  return map;
}
 
std::vector<FieldState> DifferentiablePhysics::getFieldStatesAndParamStates() const
{
  std::vector<FieldState> fields;
  fields.insert(fields.end(), field_states_.begin(), field_states_.end());
  fields.insert(fields.end(), field_params_.begin(), field_params_.end());
  return fields;
}
 
FieldState DifferentiablePhysics::getShapeDispFieldState() const { return *field_shape_displacement_; }
 
void DifferentiablePhysics::initializeReactionStates()
{
  reaction_states_.clear();
  reaction_states_.reserve(reaction_infos_.size());
  for (const auto& reaction_info : reaction_infos_) {
    auto reaction = std::make_shared<FiniteElementDual>(*reaction_info.space, reaction_info.name);
    reaction_states_.push_back(createReactionState(*checkpointer_, reaction));
  }
}
 
}  // namespace smith