colvaratoms.h

#ifndef COLVARATOMS_SOA_H
#define COLVARATOMS_SOA_H
 
#include "colvaratoms_gpu.h"
#include "colvarmodule.h"
#include "colvardeps.h"
#include "colvar_rotation_derivative.h"
#include <unordered_map>
#include <mutex>
 
class cvm::atom_group: public colvardeps {
public:
  static ag_vector_real_t pos_aos_to_soa(const std::vector<cvm::atom_pos>& aos_in);
  atom_group();
  atom_group(char const *key_in);
  ~atom_group() override;
  int init();
  int setup();
  int init_dependencies() override;
  int proxy_gpu_buffers_reallocated() override;
  virtual const std::vector<feature *> &features() const override { return ag_features; }
  std::vector<feature *> &modify_features() override { return ag_features; }
  static void delete_features()
  {
    for (size_t i = 0; i < ag_features.size(); i++) {
      delete ag_features[i];
    }
    ag_features.clear();
  }
  struct simple_atom {
    int proxy_index;
    int id;
    cvm::real mass;
    cvm::real charge;
    cvm::atom_pos pos;
    cvm::rvector vel;
    cvm::rvector total_force;
    cvm::rvector grad;
  };
  static simple_atom init_atom_from_proxy(
    colvarproxy* const p,
    cvm::residue_id const &residue,
    std::string const     &atom_name,
    std::string const     &segment_id);
  static simple_atom init_atom_from_proxy(
    colvarproxy* const p,
    int atom_number);
  static simple_atom init_atom_from_proxy(
    colvarproxy* const p,
    const simple_atom& atom);
  class atom_modifier {
  private:
    cvm::atom_group* m_ag;
    std::vector<int> m_atoms_ids;
    std::vector<simple_atom> m_atoms;
    std::unordered_map<int, int> m_atoms_ids_count;
    cvm::real m_total_mass;
    cvm::real m_total_charge;
    void update_from_soa();
    void sync_to_soa() const;
  public:
    using atom_iter = decltype(m_atoms)::iterator;
    using const_atom_iter = decltype(m_atoms)::const_iterator;
    atom_modifier(cvm::atom_group* ag);
    ~atom_modifier();
    inline simple_atom & operator [] (size_t const i){return m_atoms[i];}
    inline simple_atom const & operator [] (size_t const i) const {return m_atoms[i];}
    inline atom_iter begin(){return m_atoms.begin();}
    inline const_atom_iter begin() const {return m_atoms.begin();}
    inline atom_iter end() {return m_atoms.end();}
    inline const_atom_iter end() const {return m_atoms.end();}
    inline size_t size() const {return m_atoms.size();}
 
    int add_atom(simple_atom const &a);
    int add_atom_numbers(std::string const &numbers_conf);
    int add_atoms_of_group(const atom_group *ag);
    int add_index_group(std::string const &index_group_name, bool silent = false);
    int add_atom_numbers_range(std::string const &range_conf);
    int add_atom_name_residue_range(std::string const &psf_segid,
                                    std::string const &range_conf);
    int add_atom_id(int aid);
 
    int remove_atom(atom_iter ai);
  };
  atom_modifier get_atom_modifier() {
    return atom_modifier(this);
  }
  void clear_soa();
  int parse(std::string const &conf);
  int parse_fitting_options(std::string const &group_conf);
  int set_dummy();
  int set_dummy_pos(cvm::atom_pos const &pos);
  void update_total_mass();
  void update_total_charge();
  void print_properties(std::string const &colvar_name, int i, int j);
  inline std::vector<int> const &ids() const
  {
    return atoms_ids;
  }
  std::string const print_atom_ids() const;
  int create_sorted_ids();
  inline std::vector<int> const &sorted_ids() const
  {
    return sorted_atoms_ids;
  }
  inline std::vector<int> const &sorted_ids_map() const
  {
    return sorted_atoms_ids_map;
  }
  static int overlap(const atom_group &g1, const atom_group &g2);
  void read_positions();
  void calc_apply_roto_translation();
  void setup_rotation_derivative();
  void center_ref_pos();
  void rotate(const cvm::rmatrix& rot_mat);
  void apply_translation(cvm::rvector const &t);
  void read_velocities();
  void read_total_forces();
  int calc_required_properties();
  int calc_center_of_geometry();
  cvm::ag_vector_real_t positions() const;
  inline cvm::atom_pos center_of_geometry() const
  {
    return cog;
  }
  int calc_center_of_mass();
  inline cvm::atom_pos center_of_mass() const
  {
    return com;
  }
  inline cvm::rvector center_of_mass_scalar_gradient() const
  {
    return scalar_com_gradient;
  }
  int positions_shifted(cvm::rvector const &shift, cvm::ag_vector_real_t& out) const;
  std::vector<cvm::real> velocities() const;
  int calc_dipole(cvm::atom_pos const &dipole_center);
  inline cvm::rvector dipole() const
  {
    return dip;
  }
  std::vector<cvm::real> total_forces() const;
  cvm::rvector total_force() const;
  void set_weighted_gradient(cvm::rvector const &grad);
  void calc_fit_gradients();
  template <bool B_ag_center, bool B_ag_rotate,
            typename main_force_accessor_T,
            typename fitting_force_accessor_T>
  void calc_fit_forces_impl(
    main_force_accessor_T accessor_main,
    fitting_force_accessor_T accessor_fitting) const;
 
  template <typename main_force_accessor_T, typename fitting_force_accessor_T>
  void calc_fit_forces(
    main_force_accessor_T accessor_main,
    fitting_force_accessor_T accessor_fitting) const;
  void apply_colvar_force(cvm::real const &force);
  void apply_force(cvm::rvector const &force);
  void do_feature_side_effects(int id) override;
  size_t size() const {return num_atoms;}
  void reset_atoms_data();
  void set_ref_pos_from_aos(const std::vector<cvm::atom_pos>& pos_aos);
  inline int& id(size_t i) {return atoms_ids[i];}
  inline const int& id(size_t i) const {return atoms_ids[i];}
 
  inline cvm::real& pos_x(size_t i) {return atoms_pos[i];}
  inline cvm::real& pos_y(size_t i) {return atoms_pos[i + num_atoms];}
  inline cvm::real& pos_z(size_t i) {return atoms_pos[i + 2 * num_atoms];}
  inline const cvm::real& pos_x(size_t i) const {return atoms_pos[i];}
  inline const cvm::real& pos_y(size_t i) const {return atoms_pos[i + num_atoms];}
  inline const cvm::real& pos_z(size_t i) const {return atoms_pos[i + 2 * num_atoms];}
 
  inline cvm::real& vel_x(size_t i) {return atoms_vel[i];}
  inline cvm::real& vel_y(size_t i) {return atoms_vel[i + num_atoms];}
  inline cvm::real& vel_z(size_t i) {return atoms_vel[i + 2 * num_atoms];}
  inline const cvm::real& vel_x(size_t i) const {return atoms_vel[i];}
  inline const cvm::real& vel_y(size_t i) const {return atoms_vel[i + num_atoms];}
  inline const cvm::real& vel_z(size_t i) const {return atoms_vel[i + 2 * num_atoms];}
 
  inline cvm::real& grad_x(size_t i) {return atoms_grad[i];}
  inline cvm::real& grad_y(size_t i) {return atoms_grad[i + num_atoms];}
  inline cvm::real& grad_z(size_t i) {return atoms_grad[i + 2 * num_atoms];}
  inline const cvm::real& grad_x(size_t i) const {return atoms_grad[i];}
  inline const cvm::real& grad_y(size_t i) const {return atoms_grad[i + num_atoms];}
  inline const cvm::real& grad_z(size_t i) const {return atoms_grad[i + 2 * num_atoms];}
 
  inline cvm::real& total_force_x(size_t i) {return atoms_total_force[i];}
  inline cvm::real& total_force_y(size_t i) {return atoms_total_force[i + num_atoms];}
  inline cvm::real& total_force_z(size_t i) {return atoms_total_force[i + 2 * num_atoms];}
  inline const cvm::real& total_force_x(size_t i) const {return atoms_total_force[i];}
  inline const cvm::real& total_force_y(size_t i) const {return atoms_total_force[i + num_atoms];}
  inline const cvm::real& total_force_z(size_t i) const {return atoms_total_force[i + 2 * num_atoms];}
 
  inline cvm::real& ref_pos_x(size_t i) {return ref_pos[i];}
  inline cvm::real& ref_pos_y(size_t i) {return ref_pos[i + num_ref_pos];}
  inline cvm::real& ref_pos_z(size_t i) {return ref_pos[i + 2 * num_ref_pos];}
  inline const cvm::real& ref_pos_x(size_t i) const {return ref_pos[i];}
  inline const cvm::real& ref_pos_y(size_t i) const {return ref_pos[i + num_ref_pos];}
  inline const cvm::real& ref_pos_z(size_t i) const {return ref_pos[i + 2 * num_ref_pos];}
 
  inline cvm::real& pos_unrotated_x(size_t i) {return atoms_pos_unrotated[i];}
  inline cvm::real& pos_unrotated_y(size_t i) {return atoms_pos_unrotated[i + num_atoms];}
  inline cvm::real& pos_unrotated_z(size_t i) {return atoms_pos_unrotated[i + 2 * num_atoms];}
  inline const cvm::real& pos_unrotated_x(size_t i) const {return atoms_pos_unrotated[i];}
  inline const cvm::real& pos_unrotated_y(size_t i) const {return atoms_pos_unrotated[i + num_atoms];}
  inline const cvm::real& pos_unrotated_z(size_t i) const {return atoms_pos_unrotated[i + 2 * num_atoms];}
 
  inline cvm::real& mass(size_t i) {return atoms_mass[i];}
  inline const cvm::real& mass(size_t i) const {return atoms_mass[i];}
 
  inline cvm::real& weight(size_t i) {return atoms_weight[i];}
  inline const cvm::real& weight(size_t i) const {return atoms_weight[i];}
 
  inline cvm::real& charge(size_t i) {return atoms_charge[i];}
  inline const cvm::real& charge(size_t i) const {return atoms_charge[i];}
 
  inline cvm::real& fit_gradients_x(size_t i) {
    atom_group *group_for_fit = fitting_group ? fitting_group : this;
    return group_for_fit->fit_gradients[i];}
  inline cvm::real& fit_gradients_y(size_t i) {
    atom_group *group_for_fit = fitting_group ? fitting_group : this;
    return group_for_fit->fit_gradients[i + group_for_fit->size()];}
  inline cvm::real& fit_gradients_z(size_t i) {
    atom_group *group_for_fit = fitting_group ? fitting_group : this;
    return group_for_fit->fit_gradients[i + 2 * group_for_fit->size()];}
  inline const cvm::real& fit_gradients_x(size_t i) const {
    const atom_group *group_for_fit = fitting_group ? fitting_group : this;
    return group_for_fit->fit_gradients[i];}
  inline const cvm::real& fit_gradients_y(size_t i) const {
    const atom_group *group_for_fit = fitting_group ? fitting_group : this;
    return group_for_fit->fit_gradients[i + group_for_fit->size()];}
  inline const cvm::real& fit_gradients_z(size_t i) const {
    const atom_group *group_for_fit = fitting_group ? fitting_group : this;
    return group_for_fit->fit_gradients[i + 2 * group_for_fit->size()];}
 
  inline cvm::real& group_forces_x(size_t i) {return group_forces[i];}
  inline cvm::real& group_forces_y(size_t i) {return group_forces[i + num_atoms];}
  inline cvm::real& group_forces_z(size_t i) {return group_forces[i + 2 * num_atoms];}
  inline const cvm::real& group_forces_x(size_t i) const {return group_forces[i];}
  inline const cvm::real& group_forces_y(size_t i) const {return group_forces[i + num_atoms];}
  inline const cvm::real& group_forces_z(size_t i) const {return group_forces[i + 2 * num_atoms];}
 
  inline simple_atom operator[](size_t i) const {
    return simple_atom{
      /*.proxy_index = */atoms_index[i],
      /*.id = */atoms_ids[i],
      /*.mass = */atoms_mass[i],
      /*.charge = */atoms_charge[i],
      /*.pos = */{pos_x(i), pos_y(i), pos_z(i)},
      /*.vel = */{vel_x(i), vel_y(i), vel_z(i)},
      /*.total_force = */{total_force_x(i), total_force_y(i), total_force_z(i)},
      /*.grad = */{grad_x(i), grad_y(i), grad_z(i)}
    };
  }
 
  class group_force_object {
  public:
    group_force_object(cvm::atom_group* ag);
    ~group_force_object();
    void add_atom_force(size_t i, const cvm::rvector& force);
  private:
    cvm::atom_group* m_ag;
    cvm::atom_group* m_group_for_fit;
    bool m_has_fitting_force;
    void apply_force_with_fitting_group();
  };
  group_force_object get_group_force_object();
public:
  std::string name;
  bool b_dummy;
  // TODO Make this field part of the data structures that link a group to a CVC
  std::string key;
  cvm::atom_group *fitting_group;
  cvm::rotation rot;
  bool noforce;
  bool b_user_defined_fit;
  cvm::ag_vector_real_t fit_gradients;
  cvm::real total_mass;
  cvm::real total_charge;
  rotation_derivative* rot_deriv;
  static std::vector<feature *> ag_features;
private:
  size_t num_atoms;
  std::vector<int> atoms_index;
  cvm::ag_vector_real_t atoms_pos;
  std::vector<cvm::real> atoms_charge;
  std::vector<cvm::real> atoms_vel;
  std::vector<cvm::real> atoms_mass;
  cvm::ag_vector_real_t atoms_grad;
  std::vector<cvm::real> atoms_total_force;
  std::vector<cvm::real> atoms_weight;
  std::vector<int> atoms_ids;
  std::vector<int> sorted_atoms_ids;
  std::vector<int> sorted_atoms_ids_map;
  cvm::atom_pos dummy_atom_pos;
  int index;
  cvm::ag_vector_real_t group_forces;
  cvm::ag_vector_real_t ref_pos;
  size_t num_ref_pos; // TODO: Do I really need this?
  cvm::atom_pos ref_pos_cog;
  cvm::atom_pos cog;
  cvm::atom_pos cog_orig;
  cvm::ag_vector_real_t atoms_pos_unrotated;
  cvm::atom_pos com;
  // TODO for scalable calculations of more complex variables (e.g. rotation),
  // use a colvarvalue of vectors to hold the entire derivative
  cvm::rvector scalar_com_gradient;
  cvm::rvector dip;
  std::mutex modify_lock;
#if defined(COLVARS_CUDA) || defined (COLVARS_HIP)
  std::unique_ptr<colvars_gpu::colvaratoms_gpu> gpu_atom_group;
  friend class colvars_gpu::colvaratoms_gpu;
public:
  std::unique_ptr<colvars_gpu::colvaratoms_gpu>& get_gpu_atom_group() {
    return gpu_atom_group;
  }
#elif defined (COLVARS_SYCL)
  // TODO
#endif
};
 
#endif // COLVARATOMS_SOA_H