Source code for prolif.ifp

"""
Storing interactions --- :mod:`prolif.ifp`
==========================================
"""

from collections import UserDict
from collections.abc import Iterator
from typing import TYPE_CHECKING, NamedTuple, Union, overload

from prolif.residue import ResidueId

if TYPE_CHECKING:
    from prolif.typeshed import IFPData


[docs]class InteractionData(NamedTuple): ligand: ResidueId protein: ResidueId interaction: str metadata: dict
[docs]class IFP(UserDict[tuple[ResidueId, ResidueId], "IFPData"]): """Mapping between residue pairs and interaction fingerprint. Notes ----- This class provides an easy way to access interaction data from the :attr:`~prolif.fingerprint.Fingerprint.ifp` dictionary. This class is a dictionary formatted as: .. code-block:: text { tuple[<residue_id>, <residue_id>]: { <interaction name>: tuple[{ "indices": { "ligand": tuple[int, ...], "protein": tuple[int, ...] }, "parent_indices": { "ligand": tuple[int, ...], "protein": tuple[int, ...] }, <other metadata>: <value> }, ...] } } Here ``<residue_id>`` corresponds to a :class:`~prolif.residue.ResidueId` object. For convenience, one can directly use strings rather than ``ResidueId`` objects when indexing the IFP, e.g. ``ifp[("LIG1.G", "ASP129.A")]``. You can also use a single ``ResidueId`` or string to return a filtered IFP that only contains interactions with the specified residue, e.g. ``ifp["ASP129.A"]``. """ @overload def __getitem__( self, key: tuple[ResidueId, ResidueId] | tuple[str, str] ) -> "IFPData": ... @overload def __getitem__(self, key: ResidueId | str) -> "IFP": ... def __getitem__( self, key: tuple[ResidueId, ResidueId] | tuple[str, str] | ResidueId | str ) -> Union["IFPData", "IFP"]: try: return self.data[key] # type: ignore[index] except KeyError as exc: if isinstance(key, tuple) and len(key) == 2: lig_res, prot_res = key if isinstance(lig_res, ResidueId) and isinstance(prot_res, ResidueId): raise exc from None key = ( ResidueId.from_string(lig_res), # type: ignore[arg-type] ResidueId.from_string(prot_res), # type: ignore[arg-type] ) return self.data[key] if isinstance(key, str): key = ResidueId.from_string(key) if isinstance(key, ResidueId): return IFP( { residue_tuple: interactions for residue_tuple, interactions in self.data.items() if key in residue_tuple }, ) raise KeyError( f"{key} does not correspond to a valid IFP key: it must be a tuple of " "either ResidueId or residue string. If you need to filter the IFP, a " "single ResidueId or residue string can also be used.", )
[docs] def interactions(self) -> Iterator[InteractionData]: """Yields all interactions data as an :class:`InteractionData` namedtuple. .. versionadded:: 2.1.0 """ for (ligand_resid, protein_resid), ifp_dict in self.data.items(): for int_name, metadata_tuple in ifp_dict.items(): for metadata in metadata_tuple: yield InteractionData( ligand=ligand_resid, protein=protein_resid, interaction=int_name, metadata=metadata, )