coordinateframe.h

// -*-c++-*-
#ifndef STORMM_COORDINATE_FRAME_H
#define STORMM_COORDINATE_FRAME_H
 
#include <typeinfo>
#include <typeindex>
#include <sys/types.h>
#include "copyright.h"
#include "Accelerator/hybrid.h"
#include "Accelerator/hybrid_util.h"
#include "Constants/behavior.h"
#include "Parsing/textfile.h"
#include "Topology/atomgraph.h"
#include "phasespace.h"
#include "trajectory_enumerators.h"
 
namespace stormm {
namespace trajectory {
 
using card::default_hpc_format;
using card::Hybrid;
using card::HybridFormat;
using card::HybridKind;
using card::HybridTargetLevel;
using constants::ExceptionResponse;
using parse::TextFile;
using topology::AtomGraph;

struct CoordinateFrameWriter {

  CoordinateFrameWriter(int natom_in, UnitCellType unit_cell_in, double* xcrd_in, double* ycrd_in,
                        double* zcrd_in, double* umat_in, double* invu_in, double* boxdim_in);
  CoordinateFrameWriter(PhaseSpace *ps, HybridTargetLevel tier = HybridTargetLevel::HOST);
  CoordinateFrameWriter(PhaseSpace *ps, TrajectoryKind kind,
                        HybridTargetLevel tier = HybridTargetLevel::HOST);
  CoordinateFrameWriter(PhaseSpace *ps, TrajectoryKind kind, CoordinateCycle orientation,
                        HybridTargetLevel tier = HybridTargetLevel::HOST);

  CoordinateFrameWriter(const CoordinateFrameWriter &original) = default;
  CoordinateFrameWriter(CoordinateFrameWriter &&original) = default;
  
  const int natom;               
  const UnitCellType unit_cell;  
  double* xcrd;                  
  double* ycrd;                  
  double* zcrd;                  
  double* umat;                  
  double* invu;                  
  double* boxdim;                
};

struct CoordinateFrameReader {

  CoordinateFrameReader(int natom_in, UnitCellType unit_cell_in, const double* xcrd_in,
                        const double* ycrd_in, const double* zcrd_in, const double* umat_in,
                        const double* invu_in, const double* boxdim_in);
  CoordinateFrameReader(const CoordinateFrameWriter &cfw);
  CoordinateFrameReader(const PhaseSpace &ps, HybridTargetLevel tier = HybridTargetLevel::HOST);
  CoordinateFrameReader(const PhaseSpace *ps, HybridTargetLevel tier = HybridTargetLevel::HOST);
  CoordinateFrameReader(const PhaseSpace *ps, TrajectoryKind kind, CoordinateCycle orientation,
                        HybridTargetLevel tier = HybridTargetLevel::HOST);
  CoordinateFrameReader(const PhaseSpace &ps, TrajectoryKind kind, CoordinateCycle orientation,
                        HybridTargetLevel tier = HybridTargetLevel::HOST);
  CoordinateFrameReader(const PhaseSpace *ps, TrajectoryKind kind,
                        HybridTargetLevel tier = HybridTargetLevel::HOST);
  CoordinateFrameReader(const PhaseSpace &ps, TrajectoryKind kind,
                        HybridTargetLevel tier = HybridTargetLevel::HOST);

  CoordinateFrameReader(const CoordinateFrameReader &original) = default;
  CoordinateFrameReader(CoordinateFrameReader &&original) = default;
  
  const int natom;               
  const UnitCellType unit_cell;  
  const double* xcrd;            
  const double* ycrd;            
  const double* zcrd;            
  const double* umat;            
  const double* invu;            
  const double* boxdim;          
};

class CoordinateFrame {
public:

  CoordinateFrame(int natom_in = 0, UnitCellType unit_cell_in = UnitCellType::NONE,
                  HybridFormat format_in = default_hpc_format);
 
  CoordinateFrame(int natom_in, UnitCellType unit_cell_in, const double* xcrd_in,
                  const double* ycrd_in, const double* zcrd_in, const double* umat_in = nullptr,
                  const double* invu_in = nullptr, const double* boxdim_in = nullptr,
                  HybridFormat format_in = default_hpc_format);
  
  CoordinateFrame(const std::string &file_name_in,
                  CoordinateFileKind file_kind = CoordinateFileKind::UNKNOWN,
                  int frame_number_in = 0, HybridFormat format_in = default_hpc_format);
 
  CoordinateFrame(const TextFile &tf, CoordinateFileKind file_kind  = CoordinateFileKind::UNKNOWN,
                  int frame_number_in = 0, HybridFormat format_in = default_hpc_format);
 
  CoordinateFrame(const PhaseSpace *ps);
  
  CoordinateFrame(const PhaseSpace &ps);

  CoordinateFrame(const CoordinateFrame &original);
  CoordinateFrame(const CoordinateFrame &original, HybridFormat format_in);

  CoordinateFrame& operator=(const CoordinateFrame &other);

  CoordinateFrame(CoordinateFrame &&original);

  CoordinateFrame& operator=(CoordinateFrame &&other);

  void buildFromFile(const std::string &file_name_in, const CoordinateFileKind file_kind,
                     int frame_number = 0);
  
  void buildFromFile(const TextFile &tf, const CoordinateFileKind file_kind,
                     int frame_number = 0);

  template <typename T>
  void fill(const T* xcrd, const T* ycrd, const T* zcrd, int scale_bits = 0,
            const double* box_dims = nullptr);
 
  template <typename T>
  void fill(const std::vector<T> &xcrd, const std::vector<T> &ycrd, const std::vector<T> &zcrd,
            int scale_bits = 0, const std::vector<double> &box_dims = {});

  HybridFormat getFormat() const;
  
  std::string getFileName() const;

  int getFrameNumber() const;
  
  int getAtomCount() const;

  UnitCellType getUnitCellType() const;

  std::vector<double>
  getInterlacedCoordinates(HybridTargetLevel tier = HybridTargetLevel::HOST) const;
 
  std::vector<double>
  getInterlacedCoordinates(int low_index, int high_index,
                           HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  std::vector<double> getBoxSpaceTransform(HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  std::vector<double> getInverseTransform(HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  std::vector<double> getBoxDimensions(HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  const Hybrid<double>* getCoordinateHandle(CartesianDimension dim) const;
  Hybrid<double>* getCoordinateHandle(CartesianDimension dim);

  const Hybrid<double>* getBoxTransformHandle() const;
  Hybrid<double>* getBoxTransformHandle();

  const Hybrid<double>* getInverseTransformHandle() const;
  Hybrid<double>* getInverseTransformHandle();

  const Hybrid<double>* getBoxDimensionsHandle() const;
  Hybrid<double>* getBoxDimensionsHandle();

  const Hybrid<double>* getStorageHandle() const;
  Hybrid<double>* getStorageHandle();
  
  const CoordinateFrame* getSelfPointer() const;

  void exportToFile(const std::string &file_name,
                    CoordinateFileKind output_kind = CoordinateFileKind::AMBER_CRD,
                    PrintSituation expectation = PrintSituation::UNKNOWN,
                    HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  
  const CoordinateFrameReader data(HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  CoordinateFrameWriter data(HybridTargetLevel tier = HybridTargetLevel::HOST);
 
#ifdef STORMM_USE_HPC
  void upload();

  void download();

  const CoordinateFrameReader deviceViewToHostData() const;
  CoordinateFrameWriter deviceViewToHostData();
#endif

  void setFrameNumber(int frame_number_in);
  
private:
  HybridFormat format;                
  std::string file_name;              
  int frame_number;                   
  int atom_count;                     
  UnitCellType unit_cell;             
  Hybrid<double> x_coordinates;       
  Hybrid<double> y_coordinates;       
  Hybrid<double> z_coordinates;       
  Hybrid<double> box_space_transform; 
  Hybrid<double> inverse_transform;   
  Hybrid<double> box_dimensions;      

  Hybrid<double> storage;

  void allocate();
};

static const size_t coordinateframe_type_index =
  std::type_index(typeid(CoordinateFrame)).hash_code();
  
std::vector<CoordinateFrame> getSelectedFrames(const TextFile &tf, CoordinateFileKind kind,
                                               int atom_count, UnitCellType unit_cell,
                                               const std::vector<int> &frame_numbers);
 
std::vector<CoordinateFrame> getSelectedFrames(const std::string &file_name, int atom_count,
                                               UnitCellType unit_cell,
                                               const std::vector<int> &frame_numbers);

std::vector<CoordinateFrame> getAllFrames(const TextFile &tf, const int atom_count,
                                          const UnitCellType unit_cell,
                                          const ExceptionResponse policy);
 
std::vector<CoordinateFrame> getAllFrames(const std::string &file_name, const int atom_count,
                                          const UnitCellType unit_cell,
                                          const ExceptionResponse policy);
 
} // namespace trajectory
} // namespace stormm
 
#include "coordinateframe.tpp"
 
// As with common types and STORMM vector types, define the type indices for general use in the
// STORMM namespace.
namespace stormm {
using trajectory::coordinateframe_type_index;
} // namespace stormm
 
#endif