Source code for schrodinger.application.matsci.pbcutils

"""
Utility functions and classes for pbcs

Copyright Schrodinger, LLC. All rights reserved.
"""

from schrodinger.application.desmond import constants as dconst
from schrodinger.application.matsci import msutils
from schrodinger.application.matsci import msprops
from schrodinger.infra import mm
from schrodinger.structutils import analyze


def get_disorder_groups(struct):
    """
    Get atoms that have s_cif_disorder_group set and occupancies are not 0 or 1.

    :param structure.Structure: Input structure

    :rtype: dict or None
    :return: Dict with keys disorder group labels, values list of atom indices
        or None if nothing is present in the structure
    """

    if not msutils.has_atom_property(struct, mm.CIF_DISORDER_GROUP):
        return

    groups = {}
    occupancies = {}

    # asl "?*" search for non-empty string
    for idx in analyze.evaluate_asl(struct,
                                    'atom.%s ?*' % mm.CIF_DISORDER_GROUP):
        atom = struct.atom[idx]
        group = atom.property[mm.CIF_DISORDER_GROUP]
        groups.setdefault(group, []).append(idx)

        # Keep only one non-zero occupancy (if present) for each group
        occupancies[group] = (occupancies.get(group) or
                              atom.property[mm.M2IO_DATA_PDB_OCCUPANCY])

    # Delete groups with occupancy 0 or 1
    for group, occupancy in occupancies.items():
        if occupancy in [0, 1]:
            groups.pop(group)

    return groups if groups else None


def remove_pbc_props(struct):
    """
    Remove the pbc properties of the given structure.

    :param schrodinger.structure.Structure struct: crystal structure from which
        all the pbc properties are to be removed
    """

    DESMOND_BOX_SHAPE = 's_m_BC_Box_Shape'
    DESMOND_BOX_A = 'r_m_BC_Box_A'
    DESMOND_BOX_B = 'r_m_BC_Box_B'
    DESMOND_BOX_C = 'r_m_BC_Box_C'
    DESMOND_BOX_ALPHA = 'r_m_BC_Box_Alpha'
    DESMOND_BOX_BETA = 'r_m_BC_Box_Beta'
    DESMOND_BOX_GAMMA = 'r_m_BC_Box_Gamma'
    LATTICE_A = 's_matsci_Lattice_Vector_A/(3*Ang.)'
    LATTICE_B = 's_matsci_Lattice_Vector_B/(3*Ang.)'
    LATTICE_C = 's_matsci_Lattice_Vector_C/(3*Ang.)'

    props_to_remove = [
        mm.M2IO_PDB_CRYSTAL_SPACE_GROUP, mm.M2IO_PDB_CRYSTAL_Z,
        msprops.UNIT_CELL_VOLUME_KEY, msprops.UNIT_CELL_DENSITY_KEY,
        msprops.UNIT_CELL_FORMULA_KEY, msprops.SPACE_GROUP_ID_KEY,
        msprops.PBC_POSITION_KEY, DESMOND_BOX_SHAPE, DESMOND_BOX_A,
        DESMOND_BOX_B, DESMOND_BOX_C, DESMOND_BOX_ALPHA, DESMOND_BOX_BETA,
        DESMOND_BOX_GAMMA, LATTICE_A, LATTICE_B, LATTICE_C
    ] + list(dconst.SIM_BOX) + list(msprops.PDB_BOX_PROPS)

    msutils.remove_properties(struct, props=props_to_remove)