Synthyra/ESMFold2
0505
1from __future__ import annotations
2
3import io
4import itertools
5import random
6import re
7import warnings
8from dataclasses import asdict, dataclass, replace
9from functools import cached_property
10from pathlib import Path
11from subprocess import check_output
12from tempfile import TemporaryDirectory
13from typing import Any, Iterable, Sequence
14
15import biotite.structure as bs
16import brotli
17import msgpack
18import msgpack_numpy
19import numpy as np
20import torch
21from biotite.database import rcsb
22from biotite.file import InvalidFileError
23from biotite.structure.io.pdb import PDBFile
24from biotite.structure.io.pdbx import CIFCategory, CIFColumn, CIFData, CIFFile
25from biotite.structure.io.pdbx import set_structure as set_structure_pdbx
26from biotite.structure.io.pdbx.convert import _get_transformations, get_structure
27from biotite.structure.util import matrix_rotate
28from scipy.spatial import KDTree
29
30from . import esmfold2_residue_constants as residue_constants
31from .esmfold2_misc import slice_python_object_as_numpy
32from .esmfold2_affine3d import Affine3D
33from .esmfold2_aligner import Aligner
34from .esmfold2_atom_indexer import AtomIndexer
35from .esmfold2_metrics import compute_gdt_ts, compute_lddt_ca
36from .esmfold2_mmcif_parsing import MmcifWrapper, NoProteinError
37from .esmfold2_protein_chain import (
38 ProteinChain,
39 _str_key_to_int_key,
40 chain_to_ndarray,
41 index_by_atom_name,
42 infer_CB,
43)
44from .esmfold2_utils_types import PathOrBuffer
45
46msgpack_numpy.patch()
47
48SINGLE_LETTER_CHAIN_IDS = (
49 "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
50)
51
52
53def _parse_operation_expression(expression):
54 """
55 Get successive operation steps (IDs) for the given
56 ``oper_expression``.
57 Form the cartesian product, if necessary.
58 Copied from biotite and fixed a bug
59 """
60 # Split groups by parentheses:
61 # use the opening parenthesis as delimiter
62 # and just remove the closing parenthesis
63 expressions_per_step = expression.replace(")", "").split("(")
64 expressions_per_step = [e for e in expressions_per_step if len(e) > 0]
65 # Important: Operations are applied from right to left
66 expressions_per_step.reverse()
67
68 operations = []
69 for expr in expressions_per_step:
70 cur_expr = expr.split(",")
71 cur_op = []
72 # Deal with e='1-10,20-30,40-50' type expressions
73 for e in cur_expr:
74 if "-" in e:
75 first, last = e.split("-")
76 cur_op.extend(str(id) for id in range(int(first), int(last) + 1))
77 else:
78 cur_op.append(e)
79 operations.append(cur_op)
80
81 # Cartesian product of operations
82 return list(itertools.product(*operations))
83
84
85def _apply_transformations_fast(chains, transformation_dict, operations):
86 """
87 Get subassembly by applying the given operations to the input
88 structure containing affected asym IDs.
89 """
90 # Additional first dimesion for 'structure.repeat()'
91 results = []
92
93 # Apply corresponding transformation for each copy in the assembly
94 for c in chains:
95 for operation in operations:
96 coord = c.atom37_positions.copy()
97 # Execute for each transformation step
98 # in the operation expression
99 for op_step in operation:
100 T = transformation_dict[op_step]
101 # Rotate
102 coord = matrix_rotate(coord, T.rotation)
103 # Translate
104 coord += T.target_translation
105 new_chain = replace(c, atom37_positions=coord)
106 results.append(new_chain)
107
108 return results
109
110
111@dataclass
112class ProteinComplexMetadata:
113 entity_lookup: dict[int, int]
114 chain_lookup: dict[int, str]
115 mmcif: MmcifWrapper | None = None
116 # This is a dictionary that maps assembly ids to the list of unique chains
117 # in that assembly. Allows for usage of `switch_assembly`.
118 assembly_composition: dict[str, list[str]] | None = None
119
120
121@dataclass
122class DockQSingleScore:
123 native_chains: tuple[str, str]
124 DockQ: float
125 interface_rms: float
126 ligand_rms: float
127 fnat: float
128 fnonnat: float
129 clashes: float
130 F1: float
131 DockQ_F1: float
132
133
134@dataclass
135class DockQResult:
136 total_dockq: float
137 native_interfaces: int
138 chain_mapping: dict[str, str]
139 interfaces: dict[tuple[str, str], DockQSingleScore]
140 # zip(aligned.chain_iter(), native.chain_iter()) gives you the pairing
141 # aligned.rmsd(native) should give you a low rmsd irrespective of shuffling
142 aligned: ProteinComplex
143 aligned_rmsd: float
144
145
146@dataclass(frozen=True)
147class ProteinComplex:
148 """Dataclass with atom37 representation of an entire protein complex."""
149
150 id: str
151 sequence: str
152 entity_id: np.ndarray # entities map to unique sequences
153 chain_id: np.ndarray # multiple chains might share an entity id
154 sym_id: np.ndarray # complexes might be copies of the same chain
155 residue_index: np.ndarray
156 insertion_code: np.ndarray
157 atom37_positions: np.ndarray
158 atom37_mask: np.ndarray
159 confidence: np.ndarray
160 # This metadata is parsed from the MMCIF file. For synthetic data, we do a best effort.
161 metadata: ProteinComplexMetadata
162 atom37_confidence: np.ndarray | None = None # [L, 37] per-atom pLDDT
163
164 def __post_init__(self):
165 l = len(self.sequence)
166 assert self.atom37_positions.shape[0] == l, (self.atom37_positions.shape, l)
167 assert self.atom37_mask.shape[0] == l, (self.atom37_mask.shape, l)
168 assert self.residue_index.shape[0] == l, (self.residue_index.shape, l)
169 assert self.insertion_code.shape[0] == l, (self.insertion_code.shape, l)
170 assert self.confidence.shape[0] == l, (self.confidence.shape, l)
171 assert self.entity_id.shape[0] == l, (self.entity_id.shape, l)
172 assert self.chain_id.shape[0] == l, (self.chain_id.shape, l)
173 assert self.sym_id.shape[0] == l, (self.sym_id.shape, l)
174 if self.atom37_confidence is not None:
175 assert self.atom37_confidence.shape == self.atom37_mask.shape, (
176 self.atom37_confidence.shape,
177 self.atom37_mask.shape,
178 )
179
180 def __getitem__(self, idx: int | list[int] | slice | np.ndarray):
181 """This function slices protein complexes without consideration of chain breaks
182 NOTE: When slicing with a boolean mask, it's possible that the output array won't
183 be the expected length. This is because we do our best to preserve chainbreak tokens.
184 """
185
186 if isinstance(idx, int):
187 idx = [idx]
188 if isinstance(idx, list):
189 raise ValueError(
190 "ProteinComplex doesn't supports indexing with lists of indices"
191 )
192
193 if isinstance(idx, np.ndarray):
194 is_chainbreak = np.asarray([s == "|" for s in self.sequence])
195 idx = idx.astype(bool) | is_chainbreak
196
197 complex = self._unsafe_slice(idx)
198 if len(complex) == 0:
199 return complex
200
201 # detect runs of chainbreaks by searching for instances of '||' in complex.sequence
202 chainbreak_runs = np.asarray(
203 [
204 complex.sequence[i : i + 2] == "||"
205 for i in range(len(complex.sequence) - 1)
206 ]
207 + [complex.sequence[-1] == "|"]
208 )
209 # We should remove as many chainbreaks as possible from the start of the sequence
210 for i in range(len(chainbreak_runs)):
211 if complex.sequence[i] == "|":
212 chainbreak_runs[i] = True
213 else:
214 break
215 complex = complex._unsafe_slice(~chainbreak_runs)
216 return complex
217
218 def _unsafe_slice(self, idx: int | list[int] | slice | np.ndarray):
219 sequence = slice_python_object_as_numpy(self.sequence, idx)
220 return replace(
221 self,
222 sequence=sequence,
223 entity_id=self.entity_id[..., idx],
224 chain_id=self.chain_id[..., idx],
225 sym_id=self.sym_id[..., idx],
226 residue_index=self.residue_index[..., idx],
227 insertion_code=self.insertion_code[..., idx],
228 atom37_positions=self.atom37_positions[..., idx, :, :],
229 atom37_mask=self.atom37_mask[..., idx, :],
230 confidence=self.confidence[..., idx],
231 atom37_confidence=self.atom37_confidence[..., idx, :]
232 if self.atom37_confidence is not None
233 else None,
234 )
235
236 def __len__(self):
237 return len(self.sequence)
238
239 @property
240 def num_chains(self):
241 return len(self.chain_boundaries)
242
243 @cached_property
244 def atoms(self) -> AtomIndexer:
245 return AtomIndexer(self, property="atom37_positions", dim=-2)
246
247 @cached_property
248 def atom_mask(self) -> AtomIndexer:
249 return AtomIndexer(self, property="atom37_mask", dim=-1)
250
251 @cached_property
252 def chain_lengths(self) -> np.ndarray:
253 return np.diff(self.chain_boundaries, axis=1).flatten()
254
255 @cached_property
256 def chain_boundaries(self) -> list[tuple[int, int]]:
257 cb = [-1]
258 for i, s in enumerate(self.sequence):
259 if s == "|":
260 cb.append(i)
261 cb.append(len(self))
262 return [(cb[i] + 1, cb[i + 1]) for i in range(len(cb) - 1)]
263
264 def get_chain_by_index(self, index: int) -> ProteinChain:
265 try:
266 start, end = self.chain_boundaries[index]
267 return self[start:end].as_chain()
268 except IndexError:
269 raise IndexError(f"Chain index {index} out of bounds")
270
271 def get_chain_by_id(
272 self, chain_id: str, sample_chain_if_duplicate: bool = True
273 ) -> ProteinChain:
274 valid_indices = [
275 index
276 for index, id_of_index in self.metadata.chain_lookup.items()
277 if id_of_index == chain_id
278 ]
279 if not valid_indices:
280 raise KeyError(f"Chain ID {chain_id} not found")
281 if sample_chain_if_duplicate:
282 index_to_return = random.choice(valid_indices)
283 return self.get_chain_by_index(index_to_return)
284 else:
285 if len(valid_indices) > 1:
286 raise ValueError(f"Multiple chains with chain ID {chain_id} found")
287 return self.get_chain_by_index(valid_indices[0])
288
289 def chain_iter(self) -> Iterable[ProteinChain]:
290 for start, end in self.chain_boundaries:
291 c = self[start:end]
292 yield c.as_chain()
293
294 def as_chain(self, force_conversion: bool = False) -> ProteinChain:
295 """Convert the ProteinComplex to a ProteinChain.
296
297 Args:
298 force_conversion (bool): Forces the conversion into a protein chain even if the complex has multiple chains.
299 The purpose of this is to use ProteinChain specific functions (like cbeta_contacts).
300
301 """
302 if not force_conversion:
303 assert len(np.unique(self.chain_id)) == 1, f"{self.id}"
304 assert len(np.unique(self.entity_id)) == 1, f"{self.id}"
305 if self.chain_id[0] not in self.metadata.chain_lookup:
306 warnings.warn("Chain ID not found in metadata, using 'A' as default")
307 if self.entity_id[0] not in self.metadata.entity_lookup:
308 warnings.warn("Entity ID not found in metadata, using None as default")
309 chain_id = self.metadata.chain_lookup.get(self.chain_id[0], "A")
310 entity_id = self.metadata.entity_lookup.get(self.entity_id[0], None)
311 else:
312 chain_id = "A"
313 entity_id = None
314
315 return ProteinChain(
316 id=self.id,
317 sequence=self.sequence,
318 chain_id=chain_id,
319 entity_id=entity_id,
320 atom37_positions=self.atom37_positions,
321 atom37_mask=self.atom37_mask,
322 residue_index=self.residue_index,
323 insertion_code=self.insertion_code,
324 confidence=self.confidence,
325 mmcif=self.metadata.mmcif,
326 atom37_confidence=self.atom37_confidence,
327 )
328
329 @classmethod
330 def from_pdb(
331 cls, path: PathOrBuffer, id: str | None = None, is_predicted: bool = False
332 ) -> "ProteinComplex":
333 atom_array = PDBFile.read(path).get_structure(
334 model=1, extra_fields=["b_factor"]
335 )
336
337 chains = []
338 for chain in bs.chain_iter(atom_array):
339 chain = chain[~chain.hetero]
340 if len(chain) == 0:
341 continue
342 chains.append(ProteinChain.from_atomarray(chain, id, is_predicted))
343 return ProteinComplex.from_chains(chains)
344
345 def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True):
346 atom_array = None
347 for chain in self.chain_iter():
348 carr = (
349 chain.atom_array
350 if include_insertions
351 else chain.atom_array_no_insertions
352 )
353 atom_array = carr if atom_array is None else atom_array + carr
354 f = PDBFile()
355 f.set_structure(atom_array)
356 f.write(path)
357
358 def to_pdb_string(self, include_insertions: bool = True) -> str:
359 buf = io.StringIO()
360 self.to_pdb(buf, include_insertions=include_insertions)
361 buf.seek(0)
362 return buf.read()
363
364 def normalize_chain_ids_for_pdb(self):
365 # Since PDB files have 1-letter chain IDs and don't support the idea of a symmetric index,
366 # we can normalize it instead which might be necessary for DockQ and to_pdb.
367 ids = SINGLE_LETTER_CHAIN_IDS
368 chains = []
369 for i, chain in enumerate(self.chain_iter()):
370 chain = replace(chain, chain_id=ids[i])
371 if i > len(ids):
372 raise RuntimeError("Too many chains to write to PDB file")
373 chains.append(chain)
374
375 return ProteinComplex.from_chains(chains)
376
377 def find_assembly_ids_with_chain(self, id: str) -> list[str]:
378 good_chains = []
379 if (comp := self.metadata.assembly_composition) is not None:
380 for assembly_id, chain_ids in comp.items():
381 if id in chain_ids:
382 good_chains.append(assembly_id)
383 else:
384 raise ValueError(
385 "Cannot switch assemblies on this ProteinComplex, you must create the assembly from mmcif to support this"
386 )
387 return good_chains
388
389 def switch_assembly(self, id: str):
390 assert self.metadata.mmcif is not None
391 return get_assembly_fast(self.metadata.mmcif, assembly_id=id)
392
393 def state_dict(self, backbone_only=False, json_serializable=False):
394 """This state dict is optimized for storage, so it turns things to fp16 whenever
395 possible. Note that we also only support int32 residue indices, I'm hoping we don't
396 need more than 2**32 residues..."""
397 dct = {k: v for k, v in vars(self).items()}
398 if backbone_only:
399 dct["atom37_mask"][:, 3:] = False
400 dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]]
401 if dct.get("atom37_confidence") is not None:
402 dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]]
403 else:
404 dct.pop("atom37_confidence", None)
405 for k, v in dct.items():
406 if isinstance(v, np.ndarray):
407 match v.dtype:
408 case np.int64:
409 dct[k] = v.astype(np.int32)
410 case np.float64 | np.float32:
411 dct[k] = v.astype(np.float16)
412 case _:
413 pass
414 if json_serializable:
415 dct[k] = v.tolist()
416 elif isinstance(v, ProteinComplexMetadata):
417 dct[k] = asdict(v)
418 dct["metadata"]["mmcif"] = None
419 # These can be populated with non-serializable objects and are not needed for reconstruction
420 dct.pop("atoms", None)
421 dct.pop("atom_mask", None)
422 dct.pop("per_chain_kd_trees", None)
423 return dct
424
425 def to_blob(self, backbone_only=False) -> bytes:
426 return brotli.compress(msgpack.dumps(self.state_dict(backbone_only)), quality=5)
427
428 @classmethod
429 def from_state_dict(cls, dct):
430 # Note: assembly_composition is *supposed* to have string keys.
431 dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"])
432
433 for k, v in dct.items():
434 if isinstance(v, list):
435 dct[k] = np.array(v)
436
437 atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan)
438 atom37[dct["atom37_mask"]] = dct["atom37_positions"]
439 dct["atom37_positions"] = atom37
440 if "atom37_confidence" in dct:
441 atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32)
442 atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"]
443 dct["atom37_confidence"] = atom37_conf
444 dct = {
445 k: (
446 v.astype(np.float32)
447 if k in ["atom37_positions", "confidence", "atom37_confidence"]
448 else v
449 )
450 for k, v in dct.items()
451 }
452 if "chain_boundaries" in dct:
453 del dct["chain_boundaries"]
454 if "chain_boundaries" in dct["metadata"]:
455 del dct["metadata"]["chain_boundaries"]
456 dct["metadata"] = ProteinComplexMetadata(**dct["metadata"])
457 return cls(**dct)
458
459 @classmethod
460 def from_blob(cls, input: Path | str | io.BytesIO | bytes):
461 """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory
462 of chains from 52G/1M chains to 20G/1M chains, I think this is a good first
463 shot at compressing and dumping chains to disk. I'm sure there's better ways."""
464 match input:
465 case Path() | str():
466 bytes = Path(input).read_bytes()
467 case io.BytesIO():
468 bytes = input.getvalue()
469 case _:
470 bytes = input
471 return cls.from_state_dict(
472 msgpack.loads(brotli.decompress(bytes), strict_map_key=False)
473 )
474
475 @classmethod
476 def from_rcsb(cls, pdb_id: str, keep_source: bool = False) -> ProteinComplex:
477 f: io.StringIO = rcsb.fetch(pdb_id, "cif") # type: ignore
478 return cls.from_mmcif(f, id=pdb_id, keep_source=keep_source, is_predicted=False)
479
480 @classmethod
481 def from_mmcif(
482 cls,
483 path: PathOrBuffer,
484 id: str | None = None,
485 assembly_id: str | None = None,
486 is_predicted: bool = False,
487 keep_source: bool = False,
488 ):
489 """Return a ProteinComplex object from an mmcif file.
490 TODO(@zeming): there's actually multiple complexes per file, but for ease of implementation,
491 we only consider the first defined complex!
492
493 Args:
494 path (str | Path | io.TextIO): Path or buffer to read mmcif file from. Should be uncompressed.
495 id (str, optional): String identifier to assign to structure. Will attempt to infer otherwise.
496 is_predicted (bool): If True, reads b factor as the confidence readout. Default: False.
497 chain_id (str, optional): Select a chain corresponding to (author) chain id.
498 """
499 mmcif = MmcifWrapper.read(path, id)
500 return get_assembly_fast(mmcif, assembly_id=assembly_id)
501
502 @classmethod
503 def from_chains(
504 cls,
505 chains: Sequence[ProteinChain],
506 mmcif: MmcifWrapper | None = None,
507 all_assembly_metadata_dictionary: dict[str, list[str]] | None = None,
508 ):
509 if not chains:
510 raise ValueError(
511 "Cannot create a ProteinComplex from an empty list of chains"
512 )
513
514 # TODO(roshan): Make a proper protein complex class
515 def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray):
516 full_array = []
517 for array in arrays:
518 full_array.append(array)
519 full_array.append(sep)
520 full_array = full_array[:-1]
521 return np.concatenate(full_array, 0)
522
523 sep_tokens = {
524 "residue_index": np.array([-1]),
525 "insertion_code": np.array([""]),
526 "atom37_positions": np.full([1, 37, 3], np.nan),
527 "atom37_mask": np.zeros([1, 37], dtype=bool),
528 "confidence": np.array([0]),
529 }
530
531 any_has_atom37_conf = any(c.atom37_confidence is not None for c in chains)
532 if any_has_atom37_conf:
533 sep_tokens["atom37_confidence"] = np.full([1, 37], np.nan, dtype=np.float32)
534
535 def _get_chain_attr(chain: ProteinChain, name: str) -> np.ndarray:
536 val = getattr(chain, name)
537 if val is None and name == "atom37_confidence":
538 return np.full([len(chain), 37], np.nan, dtype=np.float32)
539 return val
540
541 array_args: dict[str, np.ndarray] = {
542 name: join_arrays([_get_chain_attr(chain, name) for chain in chains], sep)
543 for name, sep in sep_tokens.items()
544 }
545
546 multimer_arrays = []
547 chain2num_max = -1
548 chain2num = {}
549 ent2num_max = -1
550 ent2num = {}
551 total_index = 0
552 for i, c in enumerate(chains):
553 num_res = c.residue_index.shape[0]
554 if c.chain_id not in chain2num:
555 chain2num[c.chain_id] = (chain2num_max := chain2num_max + 1)
556 chain_id_array = np.full([num_res], chain2num[c.chain_id], dtype=np.int64)
557
558 if c.entity_id is None:
559 entity_num = (ent2num_max := ent2num_max + 1)
560 else:
561 if c.entity_id not in ent2num:
562 ent2num[c.entity_id] = (ent2num_max := ent2num_max + 1)
563 entity_num = ent2num[c.entity_id]
564 entity_id_array = np.full([num_res], entity_num, dtype=np.int64)
565
566 sym_id_array = np.full([num_res], i, dtype=np.int64)
567
568 multimer_arrays.append(
569 {
570 "chain_id": chain_id_array,
571 "entity_id": entity_id_array,
572 "sym_id": sym_id_array,
573 }
574 )
575
576 total_index += num_res + 1
577
578 sep = np.array([-1])
579 update = {
580 name: join_arrays([dct[name] for dct in multimer_arrays], sep=sep)
581 for name in ["chain_id", "entity_id", "sym_id"]
582 }
583 array_args.update(update)
584
585 metadata = ProteinComplexMetadata(
586 mmcif=mmcif,
587 chain_lookup={v: k for k, v in chain2num.items()},
588 entity_lookup={v: k for k, v in ent2num.items()},
589 assembly_composition=all_assembly_metadata_dictionary,
590 )
591
592 return cls(
593 id=chains[0].id,
594 sequence=residue_constants.CHAIN_BREAK_TOKEN.join(
595 chain.sequence for chain in chains
596 ),
597 metadata=metadata,
598 **array_args,
599 )
600
601 def infer_oxygen(self) -> ProteinComplex:
602 """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided."""
603 O_missing_indices = np.argwhere(
604 ~np.isfinite(self.atoms["O"]).all(axis=1)
605 ).squeeze()
606
607 O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32)
608 N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1)
609 N = torch.roll(N, -3)
610 N[..., -1, :] = torch.nan
611
612 # Get the frame defined by the CA-C-N atom
613 frames = Affine3D.from_graham_schmidt(CA, C, N)
614 O = frames.apply(O_vector)
615 atom37_positions = self.atom37_positions.copy()
616 atom37_mask = self.atom37_mask.copy()
617
618 atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = O[
619 O_missing_indices
620 ].numpy()
621 atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan(
622 atom37_positions[O_missing_indices, residue_constants.atom_order["O"]]
623 ).any(-1)
624 new_chain = replace(
625 self, atom37_positions=atom37_positions, atom37_mask=atom37_mask
626 )
627 return new_chain
628
629 def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinComplex:
630 """Return a new chain with inferred CB atoms at all residues except GLY.
631
632 Args:
633 infer_cbeta_for_glycine (bool): If True, infers a beta carbon for glycine
634 residues, even though that residue doesn't have one. Default off.
635
636 NOTE(rverkuil): The reason for having this switch in the first place
637 is that sometimes we want a (inferred) CB coordinate for every residue,
638 for example for making a pairwise distance matrix, or doing an RMSD
639 calculation between two designs for a given structural template, w/
640 CB atoms.
641 """
642 atom37_positions = self.atom37_positions.copy()
643 atom37_mask = self.atom37_mask.copy()
644
645 N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0)
646 # See usage in trDesign codebase.
647 # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140
648 inferred_cbeta_positions = infer_CB(C, N, CA, 1.522, 1.927, -2.143)
649 if not infer_cbeta_for_glycine:
650 inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan
651
652 atom37_positions[:, residue_constants.atom_order["CB"]] = (
653 inferred_cbeta_positions
654 )
655 atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan(
656 atom37_positions[:, residue_constants.atom_order["CB"]]
657 ).any(-1)
658 new_chain = replace(
659 self, atom37_positions=atom37_positions, atom37_mask=atom37_mask
660 )
661 return new_chain
662
663 @classmethod
664 def from_open_source(cls, pc: ProteinComplex):
665 # TODO(@zeming): deprecated, should delete
666 return pc
667
668 @classmethod
669 def concat(cls, objs: list[ProteinComplex]) -> ProteinComplex:
670 pdb_ids = [obj.id for obj in objs]
671 if len(set(pdb_ids)) > 1:
672 raise RuntimeError(
673 "Concatention of protein complexes across different PDB ids is unsupported"
674 )
675 return ProteinComplex.from_chains(
676 list(itertools.chain.from_iterable(obj.chain_iter() for obj in objs))
677 )
678
679 def _sanity_check_complexes_are_comparable(self, other: ProteinComplex):
680 assert len(self) == len(other), "Protein complexes must have the same length"
681 assert len(list(self.chain_iter())) == len(
682 list(other.chain_iter())
683 ), "Protein complexes must have the same number of chains"
684
685 def rmsd(
686 self,
687 target: ProteinComplex,
688 also_check_reflection: bool = False,
689 mobile_inds: list[int] | np.ndarray | None = None,
690 target_inds: list[int] | np.ndarray | None = None,
691 only_compute_backbone_rmsd: bool = False,
692 compute_chain_assignment: bool = True,
693 ):
694 """
695 Compute the RMSD between this protein chain and another.
696
697 Args:
698 target (ProteinComplex): The target (other) protein complex to compare to.
699 also_check_reflection (bool, optional): If True, also check if the reflection of the mobile atoms has a lower RMSD.
700 mobile_inds (list[int], optional): The indices of the mobile atoms to align. These are NOT residue indices
701 target_inds (list[int], optional): The indices of the target atoms to align. These are NOT residue indices
702 only_compute_backbone_rmsd (bool, optional): If True, only compute the RMSD of the backbone atoms.
703 """
704 if compute_chain_assignment:
705 aligned = self.dockq(target).aligned
706 else:
707 aligned = self
708
709 aligner = Aligner(
710 aligned if mobile_inds is None else aligned[mobile_inds],
711 target if target_inds is None else target[target_inds],
712 only_compute_backbone_rmsd,
713 )
714 avg_rmsd = aligner.rmsd
715
716 if not also_check_reflection:
717 return avg_rmsd
718
719 aligner = Aligner(
720 aligned if mobile_inds is None else aligned[mobile_inds],
721 target if target_inds is None else target[target_inds],
722 only_compute_backbone_rmsd,
723 use_reflection=True,
724 )
725 avg_rmsd_neg = aligner.rmsd
726
727 return min(avg_rmsd, avg_rmsd_neg)
728
729 def lddt_ca(
730 self,
731 target: ProteinComplex,
732 mobile_inds: list[int] | np.ndarray | None = None,
733 target_inds: list[int] | np.ndarray | None = None,
734 compute_chain_assignment: bool = True,
735 **kwargs,
736 ) -> float | np.ndarray:
737 """Compute the LDDT between this protein complex and another.
738
739 Arguments:
740 target (ProteinComplex): The other protein complex to compare to.
741 mobile_inds (list[int], np.ndarray, optional): The indices of the mobile atoms to align. These are NOT residue indices
742 target_inds (list[int], np.ndarray, optional): The indices of the target atoms to align. These are NOT residue indices
743
744 Returns:
745 float | np.ndarray: The LDDT score between the two protein chains, either
746 a single float or per-residue LDDT scores if `per_residue` is True.
747 """
748 if compute_chain_assignment:
749 aligned = self.dockq(target).aligned
750 else:
751 aligned = self
752 lddt = compute_lddt_ca(
753 torch.tensor(aligned.atom37_positions[mobile_inds]).unsqueeze(0),
754 torch.tensor(target.atom37_positions[target_inds]).unsqueeze(0),
755 torch.tensor(aligned.atom37_mask[mobile_inds]).unsqueeze(0),
756 **kwargs,
757 )
758 return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten()
759
760 def gdt_ts(
761 self,
762 target: ProteinComplex,
763 mobile_inds: list[int] | np.ndarray | None = None,
764 target_inds: list[int] | np.ndarray | None = None,
765 compute_chain_assignment: bool = True,
766 **kwargs,
767 ) -> float | np.ndarray:
768 """Compute the GDT_TS between this protein complex and another.
769
770 Arguments:
771 target (ProteinComplex): The other protein complex to compare to.
772 mobile_inds (list[int], np.ndarray, optional): The indices of the mobile atoms to align. These are NOT residue indices
773 target_inds (list[int], np.ndarray, optional): The indices of the target atoms to align. These are NOT residue indices
774
775 Returns:
776 float: The GDT_TS score between the two protein chains.
777 """
778 if compute_chain_assignment:
779 aligned = self.dockq(target).aligned
780 else:
781 aligned = self
782 gdt_ts = compute_gdt_ts(
783 mobile=torch.tensor(
784 index_by_atom_name(aligned.atom37_positions[mobile_inds], "CA"),
785 dtype=torch.float32,
786 ).unsqueeze(0),
787 target=torch.tensor(
788 index_by_atom_name(target.atom37_positions[target_inds], "CA"),
789 dtype=torch.float32,
790 ).unsqueeze(0),
791 atom_exists_mask=torch.tensor(
792 index_by_atom_name(aligned.atom37_mask[mobile_inds], "CA", dim=-1)
793 & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1)
794 ).unsqueeze(0),
795 **kwargs,
796 )
797 return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten()
798
799 def dockq(self, native: ProteinComplex):
800 # This function uses dockqv2 to compute the DockQ score. Because it does a mapping
801 # over all possible chains, it's quite slow. Be careful not to use this in an inference loop
802 # or something that requires fast scoring. It defaults to 8 CPUs.
803 #
804 # TODO(@zeming): Because we haven't properly implemented protein complexes for mmcif,
805 # if your protein has multi-letter or repeated chain IDs, this will fail. Please call
806 # pc = pc.normalize_chain_ids_for_pdb() before calling this function in that case (limit is 62 chains)
807
808 try:
809 pass
810 except BaseException:
811 raise RuntimeError(
812 "DockQ is not installed. Please update your environment."
813 )
814 self._sanity_check_complexes_are_comparable(native)
815
816 def sanity_check_chain_ids(pc: ProteinComplex):
817 ids = []
818 for i, chain in enumerate(pc.chain_iter()):
819 if i > len(SINGLE_LETTER_CHAIN_IDS):
820 raise ValueError("Too many chains to write to PDB file")
821 if len(chain.chain_id) > 1:
822 raise ValueError(
823 "We only supports single letter chain IDs for DockQ"
824 )
825 ids.append(chain.chain_id)
826 if len(set(ids)) != len(ids):
827 raise ValueError(f"Duplicate chain IDs in protein complex: {ids}")
828 return ids
829
830 sanity_check_chain_ids(self)
831 sanity_check_chain_ids(native)
832
833 with TemporaryDirectory() as tdir:
834 dir = Path(tdir)
835 self.to_pdb(dir / "self.pdb")
836 native.to_pdb(dir / "native.pdb")
837
838 output = check_output(["DockQ", dir / "self.pdb", dir / "native.pdb"])
839 lines = output.decode().split("\n")
840
841 # Remove the header comments
842 start_index = next(
843 i for i, line in enumerate(lines) if line.startswith("Model")
844 )
845 lines = lines[start_index:]
846
847 result = {}
848 interfaces = []
849 current_interface: dict = {}
850
851 for line in lines:
852 line = line.strip()
853 if not line:
854 continue
855
856 if line.startswith("Model :"):
857 pass # Tmp pdb file location, it's useless...
858 elif line.startswith("Native :"):
859 pass # Tmp pdb file location, it's useless...
860 elif line.startswith("Total DockQ"):
861 total_dockq_match = re.search(
862 r"Total DockQ over (\d+) native interfaces: ([\d.]+) with (.*) model:native mapping",
863 line,
864 )
865 if total_dockq_match:
866 result["value"] = float(total_dockq_match.group(2))
867 result["native interfaces"] = int(total_dockq_match.group(1))
868 native_chains, self_chains = total_dockq_match.group(3).split(":")
869 result["mapping"] = dict(zip(native_chains, self_chains))
870 else:
871 raise RuntimeError(
872 "Failed to parse DockQ output, maybe your DockQ version is wrong?"
873 )
874 elif line.startswith("Native chains:"):
875 if current_interface:
876 interfaces.append(current_interface)
877 current_interface = {
878 "Native chains": line.split(":")[1].strip().split(", ")
879 }
880 elif line.startswith("Model chains:"):
881 current_interface["Model chains"] = (
882 line.split(":")[1].strip().split(", ")
883 )
884 elif ":" in line:
885 key, value = line.split(":", 1)
886 current_interface[key.strip()] = float(value.strip())
887
888 if current_interface:
889 interfaces.append(current_interface)
890
891 def parse_dict(d: dict[str, Any]) -> DockQSingleScore:
892 return DockQSingleScore(
893 native_chains=tuple(d["Native chains"]), # type: ignore
894 DockQ=float(d["DockQ"]),
895 interface_rms=float(d["irms"]),
896 ligand_rms=float(d["Lrms"]), # Note the capitalization difference
897 fnat=float(d["fnat"]),
898 fnonnat=float(d["fnonnat"]),
899 clashes=float(d["clashes"]),
900 F1=float(d["F1"]),
901 DockQ_F1=float(d["DockQ_F1"]),
902 )
903
904 inv_mapping = {v: k for k, v in result["mapping"].items()}
905
906 self_chain_map = {c.chain_id: c for c in self.chain_iter()}
907 realigned = []
908 for chain in native.chain_iter():
909 realigned.append(self_chain_map[inv_mapping[chain.chain_id]])
910
911 realigned = ProteinComplex.from_chains(realigned)
912 aligner = Aligner(realigned, native)
913 realigned = aligner.apply(realigned)
914
915 result = DockQResult(
916 total_dockq=result["value"],
917 native_interfaces=result["native interfaces"],
918 chain_mapping=result["mapping"],
919 interfaces={
920 (i["Model chains"][0], i["Model chains"][1]): parse_dict(i)
921 for i in interfaces
922 },
923 aligned=realigned,
924 aligned_rmsd=aligner.rmsd,
925 )
926
927 return result
928
929 @cached_property
930 def per_chain_kd_trees(self):
931 # Iterate over chains, build KDTree for each chain
932 kdtrees = []
933
934 CA = self.atoms["CA"]
935
936 for start, end in self.chain_boundaries:
937 chain_CA = CA[start:end]
938 chain_CA = chain_CA[np.isfinite(chain_CA).all(axis=-1)]
939 kdtrees.append(KDTree(chain_CA))
940
941 return kdtrees
942
943 def chain_adjacency(self, cutoff: float = 8.0) -> np.ndarray:
944 # Compute adjacency matrix for protein complex
945 num_chains = self.num_chains
946 adjacency = np.zeros((num_chains, num_chains), dtype=bool)
947 for (i, kdtree), (j, kdtree2) in itertools.combinations(
948 enumerate(self.per_chain_kd_trees), 2
949 ):
950 adj = kdtree.query_ball_tree(kdtree2, cutoff)
951 any_is_adjacent = any(len(a) > 0 for a in adj)
952 adjacency[i, j] = any_is_adjacent
953 adjacency[j, i] = any_is_adjacent
954 return adjacency
955
956 def chain_adjacency_by_index(self, index: int, cutoff: float = 8.0) -> np.ndarray:
957 num_chains = len(self.chain_boundaries)
958 adjacency = np.zeros(num_chains, dtype=bool)
959 for i, kdtree in enumerate(self.per_chain_kd_trees):
960 if i == index:
961 continue
962 adj = kdtree.query_ball_tree(self.per_chain_kd_trees[index], cutoff)
963 adjacency[i] = any(len(a) > 0 for a in adj)
964 return adjacency
965
966 def add_prefix_to_chain_ids(self, prefix: str) -> ProteinComplex:
967 """Rename all chains in the complex with a given prefix.
968
969 Args:
970 prefix (str): The prefix to use for the new chain IDs. Each chain will be
971 named as "{prefix}_{chain_id}".
972
973 Returns:
974 ProteinComplex: A new protein complex with renamed chains.
975 """
976 new_chains = []
977 for chain in self.chain_iter():
978 # Create new chain with updated chain_id
979 new_chain = replace(chain, chain_id=f"{prefix}_{chain.chain_id}")
980 new_chains.append(new_chain)
981 return ProteinComplex.from_chains(new_chains)
982
983 def sasa(self, by_residue: bool = True):
984 chain = self.as_chain(force_conversion=True)
985 return chain.sasa(by_residue=by_residue)
986
987 def to_mmcif_string(self) -> str:
988 """Convert the ProteinComplex to mmCIF format.
989
990 Returns:
991 str: The mmCIF content as a string.
992 """
993 # Convert the ProteinComplex to a biotite AtomArray
994 # Collect all atoms from all chains
995 all_atoms = []
996 for chain in self.chain_iter():
997 chain_atom_array = chain.atom_array
998 # Convert AtomArray to list of atoms and add to collection
999 all_atoms.extend(chain_atom_array)
1000
1001 # Create combined AtomArray from all atoms
1002 if not all_atoms:
1003 raise ValueError("No atoms found in protein complex")
1004
1005 atom_array = bs.array(all_atoms)
1006
1007 # Create CIF file
1008 f = CIFFile()
1009 set_structure_pdbx(f, atom_array, data_block=self.id)
1010
1011 # Add entity information for proper mmCIF structure
1012 self._add_entity_information(f)
1013
1014 # Write to string
1015 output = io.StringIO()
1016 f.write(output)
1017 return output.getvalue()
1018
1019 def _add_entity_information(self, cif_file: CIFFile) -> None:
1020 """Add entity, entity_poly, and struct_asym sections to CIF file."""
1021
1022 # Group chains by sequence to create unique entities
1023 entity_map = {} # sequence -> entity_id
1024 chain_to_entity = {} # chain_id -> entity_id
1025 entity_sequences = {} # entity_id -> sequence
1026 entity_id_counter = 1
1027
1028 for chain in self.chain_iter():
1029 sequence = chain.sequence
1030 if sequence not in entity_map:
1031 entity_map[sequence] = entity_id_counter
1032 entity_sequences[entity_id_counter] = sequence
1033 entity_id_counter += 1
1034 chain_to_entity[chain.chain_id] = entity_map[sequence]
1035
1036 # Create _entity section
1037 entity_ids = []
1038 entity_types = []
1039 entity_descriptions = []
1040
1041 for entity_id in sorted(entity_sequences.keys()):
1042 entity_ids.append(str(entity_id))
1043 entity_types.append("polymer")
1044 entity_descriptions.append(f"Protein chain (entity {entity_id})")
1045
1046 cif_file.block["entity"] = CIFCategory(
1047 name="entity",
1048 columns={
1049 "id": CIFColumn(
1050 data=CIFData(array=np.array(entity_ids), dtype=np.str_)
1051 ),
1052 "type": CIFColumn(
1053 data=CIFData(array=np.array(entity_types), dtype=np.str_)
1054 ),
1055 "pdbx_description": CIFColumn(
1056 data=CIFData(array=np.array(entity_descriptions), dtype=np.str_)
1057 ),
1058 },
1059 )
1060
1061 # Create _entity_poly section
1062 poly_entity_ids = []
1063 poly_types = []
1064 poly_nstd_linkages = []
1065 poly_sequences = []
1066
1067 for entity_id in sorted(entity_sequences.keys()):
1068 poly_entity_ids.append(str(entity_id))
1069 poly_types.append("polypeptide(L)")
1070 poly_nstd_linkages.append("no")
1071 poly_sequences.append(entity_sequences[entity_id])
1072
1073 cif_file.block["entity_poly"] = CIFCategory(
1074 name="entity_poly",
1075 columns={
1076 "entity_id": CIFColumn(
1077 data=CIFData(array=np.array(poly_entity_ids), dtype=np.str_)
1078 ),
1079 "type": CIFColumn(
1080 data=CIFData(array=np.array(poly_types), dtype=np.str_)
1081 ),
1082 "nstd_linkage": CIFColumn(
1083 data=CIFData(array=np.array(poly_nstd_linkages), dtype=np.str_)
1084 ),
1085 "pdbx_seq_one_letter_code": CIFColumn(
1086 data=CIFData(array=np.array(poly_sequences), dtype=np.str_)
1087 ),
1088 },
1089 )
1090
1091 # Create _struct_asym section
1092 asym_ids = []
1093 asym_entity_ids = []
1094 asym_details = []
1095
1096 for chain in self.chain_iter():
1097 asym_ids.append(chain.chain_id)
1098 asym_entity_ids.append(str(chain_to_entity[chain.chain_id]))
1099 asym_details.append("")
1100
1101 cif_file.block["struct_asym"] = CIFCategory(
1102 name="struct_asym",
1103 columns={
1104 "id": CIFColumn(data=CIFData(array=np.array(asym_ids), dtype=np.str_)),
1105 "entity_id": CIFColumn(
1106 data=CIFData(array=np.array(asym_entity_ids), dtype=np.str_)
1107 ),
1108 "details": CIFColumn(
1109 data=CIFData(array=np.array(asym_details), dtype=np.str_)
1110 ),
1111 },
1112 )
1113
1114
1115def get_assembly_fast(
1116 mmcif: MmcifWrapper,
1117 assembly_id=None,
1118 model=None,
1119 data_block=None,
1120 altloc="first",
1121 use_author_fields=True,
1122):
1123 pdbx_file = mmcif.raw
1124 if pdbx_file is None:
1125 raise InvalidFileError("No mmCIF data loaded")
1126 assembly_gen_category = pdbx_file.block["pdbx_struct_assembly_gen"]
1127 if assembly_gen_category is None:
1128 raise InvalidFileError("File has no 'pdbx_struct_assembly_gen' category")
1129
1130 struct_oper_category = pdbx_file.block["pdbx_struct_oper_list"]
1131 if struct_oper_category is None:
1132 raise InvalidFileError("File has no 'pdbx_struct_oper_list' category")
1133
1134 if assembly_id is None:
1135 assembly_id = assembly_gen_category["assembly_id"].data.array[0]
1136 elif assembly_id not in assembly_gen_category["assembly_id"].data.array:
1137 raise KeyError(f"File has no Assembly ID '{assembly_id}'")
1138
1139 ### Calculate all possible transformations
1140 transformations = _get_transformations(struct_oper_category)
1141
1142 ### Get structure according to additional parameters
1143 structure = get_structure(
1144 pdbx_file, model, data_block, altloc, ["label_asym_id"], use_author_fields
1145 )[0] # type: ignore
1146 # TODO(@zeming) This line will remove all non-protein structural elements,
1147 # we should remove this when we want to parse these too.
1148 structure: bs.AtomArray = structure[
1149 bs.filter_amino_acids(structure) & ~structure.hetero # type: ignore
1150 ]
1151 if len(structure) == 0:
1152 raise NoProteinError
1153 unique_asym_ids = np.unique(structure.label_asym_id) # type: ignore
1154 asym2chain = {}
1155 asym2auth = {}
1156 for asym_id in unique_asym_ids:
1157 sub_structure: bs.AtomArray = structure[structure.label_asym_id == asym_id] # type: ignore
1158 chain_id: str = sub_structure[0].chain_id # type: ignore
1159 (
1160 sequence,
1161 atom_positions,
1162 atom_mask,
1163 residue_index,
1164 insertion_code,
1165 confidence,
1166 entity_id,
1167 ) = chain_to_ndarray(sub_structure, mmcif, chain_id, False)
1168
1169 asym2chain[asym_id] = ProteinChain(
1170 id=mmcif.id or "unknown",
1171 sequence=sequence,
1172 chain_id=chain_id,
1173 entity_id=entity_id,
1174 atom37_positions=atom_positions,
1175 atom37_mask=atom_mask,
1176 residue_index=residue_index,
1177 insertion_code=insertion_code,
1178 confidence=confidence,
1179 mmcif=None,
1180 )
1181 asym2auth[asym_id] = chain_id
1182
1183 ### Get transformations and apply them to the affected asym IDs
1184 assembly = []
1185 assembly_id_dict: dict[str, list[str]] = {}
1186
1187 # Process the target assembly ID
1188 for aid, op_expr, asym_id_expr in zip(
1189 assembly_gen_category["assembly_id"].data.array,
1190 assembly_gen_category["oper_expression"].data.array,
1191 assembly_gen_category["asym_id_list"].data.array,
1192 ):
1193 if aid == assembly_id:
1194 # Parse operations and asym IDs for this specific entry
1195 operations = _parse_operation_expression(op_expr)
1196 asym_ids = asym_id_expr.split(",")
1197
1198 # Filter affected asym IDs to only protein chains, preserving order
1199 sub_structures = [
1200 asym2chain[asym_id] for asym_id in asym_ids if asym_id in asym2chain