Source code for schrodinger.trajectory.trajectory_gui_dir.secondary_structure_data

from schrodinger.infra import mm


class SecondaryStructureData:
    """
    This class provides abiity to assign secondary structure
    property to current frame structure.
    For optimization purpose, it saves original trajectory residue
    and chain information and does quick assignment of those properties
    to each frame structure.
    """

    ALGORITHM_TYPE = 1

    def __init__(self):
        self.clearData()

    def clearData(self):
        self._initialized = False
        self._ref_ct_atom_total = None
        self._residue_index = None
        self._first_residue_in_chain = None
        self._nchains = None
        self._previous_frame_ss = None

    def _initializeData(self, ref_ct):
        """
        Store residue information from reference trajectory ct.

        :type ref_ct: structure.Structure
        :param ref_ct: Reference structure used to get residue information.
        """

        try:
            (self._residue_index, self._first_residue_in_chain,
             self._nchains) = mm.mmss_get_residue_index(ref_ct.handle)
        except mm.MmException:
            return
        self.setFrameSSData(ref_ct)
        self._ref_ct_atom_total = ref_ct.atom_total
        self._initialized = True

    def assign(self, ref_ct, working_ct, varying_atoms_frame):
        """
        Quickly assign secondary structure properties to the given working ct.

        If data is not pre-populated, then it does full assignment.
        In case of varying atoms trajectory, we always do full assignment.

        :type ref_ct: structure.Structure
        :param ref_ct: Reference structure used to get residue information.

        :type working_ct: structure.Structure
        :param working_ct: Structure in which residue information are set.

        :type varying_atoms_frame: bool
        :param varying_atoms_frame: True if ct is varying atoms frame
            trajectory.
        """

        fixed_atoms_frame = not varying_atoms_frame
        if (not self._initialized) and fixed_atoms_frame:
            self._initializeData(ref_ct)
        if self._initialized and fixed_atoms_frame:
            mm.mmss_assign_quick(working_ct.handle, self.ALGORITHM_TYPE,
                                 self._residue_index,
                                 self._first_residue_in_chain, self._nchains)
        else:
            mm.mmss_assign(working_ct.handle, self.ALGORITHM_TYPE)

        changed = False
        if self._previous_frame_ss and fixed_atoms_frame:
            for idx, old_value in enumerate(self._previous_frame_ss, 1):
                new_value = mm.mmct_atom_get_secondary_struct(working_ct, idx)
                if new_value != old_value:
                    changed = True
                    break
        else:
            changed = True
        if changed and fixed_atoms_frame:
            self.setFrameSSData(working_ct)

    def setFrameSSData(self, ct):
        """
        Store given frame ct's SS data.
        """

        atom_total = ct.atom_total
        self._previous_frame_ss = (mm.mmct_atom_get_secondary_struct(ct, at)
                                   for at in range(1, atom_total + 1))