state_manager.h

Go to the documentation of this file.

/*
  This file is part of the Ristra wonton project.
  Please see the license file at the root of this repository, or at:
  https://github.com/laristra/wonton/blob/master/LICENSE
*/

#ifndef WONTON_STATE_STATE_MANAGER_H_
#define WONTON_STATE_STATE_MANAGER_H_

#include <cassert>
#include <string>
#include <iostream>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <memory>
#include <stdexcept>
#include <iterator>
#include <algorithm>

#include "wonton/support/wonton.h"
#include "wonton/state/state_vector_base.h"
#include "wonton/state/state_vector_uni.h"
#include "wonton/state/state_vector_multi.h"

namespace Wonton {

template <class MeshWrapper>
class StateManager {
 public:
StateManager(const MeshWrapper& mesh,
             std::unordered_map<std::string, int> names = {},
             std::unordered_map<int, std::vector<int>> material_cells = {})
    :mesh_(mesh) {
    add_material_names(names);
    add_material_cells(material_cells);
  }


  // state manager level API functions and material names


  void add_material_names( const std::unordered_map<std::string, int>& names) {
    for (auto& kv : names) {

      // this will update, not overwrite
      material_ids_[kv.first]=kv.second;

      // likewise
      material_names_[kv.second] = kv.first;
    }
  }


  std::vector<std::string> const names() const {
    std::vector<std::string> keys;
    for (auto& kv : state_vectors_) {
      keys.push_back(kv.first);
    }
    return keys;
  }


  std::string material_name(int m) const {return material_names_.at(m); }


  int get_material_id(std::string const& name) const {
    return material_ids_.at(name);}


  void add_material_cells(std::unordered_map<int, std::vector<int>> cells) {
    // make sure all materials are known by index
    for (auto &kv : cells) {
      if (material_names_.find(kv.first) == material_names_.end()) {
        throw std::runtime_error(
            "Tried to add cells for an unknown material: "
            + std::to_string(kv.first));
      }

      // create the inverse map
      for (int cell : kv.second) {
        // add the material to the cell
        cell_materials_[cell].insert(kv.first);
      }

      // compute the cell index in material inverse map
      for (unsigned i = 0; i < kv.second.size(); i++)
        cell_index_in_mat_[kv.first][kv.second[i]]=i;
    }

    // set the cells of the material
    material_cells_ = cells;

  }


  std::unordered_map<int, std::vector<int>> const& get_material_cells() const {
    return material_cells_;
  }

    std::unordered_map<int, std::unordered_set<int>> const& get_cell_materials() const {
    return cell_materials_;
  }

  std::unordered_map<int, int> const get_material_shape() const {

    std::unordered_map<int, int> shape;
    for (auto& kv : material_cells_) shape[kv.first] = kv.second.size();

    return shape;
  }


  int num_materials() const {return material_names_.size();}


  // state vector metadata



  Entity_kind get_entity(std::string const& name) const {
    auto v=state_vectors_.find(name);
    return (v != state_vectors_.end()) ? v->second->get_kind() :  Entity_kind::UNKNOWN_KIND;  
  }


  Field_type field_type(Entity_kind kind, std::string const& name) const {
    return state_vectors_.at(name)->get_type();
  }


  int get_data_size(Entity_kind on_what, std::string const& name) const {
    return std::static_pointer_cast<StateVectorUni<>>(
        state_vectors_.at(name))->get_data().size();
  }

  const std::type_info& get_data_type(std::string var_name) {

    // get a shared state vector base class pointer and cast it to whatever
    std::shared_ptr<StateVectorBase> pv = get<StateVectorBase>(var_name);

        // define the data type
        const std::type_info & info = pv->data_type();

        // get the data type
    return info;
  }



  // methods to register or access a complete state vector


  void add(std::shared_ptr<StateVectorBase> sv) {
    // create the key
    std::string key{sv->get_name()};

    // if the map entry for this key already exists, throw an error
    if (state_vectors_.find(key) != state_vectors_.end())
      throw std::runtime_error("Field " + sv->get_name() +
                               " already exists in the state manager");

    // copies the shared pointer
    state_vectors_[key] = sv;
  }


  template <class T>
      void add(std::shared_ptr<StateVectorUni<T>> sv) {
    // create the key
    std::string key{sv->get_name()};

    // if the map entry for this key already exists, throw an error
    if (state_vectors_.find(key) != state_vectors_.end())
      throw std::runtime_error("Field " + sv->get_name() +
                               " already exists in the state manager");

    // check that the size is correct
    if ((int) sv->get_data().size() != mesh_.num_entities(sv->get_kind()))
      throw std::runtime_error(
            "The added data did not have the same number of elements (" +
            std::to_string(sv->get_data().size()) +") as the number of cells ("+
            std::to_string(mesh_.num_owned_cells())+")");

    // copies the shared pointer
    state_vectors_[key] = sv;
  }


  template <class T>
      void add(std::shared_ptr<StateVectorMulti<T>> sv) {
    // create the key
    std::string key{sv->get_name()};

    // if the map entry for this key already exists, throw an error
    // Note: we shouldn't use (state_vectors_[key] != nullptr) as []creates the
    // key even if the key doesn't exist
    if (state_vectors_.find(key) != state_vectors_.end())
      throw std::runtime_error("Field " + sv->get_name() +
                               " already exists in the state manager");

    // check that the size is correct
    if (!shape_is_good<T>(sv)) throw std::runtime_error(
            "The shape of the data was not the same as the shape of the"
            "material cells");

    // copies the shared pointer
    state_vectors_[key] = sv;
  }


  std::shared_ptr<StateVectorBase> get(std::string name) {

    // create the key
    std::string key{name};

    if (state_vectors_.find(key) == state_vectors_.end()) {
      return nullptr;
    }

    // do the return all in one step
    return state_vectors_[key];

  }


  std::shared_ptr<StateVectorBase const> get(std::string name) const {

    // create the key
    std::string key{name};

    if (state_vectors_.find(key) == state_vectors_.end()) {
      return nullptr;
    }

    // do the return all in one step
    return state_vectors_.at(key);

  }


  template <typename T>
      std::shared_ptr<T> get(std::string name) {
    // create the key
    std::string key{name};

    if (state_vectors_.find(key) == state_vectors_.end()) {
      return nullptr;
    }

    // do the return all in one step
    return std::dynamic_pointer_cast<T>(state_vectors_[key]);

  }


  template <typename T>
      std::shared_ptr<T> get(std::string name) const {
    // create the key
    std::string key{name};

    if (state_vectors_.find(key) == state_vectors_.end()) {
      return nullptr;
    }

    // do the return all in one step
    return std::dynamic_pointer_cast<T>(state_vectors_.at(key));
  }



  // material dominant accessor methods



  int num_material_cells(int m) const {return material_cells_.at(m).size();}


  std::vector<int> const& get_material_cells(int m) const {
    return material_cells_.at(m);
  }


  void mat_get_cells(int m, std::vector<int> *matcells) const {
    // the "this" is for a gcc compiler error about assert and templates
    assert(this->material_cells_.find(m) != this->material_cells_.end());
    matcells->clear();
    *matcells = material_cells_.at(m);
  }


  template <class T>
      void mat_get_celldata(std::string const& name, int m,
                            T const **data) const {
    *data = get<StateVectorMulti<T>>(name)->get_data(m).data();
  }


  int cell_index_in_material(int c, int m) const {
    auto const& cell_index_map = cell_index_in_mat_.at(m);
    auto const it = cell_index_map.find(c);
    return it == cell_index_map.end() ? -1 : it->second;
  }


  template <class T>
      void mat_get_celldata(std::string const& name, int m,
                            T **data) {

    // get the correct row of the ragged right data structure
    // NOTE: (this bit me) the &, the local reference needs to refer to the
    // actual data in the state manager, not a copy
    std::vector<T>& material_data = get<StateVectorMulti<T>>(name)->get_data(m);

    int n = material_cells_[m].size();

    // if the space in the row isn't already allocated, fix this
    if (material_data.size() != unsigned(n))
      material_data.resize(n);

    *data = material_data.data();
  }



  // cell dominant accessor methods



  int cell_get_num_mats(int c) const {
    // if a mesh field problem, then there will be no cell_materials_ defined
    if (cell_materials_.size() == 0) return 0;

    // a true multimaterial field
    return cell_materials_.at(c).size();
  }


  std::unordered_set<int> get_cell_materials(int c) const {
    return cell_materials_.at(c);
  }

  void cell_get_mats(int c, std::vector<int> *cellmats) const {
    cellmats->clear();

    // if a mesh field problem, then there will be no cell_materials_ defined
    if (cell_materials_.size() == 0) return;

    auto& mats = cell_materials_.at(c);
    *cellmats = std::vector<int>{mats.begin(), mats.end()};
  }



  // single material accessor methods



  template <class T>
      void mesh_get_data(Entity_kind on_what, std::string const& name,
                         T const **data) const {
    *data = get<StateVectorUni<T>>(name)->get_data().data();
  }


  template <class T>
      void mesh_get_data(Entity_kind on_what, std::string const& name,
                         T **data) {
    *data = get<StateVectorUni<T>>(name)->get_data().data();
  }



  // material dominant mutator methods


  void mat_add_cells(int m, std::vector<int> const& newcells) {
    // assign the cell ids to the material (copy constructor
    material_cells_[m] = newcells;

    // need to update cell_materials_ inverse map
    for (int c : newcells)
      cell_materials_[c].insert(m);

    int i = 0;
    for (auto const& c : material_cells_[m])
      cell_index_in_mat_[m][c] = i++;
  }


  template <class T>
      void mat_add_celldata(std::string const& name, int m,
                            T const * values) {
    // get the number of cells for this material
    // should use api, but I'm going to change it
    int ncells = material_cells_[m].size();

    // could use a std::copy, but I'll get to it later
    // need a reference, because we are modifying the data in place
    std::vector<T>& data = get<StateVectorMulti<T>>(name)->get_data(m);
    data.clear();
    for (int i = 0; i < ncells; ++i) data.push_back(*(values+i));
  }


  void add_material(
      std::string const& matname, std::vector<int> const& matcells) {
    std::cout << "in add_material: " << matname << std::endl;
    throw std::runtime_error("StateManager::add_material is not implemented");
  }




  // utility functions

  template <class T = double>
  bool shape_is_good(const std::shared_ptr<StateVectorMulti<T>> sv) {

    // get the data shape
    std::unordered_map<int, int> shape{get_material_shape()};

    // if we don't have cells yet, and don't know the shape, assume the
    // user knows what they are doing and the shape is good.
    if (shape.size() == 0) return true;

    // get the actual data out of the state vector
    const auto& data = sv->get_data();

      // this one is tricky, but if we build up the state vector material
      // we may start with no data, so for the moment, let's pass this
      if (data.size() == 0) return true;

    // check first that there are the correct number of materials
    if (shape.size() != data.size()) return false;

    for (const auto& kv : data) {
      if (shape.find(kv.first) == shape.end() ||
        shape[kv.first] != (int) kv.second.size()
      ) return false;
    }

    return true;
  }

  void clear_material_cells() {
    material_cells_.clear();
    cell_materials_.clear();
  }

 protected:

  // underlying mesh reference
  const MeshWrapper& mesh_;

  // state vectors
  std::unordered_map<std::string, std::shared_ptr<StateVectorBase>>
      state_vectors_;

  // material name to id
  std::unordered_map<std::string, int> material_ids_;

  // material id to name
  std::unordered_map<int, std::string> material_names_;

  // cell ids for a material
  std::unordered_map<int, std::vector<int>> material_cells_;

  // cell indices for a material
  std::unordered_map<int, std::unordered_map<int,int>> cell_index_in_mat_;

  // material id's for a cell
  std::unordered_map<int, std::unordered_set<int>> cell_materials_;

    // clear all the protected data structures
  void clear() {
    state_vectors_.clear();
    material_ids_.clear();
    material_names_.clear();
    material_cells_.clear();
    cell_materials_.clear();
  }

};

}  // namespace Wonton

#endif  // WONTON_STATE_STATE_MANAGER_H_