from past.utils import old_div
import numpy
degrees_per_rad = old_div(180., numpy.pi)
[docs]class PBCMeasureMananger:
    """
    A class to measure distance, angle, dihedral and planar angle under
    periodic boundary conditions.
    """
[docs]    def __init__(self, frame):
        """
        :type frame: _DesmondFrame instance
        :param frame:
            The trajectory frame on which the measurement will be performed.
        """
        self._frame = frame 
    def __call__(self,
                 atom1,
                 atom2,
                 atom3=None,
                 atom4=None,
                 atom5=None,
                 atom6=None,
                 minangle=False):
        """
        Measure the distance, angle, dihedral angle, or angle between planes
        of the provided atoms under periodic boundary condition.
        If atom1 and atom2 are provided, return the distance between them.
        If atom1, atom2, and atom3 are provided, return the angle between
        them.
        If atoms 1-4 are provided, return the dihedral angle between them.
        If atoms 1-6 are provided, return the angle between the planes
        defined by atoms 1-3 and atoms 4-6.
        Each of the atom arguments can be integer atom indices or
        structure._StructureAtom objects.
        Parameters
        :type minangle: bool
        :param minangle:
                This applies to the planar angle calculation and if True
                restricts the angle to <= 90.0 degrees.  That is, it treats the
                order of atoms defining a plane as unimportant, and the
                directionality of the plane normals is ignored.
        """
        if atom3 is None:
            return self._distance(atom1, atom2)
        elif atom4 is None:
            return self._angle(atom1, atom2, atom3)
        elif atom5 is None:
            return self._dihedral(atom1, atom2, atom3, atom4)
        else:
            return self._planarAngle(atom1, atom2, atom3, atom4, atom5, atom6,
                                     minangle)
    def _getPosition(self, atom):
        index = int(atom._index)
        pos = self._frame.position[index - 1]
        pos.shape = 1, 3  # getMinimalDifference only works for array with shape (N, 3)
        return pos
    def _distance(self, atom1, atom2):
        pos1 = self._getPosition(atom1)
        pos2 = self._getPosition(atom2)
        dist = self.calcDistance(pos1, pos2)
        return dist[0]
[docs]    def calcDistance(self, pos1, pos2):
        diff = self._frame.getMinimalDifference(pos1, pos2)
        return numpy.sqrt(numpy.sum(diff**2, axis=1)) 
    def _angle(self, atom1, atom2, atom3):
        pos1 = self._getPosition(atom1)
        pos2 = self._getPosition(atom2)
        pos3 = self._getPosition(atom3)
        theta = self.calcAngle(pos1, pos2, pos3)
        return theta[0]
[docs]    def calcAngle(self, pos1, pos2, pos3):
        r21 = self._frame.getMinimalDifference(pos1, pos2)
        r23 = self._frame.getMinimalDifference(pos3, pos2)
        norm = numpy.sqrt(numpy.sum(r21**2, axis=1) * numpy.sum(r23**2, axis=1))
        cos_theta = old_div(numpy.sum(r21 * r23, axis=1), norm)
        # FIXME: is the isnan check sufficient?
        # handle problem when pos1 or pos3 == pos2; make theta = 0.0 in this
        # case
        if numpy.any(numpy.isnan(cos_theta)):
            cos_theta[numpy.isnan(cos_theta)] = 1.0
        # handle numerical roundoff issue where abs(cos_theta) may be
        # greater than 1
        if numpy.any(numpy.abs(cos_theta) > 1.0):
            cos_theta[cos_theta > 1.0] = 1.0
            cos_theta[cos_theta < -1.0] = -1.0
        theta = numpy.arccos(cos_theta) * degrees_per_rad
        return theta 
    def _dihedral(self, atom1, atom2, atom3, atom4):
        pos1 = self._getPosition(atom1)
        pos2 = self._getPosition(atom2)
        pos3 = self._getPosition(atom3)
        pos4 = self._getPosition(atom4)
        theta = self.calcDihedral(pos1, pos2, pos3, pos4)
        return theta[0]
[docs]    def calcDihedral(self, pos1, pos2, pos3, pos4):
        global degrees_per_rad
        r12 = self._frame.getMinimalDifference(pos2, pos1)
        r32 = self._frame.getMinimalDifference(pos2, pos3)
        r42 = self._frame.getMinimalDifference(pos2, pos4)
        c1232 = numpy.cross(r12, r32)
        c4232 = numpy.cross(r42, r32)
        norm = numpy.sqrt(
            numpy.sum(c1232**2, axis=1) * numpy.sum(c4232**2, axis=1))
        cos_theta = old_div(numpy.sum(c1232 * c4232, axis=1), norm)
        # FIXME: is the isnan check sufficient?
        # handle problem when pos1 or pos3 == pos2; make theta = 0.0 in this
        # case
        if numpy.any(numpy.isnan(cos_theta)):
            cos_theta[numpy.isnan(cos_theta)] = 1.0
        # handle numerical roundoff issue where abs(cos_theta) may be
        # greater than 1
        if numpy.any(numpy.abs(cos_theta) > 1.0):
            cos_theta[cos_theta > 1.0] = 1.0
            cos_theta[cos_theta < -1.0] = -1.0
        theta = numpy.arccos(cos_theta) * degrees_per_rad
        # Get the proper sign
        sign_check = numpy.sum(r12 * c4232, axis=1)
        theta[sign_check < 0.0] *= -1.0
        return theta 
    def _planarAngle(self, atom1, atom2, atom3, atom4, atom5, atom6):
        pos1 = self._getPosition(atom1)
        pos2 = self._getPosition(atom2)
        pos3 = self._getPosition(atom3)
        pos4 = self._getPosition(atom4)
        pos5 = self._getPosition(atom5)
        pos6 = self._getPosition(atom6)
        theta = self.calcPlanarAngle(pos1, pos2, pos3, pos4, pos5, pos6)
        return theta[0]
[docs]    def calcPlanarAngle(self,
                        pos1,
                        pos2,
                        pos3,
                        pos4,
                        pos5,
                        pos6,
                        minangle=False):
        global degrees_per_rad
        # First we find the plane equations
        # Plane 1 (atomsels 1, 2, 3)
        r12 = self._frame.getMinimalDifference(pos2, pos1)
        r32 = self._frame.getMinimalDifference(pos2, pos3)
        c1232 = numpy.cross(r12, r32)[0]  # = the plane normal
        # Plane 2 (atomsels 4, 5, 6)
        r45 = self._frame.getMinimalDifference(pos5, pos4)
        r65 = self._frame.getMinimalDifference(pos5, pos6)
        c4565 = numpy.cross(r45, r65)[0]  # = the plane normal
        # Now we find the angle between the planes
        mag_c1232 = numpy.sqrt(numpy.sum(c1232**2, axis=0))
        mag_c4565 = numpy.sqrt(numpy.sum(c4565**2, axis=0))
        if mag_c1232 == 0 or mag_c4565 == 0:
            theta = 0.0
        else:
            c1232dotc4565 = numpy.dot(c1232, c4565)
            cos_theta = old_div(numpy.dot(c1232, c4565),
                                (mag_c1232 * mag_c4565))
            if cos_theta > 1.0:
                cos_theta = 1.0
            elif cos_theta < -1.0:
                cos_theta = -1.0
            theta = numpy.arccos(cos_theta) * degrees_per_rad
            if minangle and theta > 90.0:
                theta = 180 - theta  # Normalize to between 0 and 90
        return theta