Source code for prody.measure.transform

# -*- coding: utf-8 -*-
""" This module defines a class for identifying contacts."""

import numpy as np

from prody import LOGGER
from prody.atomic import AtomPointer, AtomMap
from prody.utilities import importLA, checkWeights

from .measure import calcCenter

linalg = importLA()

__all__ = ['Transformation', 'applyTransformation', 'alignCoordsets',
           'calcRMSD', 'calcTransformation', 'superpose',
           'moveAtoms', 'wrapAtoms',
           'printRMSD']


class Transformation(object):

    """A class for storing a transformation matrix."""

    __slots__ = ['_matrix']

    def __init__(self, *args):
        """Either 4x4 transformation *matrix*, or *rotation* matrix and
        *translation* vector must be provided at instantiation."""

        nargs = len(args)
        if nargs == 1:
            self.setMatrix(args[0])
        elif nargs == 2:
            self._matrix = np.eye(4)
            self.setRotation(args[0])
            self.setTranslation(args[1])
        else:
            raise ValueError('one or two array must be provided as arguments')

    def getRotation(self):
        """Returns rotation matrix."""

        if self._matrix is not None:
            return self._matrix[:3, :3]

    def setRotation(self, rotation):
        """Set rotation matrix."""

        if not isinstance(rotation, np.ndarray):
            raise TypeError('rotation must be an ndarray')
        elif rotation.shape != (3, 3):
            raise ValueError('rotation must be a 3x3 array')
        self._matrix[:3, :3] = rotation

    def getTranslation(self):
        """Returns translation vector."""

        if self._matrix is not None:
            return self._matrix[:3, 3]

    def setTranslation(self, translation):
        """Set translation vector."""

        if not isinstance(translation, np.ndarray):
            raise TypeError('translation must be an ndarray')
        elif translation.shape != (3,):
            raise ValueError('translation must be an ndarray of length 3')
        self._matrix[:3, 3] = translation

    def getMatrix(self):
        """Returns a copy of the 4x4 transformation matrix whose top left is
        rotation matrix and last column is translation vector."""

        if self._matrix is not None:
            return self._matrix.copy()

    def setMatrix(self, matrix):

        if not isinstance(matrix, np.ndarray):
            raise TypeError('matrix must be an ndarray')
        elif matrix.shape != (4, 4):
            raise ValueError('matrix must have shape (4,4)')
        self._matrix = matrix

    def apply(self, atoms):
        """Apply transformation to *atoms*, see :func:`applyTransformation`
        for details."""

        return applyTransformation(self, atoms)


def calcTransformation(mobile, target, weights=None):
    """Returns a :class:`Transformation` instance which, when applied to the
    atoms in *mobile*, minimizes the weighted RMSD between *mobile* and
    *target*.  *mobile* and *target* may be NumPy coordinate arrays, or
    :class:`.Atomic` instances, e.g. :class:`.AtomGroup`, :class:`.Chain`,
    or :class:`.Selection`."""

    if not isinstance(mobile, np.ndarray):
        try:
            mob = mobile._getCoords()
        except AttributeError:
            raise TypeError('mobile must be a numpy array or an object '
                            'with getCoords method')
    else:
        mob = mobile
    if not isinstance(target, np.ndarray):
        try:
            tar = target._getCoords()
        except AttributeError:
            raise TypeError('target must be a numpy array or an object '
                            'with getCoords method')
    else:
        tar = target

    if mob.shape != tar.shape:
        raise ValueError('reference and target coordinate arrays '
                         'must have same number of atoms')

    if mob.shape[1] != 3:
        raise ValueError('reference and target must be coordinate arrays')
    
    if weights is None:
        if isinstance(mobile, AtomMap):
            LOGGER.warn('mobile is an AtomMap instance, consider assign weights=mobile.getFlags("mapped") '
                        'if there are dummy atoms in mobile')

        if isinstance(target, AtomMap):
            LOGGER.warn('target is an AtomMap instance, consider assign weights=target.getFlags("mapped") '
                        'if there are dummy atoms in target')

    if weights is not None:
        weights = checkWeights(weights, mob.shape[0])

    return Transformation(*getTransformation(mob, tar, weights))


def getTransformation(mob, tar, weights=None):

    if weights is None:
        mob_com = mob.mean(0)
        tar_com = tar.mean(0)
        mob = mob - mob_com
        tar = tar - tar_com
        matrix = np.dot(mob.T, tar)
    else:
        weights_sum = weights.sum()
        weights_dot = np.dot(weights.T, weights)
        mob_com = (mob * weights).sum(axis=0) / weights_sum
        tar_com = (tar * weights).sum(axis=0) / weights_sum
        mob = mob - mob_com
        tar = tar - tar_com
        matrix = np.dot((mob * weights).T, (tar * weights)) / weights_dot

    U, _, Vh = linalg.svd(matrix)
    d = np.sign(linalg.det(np.dot(U, Vh)))
    Id = np.array([[1, 0, 0],
                   [0, 1, 0],
                   [0, 0, d]])
    rotation = np.dot(Vh.T, np.dot(Id, U.T))

    return rotation, tar_com - np.dot(mob_com, rotation.T)


def applyTransformation(transformation, atoms):
    """Returns *atoms* after applying *transformation*.  If *atoms*
    is a :class:`.Atomic` instance, it will be returned after
    *transformation* is applied to its active coordinate set.  If
    *atoms* is an :class:`.AtomPointer` instance, *transformation*
    will be applied to the corresponding coordinate set in the
    associated :class:`.AtomGroup`."""

    coords = None
    ag = None
    if isinstance(atoms, np.ndarray):
        atoms = np.atleast_2d(atoms)
        if atoms.shape[1] != 3:
            raise ValueError('atoms must be a 3-d coordinate array')
        coords = atoms
        atoms = None
    else:
        if isinstance(atoms, AtomPointer):
            ag = atoms.getAtomGroup()
            acsi = ag.getACSIndex()
            ag.setACSIndex(atoms.getACSIndex())
            coords = ag._getCoords()
        else:
            try:
                coords = atoms._getCoords()
            except AttributeError:
                raise TypeError('atoms must be a Atomic instance')

    if atoms is None:
        return _applyTransformation(transformation, coords)
    else:
        if ag is None:
            atoms.setCoords(_applyTransformation(transformation, coords))
        else:
            ag.setCoords(_applyTransformation(transformation, coords))
            ag.setACSIndex(acsi)
        return atoms


def _applyTransformation(t, coords):

    return np.dot(coords, t.getRotation().T) + t.getTranslation()


def superpose(mobile, target, weights=None):
    """Returns *mobile*, after its RMSD minimizing superposition onto *target*,
    and the transformation that minimizes the RMSD."""

    t = calcTransformation(mobile, target, weights)
    result = applyTransformation(t, mobile)
    return (result, t)


def moveAtoms(atoms, **kwargs):
    """Move *atoms* *to* a new location or *by* an offset.  This method will
    change the active coordinate set of the *atoms*.  Note that only one of
    *to* or *by* keyword arguments is expected.

    Move protein so that its centroid is at the origin, ``[0., 0., 0.]``:

    .. ipython:: python

       from prody import *
       from numpy import ones, zeros
       protein = parsePDB('1ubi')
       calcCenter(protein).round(3)
       moveAtoms(protein, to=zeros(3))
       calcCenter(protein).round(3)

    Move protein so that its mass center is at the origin:

    .. ipython:: python

       protein.setMasses(ones(len(protein)))
       protein.carbon.setMasses(12)
       protein.nitrogen.setMasses(14)
       protein.oxygen.setMasses(16)
       moveAtoms(protein, to=zeros(3), weights=protein.getMasses())
       calcCenter(protein, weights=protein.getMasses()).round(3)

    Move protein so that centroid of Cα atoms is at the origin:

    .. ipython:: python

       moveAtoms(protein.ca, to=zeros(3), ag=True)
       calcCenter(protein).round(3)
       calcCenter(protein.ca).round(3)

    Move protein by 10 A along each direction:

    .. ipython:: python

       moveAtoms(protein, by=ones(3) * 10)
       calcCenter(protein).round(3)
       calcCenter(protein.ca).round(3)


    :arg by: an offset array with shape ``([1,] 3)`` or ``(n_atoms, 3)`` or
        a transformation matrix with shape ``(4, 4)``
    :type by: :class:`numpy.ndarray`

    :arg to: a point array with shape ``([1,] 3)``
    :type to: :class:`numpy.ndarray`

    :arg ag: when *atoms* is a :class:`.AtomSubset`, apply translation vector
        (*to*) or transformation matrix to the :class:`.AtomGroup`,
        default is **False**
    :type ag: bool

    :arg weights: array of atomic weights with shape ``(n_atoms[, 1])``
    :type weights: :class:`numpy.ndarray`

    When *to* argument is passed, :func:`.calcCenter` function is used to
    calculate centroid or mass center."""

    point = kwargs.pop('to', None)
    if point is None:
        offset = kwargs.pop('by', None)
        if offset is None:
            raise TypeError('moveAtoms() expects one of {0} or {1} '
                            'arguments'.format(repr('to'), repr('by')))

        try:
            shape = offset.shape
        except AttributeError:
            raise TypeError('by must be a numpy array')

        if shape == (4, 4):
            if kwargs.pop('ag', False):
                try:
                    atoms = atoms.getAtomGroup()
                except AttributeError:
                    # itself must be an AtomGroup
                    pass
            try:
                coords = atoms._getCoords()
            except AttributeError:
                try:
                    coords = atoms.getCoords()
                except AttributeError:
                    raise TypeError('atoms must be an Atomic instance')

            coords = np.dot(coords, offset[:3, :3])
            coords += offset[3, :3]

            atoms.setCoords(coords)
        else:

            try:
                coords = atoms._getCoords()
            except AttributeError:
                try:
                    coords = atoms.getCoords()
                except AttributeError:
                    raise TypeError('atoms must be an Atomic instance')

            try:
                natoms = atoms.numAtoms()
            except AttributeError:
                raise TypeError('atoms must be an Atomic instance')
            if shape == (3,) or shape == (1, 3) or shape == (natoms, 3):
                atoms.setCoords(coords + offset)
            else:
                raise ValueError('by.shape is not valid')

    else:
        try:
            shape = point.shape
        except AttributeError:
            raise TypeError('to must be a numpy array')
        if shape != (3,) and shape != (1, 3):
            raise ValueError('to.shape must be ([1,] 3)')

        center = calcCenter(atoms, weights=kwargs.pop('weights', None))
        offset = point - center
        if kwargs.pop('ag', False):
            try:
                atoms = atoms.getAtomGroup()
            except AttributeError:
                # itself must be an AtomGroup
                pass

        try:
            coords = atoms._getCoords()
        except AttributeError:
            try:
                coords = atoms.getCoords()
            except AttributeError:
                raise TypeError('atoms must be an Atomic instance')
        atoms.setCoords(coords + offset)


def calcRMSD(reference, target=None, weights=None):
    """Returns root-mean-square deviation (RMSD) between reference and target
    coordinates.

    .. ipython:: python

       ens = loadEnsemble('p38_X-ray.ens.npz')
       ens.getRMSDs()
       calcRMSD(ens)
       calcRMSD(ens.getCoords(), ens.getCoordsets(), ens.getWeights())"""

    if isinstance(reference, np.ndarray):
        ref = reference
    else:
        try:
            ref = reference._getCoords()
        except AttributeError:
            raise TypeError('reference must be a numpy array or an object '
                            'with getCoords method')
        if target is None:
            try:
                target = reference._getCoordsets()
            except AttributeError:
                pass
        if weights is None:
            try:
                weights = reference._getWeights()
            except AttributeError:
                pass
    if ref.ndim != 2 or ref.shape[1] != 3:
        raise ValueError('reference must have shape (n_atoms, 3)')

    if isinstance(target, np.ndarray):
        tar = target
    else:
        try:
            tar = target._getCoords()
        except AttributeError:
            raise TypeError('target must be a numpy array or an object '
                            'with getCoords method')
    if tar.ndim not in (2, 3) or tar.shape[-1] != 3:
        raise ValueError('target must have shape ([n_confs,] n_atoms, 3)')

    if ref.shape != tar.shape[-2:]:
        raise ValueError('reference and target arrays must have the same '
                         'number of atoms')

    if weights is not None:
        n_atoms = len(ref)
        n_csets = 1 if tar.ndim == 2 else len(tar)
        weights = checkWeights(weights, n_atoms, n_csets)
    return getRMSD(ref, tar, weights)


def getRMSD(ref, tar, weights=None):
    if weights is None:
        divByN = 1.0 / ref.shape[0]
        if tar.ndim == 2:
            return np.sqrt(((ref-tar) ** 2).sum() * divByN)
        else:
            rmsd = np.zeros(len(tar))
            for i, t in enumerate(tar):
                rmsd[i] = ((ref-t) ** 2).sum()
            return np.sqrt(rmsd * divByN)
    else:
        if tar.ndim == 2:
            return np.sqrt((((ref-tar) ** 2) * weights).sum() *
                           (1. / weights.sum()))
        else:
            rmsd = np.zeros(len(tar))
            if weights.ndim == 2:
                for i, t in enumerate(tar):
                    rmsd[i] = (((ref-t) ** 2) * weights).sum()
                return np.sqrt(rmsd * (1. / weights.sum()))
            else:
                for i, t in enumerate(tar):
                    rmsd[i] = (((ref-t) ** 2) * weights[i]).sum()
                return np.sqrt(rmsd / weights.sum(1).flatten())


def printRMSD(reference, target=None, weights=None, log=True, msg=None):
    """Print RMSD to the screen.  If *target* has multiple coordinate sets,
    minimum, maximum and mean RMSD values are printed.  If *log* is **True**
    (default), RMSD is written to the standard error using package logger,
    otherwise standard output is used.  When *msg* string is given, it is
    printed before the RMSD value.  See also :func:`calcRMSD` function. """

    if log:
        write = LOGGER.info
    else:
        import sys
        write = lambda line: sys.stdout.write(line + '\n')
    msg = msg or ''
    if msg and msg[-1] != ' ':
        msg += ' '
    rmsd = calcRMSD(reference, target, weights)
    if isinstance(rmsd, np.ndarray) and len(rmsd) > 1:
        write(msg + 'RMSD: min={0:.2f}, max={1:.2f}, mean={2:.2f}'
                    .format(rmsd.min(), rmsd.max(), rmsd.mean()))
    else:
        if isinstance(rmsd, np.ndarray):
            rmsd = rmsd[0]
        write(msg + 'RMSD: {0:.2f}'.format(rmsd))


def alignCoordsets(atoms, weights=None):
    """Returns *atoms* after superposing coordinate sets onto its active
    coordinate set.  Transformations will be calculated for *atoms* and
    applied to its :class:`.AtomGroup`, when applicable.  Optionally,
    atomic *weights* can be passed for weighted superposition."""

    try:
        acsi, n_csets = atoms.getACSIndex(), atoms.numCoordsets()
    except AttributeError:
        raise TypeError('atoms must have type Atomic, not {0}'
                        .format(type(atoms)))
        if n_csets < 2:
            LOGGER.warning('{0} contains fewer than two coordinate sets, '
                           'alignment was not performed.'.format(str(atoms)))
            return

    try:
        ag = atoms.getAtomGroup()
    except AttributeError:
        ag = atoms
    agacsi = ag.getACSIndex()

    tar = atoms._getCoords()
    for i in range(n_csets):
        if i == acsi:
            continue
        atoms.setACSIndex(i)
        ag.setACSIndex(i)
        calcTransformation(atoms, tar, weights).apply(ag)
    atoms.setACSIndex(acsi)
    ag.setACSIndex(agacsi)
    return atoms


def wrapAtoms(frame, unitcell=None, center=np.array([0., 0., 0.])):
    """Wrap atoms into an image of the system simulated under periodic boundary
    conditions. When *frame* is a :class:`.Frame`, unitcell information will be
    retrieved automatically.

    .. note::
       This function will wrap all atoms into the specified periodic image, so
       covalent bonds will be broken.

    :arg frame: a frame instance or a coordinate set
    :type frame: :class:`.Frame`, :class:`.AtomGroup`, :class:`numpy.ndarray`

    :arg unitcell: orthorhombic unitcell array with shape (3,)
    :type unitcell: :class:`numpy.ndarray`

    :arg center: coordinates of the center of the wrapping cell, default is
        the origin of the Cartesian coordinate system
    :type center: :class:`numpy.ndarray`"""

    try:
        coords = frame._getCoords()
    except AttributeError:
        coords = frame
    else:
        try:
            frame.getAtomGroup()
        except AttributeError:
            pass
        else:
            raise TypeError('frame must be a Frame, AtomGroup, or numpy array,'
                            ' not a ' + str(type(frame)))

    if unitcell is None:
        try:
            unitcell = frame.getUnitcell()[:3]
        except AttributeError:
            raise TypeError('unitcell information must be provided')

    half = unitcell / 2
    ucmin = center - half
    ucmax = center + half
    for axis in range(3):
        xyz = coords[:, axis]
        which = (xyz < ucmin[axis]).nonzero()[0]
        while len(which):
            coords[which, axis] += unitcell[axis]
            which = which[coords[which, axis] < ucmin[axis]]
        which = (xyz > ucmax[axis]).nonzero()[0]
        while len(which):
            coords[which, axis] -= unitcell[axis]
            which = which[coords[which, axis] > ucmax[axis]]
    return frame