from __future__ import annotations import io import itertools import random import re import warnings from dataclasses import asdict, dataclass, replace from functools import cached_property from pathlib import Path from subprocess import check_output from tempfile import TemporaryDirectory from typing import Any, Iterable, Sequence import biotite.structure as bs import brotli import msgpack import msgpack_numpy import numpy as np import torch from biotite.database import rcsb from biotite.file import InvalidFileError from biotite.structure.io.pdb import PDBFile from biotite.structure.io.pdbx import CIFCategory, CIFColumn, CIFData, CIFFile from biotite.structure.io.pdbx import set_structure as set_structure_pdbx from biotite.structure.io.pdbx.convert import _get_transformations, get_structure from biotite.structure.util import matrix_rotate from scipy.spatial import KDTree from . import esmfold2_residue_constants as residue_constants from .esmfold2_misc import slice_python_object_as_numpy from .esmfold2_affine3d import Affine3D from .esmfold2_aligner import Aligner from .esmfold2_atom_indexer import AtomIndexer from .esmfold2_metrics import compute_gdt_ts, compute_lddt_ca from .esmfold2_mmcif_parsing import MmcifWrapper, NoProteinError from .esmfold2_protein_chain import ( ProteinChain, _str_key_to_int_key, chain_to_ndarray, index_by_atom_name, infer_CB, ) from .esmfold2_utils_types import PathOrBuffer msgpack_numpy.patch() SINGLE_LETTER_CHAIN_IDS = ( "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" ) def _parse_operation_expression(expression): """ Get successive operation steps (IDs) for the given ``oper_expression``. Form the cartesian product, if necessary. Copied from biotite and fixed a bug """ # Split groups by parentheses: # use the opening parenthesis as delimiter # and just remove the closing parenthesis expressions_per_step = expression.replace(")", "").split("(") expressions_per_step = [e for e in expressions_per_step if len(e) > 0] # Important: Operations are applied from right to left expressions_per_step.reverse() operations = [] for expr in expressions_per_step: cur_expr = expr.split(",") cur_op = [] # Deal with e='1-10,20-30,40-50' type expressions for e in cur_expr: if "-" in e: first, last = e.split("-") cur_op.extend(str(id) for id in range(int(first), int(last) + 1)) else: cur_op.append(e) operations.append(cur_op) # Cartesian product of operations return list(itertools.product(*operations)) def _apply_transformations_fast(chains, transformation_dict, operations): """ Get subassembly by applying the given operations to the input structure containing affected asym IDs. """ # Additional first dimesion for 'structure.repeat()' results = [] # Apply corresponding transformation for each copy in the assembly for c in chains: for operation in operations: coord = c.atom37_positions.copy() # Execute for each transformation step # in the operation expression for op_step in operation: T = transformation_dict[op_step] # Rotate coord = matrix_rotate(coord, T.rotation) # Translate coord += T.target_translation new_chain = replace(c, atom37_positions=coord) results.append(new_chain) return results @dataclass class ProteinComplexMetadata: entity_lookup: dict[int, int] chain_lookup: dict[int, str] mmcif: MmcifWrapper | None = None # This is a dictionary that maps assembly ids to the list of unique chains # in that assembly. Allows for usage of `switch_assembly`. assembly_composition: dict[str, list[str]] | None = None @dataclass class DockQSingleScore: native_chains: tuple[str, str] DockQ: float interface_rms: float ligand_rms: float fnat: float fnonnat: float clashes: float F1: float DockQ_F1: float @dataclass class DockQResult: total_dockq: float native_interfaces: int chain_mapping: dict[str, str] interfaces: dict[tuple[str, str], DockQSingleScore] # zip(aligned.chain_iter(), native.chain_iter()) gives you the pairing # aligned.rmsd(native) should give you a low rmsd irrespective of shuffling aligned: ProteinComplex aligned_rmsd: float @dataclass(frozen=True) class ProteinComplex: """Dataclass with atom37 representation of an entire protein complex.""" id: str sequence: str entity_id: np.ndarray # entities map to unique sequences chain_id: np.ndarray # multiple chains might share an entity id sym_id: np.ndarray # complexes might be copies of the same chain residue_index: np.ndarray insertion_code: np.ndarray atom37_positions: np.ndarray atom37_mask: np.ndarray confidence: np.ndarray # This metadata is parsed from the MMCIF file. For synthetic data, we do a best effort. metadata: ProteinComplexMetadata atom37_confidence: np.ndarray | None = None # [L, 37] per-atom pLDDT def __post_init__(self): l = len(self.sequence) assert self.atom37_positions.shape[0] == l, (self.atom37_positions.shape, l) assert self.atom37_mask.shape[0] == l, (self.atom37_mask.shape, l) assert self.residue_index.shape[0] == l, (self.residue_index.shape, l) assert self.insertion_code.shape[0] == l, (self.insertion_code.shape, l) assert self.confidence.shape[0] == l, (self.confidence.shape, l) assert self.entity_id.shape[0] == l, (self.entity_id.shape, l) assert self.chain_id.shape[0] == l, (self.chain_id.shape, l) assert self.sym_id.shape[0] == l, (self.sym_id.shape, l) if self.atom37_confidence is not None: assert self.atom37_confidence.shape == self.atom37_mask.shape, ( self.atom37_confidence.shape, self.atom37_mask.shape, ) def __getitem__(self, idx: int | list[int] | slice | np.ndarray): """This function slices protein complexes without consideration of chain breaks NOTE: When slicing with a boolean mask, it's possible that the output array won't be the expected length. This is because we do our best to preserve chainbreak tokens. """ if isinstance(idx, int): idx = [idx] if isinstance(idx, list): raise ValueError( "ProteinComplex doesn't supports indexing with lists of indices" ) if isinstance(idx, np.ndarray): is_chainbreak = np.asarray([s == "|" for s in self.sequence]) idx = idx.astype(bool) | is_chainbreak complex = self._unsafe_slice(idx) if len(complex) == 0: return complex # detect runs of chainbreaks by searching for instances of '||' in complex.sequence chainbreak_runs = np.asarray( [ complex.sequence[i : i + 2] == "||" for i in range(len(complex.sequence) - 1) ] + [complex.sequence[-1] == "|"] ) # We should remove as many chainbreaks as possible from the start of the sequence for i in range(len(chainbreak_runs)): if complex.sequence[i] == "|": chainbreak_runs[i] = True else: break complex = complex._unsafe_slice(~chainbreak_runs) return complex def _unsafe_slice(self, idx: int | list[int] | slice | np.ndarray): sequence = slice_python_object_as_numpy(self.sequence, idx) return replace( self, sequence=sequence, entity_id=self.entity_id[..., idx], chain_id=self.chain_id[..., idx], sym_id=self.sym_id[..., idx], residue_index=self.residue_index[..., idx], insertion_code=self.insertion_code[..., idx], atom37_positions=self.atom37_positions[..., idx, :, :], atom37_mask=self.atom37_mask[..., idx, :], confidence=self.confidence[..., idx], atom37_confidence=self.atom37_confidence[..., idx, :] if self.atom37_confidence is not None else None, ) def __len__(self): return len(self.sequence) @property def num_chains(self): return len(self.chain_boundaries) @cached_property def atoms(self) -> AtomIndexer: return AtomIndexer(self, property="atom37_positions", dim=-2) @cached_property def atom_mask(self) -> AtomIndexer: return AtomIndexer(self, property="atom37_mask", dim=-1) @cached_property def chain_lengths(self) -> np.ndarray: return np.diff(self.chain_boundaries, axis=1).flatten() @cached_property def chain_boundaries(self) -> list[tuple[int, int]]: cb = [-1] for i, s in enumerate(self.sequence): if s == "|": cb.append(i) cb.append(len(self)) return [(cb[i] + 1, cb[i + 1]) for i in range(len(cb) - 1)] def get_chain_by_index(self, index: int) -> ProteinChain: try: start, end = self.chain_boundaries[index] return self[start:end].as_chain() except IndexError: raise IndexError(f"Chain index {index} out of bounds") def get_chain_by_id( self, chain_id: str, sample_chain_if_duplicate: bool = True ) -> ProteinChain: valid_indices = [ index for index, id_of_index in self.metadata.chain_lookup.items() if id_of_index == chain_id ] if not valid_indices: raise KeyError(f"Chain ID {chain_id} not found") if sample_chain_if_duplicate: index_to_return = random.choice(valid_indices) return self.get_chain_by_index(index_to_return) else: if len(valid_indices) > 1: raise ValueError(f"Multiple chains with chain ID {chain_id} found") return self.get_chain_by_index(valid_indices[0]) def chain_iter(self) -> Iterable[ProteinChain]: for start, end in self.chain_boundaries: c = self[start:end] yield c.as_chain() def as_chain(self, force_conversion: bool = False) -> ProteinChain: """Convert the ProteinComplex to a ProteinChain. Args: force_conversion (bool): Forces the conversion into a protein chain even if the complex has multiple chains. The purpose of this is to use ProteinChain specific functions (like cbeta_contacts). """ if not force_conversion: assert len(np.unique(self.chain_id)) == 1, f"{self.id}" assert len(np.unique(self.entity_id)) == 1, f"{self.id}" if self.chain_id[0] not in self.metadata.chain_lookup: warnings.warn("Chain ID not found in metadata, using 'A' as default") if self.entity_id[0] not in self.metadata.entity_lookup: warnings.warn("Entity ID not found in metadata, using None as default") chain_id = self.metadata.chain_lookup.get(self.chain_id[0], "A") entity_id = self.metadata.entity_lookup.get(self.entity_id[0], None) else: chain_id = "A" entity_id = None return ProteinChain( id=self.id, sequence=self.sequence, chain_id=chain_id, entity_id=entity_id, atom37_positions=self.atom37_positions, atom37_mask=self.atom37_mask, residue_index=self.residue_index, insertion_code=self.insertion_code, confidence=self.confidence, mmcif=self.metadata.mmcif, atom37_confidence=self.atom37_confidence, ) @classmethod def from_pdb( cls, path: PathOrBuffer, id: str | None = None, is_predicted: bool = False ) -> "ProteinComplex": atom_array = PDBFile.read(path).get_structure( model=1, extra_fields=["b_factor"] ) chains = [] for chain in bs.chain_iter(atom_array): chain = chain[~chain.hetero] if len(chain) == 0: continue chains.append(ProteinChain.from_atomarray(chain, id, is_predicted)) return ProteinComplex.from_chains(chains) def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True): atom_array = None for chain in self.chain_iter(): carr = ( chain.atom_array if include_insertions else chain.atom_array_no_insertions ) atom_array = carr if atom_array is None else atom_array + carr f = PDBFile() f.set_structure(atom_array) f.write(path) def to_pdb_string(self, include_insertions: bool = True) -> str: buf = io.StringIO() self.to_pdb(buf, include_insertions=include_insertions) buf.seek(0) return buf.read() def normalize_chain_ids_for_pdb(self): # Since PDB files have 1-letter chain IDs and don't support the idea of a symmetric index, # we can normalize it instead which might be necessary for DockQ and to_pdb. ids = SINGLE_LETTER_CHAIN_IDS chains = [] for i, chain in enumerate(self.chain_iter()): chain = replace(chain, chain_id=ids[i]) if i > len(ids): raise RuntimeError("Too many chains to write to PDB file") chains.append(chain) return ProteinComplex.from_chains(chains) def find_assembly_ids_with_chain(self, id: str) -> list[str]: good_chains = [] if (comp := self.metadata.assembly_composition) is not None: for assembly_id, chain_ids in comp.items(): if id in chain_ids: good_chains.append(assembly_id) else: raise ValueError( "Cannot switch assemblies on this ProteinComplex, you must create the assembly from mmcif to support this" ) return good_chains def switch_assembly(self, id: str): assert self.metadata.mmcif is not None return get_assembly_fast(self.metadata.mmcif, assembly_id=id) def state_dict(self, backbone_only=False, json_serializable=False): """This state dict is optimized for storage, so it turns things to fp16 whenever possible. Note that we also only support int32 residue indices, I'm hoping we don't need more than 2**32 residues...""" dct = {k: v for k, v in vars(self).items()} if backbone_only: dct["atom37_mask"][:, 3:] = False dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]] if dct.get("atom37_confidence") is not None: dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]] else: dct.pop("atom37_confidence", None) for k, v in dct.items(): if isinstance(v, np.ndarray): match v.dtype: case np.int64: dct[k] = v.astype(np.int32) case np.float64 | np.float32: dct[k] = v.astype(np.float16) case _: pass if json_serializable: dct[k] = v.tolist() elif isinstance(v, ProteinComplexMetadata): dct[k] = asdict(v) dct["metadata"]["mmcif"] = None # These can be populated with non-serializable objects and are not needed for reconstruction dct.pop("atoms", None) dct.pop("atom_mask", None) dct.pop("per_chain_kd_trees", None) return dct def to_blob(self, backbone_only=False) -> bytes: return brotli.compress(msgpack.dumps(self.state_dict(backbone_only)), quality=5) @classmethod def from_state_dict(cls, dct): # Note: assembly_composition is *supposed* to have string keys. dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"]) for k, v in dct.items(): if isinstance(v, list): dct[k] = np.array(v) atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan) atom37[dct["atom37_mask"]] = dct["atom37_positions"] dct["atom37_positions"] = atom37 if "atom37_confidence" in dct: atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32) atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"] dct["atom37_confidence"] = atom37_conf dct = { k: ( v.astype(np.float32) if k in ["atom37_positions", "confidence", "atom37_confidence"] else v ) for k, v in dct.items() } if "chain_boundaries" in dct: del dct["chain_boundaries"] if "chain_boundaries" in dct["metadata"]: del dct["metadata"]["chain_boundaries"] dct["metadata"] = ProteinComplexMetadata(**dct["metadata"]) return cls(**dct) @classmethod def from_blob(cls, input: Path | str | io.BytesIO | bytes): """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory of chains from 52G/1M chains to 20G/1M chains, I think this is a good first shot at compressing and dumping chains to disk. I'm sure there's better ways.""" match input: case Path() | str(): bytes = Path(input).read_bytes() case io.BytesIO(): bytes = input.getvalue() case _: bytes = input return cls.from_state_dict( msgpack.loads(brotli.decompress(bytes), strict_map_key=False) ) @classmethod def from_rcsb(cls, pdb_id: str, keep_source: bool = False) -> ProteinComplex: f: io.StringIO = rcsb.fetch(pdb_id, "cif") # type: ignore return cls.from_mmcif(f, id=pdb_id, keep_source=keep_source, is_predicted=False) @classmethod def from_mmcif( cls, path: PathOrBuffer, id: str | None = None, assembly_id: str | None = None, is_predicted: bool = False, keep_source: bool = False, ): """Return a ProteinComplex object from an mmcif file. TODO(@zeming): there's actually multiple complexes per file, but for ease of implementation, we only consider the first defined complex! Args: path (str | Path | io.TextIO): Path or buffer to read mmcif file from. Should be uncompressed. id (str, optional): String identifier to assign to structure. Will attempt to infer otherwise. is_predicted (bool): If True, reads b factor as the confidence readout. Default: False. chain_id (str, optional): Select a chain corresponding to (author) chain id. """ mmcif = MmcifWrapper.read(path, id) return get_assembly_fast(mmcif, assembly_id=assembly_id) @classmethod def from_chains( cls, chains: Sequence[ProteinChain], mmcif: MmcifWrapper | None = None, all_assembly_metadata_dictionary: dict[str, list[str]] | None = None, ): if not chains: raise ValueError( "Cannot create a ProteinComplex from an empty list of chains" ) # TODO(roshan): Make a proper protein complex class def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray): full_array = [] for array in arrays: full_array.append(array) full_array.append(sep) full_array = full_array[:-1] return np.concatenate(full_array, 0) sep_tokens = { "residue_index": np.array([-1]), "insertion_code": np.array([""]), "atom37_positions": np.full([1, 37, 3], np.nan), "atom37_mask": np.zeros([1, 37], dtype=bool), "confidence": np.array([0]), } any_has_atom37_conf = any(c.atom37_confidence is not None for c in chains) if any_has_atom37_conf: sep_tokens["atom37_confidence"] = np.full([1, 37], np.nan, dtype=np.float32) def _get_chain_attr(chain: ProteinChain, name: str) -> np.ndarray: val = getattr(chain, name) if val is None and name == "atom37_confidence": return np.full([len(chain), 37], np.nan, dtype=np.float32) return val array_args: dict[str, np.ndarray] = { name: join_arrays([_get_chain_attr(chain, name) for chain in chains], sep) for name, sep in sep_tokens.items() } multimer_arrays = [] chain2num_max = -1 chain2num = {} ent2num_max = -1 ent2num = {} total_index = 0 for i, c in enumerate(chains): num_res = c.residue_index.shape[0] if c.chain_id not in chain2num: chain2num[c.chain_id] = (chain2num_max := chain2num_max + 1) chain_id_array = np.full([num_res], chain2num[c.chain_id], dtype=np.int64) if c.entity_id is None: entity_num = (ent2num_max := ent2num_max + 1) else: if c.entity_id not in ent2num: ent2num[c.entity_id] = (ent2num_max := ent2num_max + 1) entity_num = ent2num[c.entity_id] entity_id_array = np.full([num_res], entity_num, dtype=np.int64) sym_id_array = np.full([num_res], i, dtype=np.int64) multimer_arrays.append( { "chain_id": chain_id_array, "entity_id": entity_id_array, "sym_id": sym_id_array, } ) total_index += num_res + 1 sep = np.array([-1]) update = { name: join_arrays([dct[name] for dct in multimer_arrays], sep=sep) for name in ["chain_id", "entity_id", "sym_id"] } array_args.update(update) metadata = ProteinComplexMetadata( mmcif=mmcif, chain_lookup={v: k for k, v in chain2num.items()}, entity_lookup={v: k for k, v in ent2num.items()}, assembly_composition=all_assembly_metadata_dictionary, ) return cls( id=chains[0].id, sequence=residue_constants.CHAIN_BREAK_TOKEN.join( chain.sequence for chain in chains ), metadata=metadata, **array_args, ) def infer_oxygen(self) -> ProteinComplex: """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided.""" O_missing_indices = np.argwhere( ~np.isfinite(self.atoms["O"]).all(axis=1) ).squeeze() O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32) N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1) N = torch.roll(N, -3) N[..., -1, :] = torch.nan # Get the frame defined by the CA-C-N atom frames = Affine3D.from_graham_schmidt(CA, C, N) O = frames.apply(O_vector) atom37_positions = self.atom37_positions.copy() atom37_mask = self.atom37_mask.copy() atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = O[ O_missing_indices ].numpy() atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan( atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] ).any(-1) new_chain = replace( self, atom37_positions=atom37_positions, atom37_mask=atom37_mask ) return new_chain def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinComplex: """Return a new chain with inferred CB atoms at all residues except GLY. Args: infer_cbeta_for_glycine (bool): If True, infers a beta carbon for glycine residues, even though that residue doesn't have one. Default off. NOTE(rverkuil): The reason for having this switch in the first place is that sometimes we want a (inferred) CB coordinate for every residue, for example for making a pairwise distance matrix, or doing an RMSD calculation between two designs for a given structural template, w/ CB atoms. """ atom37_positions = self.atom37_positions.copy() atom37_mask = self.atom37_mask.copy() N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0) # See usage in trDesign codebase. # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140 inferred_cbeta_positions = infer_CB(C, N, CA, 1.522, 1.927, -2.143) if not infer_cbeta_for_glycine: inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan atom37_positions[:, residue_constants.atom_order["CB"]] = ( inferred_cbeta_positions ) atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan( atom37_positions[:, residue_constants.atom_order["CB"]] ).any(-1) new_chain = replace( self, atom37_positions=atom37_positions, atom37_mask=atom37_mask ) return new_chain @classmethod def from_open_source(cls, pc: ProteinComplex): # TODO(@zeming): deprecated, should delete return pc @classmethod def concat(cls, objs: list[ProteinComplex]) -> ProteinComplex: pdb_ids = [obj.id for obj in objs] if len(set(pdb_ids)) > 1: raise RuntimeError( "Concatention of protein complexes across different PDB ids is unsupported" ) return ProteinComplex.from_chains( list(itertools.chain.from_iterable(obj.chain_iter() for obj in objs)) ) def _sanity_check_complexes_are_comparable(self, other: ProteinComplex): assert len(self) == len(other), "Protein complexes must have the same length" assert len(list(self.chain_iter())) == len( list(other.chain_iter()) ), "Protein complexes must have the same number of chains" def rmsd( self, target: ProteinComplex, also_check_reflection: bool = False, mobile_inds: list[int] | np.ndarray | None = None, target_inds: list[int] | np.ndarray | None = None, only_compute_backbone_rmsd: bool = False, compute_chain_assignment: bool = True, ): """ Compute the RMSD between this protein chain and another. Args: target (ProteinComplex): The target (other) protein complex to compare to. also_check_reflection (bool, optional): If True, also check if the reflection of the mobile atoms has a lower RMSD. mobile_inds (list[int], optional): The indices of the mobile atoms to align. These are NOT residue indices target_inds (list[int], optional): The indices of the target atoms to align. These are NOT residue indices only_compute_backbone_rmsd (bool, optional): If True, only compute the RMSD of the backbone atoms. """ if compute_chain_assignment: aligned = self.dockq(target).aligned else: aligned = self aligner = Aligner( aligned if mobile_inds is None else aligned[mobile_inds], target if target_inds is None else target[target_inds], only_compute_backbone_rmsd, ) avg_rmsd = aligner.rmsd if not also_check_reflection: return avg_rmsd aligner = Aligner( aligned if mobile_inds is None else aligned[mobile_inds], target if target_inds is None else target[target_inds], only_compute_backbone_rmsd, use_reflection=True, ) avg_rmsd_neg = aligner.rmsd return min(avg_rmsd, avg_rmsd_neg) def lddt_ca( self, target: ProteinComplex, mobile_inds: list[int] | np.ndarray | None = None, target_inds: list[int] | np.ndarray | None = None, compute_chain_assignment: bool = True, **kwargs, ) -> float | np.ndarray: """Compute the LDDT between this protein complex and another. Arguments: target (ProteinComplex): The other protein complex to compare to. mobile_inds (list[int], np.ndarray, optional): The indices of the mobile atoms to align. These are NOT residue indices target_inds (list[int], np.ndarray, optional): The indices of the target atoms to align. These are NOT residue indices Returns: float | np.ndarray: The LDDT score between the two protein chains, either a single float or per-residue LDDT scores if `per_residue` is True. """ if compute_chain_assignment: aligned = self.dockq(target).aligned else: aligned = self lddt = compute_lddt_ca( torch.tensor(aligned.atom37_positions[mobile_inds]).unsqueeze(0), torch.tensor(target.atom37_positions[target_inds]).unsqueeze(0), torch.tensor(aligned.atom37_mask[mobile_inds]).unsqueeze(0), **kwargs, ) return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten() def gdt_ts( self, target: ProteinComplex, mobile_inds: list[int] | np.ndarray | None = None, target_inds: list[int] | np.ndarray | None = None, compute_chain_assignment: bool = True, **kwargs, ) -> float | np.ndarray: """Compute the GDT_TS between this protein complex and another. Arguments: target (ProteinComplex): The other protein complex to compare to. mobile_inds (list[int], np.ndarray, optional): The indices of the mobile atoms to align. These are NOT residue indices target_inds (list[int], np.ndarray, optional): The indices of the target atoms to align. These are NOT residue indices Returns: float: The GDT_TS score between the two protein chains. """ if compute_chain_assignment: aligned = self.dockq(target).aligned else: aligned = self gdt_ts = compute_gdt_ts( mobile=torch.tensor( index_by_atom_name(aligned.atom37_positions[mobile_inds], "CA"), dtype=torch.float32, ).unsqueeze(0), target=torch.tensor( index_by_atom_name(target.atom37_positions[target_inds], "CA"), dtype=torch.float32, ).unsqueeze(0), atom_exists_mask=torch.tensor( index_by_atom_name(aligned.atom37_mask[mobile_inds], "CA", dim=-1) & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1) ).unsqueeze(0), **kwargs, ) return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten() def dockq(self, native: ProteinComplex): # This function uses dockqv2 to compute the DockQ score. Because it does a mapping # over all possible chains, it's quite slow. Be careful not to use this in an inference loop # or something that requires fast scoring. It defaults to 8 CPUs. # # TODO(@zeming): Because we haven't properly implemented protein complexes for mmcif, # if your protein has multi-letter or repeated chain IDs, this will fail. Please call # pc = pc.normalize_chain_ids_for_pdb() before calling this function in that case (limit is 62 chains) try: pass except BaseException: raise RuntimeError( "DockQ is not installed. Please update your environment." ) self._sanity_check_complexes_are_comparable(native) def sanity_check_chain_ids(pc: ProteinComplex): ids = [] for i, chain in enumerate(pc.chain_iter()): if i > len(SINGLE_LETTER_CHAIN_IDS): raise ValueError("Too many chains to write to PDB file") if len(chain.chain_id) > 1: raise ValueError( "We only supports single letter chain IDs for DockQ" ) ids.append(chain.chain_id) if len(set(ids)) != len(ids): raise ValueError(f"Duplicate chain IDs in protein complex: {ids}") return ids sanity_check_chain_ids(self) sanity_check_chain_ids(native) with TemporaryDirectory() as tdir: dir = Path(tdir) self.to_pdb(dir / "self.pdb") native.to_pdb(dir / "native.pdb") output = check_output(["DockQ", dir / "self.pdb", dir / "native.pdb"]) lines = output.decode().split("\n") # Remove the header comments start_index = next( i for i, line in enumerate(lines) if line.startswith("Model") ) lines = lines[start_index:] result = {} interfaces = [] current_interface: dict = {} for line in lines: line = line.strip() if not line: continue if line.startswith("Model :"): pass # Tmp pdb file location, it's useless... elif line.startswith("Native :"): pass # Tmp pdb file location, it's useless... elif line.startswith("Total DockQ"): total_dockq_match = re.search( r"Total DockQ over (\d+) native interfaces: ([\d.]+) with (.*) model:native mapping", line, ) if total_dockq_match: result["value"] = float(total_dockq_match.group(2)) result["native interfaces"] = int(total_dockq_match.group(1)) native_chains, self_chains = total_dockq_match.group(3).split(":") result["mapping"] = dict(zip(native_chains, self_chains)) else: raise RuntimeError( "Failed to parse DockQ output, maybe your DockQ version is wrong?" ) elif line.startswith("Native chains:"): if current_interface: interfaces.append(current_interface) current_interface = { "Native chains": line.split(":")[1].strip().split(", ") } elif line.startswith("Model chains:"): current_interface["Model chains"] = ( line.split(":")[1].strip().split(", ") ) elif ":" in line: key, value = line.split(":", 1) current_interface[key.strip()] = float(value.strip()) if current_interface: interfaces.append(current_interface) def parse_dict(d: dict[str, Any]) -> DockQSingleScore: return DockQSingleScore( native_chains=tuple(d["Native chains"]), # type: ignore DockQ=float(d["DockQ"]), interface_rms=float(d["irms"]), ligand_rms=float(d["Lrms"]), # Note the capitalization difference fnat=float(d["fnat"]), fnonnat=float(d["fnonnat"]), clashes=float(d["clashes"]), F1=float(d["F1"]), DockQ_F1=float(d["DockQ_F1"]), ) inv_mapping = {v: k for k, v in result["mapping"].items()} self_chain_map = {c.chain_id: c for c in self.chain_iter()} realigned = [] for chain in native.chain_iter(): realigned.append(self_chain_map[inv_mapping[chain.chain_id]]) realigned = ProteinComplex.from_chains(realigned) aligner = Aligner(realigned, native) realigned = aligner.apply(realigned) result = DockQResult( total_dockq=result["value"], native_interfaces=result["native interfaces"], chain_mapping=result["mapping"], interfaces={ (i["Model chains"][0], i["Model chains"][1]): parse_dict(i) for i in interfaces }, aligned=realigned, aligned_rmsd=aligner.rmsd, ) return result @cached_property def per_chain_kd_trees(self): # Iterate over chains, build KDTree for each chain kdtrees = [] CA = self.atoms["CA"] for start, end in self.chain_boundaries: chain_CA = CA[start:end] chain_CA = chain_CA[np.isfinite(chain_CA).all(axis=-1)] kdtrees.append(KDTree(chain_CA)) return kdtrees def chain_adjacency(self, cutoff: float = 8.0) -> np.ndarray: # Compute adjacency matrix for protein complex num_chains = self.num_chains adjacency = np.zeros((num_chains, num_chains), dtype=bool) for (i, kdtree), (j, kdtree2) in itertools.combinations( enumerate(self.per_chain_kd_trees), 2 ): adj = kdtree.query_ball_tree(kdtree2, cutoff) any_is_adjacent = any(len(a) > 0 for a in adj) adjacency[i, j] = any_is_adjacent adjacency[j, i] = any_is_adjacent return adjacency def chain_adjacency_by_index(self, index: int, cutoff: float = 8.0) -> np.ndarray: num_chains = len(self.chain_boundaries) adjacency = np.zeros(num_chains, dtype=bool) for i, kdtree in enumerate(self.per_chain_kd_trees): if i == index: continue adj = kdtree.query_ball_tree(self.per_chain_kd_trees[index], cutoff) adjacency[i] = any(len(a) > 0 for a in adj) return adjacency def add_prefix_to_chain_ids(self, prefix: str) -> ProteinComplex: """Rename all chains in the complex with a given prefix. Args: prefix (str): The prefix to use for the new chain IDs. Each chain will be named as "{prefix}_{chain_id}". Returns: ProteinComplex: A new protein complex with renamed chains. """ new_chains = [] for chain in self.chain_iter(): # Create new chain with updated chain_id new_chain = replace(chain, chain_id=f"{prefix}_{chain.chain_id}") new_chains.append(new_chain) return ProteinComplex.from_chains(new_chains) def sasa(self, by_residue: bool = True): chain = self.as_chain(force_conversion=True) return chain.sasa(by_residue=by_residue) def to_mmcif_string(self) -> str: """Convert the ProteinComplex to mmCIF format. Returns: str: The mmCIF content as a string. """ # Convert the ProteinComplex to a biotite AtomArray # Collect all atoms from all chains all_atoms = [] for chain in self.chain_iter(): chain_atom_array = chain.atom_array # Convert AtomArray to list of atoms and add to collection all_atoms.extend(chain_atom_array) # Create combined AtomArray from all atoms if not all_atoms: raise ValueError("No atoms found in protein complex") atom_array = bs.array(all_atoms) # Create CIF file f = CIFFile() set_structure_pdbx(f, atom_array, data_block=self.id) # Add entity information for proper mmCIF structure self._add_entity_information(f) # Write to string output = io.StringIO() f.write(output) return output.getvalue() def _add_entity_information(self, cif_file: CIFFile) -> None: """Add entity, entity_poly, and struct_asym sections to CIF file.""" # Group chains by sequence to create unique entities entity_map = {} # sequence -> entity_id chain_to_entity = {} # chain_id -> entity_id entity_sequences = {} # entity_id -> sequence entity_id_counter = 1 for chain in self.chain_iter(): sequence = chain.sequence if sequence not in entity_map: entity_map[sequence] = entity_id_counter entity_sequences[entity_id_counter] = sequence entity_id_counter += 1 chain_to_entity[chain.chain_id] = entity_map[sequence] # Create _entity section entity_ids = [] entity_types = [] entity_descriptions = [] for entity_id in sorted(entity_sequences.keys()): entity_ids.append(str(entity_id)) entity_types.append("polymer") entity_descriptions.append(f"Protein chain (entity {entity_id})") cif_file.block["entity"] = CIFCategory( name="entity", columns={ "id": CIFColumn( data=CIFData(array=np.array(entity_ids), dtype=np.str_) ), "type": CIFColumn( data=CIFData(array=np.array(entity_types), dtype=np.str_) ), "pdbx_description": CIFColumn( data=CIFData(array=np.array(entity_descriptions), dtype=np.str_) ), }, ) # Create _entity_poly section poly_entity_ids = [] poly_types = [] poly_nstd_linkages = [] poly_sequences = [] for entity_id in sorted(entity_sequences.keys()): poly_entity_ids.append(str(entity_id)) poly_types.append("polypeptide(L)") poly_nstd_linkages.append("no") poly_sequences.append(entity_sequences[entity_id]) cif_file.block["entity_poly"] = CIFCategory( name="entity_poly", columns={ "entity_id": CIFColumn( data=CIFData(array=np.array(poly_entity_ids), dtype=np.str_) ), "type": CIFColumn( data=CIFData(array=np.array(poly_types), dtype=np.str_) ), "nstd_linkage": CIFColumn( data=CIFData(array=np.array(poly_nstd_linkages), dtype=np.str_) ), "pdbx_seq_one_letter_code": CIFColumn( data=CIFData(array=np.array(poly_sequences), dtype=np.str_) ), }, ) # Create _struct_asym section asym_ids = [] asym_entity_ids = [] asym_details = [] for chain in self.chain_iter(): asym_ids.append(chain.chain_id) asym_entity_ids.append(str(chain_to_entity[chain.chain_id])) asym_details.append("") cif_file.block["struct_asym"] = CIFCategory( name="struct_asym", columns={ "id": CIFColumn(data=CIFData(array=np.array(asym_ids), dtype=np.str_)), "entity_id": CIFColumn( data=CIFData(array=np.array(asym_entity_ids), dtype=np.str_) ), "details": CIFColumn( data=CIFData(array=np.array(asym_details), dtype=np.str_) ), }, ) def get_assembly_fast( mmcif: MmcifWrapper, assembly_id=None, model=None, data_block=None, altloc="first", use_author_fields=True, ): pdbx_file = mmcif.raw if pdbx_file is None: raise InvalidFileError("No mmCIF data loaded") assembly_gen_category = pdbx_file.block["pdbx_struct_assembly_gen"] if assembly_gen_category is None: raise InvalidFileError("File has no 'pdbx_struct_assembly_gen' category") struct_oper_category = pdbx_file.block["pdbx_struct_oper_list"] if struct_oper_category is None: raise InvalidFileError("File has no 'pdbx_struct_oper_list' category") if assembly_id is None: assembly_id = assembly_gen_category["assembly_id"].data.array[0] elif assembly_id not in assembly_gen_category["assembly_id"].data.array: raise KeyError(f"File has no Assembly ID '{assembly_id}'") ### Calculate all possible transformations transformations = _get_transformations(struct_oper_category) ### Get structure according to additional parameters structure = get_structure( pdbx_file, model, data_block, altloc, ["label_asym_id"], use_author_fields )[0] # type: ignore # TODO(@zeming) This line will remove all non-protein structural elements, # we should remove this when we want to parse these too. structure: bs.AtomArray = structure[ bs.filter_amino_acids(structure) & ~structure.hetero # type: ignore ] if len(structure) == 0: raise NoProteinError unique_asym_ids = np.unique(structure.label_asym_id) # type: ignore asym2chain = {} asym2auth = {} for asym_id in unique_asym_ids: sub_structure: bs.AtomArray = structure[structure.label_asym_id == asym_id] # type: ignore chain_id: str = sub_structure[0].chain_id # type: ignore ( sequence, atom_positions, atom_mask, residue_index, insertion_code, confidence, entity_id, ) = chain_to_ndarray(sub_structure, mmcif, chain_id, False) asym2chain[asym_id] = ProteinChain( id=mmcif.id or "unknown", sequence=sequence, chain_id=chain_id, entity_id=entity_id, atom37_positions=atom_positions, atom37_mask=atom_mask, residue_index=residue_index, insertion_code=insertion_code, confidence=confidence, mmcif=None, ) asym2auth[asym_id] = chain_id ### Get transformations and apply them to the affected asym IDs assembly = [] assembly_id_dict: dict[str, list[str]] = {} # Process the target assembly ID for aid, op_expr, asym_id_expr in zip( assembly_gen_category["assembly_id"].data.array, assembly_gen_category["oper_expression"].data.array, assembly_gen_category["asym_id_list"].data.array, ): if aid == assembly_id: # Parse operations and asym IDs for this specific entry operations = _parse_operation_expression(op_expr) asym_ids = asym_id_expr.split(",") # Filter affected asym IDs to only protein chains, preserving order sub_structures = [ asym2chain[asym_id] for asym_id in asym_ids if asym_id in asym2chain ] # Apply transformations sub_assembly = _apply_transformations_fast( sub_structures, transformations, operations ) assembly.extend(sub_assembly) # Build assembly_id_dict for this entry assembly_id_dict[aid] = assembly_id_dict.get(aid, []) + [ asym2auth[id_] for id_ in asym_ids if id_ in asym2auth ] if len(assembly) == 0: raise NoProteinError return ProteinComplex.from_chains(assembly, mmcif, assembly_id_dict) def protein_chain_to_protein_complex(chain: ProteinChain) -> ProteinComplex: if "|" not in chain.sequence: return ProteinComplex.from_chains([chain]) chain_breaks = np.array(list(chain.sequence)) == "|" chain_break_inds = np.where(chain_breaks)[0] chain_break_inds = np.concatenate([[0], chain_break_inds, [len(chain)]]) chain_break_inds = np.array(list(zip(chain_break_inds[:-1], chain_break_inds[1:]))) complex_chains = [] for start, end in chain_break_inds: if start != 0: start += 1 complex_chains.append(chain[start:end]) complex_chains = [ ProteinChain.from_atom37( chain.atom37_positions, sequence=chain.sequence, chain_id=SINGLE_LETTER_CHAIN_IDS[i], entity_id=i, ) for i, chain in enumerate(complex_chains) ] return ProteinComplex.from_chains(complex_chains)