atomgraph_stage.h

// -*-c++-*-
#ifndef STORMM_ATOMGRAPH_STAGE_H
#define STORMM_ATOMGRAPH_STAGE_H
 
#include <vector>
#include <string>
#include "copyright.h"
#include "DataTypes/stormm_vector_types.h"
#include "MoleculeFormat/mdlmol.h"
#include "atomgraph.h"
#include "atomgraph_abstracts.h"
#include "atomgraph_enumerators.h"
 
namespace stormm {
namespace topology {
 
using structure::MdlMol;

constexpr int default_minimum_solute_atoms = 15;
constexpr char4 wildcard_atom_type = { 'W', 'I', 'L', 'D' };
constexpr char4 unknown_atom_name = { 'U', 'N', 'K', ' ' };
constexpr char4 unknown_residue_name = { 'U', 'N', 'K', ' ' };
constexpr char4 unknown_atom_type = { 'U', 'N', 'K', ' ' };

struct ScreeningFactorTag {

  ScreeningFactorTag();
  ScreeningFactorTag(int dihe_idx_in, double elec14_in, double vdw14_in);
  ScreeningFactorTag(const char4 i_at_in, const char4 l_at_in, double elec14_in, double vdw14_in);
  ScreeningFactorTag(const char4 i_at_in, const char4 j_at_in, const char4 k_at_in,
                     const char4 l_at_in, double elec14_in, double vdw14_in);
 
  int dihe_idx;   
  char4 i_at;     
  char4 j_at;     
  char4 k_at;     
  char4 l_at;     
  double elec14;  
  double vdw14;   
};
  
class AtomGraphStage {
public:

  AtomGraphStage(int atom_count_in = 0, const std::vector<int> &residue_limits_in = {},
                 ExceptionResponse policy_in = ExceptionResponse::DIE);
  AtomGraphStage(const AtomGraph *ag_in, ExceptionResponse policy_in = ExceptionResponse::DIE);
  AtomGraphStage(const AtomGraph &ag_in, ExceptionResponse policy_in = ExceptionResponse::DIE);
  AtomGraphStage(const MdlMol *cmpd_in, ExceptionResponse policy_in = ExceptionResponse::DIE);
  AtomGraphStage(const MdlMol &cmpd_in, ExceptionResponse policy_in = ExceptionResponse::DIE);

  AtomGraphStage(const AtomGraphStage &original) = default;
  AtomGraphStage(AtomGraphStage &&original) = default;
  AtomGraphStage& operator=(const AtomGraphStage &other) = default;
  AtomGraphStage& operator=(AtomGraphStage &&other) = default;

  const std::string& getTitle() const;

  const std::string& getOutputFile() const;
  
  int getAtomCount() const;

  int getResidueCount() const;

  int getMoleculeCount() const;

  double getElectrostatic14Screening() const;
  
  double getVanDerWaals14Screening() const;
  
  AtomGraph exportTopology(int minimum_solute_size = default_minimum_solute_atoms,
                           const std::vector<char4> &solute_included_residue_names = {},
                           double hmass_repartition_factor = 1.0) const;

  void setTitle(const std::string &title_in);

  void setOutputFile(const std::string &output_file_in);

  void addUbrdParameters(double k_eq_in, double l_eq_in,
                         const char4 atom_i_type = unknown_atom_type,
                         const char4 atom_k_type = unknown_atom_type);

  void addCImpParameters(double k_eq_in, double phi_in,
                         const char4 atom_i_type = unknown_atom_type,
                         const char4 atom_j_type = unknown_atom_type,
                         const char4 atom_k_type = unknown_atom_type,
                         const char4 atom_l_type = unknown_atom_type);

  void addCmapParameters(const int dimension_in, const std::vector<double> &surf_in,
                         const char4 atom_i_type = unknown_atom_type,
                         const char4 atom_j_type = unknown_atom_type,
                         const char4 atom_k_type = unknown_atom_type,
                         const char4 atom_l_type = unknown_atom_type,
                         const char4 atom_m_type = unknown_atom_type);

  void addBondParameters(double k_eq_in, double l_eq_in,
                         const char4 atom_i_type = unknown_atom_type,
                         const char4 atom_j_type = unknown_atom_type);

  void addAnglParameters(double k_eq_in, double th_eq_in,
                         const char4 atom_i_type = unknown_atom_type,
                         const char4 atom_j_type = unknown_atom_type,
                         const char4 atom_k_type = unknown_atom_type);

  void addDiheParameters(double ampl_in, double period_in, double phi_in,
                         const char4 atom_i_type = unknown_atom_type,
                         const char4 atom_j_type = unknown_atom_type,
                         const char4 atom_k_type = unknown_atom_type,
                         const char4 atom_l_type = unknown_atom_type);

  std::vector<int> addAtoms(const std::vector<int> &z_numbers, int placement = -1,
                            const std::vector<int> &added_residue_limits = {},
                            const std::vector<char4> &added_atom_names = {},
                            const std::vector<char4> &added_atom_types = {},
                            const std::vector<char4> &added_residue_names = {},
                            const std::vector<int> &added_charge_indices = {},
                            const std::vector<int> &added_lennard_jones_indices = {},
                            const std::vector<double> &added_atomic_charges = {});

  int addAtom(int z_number, int placement = -1, const char4 atom_name = { 'U', 'N', 'K', ' ' },
              const char4 atom_type = { 'U', 'N', 'K', ' ' },
              const char4 res_name = { 'U', 'N', 'K', ' ' }, int charge_index = -1,
              int lennard_jones_index = -1, double atomic_charge = 0.0);
  
  std::vector<int> addVirtualSites(const std::vector<int> &parameter_indices,
                                   const std::vector<int> &frame1_atoms,
                                   const std::vector<int> &frame2_atoms,
                                   const std::vector<int> &frame3_atoms,
                                   const std::vector<int> &frame4_atoms);
  
  int addVirtualSite(int parameter_index, int frame1_atom, int frame2_atom, int frame3_atom = -1,
                     int frame4_atom = -1);

  void setBonds(const std::vector<int> &atom_i, const std::vector<int> &atom_j,
                const std::vector<int> &parameter_index);

  void setBond(int atom_i, int atom_j, int parameter_index);

  void setElectrostatic14Screening(double screening_factor_in);

  void setVanDerWaals14Screening(double screening_factor_in);
  
private:
 
  // The title and source can be specified.
  ExceptionResponse policy;  
  std::string title;         
  std::string output_file;   

  std::vector<Citation> force_fields;
  
  // Sizing constants follow the AtomGraph class
  int atom_count;                 
  int residue_count;              
  int molecule_count;             
 
  // Arrays of limits define the structures and composition of the molecular system.
  std::vector<int> residue_limits;       
  std::vector<int> molecule_limits;      
  std::vector<int> molecule_membership;  
  std::vector<int> molecule_contents;    
 
  // Arrays of atomic descriptors define the chemical properties of the system
  std::vector<int> atomic_numbers;     
 
  // Arrays of atom and residue names
  std::vector<char4> atom_names;     
  std::vector<char4> atom_types;     
  std::vector<char4> residue_names;  
 
  // Sizing constants on CHARMM force field valence terms
  int urey_bradley_term_count;                
  int charmm_impr_term_count;                 
  int cmap_term_count;                        
  int urey_bradley_parameter_count;           
  int charmm_impr_parameter_count;            
  int cmap_surface_count;                     
 
  // Atom indexing on CHARMM force field terms
  std::vector<int> urey_bradley_i_atoms;           
  std::vector<int> urey_bradley_k_atoms;           
  std::vector<int> urey_bradley_parameter_indices; 
  std::vector<int> charmm_impr_i_atoms;            
  std::vector<int> charmm_impr_j_atoms;            
  std::vector<int> charmm_impr_k_atoms;            
  std::vector<int> charmm_impr_l_atoms;            
  std::vector<int> charmm_impr_parameter_indices;  
  std::vector<int> cmap_i_atoms;                   
  std::vector<int> cmap_j_atoms;                   
  std::vector<int> cmap_k_atoms;                   
  std::vector<int> cmap_l_atoms;                   
  std::vector<int> cmap_m_atoms;                   
  std::vector<int> cmap_surface_indices;           
 
  // Valence parameters for CHARMM force field terms
  std::vector<int> cmap_surface_dimensions;      
  std::vector<int> cmap_surface_bounds;          
  std::vector<double> urey_bradley_stiffnesses;  
  std::vector<double> urey_bradley_equilibria;   
  std::vector<double> charmm_impr_stiffnesses;   
  std::vector<double> charmm_impr_phase_angles;  
  std::vector<double> cmap_surfaces;             
  std::vector<char4> urey_bradley_i_atom_types;  
  std::vector<char4> urey_bradley_k_atom_types;  
  std::vector<char4> charmm_impr_i_atom_types;   
  std::vector<char4> charmm_impr_j_atom_types;   
  std::vector<char4> charmm_impr_k_atom_types;   
  std::vector<char4> charmm_impr_l_atom_types;   
  std::vector<char4> cmap_i_atom_types;          
  std::vector<char4> cmap_j_atom_types;          
  std::vector<char4> cmap_k_atom_types;          
  std::vector<char4> cmap_l_atom_types;          
  std::vector<char4> cmap_m_atom_types;          
 
  // Sizing constants on common valence terms (AMBER, GROMACS, CHARMM)
  int bond_term_count;           
  int angl_term_count;           
  int dihe_term_count;           
  int bond_parameter_count;      
  int angl_parameter_count;      
  int dihe_parameter_count;      
  int attenuated_14_type_count;  
  int inferred_14_attenuations;  
 
  // Atom indexing on common valence terms
  std::vector<int> bond_i_atoms;              
  std::vector<int> bond_j_atoms;              
  std::vector<int> bond_parameter_indices;    
  std::vector<int> angl_i_atoms;              
  std::vector<int> angl_j_atoms;              
  std::vector<int> angl_k_atoms;              
  std::vector<int> angl_parameter_indices;    
  std::vector<int> dihe_i_atoms;              
  std::vector<int> dihe_j_atoms;              
  std::vector<int> dihe_k_atoms;              
  std::vector<int> dihe_l_atoms;              
  std::vector<int> dihe_parameter_indices;    
  std::vector<int> dihe14_parameter_indices;  
  std::vector<int> infr14_i_atoms;            
  std::vector<int> infr14_l_atoms;            
  std::vector<int> infr14_parameter_indices;  
  std::vector<char4> bond_modifiers;          
  std::vector<char4> angl_modifiers;          
  std::vector<char4> dihe_modifiers;          
 
  // Valence parameters for common force field terms
  std::vector<double> bond_stiffnesses;    
  std::vector<double> bond_equilibria;     
  std::vector<double> angl_stiffnesses;    
  std::vector<double> angl_equilibria;     
  std::vector<double> dihe_amplitudes;     
  std::vector<double> dihe_periodicities;  
  std::vector<double> dihe_phase_angles;   
  std::vector<char4> bond_i_atom_types;    
  std::vector<char4> bond_j_atom_types;    
  std::vector<char4> angl_i_atom_types;    
  std::vector<char4> angl_j_atom_types;    
  std::vector<char4> angl_k_atom_types;    
  std::vector<char4> dihe_i_atom_types;    
  std::vector<char4> dihe_j_atom_types;    
  std::vector<char4> dihe_k_atom_types;    
  std::vector<char4> dihe_l_atom_types;    
 
  // Information relevant to virtual site placement
  int virtual_site_count;                           
  int virtual_site_parameter_set_count;             
  std::vector<int> virtual_site_atoms;              
  std::vector<VirtualSiteKind> vs_frame_types;      
  std::vector<int> virtual_site_frame1_atoms;       
  std::vector<int> virtual_site_frame2_atoms;       
  std::vector<int> virtual_site_frame3_atoms;       
  std::vector<int> virtual_site_frame4_atoms;       
  std::vector<int> virtual_site_parameter_indices;  
  std::vector<double> virtual_site_frame_dim1;      
  std::vector<double> virtual_site_frame_dim2;      
  std::vector<double> virtual_site_frame_dim3;      
 
  // Sizing constants and settings for the non-bonded calculation
  int charge_type_count;            
  int lj_type_count;                
  int total_exclusions;             
  UnitCellType periodic_box_class;  
  double coulomb_constant;          
  double elec14_screening_factor;   
  double vdw14_screening_factor;    
  
  // Arrays of non-bonded parameters
  std::vector<int> charge_indices;          
  std::vector<int> lennard_jones_indices;   
  std::vector<int> nb11_exclusion_bounds;   
  std::vector<int> nb11_exclusion_list;     
  std::vector<int> nb12_exclusion_bounds;   
  std::vector<int> nb12_exclusion_list;     
  std::vector<int> nb13_exclusion_bounds;   
  std::vector<int> nb13_exclusion_list;     
  std::vector<int> nb14_exclusion_bounds;   
  std::vector<int> nb14_exclusion_list;     
  std::vector<double> atomic_charges;       
  std::vector<double> charge_parameters;    
  std::vector<double> lj_a_values;          
  std::vector<double> lj_b_values;          
  std::vector<double> lj_c_values;          
  std::vector<double> lj_14_a_values;       
  std::vector<double> lj_14_b_values;       
  std::vector<double> lj_14_c_values;       
  std::vector<double> lj_sigma_values;      
  std::vector<double> lj_epsilon_values;    
  std::vector<double> lj_14_sigma_values;   
  std::vector<double> lj_14_epsilon_values; 
 
  // Whereas the AtomGraph contains condensed arrays of the distinct 1:4 non-bonded screening
  // parameter pairs, the staging object must collect instructions and create the 1:4 screening
  // factors for every dihedral in the topology.  This is accomplished by accumulating a series of
  // atom type / screening factor pair records, then building the 1:4 screenings for all dihedrals
  // in a flash just before exporting a topology.  This array covers pairs of electrostatic and
  // van-der Waals screening factors for up to four distinct atom types. 
  std::vector<ScreeningFactorTag> attn14_tags;
 
  // Atom, atom type, and residue name overflow keys
  std::vector<char4> atom_overflow_names;    
  std::vector<char4> atom_overflow_types;    
  std::vector<char4> residue_overflow_names; 

  void validateAtomIndex(int atom_index) const;

  void validateFrameIndex(int set_index) const;

  std::vector<int> insertParticles(const std::vector<int> &z_numbers,
                                   const std::vector<int> &residue_homes, 
                                   const std::vector<int> &added_vs_parameter_indices = {},
                                   const std::vector<int> &added_vs_frame1_atoms = {},
                                   const std::vector<int> &added_vs_frame2_atoms = {},
                                   const std::vector<int> &added_vs_frame3_atoms = {},
                                   const std::vector<int> &added_vs_frame4_atoms = {});

  std::vector<int> addConnections(const std::vector<int> &atom_i, const std::vector<int> &atom_j);
  
  int addConnection(int atom_i, int atom_j);

  bool applyProperty(int array_size, int added_count, const std::string &desc);

  std::vector<bool>
  maskSoluteAtoms(int minimum_solute_size = default_minimum_solute_atoms,
                  const std::vector<char4> &solute_included_residue_names = {}) const;

  std::vector<int> createPrmtopDescriptors(int largest_residue_atoms,
                                           int total_exclusions) const;

  std::vector<double> computeAtomicMasses(double hmass_repartition_factor) const;

  void buildNonbonded14Screen(std::vector<double> *attn14_elec_factors,
                              std::vector<double> *attn14_vdw_factors) const;
};

std::vector<int> atomWithVirtualSiteChildren(const int atom_index,
                                             const std::vector<int> &nb11_exclusion_list,
                                             const std::vector<int> &nb11_exclusion_bounds);
 
} // namespace topology
} // namespace stormm
 
#endif