Source code for tmol.io._canonical_form

import attr
import torch
import numpy
from typing import Optional

from tmol.types import (
    Tensor,
    NDArray,
)




@attr.s(auto_attribs=True, frozen=False, slots=True)
class CanonicalForm:
    """This class holds the data that describe a (stack of) structure(s) in a poised, ready-to-use state.

    This datastructure holds the information necessary to determine the chemical
    identities of the residues in the structure(s), which may be under-determined
    from tmol's perspective by the source of the structure (e.g. OpenFold does not
    explicitly model termini). The atoms that are present are represented with
    non-NaN coordinates in the `coords` array; the order in which those atoms appear
    is given by a particular CanonicalOrdering object.

    The datastructure also holds convenience information such as author-provided
    residue labels (ints), chain labels (strings) & insertion codes (strings) as well
    as the occupancy and B-factor of each atom. These are not strictly necessary
    but are often useful when processing structures.
    """

    # n_poses x max_n_res
    chain_id: Tensor[torch.int64][:, :]
    # n_poses x max_n_res
    res_types: Tensor[torch.int64][:, :]
    # n_poses x max_n_res x max_n_canonical_atoms x 3
    coords: Tensor[torch.float32][:, :, :, 3]
    # n_poses x max_n_res
    res_labels: NDArray[int][:, :]
    # n_poses x max_n_res
    residue_insertion_codes: NDArray[object][:, :]
    # n_poses x max_n_res
    chain_labels: NDArray[object][:, :]
    # n_poses x max_n_res x max_n_canonical_atoms
    atom_occupancy: Optional[NDArray[numpy.float32][:, :, :]]
    # n_poses x max_n_res x max_n_canonical_atoms
    atom_b_factor: Optional[NDArray[numpy.float32][:, :, :]]
    # n_disulfides x 3
    disulfides: Optional[Tensor[torch.int64][:, 3]]
    # n_poses x max_n_res x 2
    res_not_connected: Optional[Tensor[torch.bool][:, :, 2]]
    # n_cyclic_closures x 3
    cyclic_bonds: Optional[Tensor[torch.int64][:, 3]] = None
    # n_covalent_bonds x 5: pose, res1, canonical atom1, res2, canonical atom2
    covalent_bonds: Optional[Tensor[torch.int64][:, 5]] = None
    # n_metals x 3: pose, metal res, index into GEOMETRY_NAMES
    metal_sites: Optional[Tensor[torch.int64][:, 3]] = None
    # n_filled_sites x 5: pose, metal res, site, donor res, donor canonical atom
    metal_coordination: Optional[Tensor[torch.int64][:, 5]] = None
    # n_poses x max_n_res: for a metal split out of a component, (res label,
    # insertion code, res name, atom name) of where it sat; None elsewhere
    metal_origins: Optional[NDArray[object][:, :]] = None
    # n_poses x max_n_res structured array of the input's RESIDUE_ANNOTATIONS
    residue_annotations: Optional[numpy.ndarray] = None
    # n_poses x max_n_res: the res_type_variant the input's protonation state
    # selects, -1 where it leaves the variant to tmol
    protonation_variants: Optional[Tensor[torch.int64][:, :]] = None

    def __iter__(self):
        return iter(attr.astuple(self, recurse=False))



    def as_dict(self):
        """Constructor keyword arguments sharing the original tensors and arrays."""
        values = attr.asdict(self, recurse=False)
        values["res_ins_codes"] = values.pop("residue_insertion_codes")
        return values