selection.h

/*
 * This file is a part of
 *
 * ============================================
 * ###   Pteros molecular modeling library  ###
 * ============================================
 *
 * https://github.com/yesint/pteros
 *
 * (C) 2009-2021, Semen Yesylevskyy
 *
 * All works, which use Pteros, should cite the following papers:
 *
 *  1.  Semen O. Yesylevskyy, "Pteros 2.0: Evolution of the fast parallel
 *      molecular analysis library for C++ and python",
 *      Journal of Computational Chemistry, 2015, 36(19), 1480–1488.
 *      doi: 10.1002/jcc.23943.
 *
 *  2.  Semen O. Yesylevskyy, "Pteros: Fast and easy to use open-source C++
 *      library for molecular analysis",
 *      Journal of Computational Chemistry, 2012, 33(19), 1632–1636.
 *      doi: 10.1002/jcc.22989.
 *
 * This is free software distributed under Artistic License:
 * http://www.opensource.org/licenses/artistic-license-2.0.php
 *
*/
 
 
#pragma once
 
#include <string>
//#include <memory>
#include <map>
#include <vector>
#include <functional>
#include <Eigen/Core>
//#include <Eigen/Geometry>
#include "pteros/core/atom_handler.h"
#include "pteros/core/system.h"
#include "pteros/core/typedefs.h"
 
namespace pteros {
 
// Forward declaration of friend classes
class System;
class SelectionParser;
 
 
class Selection {
  friend class System;
  friend class SelectionParser;
  friend class Grid_searcher;
 
  public:
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    Selection();
 
    explicit Selection(const System& sys);
 
    Selection(const System& sys, std::string str, int fr = 0);
 
    Selection(const System& sys, int ind1, int ind2);
 
    Selection(const System& sys, const std::vector<int>& ind);
 
    Selection(const System& sys,
              std::vector<int>::iterator it1,
              std::vector<int>::iterator it2);
 
    Selection(const System& sys,
              const std::function<void(const System&,int,std::vector<int>&)>& callback,
              int fr = 0);
 
    Selection(const Selection& other);
 
    virtual ~Selection();
 
    Selection& operator=(const Selection& other);
 
    bool operator==(const Selection &other) const;
 
    bool operator!=(const Selection &other) const {
        return !(*this == other);
    }
 
    AtomHandler operator[](int ind);
 
    AtomHandler operator[](const std::pair<int,int>& ind_fr);
 
    friend std::ostream& operator<<(std::ostream& os, const Selection& sel);
 
    friend Selection operator|(const Selection& sel1, const Selection& sel2);        
 
    friend Selection operator&(const Selection& sel1, const Selection& sel2);
 
    friend Selection operator-(const Selection& sel1, const Selection& sel2);
 
    Selection operator~() const;        
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    void append(const Selection& sel);
 
    void append(int ind);
 
    void append(const std::vector<Selection>& sel_vec);
 
    void remove(const Selection& sel);
 
    void remove(int ind);
 
    void invert();
 
    void set_system(const System& sys);
 
    void modify(std::string str, int fr = 0);
 
    void modify(int ind1, int ind2);
 
    void modify(const std::vector<int>& ind);
 
    void modify(std::vector<int>::iterator it1, std::vector<int>::iterator it2);
 
    void modify(const std::function<void(const System&,int,std::vector<int>&)>& callback, int fr = 0);
 
    void modify(const System& sys, std::string str, int fr = 0);
 
    void modify(const System& sys, int ind1, int ind2);
 
    void modify(const System& sys, const std::vector<int>& ind);
 
    void modify(const System& sys,
                std::vector<int>::iterator it1,
                std::vector<int>::iterator it2);
 
    void modify(const System& sys, const std::function<void(const System&,int,std::vector<int>&)>& callback, int fr = 0);
 
    void apply();
 
    void update();
 
    void clear();
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    Selection select(std::string str);
    Selection operator()(std::string str);
 
    Selection select(int ind1, int ind2);
    Selection operator()(int ind1, int ind2);
 
    Selection select(const std::vector<int>& ind);
    Selection operator()(const std::vector<int>& ind);
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    std::vector<int>::const_iterator index_begin() const { return _index.begin(); }
 
    std::vector<int>::const_iterator index_end() const { return _index.end(); }
 
    std::back_insert_iterator<std::vector<int> > index_back_inserter() {
        return std::back_inserter(_index);
    }
 
    class iterator;
 
    iterator begin();
    iterator end();
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    int get_frame() const {return frame;}
 
    void set_frame(int fr);
 
    System* get_system() const { return system; }
 
    std::string get_text() const;
 
    std::vector<int> get_index() const { return _index; }
 
    std::vector<char> get_chain(bool unique=false) const;
 
    void set_chain(const std::vector<char>& data);
 
    void set_chain(char data);
 
 
    std::vector<int> get_resid(bool unique=false) const;
 
    void set_resid(const std::vector<int>& data);
 
    void set_resid(int data);
 
 
    std::vector<std::string> get_name(bool unique=false) const;
 
    void set_name(const std::vector<std::string>& data);
 
    void set_name(std::string data);
 
 
    std::vector<std::string> get_resname(bool unique=false) const;
 
    void set_resname(const std::vector<std::string>& data);
 
    void set_resname(std::string data);
 
 
    std::vector<int> get_resindex(bool unique=false) const;
 
 
    Eigen::MatrixXf get_xyz(bool make_row_major_matrix = false) const;
 
    void set_xyz(MatrixXf_const_ref coord);
 
 
    std::vector<float> get_mass() const;
 
    void set_mass(const std::vector<float>& m);
 
    void set_mass(float data);
 
 
    std::vector<float> get_beta() const;
 
    void set_beta(const std::vector<float>& data);
 
    void set_beta(float data);
 
 
    std::vector<float> get_occupancy() const;
 
    void set_occupancy(const std::vector<float>& data);
 
    void set_occupancy(float data);
 
 
    std::vector<float> get_charge() const;
 
    void set_charge(const std::vector<float>& data);
 
    void set_charge(float data);
 
    float get_total_charge() const;
 
    std::vector<std::string> get_tag(bool unique=false) const;
 
    void set_tag(const std::vector<std::string>& data);
 
    void set_tag(std::string data);
 
 
    Eigen::MatrixXf get_vel(bool make_row_major_matrix = false) const;
 
    void set_vel(MatrixXf_const_ref data);
 
    Eigen::MatrixXf get_force(bool make_row_major_matrix = false) const;
 
    void set_force(MatrixXf_const_ref coord);
 
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    bool is_large();
 
    Eigen::Vector3f center(bool mass_weighted = false,
                           Array3i_const_ref pbc = noPBC,
                           int pbc_atom = -1) const;
 
    // Get the center of selection with user-supplied weights
    Eigen::Vector3f center(const std::vector<float>& weights,
                           Array3i_const_ref pbc = noPBC,
                           int pbc_atom = -1) const;
 
    void minmax(Vector3f_ref min, Vector3f_ref max) const;
 
    float powersasa(float probe_r = 0.14,
               std::vector<float>* area_per_atom = nullptr,
               float* total_volume = nullptr,
               std::vector<float>* volume_per_atom = nullptr) const;
 
    float sasa(float probe_r = 0.14, std::vector<float>* area_per_atom = nullptr, int n_sphere_points = 960) const;
 
    Eigen::MatrixXf average_structure(int b=0, int e=-1, bool make_row_major_matrix = false) const;
 
    Eigen::MatrixXf atom_traj(int ind, int b=0, int e=-1, bool make_row_major_matrix = false) const;
 
    void inertia(Vector3f_ref moments, Matrix3f_ref axes,
                 Array3i_const_ref pbc = noPBC,
                 int pbc_atom = -1) const;
 
    float gyration(Array3i_const_ref pbc = noPBC, int pbc_atom = -1) const;
 
    /* Computes dipole moment of selection in Debye
    If the size of selection is larger than 1/2 of the box size in
    any dimension you will get incorrect results if periodic is set to true.
    This is not checked automatically!
    In this case use one of unwrapping options first.
    */
    Eigen::Vector3f dipole(bool as_charged=false, Array3i_const_ref pbc = fullPBC, int pbc_atom = -1) const;
 
    float distance(int i, int j, Array3i_const_ref pbc = fullPBC) const;
 
    float angle(int i, int j, int k, Array3i_const_ref pbc = fullPBC) const;
 
    float dihedral(int i, int j, int k, int l, Array3i_const_ref pbc = fullPBC) const;
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    void translate(Vector3f_const_ref v);
 
    void translate_to(Vector3f_const_ref p,
                      bool mass_weighted = false,
                      Array3i_const_ref pbc = noPBC,
                      int pbc_atom = -1);
 
    void rotate(Vector3f_const_ref pivot, Vector3f_const_ref axis, float angle);
 
    void mirror(int dim, float around=0);
 
    void mirror(Vector3f_const_ref normal, Vector3f_const_ref point=Eigen::Vector3f::Zero());
 
    void wrap(Array3i_const_ref pbc = fullPBC);
 
    void unwrap(Array3i_const_ref pbc = fullPBC, int pbc_atom = -1);
 
    int unwrap_bonds(float d, Array3i_const_ref pbc = fullPBC, int pbc_atom = -1);
 
    Eigen::Affine3f principal_transform(Array3i_const_ref pbc = noPBC, int pbc_atom = -1) const;
 
    void principal_orient(Array3i_const_ref pbc = noPBC, int pbc_atom = -1);
 
    int num_residues() const;
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    float rmsd(int fr) const;
 
    float rmsd(int fr1, int fr2) const;
 
    friend float rmsd(const Selection& sel1, const Selection& sel2);
 
    friend float rmsd(const Selection& sel1, int fr1, const Selection& sel2, int fr2);
 
    friend void fit(Selection& sel1, const Selection& sel2);
 
    void fit_trajectory(int ref_frame=0, int b=0, int e=-1);
 
    friend Eigen::Affine3f fit_transform(const Selection& sel1, const Selection& sel2);
 
    Eigen::Affine3f fit_transform(int fr1, int fr2) const;
 
    void fit(int fr1, int fr2);
 
    void apply_transform(const Eigen::Affine3f& t);
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    Eigen::Vector2f non_bond_energy(float cutoff=0, bool pbc = true) const;
 
 
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    void write(std::string fname, int b=-1,int e=-1) const;
 
    void write(const std::unique_ptr<FileHandler>& handler, FileContent what,int b=-1,int e=-1) const;
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    int size() const {return _index.size();}
 
    bool text_based() const;
 
    bool coord_dependent() const;
 
    void flatten();
 
    std::string to_gromacs_ndx(std::string name) const;
 
    friend void copy_coord(const Selection& from, int from_fr, Selection& to, int to_fr);        
 
    friend void copy_coord(const Selection& from, Selection& to);
 
    int find_index(int global_index) const;
 
    std::vector<std::vector<int>> get_internal_bonds(float d, bool periodic=true) const;
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    void split_by_connectivity(float d, std::vector<Selection>& res, bool periodic=true);    
 
    void split_by_residue(std::vector<Selection>& res);
 
    void split_by_molecule(std::vector<Selection>& res);
 
    void split_by_chain(std::vector<Selection>& chains);
 
    void split_by_contiguous_index(std::vector<Selection>& parts);
 
    void split_by_contiguous_residue(std::vector<Selection>& parts);
 
    template<class F>
    void split(std::vector<Selection>& parts, F callback){
        using namespace std;
        parts.clear();
        using Ret = decltype(callback(*this,0));
        map<Ret,vector<int> > m;
        for(int i=0; i<size(); ++i){
            m[callback(*this,i)].push_back(index(i));
        }
        // Create selections
        typename map<Ret,vector<int> >::iterator it;
        for(it=m.begin();it!=m.end();it++){
            parts.emplace_back(*system, it->second.begin(), it->second.end());
        }
    }
 
 
    void each_residue(std::vector<Selection>& sel) const;
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    void dssp(std::string fname) const;
 
    void dssp(std::ostream& os) const;
 
    std::string dssp() const;
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    inline float& x(int ind)       { return system->traj[frame].coord[_index[ind]](0);  }
    inline float  x(int ind) const { return system->traj[frame].coord[_index[ind]](0); }
    inline float& x(int ind, int fr)       { return system->traj[fr].coord[_index[ind]](0); }
    inline float  x(int ind, int fr) const { return system->traj[fr].coord[_index[ind]](0); }
 
    inline float& y(int ind)       { return system->traj[frame].coord[_index[ind]](1); }
    inline float  y(int ind) const { return system->traj[frame].coord[_index[ind]](1); }
    inline float& y(int ind, int fr)       { return system->traj[fr].coord[_index[ind]](1); }
    inline float  y(int ind, int fr) const { return system->traj[fr].coord[_index[ind]](1); }
 
    inline float& z(int ind)       { return system->traj[frame].coord[_index[ind]](2); }
    inline float  z(int ind) const { return system->traj[frame].coord[_index[ind]](2); }
    inline float& z(int ind, int fr)       { return system->traj[fr].coord[_index[ind]](2); }
    inline float  z(int ind, int fr) const { return system->traj[fr].coord[_index[ind]](2); }
 
    inline       Eigen::Vector3f& xyz(int ind)       { return system->traj[frame].coord[_index[ind]]; }
    inline const Eigen::Vector3f& xyz(int ind) const { return system->traj[frame].coord[_index[ind]]; }
    inline       Eigen::Vector3f& xyz(int ind, int fr)       { return system->traj[fr].coord[_index[ind]]; }
    inline const Eigen::Vector3f& xyz(int ind, int fr) const { return system->traj[fr].coord[_index[ind]]; }
 
    inline Eigen::Vector3f* xyz_ptr(int ind)         const { return &(system->traj[frame].coord[_index[ind]]); }
    inline Eigen::Vector3f* xyz_ptr(int ind, int fr) const { return &(system->traj[fr].coord[_index[ind]]); }
 
    inline       Eigen::Vector3f& vel(int ind)       { return system->traj[frame].vel[_index[ind]]; }
    inline const Eigen::Vector3f& vel(int ind) const { return system->traj[frame].vel[_index[ind]]; }
    inline       Eigen::Vector3f& vel(int ind, int fr)       { return system->traj[fr].vel[_index[ind]]; }
    inline const Eigen::Vector3f& vel(int ind, int fr) const { return system->traj[fr].vel[_index[ind]]; }
 
    inline       Eigen::Vector3f& force(int ind)       { return system->traj[frame].force[_index[ind]]; }
    inline const Eigen::Vector3f& force(int ind) const { return system->traj[frame].force[_index[ind]]; }
    inline       Eigen::Vector3f& force(int ind, int fr)       { return system->traj[fr].force[_index[ind]]; }
    inline const Eigen::Vector3f& force(int ind, int fr) const { return system->traj[fr].force[_index[ind]]; }
 
#define DEFINE_ACCESSOR(T,prop) \
    inline T& prop(int ind){ return system->atoms[_index[ind]].prop; } \
    inline const T& prop(int ind) const { return system->atoms[_index[ind]].prop; }
 
    DEFINE_ACCESSOR(int,type)
    DEFINE_ACCESSOR(std::string,type_name)
    DEFINE_ACCESSOR(std::string,resname)
    DEFINE_ACCESSOR(char,chain)
    DEFINE_ACCESSOR(std::string,name)
    DEFINE_ACCESSOR(float,mass)
    DEFINE_ACCESSOR(float,charge)
    DEFINE_ACCESSOR(float,beta)
    DEFINE_ACCESSOR(float,occupancy)
    DEFINE_ACCESSOR(int,resid)
    DEFINE_ACCESSOR(std::string,tag)
 
    
    inline int index(int ind) const { return _index[ind]; }
 
    inline Atom& atom(int ind){ return system->atoms[_index[ind]]; }
    inline const Atom& atom(int ind) const { return system->atoms[_index[ind]]; }
 
    DEFINE_ACCESSOR(int,resindex)
    DEFINE_ACCESSOR(int,atomic_number)
 
    
    float vdw(int ind) const;
 
    std::string element_name(int ind) const;
 
    inline PeriodicBox& box() { return system->traj[frame].box; }
    inline const PeriodicBox& box() const { return system->traj[frame].box; }
 
    inline float& time() { return system->traj[frame].time; }
    inline const float& time() const { return system->traj[frame].time; }
 
 
protected:
    // Raw text of selection
    std::string sel_text;    
    // Indexes of atoms in selection
    std::vector<int> _index;
    // Pointer to target system
    System* system;
 
    // Stores current frame
    int frame;
 
    // Holds an instance of selection parser
    std::unique_ptr<SelectionParser> parser;
    void allocate_parser();
    void sort_and_remove_duplicates();    
    void process_pbc_atom(int& a) const;
    void get_local_bonds_from_topology(std::vector<std::vector<int>>& con) const;
};
 
//-----------------------------------------------------------------------
class Selection::iterator {
public:
    using value_type = AtomHandler;
    using difference_type = size_t;
    using pointer = AtomHandler*;
    using reference = AtomHandler&;
    using iterator_category = std::random_access_iterator_tag;
 
    iterator(const Selection& sel, int i){
        index_it = sel.index_begin() + i;
        proxy.set(sel,*index_it);
    }
    iterator operator++(int junk) { iterator tmp = *this; advance_proxy(); return tmp; }
    iterator& operator++() { advance_proxy(); return *this; }
    iterator& operator+(int i) {advance_proxy(i); return *this;}
    iterator& operator-(int i) {advance_proxy(-i); return *this;}
    reference operator*() { return proxy; }
    pointer   operator->() { return &proxy; }
    bool operator==(const iterator& rhs) { return index_it == rhs.index_it; }
    bool operator!=(const iterator& rhs) { return !(*this==rhs); }
 
private:
    AtomHandler proxy;
    // We have to keep an iterator for index array of parent selection
    std::vector<int>::const_iterator index_it;
 
    void advance_proxy(int n=1){
        int offset = *index_it; // save current index
        index_it += n;
        offset = *index_it-offset; // compute offset to new index
        proxy.ind = *index_it;
        proxy.atom_ptr += offset;
        proxy.coord_ptr += offset;
    }
};
 
bool check_selection_overlap(const std::vector<Selection> &sel_vec);
 
 
Eigen::Vector2f non_bond_energy(const Selection& sel1,
                                const Selection& sel2,
                                float cutoff = 0,
                                int fr = -1,
                                bool pbc = true);
 
} // namespace pteros