Program Listing for File aatopology.hpp

↰ Return to documentation for file (source/pele/aatopology.hpp)

/*
 * This is a partial c++ implementation of the tools needed to interact with
 * systems of rigid bodies.  This is not a complete reimplementation, only the
 * parts that were too slow in python were implemented here.
 *
 */
#ifndef _PELE_AATOPOLOGY_H_
#define _PELE_AATOPOLOGY_H_

#include <cmath>
#include <list>
#include <stdexcept>
#include <string>

#include "pele/array.hpp"
#include "pele/base_potential.hpp"
#include "pele/distance.hpp"
#include "pele/lowest_eig_potential.hpp"
#include "pele/matrix.hpp"
#include "pele/rotations.hpp"
#include "pele/vecn.hpp"

namespace pele {

class CoordsAdaptor {
  size_t _nrigid;
  size_t _natoms;
  pele::Array<double> _coords;
  static const size_t _ndim = 3;

 public:
  CoordsAdaptor(size_t nrigid, size_t natoms, pele::Array<double> coords)
      : _nrigid(nrigid), _natoms(natoms), _coords(coords) {
    if (coords.size() != (6 * _nrigid + 3 * _natoms)) {
      throw std::invalid_argument(std::string("coords has the wrong size ") +
                                  std::to_string(coords.size()));
    }
    if (natoms != 0)
      throw std::runtime_error("coords adaptor doesn't support free atoms yet");
  }

  pele::Array<double> get_coords() { return _coords; }

  pele::Array<double> get_rb_positions() {
    if (_nrigid == 0) {
      // return empty array
      return pele::Array<double>();
    }
    return _coords.view(0, 3 * _nrigid);
  }

  pele::Array<double> get_rb_position(size_t isite) {
    if (_nrigid == 0) {
      // return empty array
      return pele::Array<double>();
    }
    if (isite > _nrigid) {
      throw std::invalid_argument("isite must be less than nrigid");
    }
    size_t const istart = 3 * isite;
    return _coords.view(istart, istart + 3);
  }

  pele::Array<double> get_rb_rotations() {
    if (_nrigid == 0) {
      // return empty array
      return pele::Array<double>();
    }
    return _coords.view(3 * _nrigid, 6 * _nrigid);
  }

  pele::Array<double> get_rb_rotation(size_t isite) {
    if (_nrigid == 0) {
      // return empty array
      return pele::Array<double>();
    }
    if (isite > _nrigid) {
      throw std::invalid_argument("isite must be less than nrigid");
    }
    size_t const istart = 3 * _nrigid + 3 * isite;
    return _coords.view(istart, istart + 3);
  }
  pele::Array<double> get_atom_positions() {
    if (_natoms == 0) {
      // return empty array
      return pele::Array<double>();
    }

    return _coords.view(6 * _nrigid, 6 * _nrigid + 3 * _natoms);
  }
};

// forward definition of RBTopology needed for TrasnformAACluster
class RBTopology;

class TransformPolicy {
 public:
  //    void translate(self, X, d) {
  //        ''' translate the coordinates '''
  //    }
  virtual ~TransformPolicy() {}

  virtual void rotate(pele::Array<double> x,
                      pele::MatrixNM<3, 3> const &mx) = 0;

  //    def can_invert(self):
  //        ''' returns True or False if an inversion can be performed'''
  //
  //    def invert(self, X):
  //        ''' perform an inversion at the origin '''
  //
  //    def permute(self, X, perm):
  //        ''' returns the permuted coordinates '''
};

class TransformAACluster : public TransformPolicy {
 public:
  pele::RBTopology *m_topology;
  TransformAACluster(pele::RBTopology *topology) : m_topology(topology) {}
  virtual ~TransformAACluster() {}

  void rotate(pele::Array<double> x, pele::MatrixNM<3, 3> const &mx);
  //    inline void rotate(pele::Array<double> x, pele::Array<double> mx)
  //    {
  //        return rotate(x, pele::MatrixNM<3,3>(mx));
  //    }
};

class MeasureAngleAxisCluster {
 public:
  pele::RBTopology *m_topology;
  MeasureAngleAxisCluster(pele::RBTopology *topology) : m_topology(topology) {}

  void align(pele::Array<double> const x1, pele::Array<double> x2);
};

class RigidFragment {
  static const size_t _ndim = 3;
  pele::Array<double> _atom_positions;
  pele::MatrixAdapter<double> _atom_positions_matrix;
  size_t _natoms;

  std::shared_ptr<DistanceInterface> m_distance_function;

  double m_M;           // total mass of the angle axis site
  double m_W;           // sum of all weights
  pele::VecN<3> m_cog;  // center of geometry
  pele::MatrixNM<3, 3>
      m_S;  // weighted tensor of gyration S_ij = \sum m_i x_i x_j
  pele::MatrixNM<3, 3>
      m_inversion;  // matrix that applies the appropriate inversion
  bool m_can_invert;

  // a list of rotations that leave the rigid body unchanged.
  std::vector<pele::MatrixNM<3, 3>> m_symmetry_rotations;

  // the list of indices that each atom takes in the coords array
  // coords[atom_indices[i]] is the first coordinate of the i'th atom
  pele::Array<size_t> m_atom_indices;

 public:
  RigidFragment(pele::Array<double> atom_positions, Array<double> cog, double M,
                double W, Array<double> S, Array<double> inversion,
                bool can_invert,
                std::shared_ptr<DistanceInterface> distance_function)
      : _atom_positions(atom_positions.copy()),
        _atom_positions_matrix(_atom_positions, _ndim),
        _natoms(_atom_positions.size() / _ndim),
        m_distance_function(distance_function),
        m_M(M),
        m_W(W),
        m_cog(cog),
        m_S(S),
        m_inversion(inversion),
        m_can_invert(can_invert) {
    if (_atom_positions.size() == 0) {
      throw std::invalid_argument("the atom positions must not have zero size");
    }
    if (_atom_positions.size() != _natoms * _ndim) {
      throw std::invalid_argument(
          "the length of atom_positions must be divisible by 3");
    }
  }

  inline size_t natoms() const { return _natoms; }

  inline void add_symmetry_rotation(pele::Array<double> R) {
    m_symmetry_rotations.push_back(R.copy());
  }

  inline void set_atom_indices(pele::Array<size_t> atom_indices) {
    if (atom_indices.size() != natoms()) {
      throw std::runtime_error("atom_indices has the wrong size");
    }
    m_atom_indices = atom_indices.copy();
  }

  inline pele::Array<size_t> &get_atom_indices() { return m_atom_indices; }

  inline std::vector<pele::MatrixNM<3, 3>> const &get_symmetry_rotations()
      const {
    return m_symmetry_rotations;
  }

  pele::Array<double> to_atomistic(pele::Array<double> const com,
                                   pele::VecN<3> const &p);

  void transform_grad(pele::VecN<3> const &p, pele::Array<double> const g,
                      pele::VecN<3> &g_com, pele::VecN<3> &g_rot);

  void transform_grad(pele::Array<double> const p, pele::Array<double> const g,
                      pele::Array<double> g_com, pele::Array<double> g_rot) {
    pele::VecN<3> p_vec = p;
    pele::VecN<3> g_com_vec = g_com;
    pele::VecN<3> g_rot_vec = g_rot;
    transform_grad(p_vec, g, g_com_vec, g_rot_vec);
    for (size_t i = 0; i < 3; ++i) {
      g_com[i] = g_com_vec[i];
      g_rot[i] = g_rot_vec[i];
    }
  }

  inline pele::VecN<3> get_smallest_rij(pele::VecN<3> const &com1,
                                        pele::VecN<3> const &com2) const {
    pele::VecN<3> distance;
    m_distance_function->get_rij(distance.data(), com2.data(), com1.data());
    return distance;
  }

  double distance_squared(pele::VecN<3> const &com1, pele::VecN<3> const &p1,
                          pele::VecN<3> const &com2,
                          pele::VecN<3> const &p2) const;

  void distance_squared_grad(pele::VecN<3> const &com1, pele::VecN<3> const &p1,
                             pele::VecN<3> const &com2, pele::VecN<3> const &p2,
                             VecN<3> &g_M, VecN<3> &g_P) const;
};

// Inactive angle-axis topology prototype.
// class AATopology {
//};

class RBTopology {
  std::vector<RigidFragment> _sites;
  size_t _natoms_total;

 public:
  RBTopology() : _natoms_total(0) {}

  void add_site(RigidFragment const &site) {
    _sites.push_back(site);
    _natoms_total += site.natoms();
  }

  std::vector<RigidFragment> const &get_sites() const { return _sites; };

  void assign_atom_indices() {
    size_t atom = 0;
    for (auto &site : _sites) {
      pele::Array<size_t> atom_indices(site.natoms());
      for (size_t i = 0; i < atom_indices.size(); ++i)
        atom_indices[i] = atom + i;
      site.set_atom_indices(atom_indices);
      atom += site.natoms();
    }
  }

  size_t nrigid() const { return _sites.size(); }

  size_t natoms_total() const { return _natoms_total; }

  size_t number_of_non_rigid_atoms() const { return 0; }

  CoordsAdaptor get_coords_adaptor(pele::Array<double> x) const {
    return CoordsAdaptor(nrigid(), number_of_non_rigid_atoms(), x);
  }

  Array<double> to_atomistic(Array<double> rbcoords);

  void transform_gradient(pele::Array<double> rbcoords,
                          pele::Array<double> grad, pele::Array<double> rbgrad);

  pele::VecN<3> align_angle_axis_vectors(pele::VecN<3> const &p1,
                                         pele::VecN<3> const &p2in);

  void align_all_angle_axis_vectors(pele::Array<double> x1,
                                    pele::Array<double> x2);

  void align_path(std::list<pele::Array<double>> path);

  void get_zero_modes(pele::Array<double> const x,
                      std::vector<pele::Array<double>> &zev) {
    auto ca = get_coords_adaptor(x);
    pele::Array<double> v(x.size(), 0);
    auto cv = get_coords_adaptor(v);

    // get the zero eigenvectors corresponding to translation
    std::vector<pele::Array<double>> zev_t;
    pele::zero_modes_translational(zev_t, nrigid(), 3);

    for (auto const &v : zev_t) {
      cv.get_rb_positions().assign(v);
      zev.push_back(cv.get_coords().copy());
    }

    // get the zero eigenvectors corresponding to rotation
    TransformAACluster transform(this);
    double d = 1e-5;
    pele::VecN<3> v3;
    pele::Array<double> delta(x.size());

    // do rotations around the x y and z axes
    for (size_t i = 0; i < 3; ++i) {
      delta.assign(x);
      v3.assign(0.0);
      v3[i] = d;
      transform.rotate(delta, pele::aa_to_rot_mat(v3));
      align_all_angle_axis_vectors(x, delta);
      delta -= x;
      delta /= norm(delta);
      zev.push_back(delta.copy());
    }
  }

  double distance_squared(pele::Array<double> const x1,
                          pele::Array<double> const x2) const {
    double d_sq = 0;
    auto ca1 = get_coords_adaptor(x1);
    auto ca2 = get_coords_adaptor(x2);
    for (size_t isite = 0; isite < nrigid(); ++isite) {
      d_sq += _sites[isite].distance_squared(
          ca1.get_rb_position(isite), ca1.get_rb_rotation(isite),
          ca2.get_rb_position(isite), ca2.get_rb_rotation(isite));
    }
    return d_sq;
  }

  void distance_squared_grad(pele::Array<double> const x1,
                             pele::Array<double> const x2,
                             pele::Array<double> grad) const {
    if (grad.size() != x1.size()) {
      throw std::runtime_error("grad has the wrong size");
    }
    grad.assign(0);
    auto ca1 = get_coords_adaptor(x1);
    auto ca2 = get_coords_adaptor(x2);
    auto ca_spring = get_coords_adaptor(grad);

    // first distance for sites only
    for (size_t isite = 0; isite < nrigid(); ++isite) {
      pele::VecN<3> g_M, g_P;
      _sites[isite].distance_squared_grad(
          ca1.get_rb_position(isite), ca1.get_rb_rotation(isite),
          ca2.get_rb_position(isite), ca2.get_rb_rotation(isite), g_M, g_P);
      auto spring_com = ca_spring.get_rb_position(isite);
      std::copy(g_M.begin(), g_M.end(), spring_com.begin());
      auto spring_rot = ca_spring.get_rb_rotation(isite);
      std::copy(g_P.begin(), g_P.end(), spring_rot.begin());
    }
  }
};

class RBPotentialWrapper : public BasePotential {
  std::shared_ptr<BasePotential> potential_;
  std::shared_ptr<RBTopology> topology_;

 public:
  RBPotentialWrapper(std::shared_ptr<BasePotential> potential,
                     std::shared_ptr<RBTopology> top)
      : potential_(potential),
        topology_(top)

  {}

  inline double get_energy(pele::Array<double> const &rbcoords) {
    auto x = topology_->to_atomistic(rbcoords);
    return potential_->get_energy(x);
  }

  inline double get_energy_gradient(pele::Array<double> const &rbcoords,
                                    pele::Array<double> &rbgrad) {
    auto x = topology_->to_atomistic(rbcoords);
    pele::Array<double> grad_atomistic(topology_->natoms_total() * 3);
    double e = potential_->get_energy_gradient(x, grad_atomistic);
    topology_->transform_gradient(rbcoords, grad_atomistic, rbgrad);
    return e;
  }
};

}  // namespace pele

#endif