flat_state_mm_wrapper.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_FLAT_FLAT_STATE_MM_WRAPPER_H_
#define WONTON_STATE_FLAT_FLAT_STATE_MM_WRAPPER_H_

#include <vector>
#include <string>
#include <unordered_map>
#include <memory>
#include <type_traits>

#include "wonton/state/state_manager.h"
#include "wonton/support/wonton.h"
#include "wonton/state/state_vector_multi.h"

namespace Wonton {

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


  Flat_State_Wrapper & operator = (Flat_State_Wrapper const &) = delete;


  ~Flat_State_Wrapper() = default;


  template <class State_Wrapper>
  void initialize(State_Wrapper const & input,
                   std::vector<std::string> const& var_names) {

    // clear any existing state data
    StateManager<MeshWrapper>::clear();

     // get the number of materials
    int nmats = input.num_materials();

    // we may need to add volume fraction an centroid data to the list of
    // distributed fields, so define a temp variable for these fields
    std::vector<std::string> distribute_var_names(var_names.begin(), var_names.end());

    // is this a multimaterial problem, if so, add material names and cell indices
    // to the flat state. Do this only once regardless of the number of fields.
    if (nmats>0){

      // create the name to id map
      std::unordered_map<std::string, int> name_map;
      for (int i=0; i<nmats; ++i){

        // get the material name
        std::string material_name = input.material_name(i);

        // add to the temp name map
        name_map[material_name]=i;
      }

      // add the names to the state manager
      StateManager<MeshWrapper>::add_material_names(name_map);

      // copy cell indices
      for (int m=0; m<nmats; ++m){

        // get the cells for this material from the input state wrapper
        std::vector<int> mat_cells;
        input.mat_get_cells(m, &mat_cells);

        // add the cells to the flat state
        StateManager<MeshWrapper>::mat_add_cells(m, mat_cells);
      }      
      
      // if this is a multimaterial problem we will also need to distribute
      // volume fraction and potentially centroids if they are defined
      distribute_var_names.emplace_back("mat_volfracs");
      if (input.get_entity("mat_centroids") != Entity_kind::UNKNOWN_KIND)
        distribute_var_names.emplace_back("mat_centroids");

    }

    // add material fields to the flat state
    for (std::string varname : distribute_var_names) {

      // get the entity kind
      Entity_kind entity = input.get_entity(varname);

      // get the field type
      Field_type type = input.field_type(entity, varname);

      if (type==Wonton::Field_type::MULTIMATERIAL_FIELD){

        // We need a non-const state wrapper for the get_data_type function
        auto& input_nonconst = const_cast<State_Wrapper&>(input);

        // We need to create a state vector but the type is variable
        const std::type_info& data_type = input_nonconst.get_data_type(varname);

        // construct the state vector type by if statement
        // Is there a better/more idiomatic way to do this???
        if (data_type == typeid(double)){

          // create a uni state vector
          std::shared_ptr<StateVectorMulti<double>> pv =
            std::make_shared<StateVectorMulti<double>> (varname);

          // add to database (inside loop, since pv is typed diffently below and
          // pv has local scope
          StateManager<MeshWrapper>::add(pv);

          // loop over materials
          for (int m=0; m<nmats; ++m){

            // get the data from the input state wrapper
            double const* data;
            input.mat_get_celldata(varname, m, &data);

            // add the data to the flat state
            StateManager<MeshWrapper>::mat_add_celldata(varname, m, data);
            auto pv1 = StateManager<MeshWrapper>::get(varname);
            auto pv2 = std::dynamic_pointer_cast<StateVectorMulti<double>>(pv1);
          }

        } else if (data_type == typeid(Wonton::Point<2>)){

          // create a uni state vector
          std::shared_ptr<StateVectorMulti<Wonton::Point<2>>> pv =
            std::make_shared<StateVectorMulti<Wonton::Point<2>>> (varname);

          // add to database (inside loop, since pv is typed diffently below and
          // pv has local scope
          StateManager<MeshWrapper>::add(pv);

          // loop over materials
          for (int m=0; m<nmats; ++m){

            // get the data from the input state wrapper
            Wonton::Point<2> const* data;
            input.mat_get_celldata(varname, m, &data);

            // add the data to the flat state
            StateManager<MeshWrapper>::mat_add_celldata(varname, m, data);
            auto pv1 = StateManager<MeshWrapper>::get(varname);
            auto pv2 = std::dynamic_pointer_cast<StateVectorMulti<Wonton::Point<2>>>(pv1);
          }
        } else if (data_type == typeid(Wonton::Point<3>)){

          // create a uni state vector
          std::shared_ptr<StateVectorMulti<Wonton::Point<3>>> pv =
            std::make_shared<StateVectorMulti<Wonton::Point<3>>> (varname);

          // add to database (inside loop, since pv is typed diffently below and
          // pv has local scope
          StateManager<MeshWrapper>::add(pv);

          // loop over materials
          for (int m=0; m<nmats; ++m){

            // get the data from the input state wrapper
            Wonton::Point<3> const* data;
            input.mat_get_celldata(varname, m, &data);

            // add the data to the flat state
            StateManager<MeshWrapper>::mat_add_celldata(varname, m, data);
            auto pv1 = StateManager<MeshWrapper>::get(varname);
            auto pv2 = std::dynamic_pointer_cast<StateVectorMulti<Wonton::Point<3>>>(pv1);
          }
        } else {
          throw std::runtime_error("found an unknown field type!");
        }

      } else {

        // AT THE MOMENT WE ARE ASSUMING SINGLE MATERIAL FIELDS ARE ALWAYS DOUBLES
        // get pointer to data for state from input state wrapper
        double const* data;
        input.mesh_get_data(entity, varname, &data);

        // Note, this get_data_size is of the input wrapper, which is defined
        // not the base class state manager get_data_size which is not implemented
        size_t dataSize = input.get_data_size(entity, varname);

        // create a uni state vector
        std::shared_ptr<StateVectorUni<double>> pv =
          std::make_shared<StateVectorUni<double>> (varname, data,
                                                    data+dataSize, entity);
        // add to database
        StateManager<MeshWrapper>::add(pv);
      }
    }
  }


  size_t get_num_vectors() {
    return StateManager<MeshWrapper>::state_vectors_.size(); }


  std::vector<double> pack(std::string field_name) {

    // get the state vector base class shared pointer
    std::shared_ptr<StateVectorBase> pv =
        StateManager<MeshWrapper>::get(field_name);

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

    // get the field type
    const Wonton:: Field_type field_type = pv->get_type();

    if (field_type == Wonton::Field_type::MESH_FIELD){

      // this is a mesh field

      if ( data_type == typeid(double)){

          // field data is doubles
        return std::static_pointer_cast<StateVectorUni<double>>(pv)->get_data();
      }

    } else if (field_type == Wonton::Field_type::MULTIMATERIAL_FIELD){

      // this is a multimaterial field

      if ( data_type == typeid(Wonton::Point<2>)){

        // field data is wonton 2d points

        // get the ordered keys
        std::vector<int> const ids = get_material_ids();

        // get the raw mm data
        std::unordered_map<int, std::vector<Wonton::Point<2>>> mm_data =
            std::static_pointer_cast<StateVectorMulti<Wonton::Point<2>>>(pv)->get_data();

        // define the result
        std::vector<double> result;

        // loop over material ids in the mm state in sort order
        for (int id : ids){
          for (auto& d : mm_data.at(id)){
            result.emplace_back(d[0]);
            result.emplace_back(d[1]);
          }
        }

        return result;
      } else if ( data_type == typeid(Wonton::Point<3>)){

        // field data is wonton 3d points

        // get the ordered keys
        std::vector<int> const ids = get_material_ids();

        // get the raw mm data
        std::unordered_map<int, std::vector<Wonton::Point<3>>> mm_data =
            std::static_pointer_cast<StateVectorMulti<Wonton::Point<3>>>(pv)->get_data();

        // define the result
        std::vector<double> result;

        // loop over material ids in the mm state in sort order
        for (int id : ids){
          for (auto& d : mm_data.at(id)){
            result.emplace_back(d[0]);
            result.emplace_back(d[1]);
            result.emplace_back(d[2]);
          }
        }

        return result;
      } else if (data_type == typeid(double)){

        // field data is doubles

        // get the ordered keys
        std::vector<int> const ids = get_material_ids();

        // get the raw mm data
        std::unordered_map<int, std::vector<double>> mm_data =
          std::static_pointer_cast<StateVectorMulti<double>>(pv)->get_data();

        // define the result
        std::vector<double> result;

        // loop over material ids in the mm state in sort order
        for (int id : ids){
          for (auto& d : mm_data.at(id)){
            result.emplace_back(d);
          }
        }

        return result;
      }
    }
    
    throw std::runtime_error("invalid field type");
  }


  void unpack(std::string field_name, const std::vector<double> & flat_data,
    const std::vector<int> & all_material_ids={},
    const std::vector<int> & all_material_shapes={},
    const std::map<int, std::vector<int>>& distributedMaterialCellIds_={} ) {

    // get the state vector base class shared pointer
    std::shared_ptr<StateVectorBase> pv =
      StateManager<MeshWrapper>::get(field_name);

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

    // get the field type
    const Wonton:: Field_type field_type = pv->get_type();

    if (field_type == Wonton::Field_type::MESH_FIELD){

      // this is a mesh field

      if ( data_type == typeid(double)){

        // copy the mesh data into the state vector
        std::static_pointer_cast<StateVectorUni<double>>(pv)->get_data()=flat_data;
      }

    } else if (field_type == Wonton::Field_type::MULTIMATERIAL_FIELD){

      // THIS SHOULD BE MADE PRETTIER DURING A CLEANUP SPRINT. THE FOLLOWING
      // THREE CODE BLOCKS POINT<2>, POINT<3>, AND DOUBLE ARE NEARLY IDENTICAL
      // POINT<2> AND POINT<3> COULD OBVIOUSLY BE TEMPLATED, THE DOUBLE CASE
      // MAY NEED IT'S OWN FUNCTION. BUT, THE CASE SHOULD BE WRITTEN AS FUNCTIONS
      // RATHER THAN INLINE

      // this is a multimaterial field

      if ( data_type == typeid(Wonton::Point<2>)){

        // field data is 2d wonton points

        // allocate the vector data
        std::vector<Wonton::Point<2>> correctly_typed_data(flat_data.size()/2);

        // convert to the correct type
        int const num_flat_data = flat_data.size();
        for (int i = 0; i < num_flat_data; i += 2) {
          correctly_typed_data[i/2] = Wonton::Point<2>(flat_data[i], flat_data[i+1]);
        }

        // allocate the data for unpacking
        std::unordered_map<int,std::vector<Wonton::Point<2>>> material_data;


        // We need to turn the flattened material cells into a correctly shaped
        // ragged right structure for use as the material cells in the flat
        // state wrapper. Just as in the flat state mesh field (and associated
        // cell ids), we aren't removing duplicates, just concatnating by material


        // reset the running counter
        int running_counter=0;

        // loop over material ids on different nodes
        int const num_all_material_ids = all_material_ids.size();
        for (int i = 0; i < num_all_material_ids; ++i) {

          // get the current working material
          int mat_id = all_material_ids[i];

          // get the current number of material cells for this material
          int nmat_cells = all_material_shapes[i];

          // get or create a reference to the correct material cell vector
          auto& these_material_data = material_data[mat_id];

          // loop over the correct number of material cells
          for (int j=0; j<nmat_cells; ++j){
              these_material_data.push_back(correctly_typed_data[running_counter++]);
          }

        }

        // the underlying api mat_add_celldata clears the existing data
        for ( auto& kv: material_data){

            // the data needs to be merged within each material

            // material id
            int m = kv.first;

            // existing material data
            std::vector<Wonton::Point<2>> & current_material_data=kv.second;

            // vector to merge material data
            std::vector<int> distributedMaterialCellIds =distributedMaterialCellIds_.at(m);

            // merged material data
            std::vector<Wonton::Point<2>> merged_material_data(distributedMaterialCellIds.size());

            // merge the data
            int c=0;
            for (auto id: distributedMaterialCellIds)
              merged_material_data[c++]=current_material_data[id];

            // add this cell data to the state manager
            StateManager<MeshWrapper>::mat_add_celldata(field_name, m, merged_material_data.data());
        }

      } else if ( data_type == typeid(Wonton::Point<3>)){

        // field data is 2d wonton points

        // allocate the vector data
        std::vector<Wonton::Point<3>> correctly_typed_data(flat_data.size()/3);

        // convert to the correct type
        int const num_flat_data = flat_data.size();
        for (int i = 0; i < num_flat_data; i += 3) {
          correctly_typed_data[i/3] = Wonton::Point<3>(flat_data[i], flat_data[i+1], flat_data[i+2]);
        }

        // allocate the data for unpacking
        std::unordered_map<int,std::vector<Wonton::Point<3>>> material_data;


        // We need to turn the flattened material cells into a correctly shaped
        // ragged right structure for use as the material cells in the flat
        // state wrapper. Just as in the flat state mesh field (and associated
        // cell ids), we aren't removing duplicates, just concatenating by material


        // reset the running counter
        int running_counter=0;

        // loop over material ids on different nodes
        int const num_all_material_ids = all_material_ids.size();
        for (int i = 0; i < num_all_material_ids; ++i) {

          // get the current working material
          int mat_id = all_material_ids[i];

          // get the current number of material cells for this material
          int nmat_cells = all_material_shapes[i];

          // get or create a reference to the correct material cell vector
          auto& these_material_data = material_data[mat_id];

          // loop over the correct number of material cells
          for (int j=0; j<nmat_cells; ++j){
              these_material_data.push_back(correctly_typed_data[running_counter++]);
          }

        }

        // add the material indices by keys
        // the underlying api mat_add_celldata clears the existing data
        for ( auto& kv: material_data){

            // the data needs to be merged within each material

            // material id
            int m = kv.first;

            // existing material data
            std::vector<Wonton::Point<3>> & current_material_data=kv.second;

            // vector to merge material data
            std::vector<int> distributedMaterialCellIds =distributedMaterialCellIds_.at(m);

            // merged material data
            std::vector<Wonton::Point<3>> merged_material_data(distributedMaterialCellIds.size());

            // merge the data
            int c=0;
            for (auto id: distributedMaterialCellIds)
              merged_material_data[c++]=current_material_data[id];

            // add this cell data to the state manager
            StateManager<MeshWrapper>::mat_add_celldata(field_name, m, merged_material_data.data());
        }

      } else if ( data_type == typeid(double)){

        // field data is doubles

        // allocate the data for unpacking
        std::unordered_map<int,std::vector<double>> material_data;


        // We need to turn the flattened material cells into a correctly shaped
        // ragged right structure for use as the material cells in the flat
        // state wrapper. Just as in the flat state mesh field (and associated
        // cell ids), we aren't removing duplicates, just concatnating by material


        // reset the running counter
        int running_counter=0;

        // loop over material ids on different nodes
        int const num_all_material_ids = all_material_ids.size();
        for (int i = 0; i < num_all_material_ids; ++i) {

          // get the current working material
          int mat_id = all_material_ids[i];

          // get the current number of material cells for this material
          int nmat_cells = all_material_shapes[i];

          // get or create a reference to the correct material cell vector
          auto& these_material_data = material_data[mat_id];

          // loop over the correct number of material cells
          for (int j=0; j<nmat_cells; ++j){
              these_material_data.push_back(flat_data[running_counter++]);
          }

        }

        // add the material indices by keys
        // the underlying api mat_add_celldata clears the existing data
        for ( auto& kv: material_data){

            // material id
            int m = kv.first;

            // existing material data
            std::vector<double> & current_material_data=kv.second;

            // vector to merge material data
            std::vector<int> distributedMaterialCellIds =distributedMaterialCellIds_.at(m);

            // merged material data
            std::vector<double> merged_material_data(distributedMaterialCellIds.size());

            // merge the data
            int c=0;
            for (auto id: distributedMaterialCellIds)
              merged_material_data[c++]=current_material_data[id];

            // add this cell data to the state manager
            StateManager<MeshWrapper>::mat_add_celldata(field_name, m, merged_material_data.data());
        }
      }
    }
    else
      throw std::runtime_error("invalid field type");
  }

  size_t get_field_stride(std::string field_name) {

      // get the state vector base class shared pointer
    std::shared_ptr<StateVectorBase> pv =
        StateManager<MeshWrapper>::get(field_name);

    const std::type_info& data_type = pv->data_type();

    if ( data_type == typeid(double)) return 1;
    else if ( data_type == typeid(Wonton::Point<2>)) return 2;
    else if ( data_type == typeid(Wonton::Point<3>)) return 3;

    throw std::runtime_error("invalid data type");
  }

  int num_material_cells() const {
    int n=0;
    for (auto& kv : StateManager<MeshWrapper>::material_cells_){
        n += kv.second.size();
    }
    return n;
  }
    
  std::vector<int> get_material_ids() const {

    std::vector<int> material_ids;
    for (auto& kv : StateManager<MeshWrapper>::material_cells_) material_ids.emplace_back(kv.first);

    std::sort(material_ids.begin(), material_ids.end());

    return material_ids;
  }

  std::vector<int> get_material_shapes() const {

    std::vector<int> material_ids=get_material_ids();
    std::vector<int> material_shapes;

    for (auto m : material_ids)
      material_shapes.emplace_back(StateManager<MeshWrapper>::material_cells_.at(m).size());

    return material_shapes;
  }


  std::vector<int> get_material_cells() const {

    std::vector<int> material_ids=get_material_ids();
    std::vector<int> material_cells;

    for (auto m : material_ids){
      std::vector<int> this_material_cells = StateManager<MeshWrapper>::get_material_cells(m);
      material_cells.insert(material_cells.end(), this_material_cells.begin(), this_material_cells.end());
    }

    return material_cells;
  }

};  // class Flat_State_Wrapper

}  // namespace Wonton

#endif  // WONTON_STATE_FLAT_FLAT_STATE_MM_WRAPPER_H_