Source code for schrodinger.test.build_test_structures

"""
Some convenience functions to quickly and reproducibly create structures for
testing.
"""

import random

import schrodinger.protein.helix
from schrodinger import adapter
from schrodinger import structure
from schrodinger.infra import mm
from schrodinger.infra import mminit
from schrodinger.structutils import build
from schrodinger.utils import mmutil


def structure_from_smiles(smiles: str):
    """
    Create a `structure.Structure` from a SMILES string.  The structure is not
    energy minimized, so it will probably be distorted.

    :param smiles: SMILES string representation of the desired molecule.
    :rtype: structure.Structure
    :return: CT of the input SMILES string.
    """
    if mmutil.feature_flag_is_enabled(mmutil.USE_RDKIT_FOR_SMILESSTRUCTURE):
        return adapter.to_structure(smiles,
                                    adapter.Generate2DCoordinates.Enable)

    mm.mmsmilesmae_initialize_lic(mm.error_handler)
    try:
        distorted_ct = mm.mmsmiles_get_distorted_ct(smiles)
    finally:
        mm.mmsmilesmae_terminate()
    return structure.Structure(distorted_ct)


def create_fragment(resname, restype="organic", assign_atom_names=False):
    """
    Use mmfrag library to create arbitrary fragments.  See the mmfrag
    documentation for a list of possible fragments that can are available
    and which "type" they belong to.  By default the pdb atom names for all
    atoms are those set by mmfrag (usually blank), but setting
    assign_atom_names to True will set unique atom names in the manner
    described in _assignAtomNames.

    See .bld files in mmfrag for resname and restype.

    """
    with mminit.mmfrag() as mmfrag:
        with mmfrag.new(restype) as mmfrag_handle:
            mm.mmfrag_set_fragment_name(mmfrag_handle, resname)
            resname, ict = mm.mmfrag_duplicate_fragment(mmfrag_handle,
                                                        mm.MM_ACTIVE)
    ct = structure.Structure(ict)
    if assign_atom_names:
        _assign_atom_names(ct)
    return ct


def _assign_atom_names(ct):
    """
    Assign unique atom names by residue for a structure.  Atom names
    consist of the element (2 character) followed by a unique number for
    that element.  (ie C1,C2,H1)
    """
    for res in ct.residue:
        element = {}
        for atom in res.atom:
            if (atom.element not in element):
                element[atom.element] = 0
            element[atom.element] += 1
            atom.pdbname = "%2s%-2d" % (atom.element, element[atom.element])


def shuffle_ct_atoms(ct):
    """
    Randomize the atom-ordering within a ct.  This is useful for testing
    to make sure there is no atom-order dependence for a given function.
    The input ct is not changed and the reordered output ct is returned.

    :type ct: structure.Structure
    :param ct: Structure to be shuffled. Atom order of this structure will be
               preserved.
    :rtype: structure.Structure
    :return: Structure with atom order shuffled.
    """

    order_map = list(range(1, ct.atom_total + 1))
    random.shuffle(order_map)

    shuffled_ct = build.reorder_atoms(ct, order_map)

    return shuffled_ct


def build_protein_helix(sequence):
    """
    Uses schrodinger.protein.helix to create an Alpha-helix protein based on the
    input sequence.

    :ptype sequence: str
    :param sequence: String of single letter amino acid names.
    """
    return schrodinger.protein.helix.process_sequence(sequence)