restraint_apparatus.h

// -*-c++-*-
#ifndef STORMM_RESTRAINT_APPARATUS_H
#define STORMM_RESTRAINT_APPARATUS_H
 
#include "copyright.h"
#include "Accelerator/hybrid.h"
#include "Topology/atomgraph.h"
#include "bounded_restraint.h"
 
namespace stormm {
namespace restraints {
 
using card::Hybrid;
using card::HybridTargetLevel;
using topology::AtomGraph;
  
enum class RestraintStage {
  INITIAL,  
  FINAL     
};
  
template <typename T, typename T2, typename T4>
struct RestraintKit {

  explicit RestraintKit(int total_rst_in, int nposn_in, int nbond_in, int nangl_in,
                        int ndihe_in, bool time_dependence_in, const int* rposn_atoms_in,
                        const int* rbond_i_atoms_in, const int* rbond_j_atoms_in,
                        const int* rangl_i_atoms_in, const int* rangl_j_atoms_in,
                        const int* rangl_k_atoms_in, const int* rdihe_i_atoms_in,
                        const int* rdihe_j_atoms_in, const int* rdihe_k_atoms_in,
                        const int* rdihe_l_atoms_in, const int* rposn_init_step_in,
                        const int* rposn_finl_step_in, const int* rbond_init_step_in,
                        const int* rbond_finl_step_in, const int* rangl_init_step_in,
                        const int* rangl_finl_step_in, const int* rdihe_init_step_in,
                        const int* rdihe_finl_step_in, const T2* rposn_init_keq_in,
                        const T2* rposn_finl_keq_in, const T2* rposn_init_xy_in,
                        const T2* rposn_finl_xy_in, const T*  rposn_init_z_in,
                        const T*  rposn_finl_z_in, const T2* rbond_init_keq_in,
                        const T2* rbond_finl_keq_in, const T2* rangl_init_keq_in,
                        const T2* rangl_finl_keq_in, const T2* rdihe_init_keq_in,
                        const T2* rdihe_finl_keq_in, const T4* rposn_init_r_in,
                        const T4* rposn_finl_r_in, const T4* rbond_init_r_in,
                        const T4* rbond_finl_r_in, const T4* rangl_init_r_in,
                        const T4* rangl_finl_r_in, const T4* rdihe_init_r_in,
                        const T4* rdihe_finl_r_in, const AtomGraph *ag_pointer_in);

  RestraintKit(const RestraintKit &original) = default;
  RestraintKit(RestraintKit &&other) = default;
 
  const int total_rst;          
  const int nposn;              
  const int nbond;              
  const int nangl;              
  const int ndihe;              
  const bool time_dependence;   
  const int* rposn_atoms;       
  const int* rbond_i_atoms;     
  const int* rbond_j_atoms;     
  const int* rangl_i_atoms;     
  const int* rangl_j_atoms;     
  const int* rangl_k_atoms;     
  const int* rdihe_i_atoms;     
  const int* rdihe_j_atoms;     
  const int* rdihe_k_atoms;     
  const int* rdihe_l_atoms;     
  const int* rposn_init_step;   
  const int* rposn_finl_step;   
  const int* rbond_init_step;   
  const int* rbond_finl_step;   
  const int* rangl_init_step;   
  const int* rangl_finl_step;   
  const int* rdihe_init_step;   
  const int* rdihe_finl_step;   
  const T2* rposn_init_keq;     
  const T2* rposn_finl_keq;     
  const T2* rposn_init_xy;      
  const T2* rposn_finl_xy;      
  const T*  rposn_init_z;       
  const T*  rposn_finl_z;       
  const T2* rbond_init_keq;     
  const T2* rbond_finl_keq;     
  const T2* rangl_init_keq;     
  const T2* rangl_finl_keq;     
  const T2* rdihe_init_keq;     
  const T2* rdihe_finl_keq;     
  const T4* rposn_init_r;       
  const T4* rposn_finl_r;       
  const T4* rbond_init_r;       
  const T4* rbond_finl_r;       
  const T4* rangl_init_r;       
  const T4* rangl_finl_r;       
  const T4* rdihe_init_r;       
  const T4* rdihe_finl_r;       
  const AtomGraph *ag_pointer;  
};

class RestraintApparatus {
public:

  RestraintApparatus(const AtomGraph *ag_in = nullptr);
  RestraintApparatus(const std::vector<BoundedRestraint> &rbasis,
                     const AtomGraph *ag_in = nullptr);
  
  RestraintApparatus(const RestraintApparatus &original);

  RestraintApparatus(RestraintApparatus &&original);

  RestraintApparatus& operator=(const RestraintApparatus &other);
  
  RestraintApparatus& operator=(RestraintApparatus &&other);

  int getTotalRestraintCount() const;

  int getPositionalRestraintCount() const;
  
  int getDistanceRestraintCount() const;
  
  int getAngleRestraintCount() const;
  
  int getDihedralRestraintCount() const;

  bool getTimeDependence() const;

  const AtomGraph* getTopologyPointer() const;
  
  RestraintKit<double, double2, double4>
  dpData(HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  RestraintKit<float, float2, float4>
  spData(HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  const int* getApplicationStepPointer(int order, RestraintStage stage,
                                       HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  
  const double2*
  getHarmonicStiffnessPointer(int order, RestraintStage stage,
                              HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  const double4* getDisplacementPointer(int order, RestraintStage stage,
                                        HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  std::vector<BoundedRestraint> getRestraintList() const;

  const RestraintApparatus* getSelfPointer() const;
  
  void addRestraints(const std::vector<BoundedRestraint> &new_rest);
  void addRestraint(const BoundedRestraint &new_rest);
  
private:
 
  // Overall counts and items of interest
  int total_restraint_count;   
  int position_count;          
  int distance_count;          
  int angle_count;             
  int dihedral_count;          
  bool time_based_restraints;  
 
  // Integer data 
  Hybrid<int> rposn_atoms;       
  Hybrid<int> rbond_i_atoms;     
  Hybrid<int> rbond_j_atoms;     
  Hybrid<int> rangl_i_atoms;     
  Hybrid<int> rangl_j_atoms;     
  Hybrid<int> rangl_k_atoms;     
  Hybrid<int> rdihe_i_atoms;     
  Hybrid<int> rdihe_j_atoms;     
  Hybrid<int> rdihe_k_atoms;     
  Hybrid<int> rdihe_l_atoms;     
  Hybrid<int> rposn_init_step;   
  Hybrid<int> rposn_final_step;  
  Hybrid<int> rbond_init_step;   
  Hybrid<int> rbond_final_step;  
  Hybrid<int> rangl_init_step;   
  Hybrid<int> rangl_final_step;  
  Hybrid<int> rdihe_init_step;   
  Hybrid<int> rdihe_final_step;  
  Hybrid<int> int_data;          
 
  // Real data in double-precision format
  Hybrid<double2> rposn_init_keq;  
  Hybrid<double2> rposn_final_keq; 
  Hybrid<double4> rposn_init_r;    
  Hybrid<double4> rposn_final_r;   
  Hybrid<double2> rposn_init_xy;   
  Hybrid<double> rposn_init_z;     
  Hybrid<double2> rposn_final_xy;  
  Hybrid<double> rposn_final_z;    
  Hybrid<double2> rbond_init_keq;  
  Hybrid<double2> rbond_final_keq; 
  Hybrid<double4> rbond_init_r;    
  Hybrid<double4> rbond_final_r;   
  Hybrid<double2> rangl_init_keq;  
  Hybrid<double2> rangl_final_keq; 
  Hybrid<double4> rangl_init_r;    
  Hybrid<double4> rangl_final_r;   
  Hybrid<double2> rdihe_init_keq;  
  Hybrid<double2> rdihe_final_keq; 
  Hybrid<double4> rdihe_init_r;    
  Hybrid<double4> rdihe_final_r;   
  Hybrid<double> double_data;      
  Hybrid<double2> double2_data;    
  Hybrid<double4> double4_data;    
 
  // Real data in single-precision format
  Hybrid<float2> sp_rposn_init_keq;  
  Hybrid<float2> sp_rposn_final_keq; 
  Hybrid<float4> sp_rposn_init_r;    
  Hybrid<float4> sp_rposn_final_r;   
  Hybrid<float2> sp_rposn_init_xy;   
  Hybrid<float> sp_rposn_init_z;     
  Hybrid<float2> sp_rposn_final_xy;  
  Hybrid<float> sp_rposn_final_z;    
  Hybrid<float2> sp_rbond_init_keq;  
  Hybrid<float2> sp_rbond_final_keq; 
  Hybrid<float4> sp_rbond_init_r;    
  Hybrid<float4> sp_rbond_final_r;   
  Hybrid<float2> sp_rangl_init_keq;  
  Hybrid<float2> sp_rangl_final_keq; 
  Hybrid<float4> sp_rangl_init_r;    
  Hybrid<float4> sp_rangl_final_r;   
  Hybrid<float2> sp_rdihe_init_keq;  
  Hybrid<float2> sp_rdihe_final_keq; 
  Hybrid<float4> sp_rdihe_init_r;    
  Hybrid<float4> sp_rdihe_final_r;   
  Hybrid<float> float_data;          
  Hybrid<float2> float2_data;        
  Hybrid<float4> float4_data;        

  const AtomGraph *ag_pointer;

  void checkTopologyPointers(const std::vector<BoundedRestraint> &rbasis);
  
  void allocate();

  void populateInternalArrays(const std::vector<BoundedRestraint> &rbasis);

  bool assessTimeDependence(int* init_steps, int* final_steps, int nrest);
  void assessTimeDependence();
};

std::vector<RestraintApparatus>
createBlankRestraintApparatus(const std::vector<AtomGraph*> ags);
  
} // namespace restraints
} // namespace stormm
 
#include "restraint_apparatus.tpp"
 
#endif