field_store.hpp

// 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/differentiable_numerics/field_state.hpp"
#include "smith/differentiable_numerics/time_integration_rule.hpp"
#include "smith/physics/functional_weak_form.hpp"
#include "smith/physics/mesh.hpp"
 
#include <map>
#include <set>
#include <string>
#include <vector>
#include <memory>
 
namespace smith {
 
class DirichletBoundaryConditions;
class BoundaryConditionManager;
 
struct ReactionInfo {
  std::string name;                                    
  const mfem::ParFiniteElementSpace* space = nullptr;  
};
 
template <typename Space, typename Time = void*>
struct FieldType {
  using space_type = Space;  
 
  FieldType(std::string n, bool is_unknown_ = false) : name(n), is_unknown(is_unknown_) {}
  std::string name;         
  bool is_unknown = false;  
};
 
struct FieldStore {
  FieldStore(std::shared_ptr<Mesh> mesh, size_t storage_size = 50, std::string prepend_name = "");
 
  std::string prefix(const std::string& base) const;
 
  enum class TimeDerivative
  {
    VAL,   //< The value of the field.
    DOT,   
    DDOT,  
    DDDOT  
  };
 
  template <typename Space>
  void addShapeDisp(FieldType<Space>& type)
  {
    type.name = prefix(type.name);
    shape_disp_.push_back(smith::createFieldState<Space>(*graph_, Space{}, type.name, mesh_->tag()));
  }
 
  template <typename Space>
  void addParameter(FieldType<Space>& type)
  {
    type.name = prefix(type.name);
    if (to_params_index_.count(type.name)) {
      // Already registered — expected when multiple systems share the same parameter.
      // Verify the space matches by checking vdim (== Space::components).
      auto& existing = params_[to_params_index_.at(type.name)];
      SLIC_ERROR_ROOT_IF(existing.get()->space().GetVDim() != Space::components,
                         axom::fmt::format("Parameter '{}' re-registered with a different space "
                                           "(existing vdim={}, new vdim={})",
                                           type.name, existing.get()->space().GetVDim(), Space::components));
      return;
    }
    to_params_index_[type.name] = params_.size();
    params_.push_back(smith::createFieldState<Space>(*graph_, Space{}, type.name, mesh_->tag()));
  }
 
  template <typename Space>
  std::shared_ptr<DirichletBoundaryConditions> addIndependent(FieldType<Space>& type,
                                                              std::shared_ptr<TimeIntegrationRule> time_rule)
  {
    type.name = prefix(type.name);
    type.is_unknown = true;
    to_states_index_[type.name] = states_.size();
    FieldState new_field = smith::createFieldState<Space>(*graph_, Space{}, type.name, mesh_->tag());
    states_.push_back(new_field);
    is_solve_state_.push_back(true);
    auto latest_bc = addBoundaryConditions(new_field.get());
    boundary_conditions_[type.name] = latest_bc;
 
    SLIC_ERROR_IF(!time_rule, "Invalid time_rule");
 
    TimeIntegrationMapping mapping;
    mapping.primary_name = type.name;
    independent_name_to_rule_index_[type.name] = time_integration_rules_.size();
    time_integration_rules_.push_back({time_rule, mapping});
 
    return latest_bc;
  }
 
  template <typename Space>
  auto addDependent(FieldType<Space> independent_field, TimeDerivative derivative, std::string name_override = "")
  {
    std::string suffix;
    if (derivative == TimeDerivative::VAL) {
      suffix = "_old";
    } else if (derivative == TimeDerivative::DOT) {
      suffix = "_dot_old";
    } else if (derivative == TimeDerivative::DDOT) {
      suffix = "_ddot_old";
    } else {
      SLIC_ERROR("Unsupported TimeDerivative");
    }
 
    std::string name = name_override.empty() ? independent_field.name + suffix : prefix(name_override);
 
    if (independent_name_to_rule_index_.count(independent_field.name)) {
      size_t rule_idx = independent_name_to_rule_index_.at(independent_field.name);
      auto& mapping = time_integration_rules_[rule_idx].second;
      if (derivative == TimeDerivative::VAL) {
        mapping.history_name = name;
      } else if (derivative == TimeDerivative::DOT) {
        mapping.dot_name = name;
      } else if (derivative == TimeDerivative::DDOT) {
        mapping.ddot_name = name;
      }
    } else {
      SLIC_WARNING("Adding dependent time integration field for independent field '"
                   << independent_field.name << "' which has no registered TimeIntegrationRule.");
    }
 
    to_states_index_[name] = states_.size();
    states_.push_back(smith::createFieldState<Space>(*graph_, Space{}, name, mesh_->tag()));
    is_solve_state_.push_back(false);
    return FieldType<Space>(name);
  }
 
  void addWeakFormUnknownArg(std::string weak_form_name, std::string argument_name, size_t argument_index);
 
  void addWeakFormArg(std::string weak_form_name, std::string argument_name, size_t argument_index);
 
  void addWeakFormReaction(std::string weak_form_name, std::string field_name);
 
  void markWeakFormInternal(const std::string& weak_form_name);
 
  std::string getWeakFormReaction(const std::string& weak_form_name) const;
 
  template <typename... FieldTypes>
  std::vector<const mfem::ParFiniteElementSpace*> createSpaces(const std::string& weak_form_name,
                                                               const std::string& reaction_field_name,
                                                               FieldTypes... types)
  {
    addWeakFormReaction(weak_form_name, reaction_field_name);
    std::vector<const mfem::ParFiniteElementSpace*> spaces;
    size_t arg_num = 0;
    auto register_field = [&](auto type) {
      spaces.push_back(&getField(type.name).get()->space());
      addWeakFormArg(weak_form_name, type.name, arg_num);
      if (type.is_unknown) {
        addWeakFormUnknownArg(weak_form_name, type.name, arg_num);
      }
      ++arg_num;
    };
    (register_field(types), ...);
    return spaces;
  }
 
  struct TimeIntegrationMapping {
    std::string primary_name;  
    std::string history_name;  
    std::string dot_name;      
    std::string ddot_name;     
  };
 
  const std::vector<std::pair<std::shared_ptr<TimeIntegrationRule>, TimeIntegrationMapping>>& getTimeIntegrationRules()
      const;
 
  void printMap();
 
  std::vector<std::vector<size_t>> indexMap(const std::vector<std::string>& residual_names) const;
 
  std::vector<const BoundaryConditionManager*> getBoundaryConditionManagers(
      const std::vector<std::string>& weak_form_names) const;
 
  std::vector<const BoundaryConditionManager*> getBoundaryConditionManagersForFields(
      const std::vector<std::string>& field_names) const;
 
  bool hasField(const std::string& field_name) const;
 
  size_t getFieldIndex(const std::string& field_name) const;
 
  FieldState getField(const std::string& field_name) const;
 
  FieldState getParameter(const std::string& param_name) const;
 
  void setField(const std::string& field_name, FieldState updated_field);
 
  void setField(size_t index, FieldState updated_field);
 
  FieldState getShapeDisp() const;
 
  const std::vector<FieldState>& getAllFields() const;
 
  std::vector<FieldState> getStates(const std::string& weak_form_name) const;
 
  std::vector<FieldState> getStatesFromVectors(const std::string& weak_form_name,
                                               const std::vector<FieldState>& state_fields,
                                               const std::vector<FieldState>& param_fields) const;
 
  const std::vector<FieldState>& getParameterFields() const;
 
  const std::vector<FieldState>& getStateFields() const;
 
  std::vector<FieldState> getOutputFieldStates() const;
 
  std::vector<ReactionInfo> getReactionInfos() const;
 
  const std::shared_ptr<smith::Mesh>& getMesh() const;
 
  std::shared_ptr<DirichletBoundaryConditions> getBoundaryConditions(const std::string& field_name) const;
 
  const std::shared_ptr<gretl::DataStore>& graph() const;
 
 private:
  std::string resolveFieldName(const std::string& field_name) const;
 
  std::shared_ptr<Mesh> mesh_;
  std::shared_ptr<gretl::DataStore> graph_;
  std::string prepend_name_;
 
  std::vector<FieldState> shape_disp_;
  std::vector<FieldState> params_;
  std::vector<FieldState> states_;
  std::vector<bool> is_solve_state_;
 
  std::map<std::string, size_t> to_states_index_;
  std::map<std::string, size_t> to_params_index_;
 
  std::map<std::string, std::shared_ptr<DirichletBoundaryConditions>> boundary_conditions_;
 
  struct FieldLabel {
    std::string field_name;
    size_t field_index_in_residual;
  };
 
  std::shared_ptr<DirichletBoundaryConditions> addBoundaryConditions(FEFieldPtr field);
 
  std::map<std::string, std::vector<FieldLabel>> weak_form_name_to_unknown_name_index_;
 
  std::map<std::string, std::vector<size_t>> weak_form_name_to_field_indices_;
  std::map<std::string, std::vector<std::string>> weak_form_name_to_field_names_;
 
  std::vector<std::pair<std::string, std::string>> weak_form_to_test_field_;
  std::set<std::string> internal_weak_forms_;  
 
  std::vector<std::pair<std::shared_ptr<TimeIntegrationRule>, TimeIntegrationMapping>> time_integration_rules_;
  std::map<std::string, size_t> independent_name_to_rule_index_;
};
 
template <int spatial_dim, typename TestSpaceType, typename... InputSpaceTypes>
auto createWeakForm(std::string name, FieldType<TestSpaceType> test_type, FieldStore& field_store,
                    FieldType<InputSpaceTypes>... field_types)
{
  return std::make_shared<FunctionalWeakForm<spatial_dim, TestSpaceType, Parameters<InputSpaceTypes...>>>(
      name, field_store.getMesh(), field_store.getField(test_type.name).get()->space(),
      field_store.createSpaces(name, test_type.name, field_types...));
}
 
inline std::shared_ptr<FieldStore> fieldStore(std::shared_ptr<Mesh> mesh, std::size_t storage_size = 200,
                                              std::string prefix = "")
{
  return std::make_shared<FieldStore>(std::move(mesh), storage_size, std::move(prefix));
}
 
}  // namespace smith