Source code for schrodinger.application.prepwizard2.diagnostics

# Note: this file must not import any gui-dependent modules (QtGui, QtWidgets,
# maestro_ui) so the tasks can run without display dependency
from typing import List
from typing import Tuple

from schrodinger import structure
from schrodinger.application.prepwizard2 import prepare
from schrodinger.infra import mm
from schrodinger.models import parameters
from schrodinger.structutils import analyze
from schrodinger.structutils import measure

OVERLAP_CUTOFF_DIST = 0.8


class Problems(parameters.CompoundParam):
    invalid_types: List[Tuple[int, int, int]]
    missing: List[Tuple[int, int, int, bool]]
    overlapping: List[Tuple[int, int]]
    alternates: List[Tuple[int, float]]


def get_problems(st):
    return Problems(invalid_types=get_atoms_with_improper_atom_types(st),
                    missing=find_missing_res_atoms(st),
                    overlapping=get_overlapping_atoms(st),
                    alternates=find_residues_with_multiple_occupancies(st))


def get_atoms_with_improper_atom_types(st):
    """
    Returns a list of (atomnum, expected_valence, actual_valence) items.
    expected_valence is None for atom types of 150 or above.
    """

    # Ev:79670 Check atom valences:
    st.retype()
    bad_atoms_list = []
    metal_atoms = analyze.evaluate_asl(st, 'metals')
    for atom in st.atom:
        if atom.index in metal_atoms:
            continue

        actual_valence = 0
        num_zobs = 0
        for bond in atom.bond:
            actual_valence += bond.order
            if bond.order == 0:
                num_zobs += 1

        atype = atom.atom_type
        if atype >= 150:
            # These are generalized atom types, we don't know the expected
            # valence. Displaying these values to the user is unhelpful,
            # since there is no "expected" value.
            continue

        # Ev:95680 Check for the valence (except metals):
        expected_valence = mm.mmat_get_valence(atype)

        actual_charge = atom.formal_charge
        expected_charge = mm.mmat_get_formal_charge(atype)

        if actual_valence == 0 and num_zobs > 0 and expected_valence > 0:
            # This atom has zero-order bonds, ignore the valence problem.
            continue

        if actual_valence - actual_charge != expected_valence - expected_charge:
            bad_atoms_list.append((atom.index, expected_valence, actual_valence,
                                   expected_charge, actual_charge))

    return bad_atoms_list


def find_missing_res_atoms(st):
    """
    Seaches the specified structure for missing residue atoms, and
    returns a list of:
    (atomnum, num_heavy, expected_heavy, missing_only_sidechains)

    This list DOES include residues with only backbone atoms missing.
    """

    # Ev:96128 Use a new mechanism to determine which residues have missing
    # atoms. This method does NOT rely on the i_m_pdb_convert_problem
    # property, which may be out-of-date:

    missing_res = []
    for res in st.residue:
        if not res.hasMissingAtoms():
            continue

        # This function will return False if any backbone atoms are missing:
        missing_only_sidechains = prepare.does_res_have_missing_side_chains(res)

        pdbres = res.pdbres
        heavy_elements = []
        for atom in res.atom:
            if atom.atomic_number > 1:
                heavy_elements.append(atom.element)

        a = res.atom[1]
        num_heavy = len(heavy_elements)
        expected_heavy = structure._res_sizes[pdbres]

        missing_res.append(
            (a.index, num_heavy, expected_heavy, missing_only_sidechains))

    return missing_res


def get_overlapping_atoms(st):
    return measure.get_close_atoms(st, OVERLAP_CUTOFF_DIST)


def find_residues_with_multiple_occupancies(st):
    """
    Find residues that have at least one atom with multiple occupancies
    (alternate states).

    :param st: Protein structure
    :type st: `structure.Structure`

    :return: List of (Atom index, average occupancy)
    :rtype: list
    """
    residues = []
    for res in st.residue:
        occupancies = []
        for atom in res.atom:
            atom_occ = atom.property.get("r_m_pdb_occupancy")
            alt_loc = atom.property.get('s_pdb_altloc_chars')
            if atom_occ is not None and alt_loc and len(alt_loc) > 1:
                occupancies.append(atom_occ)

        if occupancies:
            average_occ = sum(occupancies) / len(occupancies)
            a = res.atom[1]
            residues.append((a.index, average_occ))
    return residues