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/differentiable_numerics/time_discretized_weak_form.hpp"
#include "smith/physics/mesh.hpp"
 
#include <map>
#include <string>
#include <vector>
#include <memory>
 
namespace smith {
 
class DirichletBoundaryConditions;
class BoundaryConditionManager;
 
template <typename Space, typename Time = void*>
struct FieldType {
  FieldType(std::string n, int unknown_index_ = -1) : name(n), unknown_index(unknown_index_) {}
  std::string name;   
  int unknown_index;  
};
 
struct FieldStore {
  FieldStore(std::shared_ptr<Mesh> mesh, size_t storage_size = 50);
 
  enum class TimeDerivative
  {
    VAL,   //< The value of the field.
    DOT,   
    DDOT,  
    DDDOT  
  };
 
  template <typename Space>
  void addShapeDisp(FieldType<Space> type)
  {
    shape_disp_.push_back(smith::createFieldState<Space>(*graph_, Space{}, type.name, mesh_->tag()));
  }
 
  template <typename Space>
  void addParameter(FieldType<Space> type)
  {
    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.unknown_index = static_cast<int>(num_unknowns_);
    to_states_index_[type.name] = states_.size();
    to_unknown_index_[type.name] = num_unknowns_;
    FieldState new_field = smith::createFieldState<Space>(*graph_, Space{}, type.name, mesh_->tag());
    states_.push_back(new_field);
    auto latest_bc = addBoundaryConditions(new_field.get());
    ++num_unknowns_;
    SLIC_ERROR_IF(num_unknowns_ != boundary_conditions_.size(),
                  "Inconcistency between num unknowns and boundary condition size");
 
    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 : 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()));
    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);
 
  size_t getNumUnknowns() const { return num_unknowns_; }
 
  void addWeakFormReaction(std::string weak_form_name, std::string field_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.unknown_index >= 0) {
        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::shared_ptr<DirichletBoundaryConditions> getBoundaryConditions(size_t unknown_index) const;
 
  size_t getFieldIndex(const std::string& field_name) const;
 
  size_t getUnknownIndex(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::shared_ptr<smith::Mesh>& getMesh() const;
 
  const std::shared_ptr<gretl::DataStore>& graph() const;
 
 private:
  std::shared_ptr<Mesh> mesh_;
  std::shared_ptr<gretl::DataStore> graph_;
 
  std::vector<FieldState> shape_disp_;
  std::vector<FieldState> params_;
  std::vector<FieldState> states_;
 
  std::map<std::string, size_t> to_states_index_;
  std::map<std::string, size_t> to_params_index_;
 
  size_t num_unknowns_ = 0;
  std::map<std::string, size_t> to_unknown_index_;
  std::vector<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::map<std::string, std::string> weak_form_to_test_field_;
 
  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<TimeDiscretizedWeakForm<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...));
}
 
}  // namespace smith