system.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 <vector>
#include <functional>
#include <memory>
#include <Eigen/Core>
#include <Eigen/Dense>
#include "pteros/core/atom.h"
#include "pteros/core/force_field.h"
#include "pteros/core/periodic_box.h"
#include "pteros/core/typedefs.h"
#include "pteros/core/atom_handler.h"
 
 
namespace pteros {
 
struct Frame {
    std::vector<Eigen::Vector3f> coord;
    std::vector<Eigen::Vector3f> vel;
    std::vector<Eigen::Vector3f> force;
    PeriodicBox box;
    float time;
 
    Frame(): time(0.0) {}
 
    bool has_vel() const { return !vel.empty(); }
    bool has_force() const { return !force.empty(); }
 
    void swap(int i, int j);
};
 
//====================================================================================
 
//Forward declarations
class Selection;
class FileHandler;
class FileContent;
class AtomHandler;
 
//====================================================================================
 
class System {
    // System and Selection are friends because they are closely integrated.
    friend class Selection;    
    // Selection_parser must access internals of Selection
    friend class SelectionParser;
    // Needs an access for constructing the system in IO handlers
    friend class SystemBuilder;
 
public:    
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    System();
 
    System(std::string fname);
 
    System(const System& other);
 
    System& operator=(const System& other);
 
    System(const Selection& sel);
 
    virtual ~System();
 
 
 
    Selection append(const System& sys);
 
    Selection append(const Selection& sel, bool current_frame=false);
 
    Selection append(const Atom& at, Vector3f_const_ref coord);
 
    Selection append(const AtomHandler& at);
 
    void rearrange(const std::vector<std::string>& sel_strings);
 
    void rearrange(const std::vector<Selection>& sel_vec);
 
    void rearrange(const std::vector<std::string>& begin_strings, const std::vector<std::string>& end_strings);
 
    void rearrange(const std::vector<Selection>& begin_vec, const std::vector<Selection>& end_vec);
 
    std::vector<std::pair<std::string,int>> rearrange_by_resname();
 
    void keep(const std::string& sel_str);
 
    void keep(const Selection& sel);
 
    void remove(const std::string& sel_str);
 
    void remove(Selection& sel);
 
    void distribute(const Selection sel, Vector3i_const_ref ncopies, Matrix3f_const_ref shift);
 
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    inline int num_frames() const { return traj.size(); }
 
    inline int num_atoms() const { return atoms.size(); }
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    Selection select(std::string str, int fr = 0);
    Selection operator()(std::string str, int fr = 0);
 
    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);
 
    Selection select(std::vector<int>::iterator it1,
                     std::vector<int>::iterator it2);
    Selection operator()(std::vector<int>::iterator it1,
                         std::vector<int>::iterator it2);
 
    Selection select(const std::function<void(const System&,int,std::vector<int>&)>& callback, int fr = 0);
    Selection operator()(const std::function<void(const System&,int,std::vector<int>&)>& callback, int fr = 0);
 
    Selection select_all() const;
    Selection operator()();
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    void load(std::string fname,
              int b=0,
              int e=-1,
              int skip = 0,
              std::function<bool(System*,int)> on_frame = 0);
 
    bool load(const std::unique_ptr<FileHandler> &handler,
              FileContent what,
              std::function<bool(System*,int)> on_frame = 0);    
 
    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;
 
 
    std::vector<std::pair<std::string, Selection> > load_gromacs_ndx(std::string fname);
 
 
    void set_filter(std::string str);
    void set_filter(int ind1, int ind2);
    void set_filter(const std::vector<int>& ind);
    void set_filter(std::vector<int>::iterator it1,
                    std::vector<int>::iterator it2);
 
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    int frame_dup(int fr);
 
    void frame_append(const Frame& fr);
 
    void frame_copy(int fr1, int fr2);
 
    void frame_delete(int b = 0, int e = -1);
 
    void frame_swap(int fr1, int fr2);
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    inline PeriodicBox& box(int fr=0){ return traj[fr].box; }
 
    inline const PeriodicBox& box(int fr=0) const { return traj[fr].box; }
 
    inline float& time(int fr=0){ return traj[fr].time; }
 
    inline const float& time(int fr=0) const { return traj[fr].time; }
 
    inline Eigen::Vector3f& xyz(int ind, int fr=0){ return traj[fr].coord[ind]; }
 
    inline const Eigen::Vector3f& xyz(int ind, int fr=0) const { return traj[fr].coord[ind]; }
 
    inline Eigen::Vector3f& vel(int ind, int fr=0){ return traj[fr].vel[ind]; }
 
    inline const Eigen::Vector3f& vel(int ind, int fr=0) const { return traj[fr].vel[ind]; }
 
    inline Eigen::Vector3f& force(int ind, int fr=0){ return traj[fr].force[ind]; }
 
    inline const Eigen::Vector3f& force(int ind, int fr=0) const { return traj[fr].force[ind]; }
 
    inline Atom& atom(int ind) { return atoms[ind]; }
 
    inline const Atom& atom(int ind) const { return atoms[ind]; }
 
    inline Frame& frame(int fr){ return traj[fr]; }
 
    inline const Frame& frame(int fr) const { return traj[fr]; }
 
 
    AtomIterator atoms_begin(){ return atoms.begin(); }
    AtomIterator atoms_end(){ return atoms.end(); }
    std::vector<Frame>::iterator traj_begin(){ return traj.begin(); }
    std::vector<Frame>::iterator traj_end(){ return traj.end(); }
 
    AtomHandler operator[](const std::pair<int,int>& ind_fr);
 
    /*
    class atom_iterator;
    atom_iterator begin(int fr);
    atom_iterator end(int fr);
    */
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    Selection atoms_dup(const std::vector<int>& ind);
 
    Selection atoms_add(const std::vector<Atom>& atm,
                   const std::vector<Eigen::Vector3f>& crd);
 
    void atoms_delete(const std::vector<int>& ind);
 
    void atom_move(int i, int j);
 
    Selection atom_clone(int source);
 
    void atom_swap(int i, int j);
 
    Selection atom_add_1h(int target, int at1, int at2, int at3, float dist=0.109, bool pbc=true);
 
    Selection atom_add_2h(int target, int at1, int at2, float dist=0.109, bool pbc=true);
 
    Selection atom_add_3h(int target, int at1, float dist=0.109, bool pbc=true);
 
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    void wrap(int fr, Array3i_const_ref pbc = fullPBC);
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    float distance(int i, int j, int fr, Array3i_const_ref pbc = fullPBC) const;
 
    float angle(int i, int j, int k, int fr, Array3i_const_ref pbc = fullPBC) const;
 
    float dihedral(int i, int j, int k, int l, int fr, Array3i_const_ref pbc = fullPBC) const;
 
 
    //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    void clear();
 
    bool force_field_ready(){return force_field.ready;}
 
    ForceField& get_force_field(){
        return force_field;
    }
 
    void assign_resindex(int start=0);
 
    void sort_by_resindex();
 
    void clear_vel();
 
    void clear_force();
 
    Eigen::Vector2f get_energy_for_list(const std::vector<Eigen::Vector2i>& pairs,
                                        const std::vector<float>& dist,
                                        std::vector<Eigen::Vector2f>* pair_en=nullptr);
 
 
 
protected:
 
    // Holds all atom attributes except the coordinates
    std::vector<Atom>  atoms;
 
    // Coordinates for any number of frames
    std::vector<Frame> traj;
 
    // Force field parameters
    ForceField force_field;
 
    // Indexes for filtering
    std::vector<int> filter;
    // Filter selection text for text-based filters
    std::string filter_text;
 
    void filter_atoms();
    void filter_coord(int fr);
};
 
} // namespace