Source code for schrodinger.protein.rotamers

"""
A rotamer library which can be applied to protein sidechains.

A Rotamers object is initialized with an atom from a Structure object.
Exceptions will be generated if there is no rotamer information for the
residue correspoding to that atom or if the residue is incomplete.
Once the Rotamers object is created it's possible to iterate through the
rotamers, to apply any of them or to restore the original geometry.

Copyright Schrodinger LLC. All rights reserved.

"""

# $Revision: 1.7 $
# $Date: 2008/12/11 16:03:58 $

#Contributors: Quentin McDonald

from schrodinger.infra import mm
from schrodinger.infra import mminit


class _Rotamer(object):
    """
    A single instance of a rotamer - represents a set of CHI angle
    values
    """

    def __init__(self, rl, idx):
        """
        Initialize the rotamer object with the Rotamer library and
        an the index of which rotamer this is:
        """
        self._rotamer_lib = rl
        self._index = idx

    def __index__(self):
        """
        Return the index - starting from 1 as the user sees them
        """
        return (self._index + 1)

    def apply(self):
        """
        Apply this rotamer to the structure - the sidechain angles will be
        set to the CHI values from the rotamer library
        """

        self._rotamer_lib._apply(self._index)

    def percentage(self):
        """
        Return the percentage for this rotamer
        """
        return mm.mmrotamer_get_population(self._rotamer_lib.resname,
                                           self._index)

    def chiAngles(self):
        """
        Return a list of the Chi angles for this rotamer.
        """
        ret_list = []

        for ichi in range(
                mm.mmrotamer_get_num_chi_angles(self._rotamer_lib.resname)):

            # SHARED-6000: see comment in Rotamers._apply
            try:
                angle = mm.mmrotamer_get_rotamer_chi(self._rotamer_lib.resname,
                                                     self._index, ichi)
            except mm.MmException as e:
                if e.rc != mm.MMROTAMER_INVALID_CHI_ANGLE:
                    raise
            else:
                ret_list.append(angle)

        return ret_list


class _RotamerIterator(object):
    """
    Container for the rotamers which are associated with a particular
    residue type
    """

    def __init__(self, rl):
        """
        Initialize the iterator with the rotamer library
        """
        self._rotamer_lib = rl

    def __len__(self):
        """
        Return the number of rotamers which are available:
        """
        return self._rotamer_lib.num_rotamers

    def __getitem__(self, rotamer_num):
        """
        Return the rotamer which is associated with rotamer_num. Index
        begins at 1
        """
        if (rotamer_num < 1 or rotamer_num > self._rotamer_lib.num_rotamers):
            raise KeyError("Rotamer key out of range (starts at 1): %d" %
                           rotamer_num)
        return _Rotamer(self._rotamer_lib, rotamer_num - 1)

    def __iter__(self):
        """
        Iterate through the available rotamers for the structure
        """
        for rot_num in range(self._rotamer_lib.num_rotamers):
            yield _Rotamer(self._rotamer_lib, rot_num)

        # Initialize the mmrotamer library via an Initializer object. This is to avoid
        # the use of a __del__ method in the Rotamers class, since there exists a
        # reference cycle in this code. The cycle is
        #   Rotamers -> _RotamerIterator -> _Rotamer -> Rotamers


_mmrotamer_initializer = mminit.Initializer([mm.mmrotamer_initialize],
                                            [mm.mmrotamer_terminate])


class Rotamers(object):
    """
    Manages a rotamer library for a given residue type.

    """

    def __init__(self, st, atom):
        """
        Initialize the Rotamer Library with respect to a given atom
        in the input structure 'st'.

        """

        # Start by setting up some basic information about this residue
        self.resname = st.atom[int(atom)].pdbres

        try:
            self.num_chi_angles = mm.mmrotamer_get_num_chi_angles(self.resname)
        except mm.MmException as e:
            if e.rc == mm.MMROTAMER_RES_NOT_FOUND:
                raise RuntimeError(
                    "There is no rotamer information for residue %s" %
                    self.resname)

        try:
            self.num_rotamers = mm.mmrotamer_get_num_rotamers(self.resname)
        except mm.MmException as e:
            if e.rc == mm.MMROTAMER_RES_NOT_FOUND:
                raise RuntimeError(
                    "There is no rotamer information for residue %s" %
                    self.resname)

        self._st = st

        res = st.atom[int(atom)].getResidue()
        res_atoms = res.atom

        # Locate the atoms in the residue
        self.chi_atoms = []
        for i in range(self.num_chi_angles):
            dihedral_atoms = self.findChiAtoms(res, i)
            self.chi_atoms.append(dihedral_atoms)

        # Make a list of original coordinates for each atom in this residue:
        self.chi_coords = {}
        for resatom in res.atom:
            self.chi_coords[resatom.index] = resatom.xyz

        # Initialize the rotamer iterator
        self._rotamers = _RotamerIterator(self)

    def findChiAtoms(self, res, chi_index):
        """
        Return a list of atom numbers for the given Chi dihedral of the residue.
        """

        at_names = mm.mmrotamer_get_chi_def(self.resname, chi_index)

        dihedral_anums = []
        for i, pdbname in enumerate(at_names):
            anum = None
            matched_atom = res.getAtomByPdbName(pdbname)
            if matched_atom:
                anum = matched_atom.index
            elif i == 3 and 'H' in pdbname:
                # The rotamer library uses hard-coded names for all atoms
                # including hydrogens (See chidef.res of the mmrotamer
                # library). For Prime/IFD we don't require hydrogens to
                # have these exact names, so this code will detect the
                # hydrogen by the name of the heavy atom that the hydrogen
                # is bonded to.
                # Hydrogen is always the last atom in the list returned by
                # mmrotamer library (for bonds that contain hydrogens).
                heavy_atom = dihedral_anums[2]
                for neighbor in self._st.atom[heavy_atom].bonded_atoms:
                    if neighbor.element == 'H':
                        anum = neighbor.index
                        break
            if anum is None:
                # Didn't find one of the sidechain atom so raise exception
                raise RuntimeError('Could not find atom "%s" of %s for chi%d' %
                                   (pdbname, res, chi_index + 1))
            dihedral_anums.append(anum)
        return dihedral_anums

    def _apply(self, idx):
        """
        A private function which is used to apply the specified rotamer
        index (idx) to the structure
        """

        for chi in range(self.num_chi_angles):
            state_angles = self.chi_atoms[chi]
            try:
                angle = mm.mmrotamer_get_rotamer_chi(self.resname, idx, chi)
            except mm.MmException as e:
                # Related to EV 74929: Ignore any CHI angles for which
                # we don't actually have rotamer data, O-H or SER for example:
                if e.rc != mm.MMROTAMER_INVALID_CHI_ANGLE:
                    raise
            self._st.adjust(angle, state_angles[0], state_angles[1],
                            state_angles[2], state_angles[3])

    def restore(self):
        """
        Restore the sidechain back to the original coordinates
        """
        for anum, xyz in self.chi_coords.items():
            self._st.atom[anum].xyz = xyz

    def _getRotamerIterator(self):
        return self._rotamers

    _doc = """
    A list of all the rotamers available for the current residue.
    """

    rotamers = property(_getRotamerIterator, doc=_doc)