state_manager.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 <optional>
#include <unordered_map>
#include <array>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
 
#include "mfem.hpp"
#include "axom/sidre/core/MFEMSidreDataCollection.hpp"
#include "smith/infrastructure/logger.hpp"
#include "smith/physics/state/finite_element_state.hpp"
#include "smith/physics/state/finite_element_dual.hpp"
#include "smith/physics/field_types.hpp"
#include "smith/numerics/functional/quadrature_data.hpp"
#include "smith/numerics/functional/domain.hpp"
#include "smith/numerics/functional/finite_element.hpp"
#include "smith/numerics/functional/geometry.hpp"
 
namespace smith {
 
class StateManager {
 public:
  static void initialize(axom::sidre::DataStore& ds, const std::string& output_directory);
 
  static bool hasState(const std::string& name) { return named_states_.find(name) != named_states_.end(); }
 
  template <typename FunctionSpace>
  static FiniteElementState newState(FunctionSpace space, const std::string& state_name, const std::string& mesh_tag)
  {
    SLIC_ERROR_ROOT_IF(!ds_, "Smith's data store was not initialized - call StateManager::initialize first");
    SLIC_ERROR_ROOT_IF(!hasMesh(mesh_tag), std::format("Mesh tag '{}' not found in the data store", mesh_tag));
    SLIC_ERROR_ROOT_IF(hasState(state_name),
                       std::format("StateManager already contains a state named '{}'", state_name));
 
    FiniteElementState state(mesh(mesh_tag), space, state_name);
 
    storeState(state);
    return state;
  }
 
  static FiniteElementState newState(const mfem::ParFiniteElementSpace& space, const std::string& state_name);
 
  static void storeState(FiniteElementState& state);
 
  template <typename T>
  static void storeQuadratureData(const std::string& mesh_tag, std::shared_ptr<QuadratureData<T>> qdata)
  {
    SLIC_ERROR_ROOT_IF(!ds_, "Smith's data store was not initialized - call StateManager::initialize first");
    SLIC_ERROR_ROOT_IF(!hasMesh(mesh_tag),
                       std::format("Smith's state manager does not have a mesh with given tag '{}'", mesh_tag));
 
    constexpr const char* qds_group_name = "quadraturedatas";
 
    // Get Sidre location for quadrature data inside data collection
    auto& datacoll = datacolls_.at(mesh_tag);
    axom::sidre::Group* bp_group = datacoll.GetBPGroup();  // mesh_datacoll
    // For each geometry type, use i to get both type and name from matching arrays
    for (std::size_t i = 0; i < detail::qdata_geometries.size(); ++i) {
      auto geom_type = detail::qdata_geometries[i];
 
      // Check if geometry type has any data
      if ((*qdata).data.find(geom_type) != (*qdata).data.end()) {
        auto geom_name = detail::qdata_geometry_names[i];
 
        // Get axom::Array of states in map
        auto states = (*qdata)[geom_type];
 
        // Get various size information
        auto num_states = static_cast<axom::IndexType>(states.size());
        SLIC_ERROR_ROOT_IF(num_states == 0, "Number of States should be more than 0 at this point.");
        auto state_size = static_cast<axom::IndexType>(sizeof(*(states.begin())));
        auto total_size = num_states * state_size;
        // Sidre treats information as an array of uint8s
        auto num_uint8s = total_size / static_cast<axom::IndexType>(sizeof(std::uint8_t));
 
        if (!is_restart_) {
          axom::sidre::Group* qdatas_group = bp_group->createGroup(qds_group_name);
 
          // Create Sidre group, store basic information, and point Sidre at the array external to Sidre
          // Note: Sidre will not own this data.
          axom::sidre::Group* geom_group = qdatas_group->createGroup(std::string(geom_name));
          geom_group->createViewScalar("num_states", num_states);
          geom_group->createViewScalar("state_size", state_size);
          geom_group->createViewScalar("total_size", total_size);
 
          // Tell Sidre where the external array is, how large it is (calculated above), and what is in it (faking
          // uint8)
          axom::sidre::View* states_view = geom_group->createView("states");
          states_view->setExternalDataPtr(axom::sidre::UINT8_ID, num_uint8s, states.data());
        } else {
          // Get Sidre group of where the states were stored.
          // Note: this data is not owned by Sidre and the array should have been created at this point but
          // the previous data has not been loaded yet into the array.
          SLIC_ERROR_ROOT_IF(!bp_group->hasGroup(qds_group_name),
                             std::format("Loaded Sidre Datastore did not have group for Quadrature Datas"));
          axom::sidre::Group* qdatas_group = bp_group->getGroup(qds_group_name);
          SLIC_ERROR_ROOT_IF(
              !qdatas_group->hasGroup(std::string(geom_name)),
              std::format("Loaded Sidre Datastore did not have group for Quadrature Data geometry type '{}'",
                          std::string(geom_name)));
          axom::sidre::Group* geom_group = qdatas_group->getGroup(std::string(geom_name));
 
          verifyQuadratureDataSize(
              geom_group, num_states, "num_states",
              "Current number of Quadrature Data States '{}' does not match value in restart '{}'.");
          verifyQuadratureDataSize(geom_group, state_size, "state_size",
                                   "Current size of Quadrature Data State '{}' does not match value in restart '{}'.");
          verifyQuadratureDataSize(
              geom_group, total_size, "total_size",
              "Current total size of Quadrature Data States '{}' does not match value in restart '{}'.");
 
          // Tell Sidre where the external array is
          SLIC_ERROR_ROOT_IF(!geom_group->hasView("states"),
                             "Loaded Quadrature Data geometry Sidre view did not have 'states'");
          axom::sidre::View* states_view = geom_group->getView("states");
          states_view->setExternalDataPtr(states.data());
 
          // TODO: swap this code for the one below after updating Axom
          // Load this set of quadrature data only
          // std::string group_name_to_load = std::format("{0}/{1}", qds_group_name, geom_name);
          // datacoll.LoadExternalData("", group_name_to_load);
        }
      }
    }
    if (is_restart_) {
      // NOTE: This call will reload all external buffers from file stored in the DataStore
      // TODO: This should be changed to load only the current material quadrature data after
      // MFEMSidreDatacollection::LoadExternalData is enhanced to allow loading the
      // external data piecemeal. After Axom PR#1555 goes in, uncomment the loadExternalData call
      // above and remove this if block.
      datacoll.LoadExternalData();
    }
  }
 
  template <typename T>
  static std::shared_ptr<QuadratureData<T>> newQuadratureDataBuffer(const std::string& mesh_tag, const Domain& domain,
                                                                    int order, int dim, T initial_state)
  {
    int Q = order + 1;
 
    std::array<uint32_t, mfem::Geometry::NUM_GEOMETRIES> elems = geometry_counts(domain);
    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));
    }
 
    auto qdata = std::make_shared<QuadratureData<T>>(elems, qpts_per_elem, initial_state);
    storeQuadratureData<T>(mesh_tag, qdata);
    return qdata;
  }
 
  static bool hasDual(const std::string& name) { return named_duals_.find(name) != named_duals_.end(); }
 
  template <typename FunctionSpace>
  static FiniteElementDual newDual(FunctionSpace space, const std::string& dual_name, const std::string& mesh_tag)
  {
    SLIC_ERROR_ROOT_IF(!ds_, "Smith's data store was not initialized - call StateManager::initialize first");
    SLIC_ERROR_ROOT_IF(!hasMesh(mesh_tag), std::format("Mesh tag '{}' not found in the data store", mesh_tag));
    SLIC_ERROR_ROOT_IF(hasDual(dual_name), std::format("StateManager already contains a dual named '{}'", dual_name));
 
    auto dual = FiniteElementDual(mesh(mesh_tag), space, dual_name);
 
    storeDual(dual);
    return dual;
  }
  static FiniteElementDual newDual(const mfem::ParFiniteElementSpace& space, const std::string& dual_name);
 
  static void storeDual(FiniteElementDual& dual);
 
  static void updateState(const FiniteElementState& state)
  {
    SLIC_ERROR_ROOT_IF(!hasState(state.name()),
                       std::format("State manager does not contain state named '{}'", state.name()));
 
    state.fillGridFunction(*named_states_[state.name()]);
  }
 
  static void updateDual(const FiniteElementDual& dual)
  {
    SLIC_ERROR_ROOT_IF(!hasDual(dual.name()),
                       std::format("State manager does not contain dual named '{}'", dual.name()));
 
    dual.space().GetRestrictionMatrix()->MultTranspose(dual, *named_duals_[dual.name()]);
  }
 
  static void save(const double t, const int cycle, const std::string& mesh_tag);
 
  static double load(const int cycle_to_load, const std::string& mesh_tag)
  {
    // FIXME: Assumes that if one DataCollection is going to be reloaded all DataCollections will be
    is_restart_ = true;
    return newDataCollection(mesh_tag, cycle_to_load);
  }
 
  static void reset()
  {
    named_states_.clear();
    named_duals_.clear();
    shape_displacements_.clear();
    shape_displacement_duals_.clear();
    datacolls_.clear();
    output_dir_.clear();
    is_restart_ = false;
    ds_ = nullptr;
  };
 
  static bool hasMesh(const std::string& mesh_tag) { return datacolls_.find(mesh_tag) != datacolls_.end(); }
 
  static mfem::ParMesh& setMesh(std::unique_ptr<mfem::ParMesh> pmesh, const std::string& mesh_tag);
 
  static mfem::ParMesh& mesh(const std::string& mesh_tag);
 
  static FiniteElementState& shapeDisplacement(const std::string& mesh_tag);
 
  static FiniteElementDual& shapeDisplacementDual(const std::string& mesh_tag);
 
  static void loadCheckpointedStates(int cycle_to_load, std::vector<FiniteElementState*> states_to_load);
 
  static FiniteElementDual& shapeDisplacementSensitivity(const std::string& mesh_tag);
 
  static std::string collectionID(const mfem::ParMesh* pmesh);
 
  static bool isRestart() { return is_restart_; }
 
  static int cycle(std::string mesh_tag);
 
  static double time(std::string mesh_tag);
 
 private:
  static void verifyQuadratureDataSize(axom::sidre::Group* group, axom::IndexType value, const char* view_name,
                                       const char* err_msg)
  {
    SLIC_ERROR_IF(!group->hasView(view_name),
                  std::format("Loaded Sidre Datastore does not have value '{}' for Quadrature Data.", view_name));
    auto prev_value = group->getView(view_name)->getData<axom::IndexType>();
    SLIC_ERROR_IF(value != prev_value, std::vformat(err_msg, std::make_format_args(value, prev_value)));
  }
 
  static double newDataCollection(const std::string& mesh_tag, const std::optional<int> cycle_to_load = {});
 
  static std::string getCollectionName(const std::string& mesh_tag) { return mesh_tag + "_datacoll"; }
 
  static void constructShapeFields(const std::string& mesh_tag);
 
  static std::unordered_map<std::string, axom::sidre::MFEMSidreDataCollection> datacolls_;
 
  static std::unordered_map<std::string, std::unique_ptr<FiniteElementState>> shape_displacements_;
  static std::unordered_map<std::string, std::unique_ptr<FiniteElementDual>> shape_displacement_duals_;
 
  static bool is_restart_;
 
  static axom::sidre::DataStore* ds_;
  static std::string output_dir_;
 
  static std::unordered_map<std::string, mfem::ParGridFunction*> named_states_;
  static std::unordered_map<std::string, mfem::ParGridFunction*> named_duals_;
};
 
void checkMesh(const mfem::ParMesh& pmesh, bool is_restart = false);
 
}  // namespace smith