phasespace.h

// -*-c++-*-
#ifndef STORMM_PHASE_SPACE_H
#define STORMM_PHASE_SPACE_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 "Math/matrix_ops.h"
#include "Topology/atomgraph.h"
#include "trajectory_enumerators.h"
#include "write_frame.h"
 
namespace stormm {
namespace trajectory {
 
using card::default_hpc_format;
using card::GpuDetails;
using card::Hybrid;
using card::HybridFormat;
using card::HybridTargetLevel;
using constants::CartesianDimension;
using stmath::computeBoxTransform;
using topology::AtomGraph;
using topology::UnitCellType;

struct PhaseSpaceWriter {

  PhaseSpaceWriter(const int natom_in, const UnitCellType unit_cell_in, double* xcrd_in,
                   double* ycrd_in, double* zcrd_in, double* umat_in, double* invu_in,
                   double* boxdim_in, double* umat_alt_in, double* invu_alt_in,
                   double* boxdim_alt_in, double* xvel_in, double* yvel_in, double* zvel_in,
                   double* xfrc_in, double* yfrc_in, double* zfrc_in, double* xalt_in,
                   double* yalt_in, double* zalt_in, double* vxalt_in, double* vyalt_in,
                   double* vzalt_in, double* fxalt_in, double* fyalt_in, double* fzalt_in);

  PhaseSpaceWriter(const PhaseSpaceWriter &original) = default;
  PhaseSpaceWriter(PhaseSpaceWriter &&original) = default;
 
  const int natom;               
  const UnitCellType unit_cell;  
  double* xcrd;                  
  double* ycrd;                  
  double* zcrd;                  
  double* umat;                  
  double* invu;                  
  double* boxdim;                
  double* umat_alt;              
  double* invu_alt;              
  double* boxdim_alt;            
  double* xvel;                  
  double* yvel;                  
  double* zvel;                  
  double* xfrc;                  
  double* yfrc;                  
  double* zfrc;                  
  double* xalt;                  
  double* yalt;                  
  double* zalt;                  
  double* vxalt;                 
  double* vyalt;                 
  double* vzalt;                 
  double* fxalt;                 
  double* fyalt;                 
  double* fzalt;                 
};

struct PhaseSpaceReader {

  PhaseSpaceReader(const int natom_in, const 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, const double* umat_alt_in,
                   const double* invu_alt_in, const double* boxdim_alt_in, const double* xvel_in,
                   const double* yvel_in, const double* zvel_in, const double* xfrc_in,
                   const double* yfrc_in, const double* zfrc_in, const double* xalt_in,
                   const double* yalt_in, const double* zalt_in, const double* vxalt_in,
                   const double* vyalt_in, const double* vzalt_in, const double* fxalt_in,
                   const double* fyalt_in, const double* fzalt_in);
 
  PhaseSpaceReader(const PhaseSpaceWriter &psw);
  
  PhaseSpaceReader(const PhaseSpaceReader &original) = default;
  PhaseSpaceReader(PhaseSpaceReader &&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;          
  const double* umat_alt;        
  const double* invu_alt;        
  const double* boxdim_alt;      
  const double* xvel;            
  const double* yvel;            
  const double* zvel;            
  const double* xfrc;            
  const double* yfrc;            
  const double* zfrc;            
  const double* xalt;            
  const double* yalt;            
  const double* zalt;            
  const double* vxalt;           
  const double* vyalt;           
  const double* vzalt;           
  const double* fxalt;           
  const double* fyalt;           
  const double* fzalt;           
};

class PhaseSpace {
public:

  PhaseSpace(int atom_count_in = 0, UnitCellType unit_cell_in = UnitCellType::NONE,
             HybridFormat format_in = default_hpc_format);
 
  PhaseSpace(const std::string &file_name_in,
             CoordinateFileKind file_kind = CoordinateFileKind::UNKNOWN, int frame_number = 0,
             HybridFormat format_in = default_hpc_format);
 
  PhaseSpace(const std::string &file_name_in, const AtomGraph &ag,
             CoordinateFileKind file_kind = CoordinateFileKind::UNKNOWN,
             int frame_number = 0, HybridFormat format_in = default_hpc_format);
 
  PhaseSpace(const std::vector<PhaseSpace*> &input_crd, const std::vector<AtomGraph*> &input_top,
             const std::vector<int> &counts);
 
  PhaseSpace(const PhaseSpace *ps_a, const AtomGraph *ag_a, const PhaseSpace *ps_b,
             const AtomGraph *ag_b);
 
  PhaseSpace(const PhaseSpace &ps_a, const AtomGraph &ag_a, const PhaseSpace &ps_b,
             const AtomGraph &ag_b);
 
  PhaseSpace(const PhaseSpace &original, const std::vector<int> &atom_list,
             HybridFormat format_in = default_hpc_format);
 
  PhaseSpace(const PhaseSpace &original, const std::vector<uint> &mask,
             HybridFormat format_in = default_hpc_format);

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

  PhaseSpace& operator=(const PhaseSpace &other);

  PhaseSpace(PhaseSpace &&original);

  PhaseSpace& operator=(PhaseSpace &&other);

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

  template <typename T>
  void fill(const T* xcrd, const T* ycrd, const T* zcrd,
            TrajectoryKind kind = TrajectoryKind::POSITIONS,
            CoordinateCycle cycle_in = CoordinateCycle::WHITE, 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,
            TrajectoryKind kind = TrajectoryKind::POSITIONS,
            CoordinateCycle cycle_in = CoordinateCycle::WHITE, int scale_bits = 0,
            const std::vector<double> &box_dims = {});

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

  int getAtomCount() const;

  UnitCellType getUnitCellType() const;

  CoordinateCycle getCyclePosition() const;  
  
  const double* getCoordinatePointer(CartesianDimension dim, TrajectoryKind kind,
                                     CoordinateCycle orientation,
                                     HybridTargetLevel tier = HybridTargetLevel::HOST) const;
 
  double* getCoordinatePointer(CartesianDimension dim, TrajectoryKind kind,
                               CoordinateCycle orientation,
                               HybridTargetLevel tier = HybridTargetLevel::HOST);
 
  const double* getCoordinatePointer(CartesianDimension dim,
                                     TrajectoryKind kind = TrajectoryKind::POSITIONS,
                                     HybridTargetLevel tier = HybridTargetLevel::HOST) const;
 
  double* getCoordinatePointer(CartesianDimension dim,
                               TrajectoryKind kind = TrajectoryKind::POSITIONS,
                               HybridTargetLevel tier = HybridTargetLevel::HOST);
 
  const double* getCoordinatePointer(CartesianDimension dim, CoordinateCycle orientation,
                                     HybridTargetLevel tier = HybridTargetLevel::HOST) const;
 
  double* getCoordinatePointer(CartesianDimension dim, CoordinateCycle orientation,
                               HybridTargetLevel tier = HybridTargetLevel::HOST);
  
  std::vector<double>
  getInterlacedCoordinates(TrajectoryKind kind = TrajectoryKind::POSITIONS,
                           HybridTargetLevel tier = HybridTargetLevel::HOST) const;
 
  std::vector<double>
  getInterlacedCoordinates(CoordinateCycle orientation,
                           TrajectoryKind kind = TrajectoryKind::POSITIONS,
                           HybridTargetLevel tier = HybridTargetLevel::HOST) const;
 
  std::vector<double>
  getInterlacedCoordinates(int low_index, int high_index,
                           TrajectoryKind kind = TrajectoryKind::POSITIONS,
                           HybridTargetLevel tier = HybridTargetLevel::HOST) const;
 
  std::vector<double>
  getInterlacedCoordinates(int low_index, int high_index, CoordinateCycle orientation,
                           TrajectoryKind kind = TrajectoryKind::POSITIONS,
                           HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  const double*
  getBoxSpaceTransformPointer(CoordinateCycle orientation,
                              HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  double* getBoxSpaceTransformPointer(CoordinateCycle orientation,
                                      HybridTargetLevel tier = HybridTargetLevel::HOST);
  const double*
  getBoxSpaceTransformPointer(HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  double* getBoxSpaceTransformPointer(HybridTargetLevel tier = HybridTargetLevel::HOST);

  const double* getInverseTransformPointer(CoordinateCycle orientation,
                                           HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  double* getInverseTransformPointer(CoordinateCycle orientation,
                                     HybridTargetLevel tier = HybridTargetLevel::HOST);
  const double* getInverseTransformPointer(HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  double* getInverseTransformPointer(HybridTargetLevel tier = HybridTargetLevel::HOST);

  const double* getBoxSizePointer(CoordinateCycle orientation,
                                  HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  double* getBoxSizePointer(CoordinateCycle orientation,
                            HybridTargetLevel tier = HybridTargetLevel::HOST);
  const double* getBoxSizePointer(HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  double* getBoxSizePointer(HybridTargetLevel tier = HybridTargetLevel::HOST);

  std::vector<double> getBoxSpaceTransform(CoordinateCycle orientation,
                                           HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  std::vector<double> getBoxSpaceTransform(HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  
  std::vector<double> getInverseTransform(CoordinateCycle orientation,
                                          HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  std::vector<double> getInverseTransform(HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  const Hybrid<double>* getCoordinateHandle(CartesianDimension dim, TrajectoryKind kind,
                                            CoordinateCycle orientation) const;
 
  const Hybrid<double>* getCoordinateHandle(CartesianDimension dim,
                                            TrajectoryKind kind = TrajectoryKind::POSITIONS) const;
 
  Hybrid<double>* getCoordinateHandle(CartesianDimension dim, TrajectoryKind kind,
                                      CoordinateCycle orientation);
 
  Hybrid<double>* getCoordinateHandle(CartesianDimension dim,
                                      TrajectoryKind kind = TrajectoryKind::POSITIONS);

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

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

  const Hybrid<double>* getBoxDimensionsHandle(CoordinateCycle orientation) const;
  const Hybrid<double>* getBoxDimensionsHandle() const;
  Hybrid<double>* getBoxDimensionsHandle(CoordinateCycle orientation);
  Hybrid<double>* getBoxDimensionsHandle();
  
  const Hybrid<double>* getStorageHandle() const;
  Hybrid<double>* getStorageHandle();

  void initializeForces(CoordinateCycle orientation);
  void initializeForces();
  
  void updateCyclePosition();
  void updateCyclePosition(CoordinateCycle time_point);
  
  void exportToFile(const std::string &file_name, double current_time = 0.0,
                    TrajectoryKind traj_kind = TrajectoryKind::POSITIONS,
                    CoordinateFileKind output_kind = CoordinateFileKind::AMBER_INPCRD,
                    PrintSituation expectation = PrintSituation::UNKNOWN) const;

  const PhaseSpace* getSelfPointer() const;

  const PhaseSpaceReader data(HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  const PhaseSpaceReader data(CoordinateCycle orientation,
                              HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  PhaseSpaceWriter data(HybridTargetLevel tier = HybridTargetLevel::HOST);
  PhaseSpaceWriter data(CoordinateCycle orientation,
                        HybridTargetLevel tier = HybridTargetLevel::HOST);
  
#ifdef STORMM_USE_HPC
  const PhaseSpaceReader deviceViewToHostData(CoordinateCycle orientation) const;
  const PhaseSpaceReader deviceViewToHostData() const;
  PhaseSpaceWriter deviceViewToHostData(CoordinateCycle orientation);
  PhaseSpaceWriter deviceViewToHostData();

  void upload();

  void uploadPositions(CoordinateCycle orientation);
  void uploadPositions();

  void uploadTransformations();

  void uploadVelocities(CoordinateCycle orientation);
  void uploadVelocities();

  void uploadForces(CoordinateCycle orientation);
  void uploadForces();

  void download();

  void downloadPositions(CoordinateCycle orientation);
  void downloadPositions();

  void downloadTransformations();

  void downloadVelocities(CoordinateCycle orientation);
  void downloadVelocities();

  void downloadForces(CoordinateCycle orientation);
  void downloadForces();
#endif
  
private:
  HybridFormat format;             
  std::string file_name;           
  int atom_count;                  
  UnitCellType unit_cell;          
  CoordinateCycle cycle_position;  
  Hybrid<double> x_coordinates;    
  Hybrid<double> y_coordinates;    
  Hybrid<double> z_coordinates;    
 
  // The actual nature of the following arrays, as well as the [x,y,z]_coordinates themselves, can
  // change based on the CoordinateCycle position.
  Hybrid<double> x_alt_coordinates;  
  Hybrid<double> y_alt_coordinates;  
  Hybrid<double> z_alt_coordinates;  
 
  // Transformation matrices are given directly after the coordinates in the overall order of data
  Hybrid<double> box_space_transform;  
  Hybrid<double> inverse_transform;    
  Hybrid<double> box_dimensions;       
 
  // Alternate transformation matrices are stored next, for consistency.  This provides a complete
  // record of the previous unit cell dimensions in the event that the trajectory needs to
  // backtrack against a rejected Monte-Carlo move, or a complete space to build the next state of
  // the system.
  Hybrid<double> alt_box_space_transform;  
  Hybrid<double> alt_inverse_transform;    
  Hybrid<double> alt_box_dimensions;       
 
  // Like coordinates, velocities and forces appear in WHITE -> BLACK blocks
  Hybrid<double> x_velocities;      
  Hybrid<double> y_velocities;      
  Hybrid<double> z_velocities;      
  Hybrid<double> x_alt_velocities;  
  Hybrid<double> y_alt_velocities;  
  Hybrid<double> z_alt_velocities;  
  Hybrid<double> x_forces;          
  Hybrid<double> y_forces;          
  Hybrid<double> z_forces;          
  Hybrid<double> x_alt_forces;      
  Hybrid<double> y_alt_forces;      
  Hybrid<double> z_alt_forces;      

  Hybrid<double> storage;

  void allocate();
};

static const size_t phasespace_type_index = std::type_index(typeid(PhaseSpace)).hash_code();

UnitCellType determineUnitCellTypeByShape(const double* inverse_transform);
UnitCellType determineUnitCellTypeByShape(const std::vector<double> &inverse_transform);
UnitCellType determineUnitCellTypeByShape(const Hybrid<double> &inverse_transform);

std::vector<double> interlaceXYZ(const double* xptr, const double* yptr, const double* zptr,
                                 int low_index, int high_index);
  
} // namespace trajectory
} // namespace stormm
 
#include "phasespace.tpp"
 
// As with common types and STORMM vector types, define the type indices for general use in the
// STORMM namespace.
namespace stormm {
using trajectory::phasespace_type_index;
} // namespace stormm
 
#endif