import gzip
import operator
from collections.abc import Mapping, Sequence, MutableMapping
from functools import reduce
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Literal, overload
import gemmi
import numpy as np
import numpy.typing as npt
import openprotein.utils.chain_id as _chain_id_utils
import openprotein.utils.cif as _cif_utils
from .chains import DNA, RNA, Ligand
from .protein import Protein
if TYPE_CHECKING:
from .template import Template
Chain = Protein | DNA | RNA | Ligand | str
def _is_chain_type(chain):
return isinstance(chain, Chain)
# TODO: deserialization note about plddt parsed per residue
[docs]
class Complex(MutableMapping):
def __init__(
self,
chains: Mapping[str, Protein | DNA | RNA | Ligand] | None = None,
name: bytes | str | None = None,
):
collected: dict[str, Protein | DNA | RNA | Ligand] = {}
if chains is not None:
for key, value in chains.items():
if not isinstance(key, str):
raise TypeError(f"chain id must be str; got {key!r}")
collected[key] = value
# preserve insertion order of the chains by not sorting the dictionary
# Note that as of python 3.7 dicts preserve insertion order (OrderedDict not needed)
self._chains = collected
self._templates: "Sequence[Protein | Complex | Template]" = ()
self.name = name
# --- NAME methods ---
@property
def name(self) -> str | None:
return self._name
@name.setter
def name(self, x: bytes | str | None) -> None:
self._name = x.decode() if isinstance(x, bytes) else x
def get_name(self) -> str | None:
return self._name
def set_name(self, x: bytes | str | None) -> "Complex":
self.name = x
return self
# --- TEMPLATE methods ---
@property
def templates(self) -> "Sequence[Protein | Complex | Template]":
"""A list of templates for guiding the structure prediction of this molecular complex."""
return self._templates
@templates.setter
def templates(self, templates: "Sequence[Protein | Complex | Template]") -> None:
self._templates = tuple(templates)
def get_templates(self) -> "Sequence[Protein | Complex | Template]":
return self.templates
def set_templates(
self, templates: "Sequence[Protein | Complex | Template]"
) -> "Complex":
self.templates = templates
return self
# --- dict-like methods and other general ops ---
[docs]
def get_chain_ids(self):
"""Returns the Chain IDs in this Complex"""
return self._chains.keys()
[docs]
def get_chains(self) -> Mapping[str, Protein | DNA | RNA | Ligand]:
"""Get all chains as a {chain_id: chain} mapping"""
return MappingProxyType(self._chains)
[docs]
def get_chain(self, chain_id: str):
"""Get the chain for this chain_id"""
return self._chains[chain_id]
[docs]
def set_chain(
self, chain_id: str, value: Protein | DNA | RNA | Ligand
) -> "Complex":
"""Add a new chain to this complex or overwrite an existing one with the same chain_id"""
self._chains[chain_id] = value
return self
[docs]
def remove_chain(self, chain_id: str):
"""Remove a chain from this complex"""
del self._chains[chain_id]
[docs]
def insert_chain(self, chain: Protein | DNA | RNA | Ligand | str):
"""Inserts a new chain into this Complex with an auto-generated chain_id. Casts strings into Proteins."""
assert _is_chain_type(chain)
id_gen = _chain_id_utils.id_generator(list(self._chains.keys()))
if isinstance(chain, str):
chain = Protein.from_expr(chain)
chain_id = next(id_gen)
self.set_chain(chain_id, chain)
return self
def __len__(self):
"""Returns the number of chains in the complex"""
return len(self._chains)
def __getitem__(self, key):
"""Retrieve the given chain"""
return self.get_chain(key)
def __setitem__(self, key, value):
"""Add a chain or overwrite a chain with the given chain_id"""
self.set_chain(key, value)
def __delitem__(self, key):
"""Remove the given chain from this complex"""
self.remove_chain(key)
def __contains__(self, chain_id):
return chain_id in self._chains
def __eq__(self, other):
if isinstance(other, Complex):
if self._chains != other._chains:
return False
if self._templates != other._templates:
return False
return True
return False
def __iter__(self):
"""Iterate the chains in this complex"""
return iter(self._chains)
[docs]
def keys(self):
return self._chains.keys()
[docs]
def values(self):
return self._chains.values()
[docs]
def items(self):
return self._chains.items()
[docs]
def update(self, **kwargs):
"""Add a collection of chains to this complex"""
overlapping_chain_ids = []
for key, value in kwargs.items():
if not isinstance(key, str):
raise TypeError(f"chain id must be str; got {key!r}")
# TODO - assert that chain objects are the right types too
# assert _is_chain_type(value)
# strictly do not allow replacing existing chain_ids
if key in self._chains:
overlapping_chain_ids.append(key)
if len(overlapping_chain_ids) > 0:
raise ValueError(
f"Trying to combine two sets of chains with overlapping chain ids: {overlapping_chain_ids}"
)
# update chains
self._chains.update(**kwargs)
[docs]
def clear(self):
"""Remove all chains and templates from this complex"""
self._chains.clear()
# also deletes the templates if you clear this complex
self._templates = ()
[docs]
def copy(self) -> "Complex":
"""Copy this Complex. Returns a deep copy of the contained Chains"""
# TODO - should this deepcopy the chain objects themselves?
chains_copy: dict[str, Protein | DNA | RNA | Ligand] = {
chain_id: chain.copy() for chain_id, chain in self._chains.items()
}
return Complex(chains=chains_copy, name=self._name)
# --- other dictionary operators like set combine ---
# python actually uses the or operator for dict combine
def __ior__(self, other) -> "Complex":
"""In-place update to this Complex using |= """
# other is either a single Chain or a Complex
# if single chain, we insert it with an autogenerated ID
# if complex, we add all of its chains
if _is_chain_type(other):
self.insert_chain(other)
else:
self.update(**other)
return self
def __or__(self, right) -> "Complex":
"""Create a combined copy: self | right"""
a = self.copy()
a |= right # handles right being Complex or single Chain
return a
def __ror__(self, left) -> "Complex":
if isinstance(left, str):
left = Protein.from_expr(expr=left)
return left | self
# NOTE - python uses the or operator (|) for set and dict unions, but we also
# support this using the and operator (&)
__rand__ = __ror__
__and__ = __or__
# -- Chain-type specific methods ---
def get_proteins(self) -> Mapping[str, Protein]:
return MappingProxyType(
{k: v for k, v in self._chains.items() if isinstance(v, Protein)}
)
def get_protein(self, chain_id: str) -> Protein:
chain = self._chains[chain_id]
assert isinstance(chain, Protein)
return chain
def get_dnas(self) -> Mapping[str, DNA]:
return MappingProxyType(
{k: v for k, v in self._chains.items() if isinstance(v, DNA)}
)
def get_dna(self, chain_id: str) -> DNA:
chain = self._chains[chain_id]
assert isinstance(chain, DNA)
return chain
def get_rnas(self) -> Mapping[str, RNA]:
return MappingProxyType(
{k: v for k, v in self._chains.items() if isinstance(v, RNA)}
)
def get_rna(self, chain_id: str) -> RNA:
chain = self._chains[chain_id]
assert isinstance(chain, RNA)
return chain
def get_ligands(self) -> Mapping[str, Ligand]:
return MappingProxyType(
{k: v for k, v in self._chains.items() if isinstance(v, Ligand)}
)
def get_ligand(self, chain_id: str) -> Ligand:
chain = self._chains[chain_id]
assert isinstance(chain, Ligand)
return chain
# --- structure comparison and geometry methods ---
@overload
def rmsd(
self,
tgt: "Complex",
backbone_only: bool | str | Sequence[str] = False,
return_transform: Literal[False] = False,
) -> float: ...
@overload
def rmsd(
self,
tgt: "Complex",
backbone_only: bool | str | Sequence[str] = False,
return_transform: Literal[True] = True,
) -> tuple[float, npt.NDArray[np.floating], npt.NDArray[np.floating]]: ...
def rmsd(
self,
tgt: "Complex",
backbone_only: bool | str | Sequence[str] = False,
return_transform: bool = False,
) -> float | tuple[float, npt.NDArray[np.floating], npt.NDArray[np.floating]]:
assert all(
isinstance(v, Protein) for v in self._chains.values()
), "rmsd supported only for Protein chains, not supported for non-protein chains"
assert all(
isinstance(v, Protein) for v in tgt._chains.values()
), "rmsd supported only for Protein chains, not supported for non-protein chains"
src_proteins, tgt_proteins = self.get_proteins(), tgt.get_proteins()
assert tgt_proteins.keys() == src_proteins.keys()
assert [len(x) for x in src_proteins.values()] == [
len(x) for x in tgt_proteins.values()
]
src_protein: Protein = reduce(operator.add, src_proteins.values())
tgt_protein: Protein = reduce(operator.add, tgt_proteins.values())
return src_protein.rmsd(
tgt_protein,
backbone_only=backbone_only,
return_transform=return_transform,
)
def transform(
self,
R: npt.NDArray[np.floating] | None = None,
t: npt.NDArray[np.floating] | None = None,
) -> "Complex":
assert all(
isinstance(v, Protein) for v in self._chains.values()
), "transform supported only for Protein chains, not supported for non-protein chains"
for protein in self.get_proteins().values():
protein.transform(R=R, t=t)
return self
def superimpose_onto(
self, tgt: "Complex", backbone_only: bool | str | Sequence[str] = False
) -> "Complex":
_, R, t = tgt.rmsd(self, backbone_only=backbone_only, return_transform=True)
return self.transform(R=R, t=t)
# --- read and write PDB and CIF files and strings ---
[docs]
def to_string(self, format: Literal["cif", "pdb"] = "cif") -> str:
"""
Serialize this Complex to a string. Note that format="pdb" may not serialize all
aspects of this object, so format="cif", the default, is preferred.
"""
if format == "cif":
return self._make_cif_string()
elif format == "pdb":
return self._make_pdb_string()
else:
raise ValueError(format)
@staticmethod
def from_filepath(
path: Path | str,
use_bfactor_as_plddt: bool | None = None,
model_idx: int = 0,
verbose: bool = True,
) -> "Complex":
path = Path(path)
if path.suffix == ".gz":
if path.name.endswith(".cif.gz"):
ext, format = ".cif.gz", "cif"
elif path.name.endswith(".pdb.gz"):
ext, format = ".pdb.gz", "pdb"
else:
raise ValueError(f"unsupported format: {path}")
with gzip.open(path, "rb") as f:
data = f.read()
else:
ext = path.suffix
format = ext.removeprefix(".")
assert format == "cif" or format == "pdb"
data = path.read_bytes()
return Complex.from_string(
filestring=data,
format=format,
use_bfactor_as_plddt=use_bfactor_as_plddt,
model_idx=model_idx,
verbose=verbose,
).set_name(path.name.removesuffix(ext))
@staticmethod
def from_string(
filestring: bytes | str,
format: Literal["pdb", "cif"],
use_bfactor_as_plddt: bool | None = None,
model_idx: int = 0,
verbose: bool = True,
) -> "Complex":
structure_block = _cif_utils.StructureCIFBlock(
filestring=filestring, format=format
)
return Complex._from_structure_block(
structure_block=structure_block,
use_bfactor_as_plddt=use_bfactor_as_plddt,
model_idx=model_idx,
verbose=verbose,
)
def _assert_valid_templates(self):
from .template import Template
for template in self.templates:
(
template if isinstance(template, Template) else Template(template)
).validate_for_target(self)
for chain_id, protein in self.get_proteins().items():
for template in protein.templates:
(
template
if isinstance(template, Template)
else Template(template, mapping=chain_id)
).validate_for_target(Complex({chain_id: protein}))
@staticmethod
def _from_structure_block(
structure_block: _cif_utils.StructureCIFBlock,
use_bfactor_as_plddt: bool | None = None,
model_idx: int = 0,
verbose: bool = True,
) -> "Complex":
block, structure = structure_block.block, structure_block.structure
model = structure[model_idx] if len(structure) > 0 else None
# Use block info directly so that we can get chains with empty struct info
subchain_ids = [x for x in block.find_loop("_struct_asym.id")]
if len(subchain_ids) == 0 and model is not None:
# Try to get actual chain IDs from the structure
subchain_ids = [subchain.subchain_id() for subchain in model.subchains()]
# collect chains
chains = {}
for subchain_id in sorted(subchain_ids):
subchain = model.get_subchain(subchain_id) if model is not None else None
# Get the entity for this chain to determine its type
if subchain is not None and len(subchain) > 0:
entity = structure.get_entity_of(subchain)
if entity is None:
raise ValueError(f"Could not find entity for chain {subchain_id}")
else:
matching_entities = [
e for e in structure.entities if subchain_id in e.subchains
]
assert len(matching_entities) == 1, (
f"expected only one entity to match chain_id={subchain_id!r}, "
f"but found {len(matching_entities)}: {matching_entities}"
)
entity = matching_entities[0]
del matching_entities
# Determine chain type based on entity type and polymer type
if (entity_type := entity.entity_type) == gemmi.EntityType.Polymer:
if structure.input_format == gemmi.CoorFormat.Pdb:
assert subchain_id.endswith("xp")
chain_id = subchain_id.removesuffix("xp")
assert chain_id not in chains
else:
chain_id = subchain_id
if (polymer_type := entity.polymer_type) in (
gemmi.PolymerType.PeptideL,
gemmi.PolymerType.PeptideD,
):
chains[chain_id] = Protein._from_structure_block(
structure_block=structure_block,
chain_id=subchain_id,
use_bfactor_as_plddt=use_bfactor_as_plddt,
model_idx=model_idx,
verbose=verbose,
)
elif polymer_type == gemmi.PolymerType.Dna:
chains[chain_id] = DNA._from_structure_block(
structure_block=structure_block,
chain_id=subchain_id,
model_idx=model_idx,
)
elif polymer_type == gemmi.PolymerType.Rna:
chains[chain_id] = RNA._from_structure_block(
structure_block=structure_block,
chain_id=subchain_id,
model_idx=model_idx,
)
else:
# if verbose:
# print(
# f"Warning: Skipping unsupported polymer type {polymer_type} for chain {subchain_id}"
# )
continue
elif entity_type == gemmi.EntityType.NonPolymer:
if structure.input_format == gemmi.CoorFormat.Pdb:
raise ValueError("ligands from pdb files not supported yet")
chain_id = subchain_id
assert (
structure.input_format != gemmi.CoorFormat.Pdb
), "ligands from pdb files not supported yet"
chains[chain_id] = Ligand._from_structure_block(
structure_block=structure_block,
chain_id=subchain_id,
model_idx=model_idx,
)
elif entity_type == gemmi.EntityType.Water:
continue
else:
# if verbose:
# print(
# f"Warning: Skipping unsupported entity type {entity_type} for chain {subchain_id}"
# )
continue
return Complex(chains=chains, name=structure.name)
def _make_cif_string(self) -> str:
structure = self._make_structure()
block = structure.make_mmcif_block(
groups=gemmi.MmcifOutputGroups(True, chem_comp=False)
)
sequence_loop, atom_loop = _cif_utils.init_loops(block=block)
for chain_id, chain in self._chains.items():
chain._append_loop_data(
chain_id=chain_id, sequence_loop=sequence_loop, atom_loop=atom_loop
)
return block.as_string()
def _make_pdb_string(self) -> str:
structure = self._make_structure()
return structure.make_pdb_string(gemmi.PdbWriteOptions(minimal=True))
def _make_structure(self) -> gemmi.Structure:
assert (
len(set(x._structure_block for x in self.get_ligands().values())) <= 1
), "can only serialize ligands if they all originate from the same structure file"
structure = gemmi.Structure()
for chain_id, chain in self._chains.items():
structure = chain._make_structure(
structure=structure,
model_idx=0,
chain_id=chain_id,
entity_name=str(len(structure.entities) + 1),
)
structure.setup_entities() # this should deduplicate polymer entities
for entity_idx, entity in enumerate(structure.entities):
entity.name = str(entity_idx + 1)
if self._name is not None:
structure.name = self._name
return structure
