Synthyra/ESMFold2
0505
1from __future__ import annotations
2
3import io
4import os
5import re
6from dataclasses import asdict, dataclass
7from pathlib import Path
8from subprocess import check_output
9from tempfile import TemporaryDirectory
10from typing import TYPE_CHECKING, Any
11
12import biotite.structure as bs
13import biotite.structure.io.pdbx as pdbx
14import brotli
15import msgpack
16import numpy as np
17import torch
18from biotite.structure.io.pdbx import (
19 CIFCategory,
20 CIFColumn,
21 CIFData,
22 CIFFile,
23 set_structure,
24)
25
26from . import esmfold2_residue_constants as residue_constants
27from .esmfold2_metrics import compute_lddt, compute_rmsd
28from .esmfold2_protein_complex import ProteinComplex, ProteinComplexMetadata
29
30
31@dataclass
32class MolecularComplexResult:
33 """Result of molecular complex folding"""
34
35 complex: MolecularComplex
36 plddt: torch.Tensor | None = None
37 ptm: float | None = None
38 iptm: float | None = None
39 pae: torch.Tensor | None = None
40 distogram: torch.Tensor | None = None
41 pair_chains_iptm: torch.Tensor | None = None
42 output_embedding_sequence: torch.Tensor | None = None
43 output_embedding_pair_pooled: torch.Tensor | None = None
44 residue_index: torch.Tensor | None = None
45 entity_id: torch.Tensor | None = None
46 sae_features: np.ndarray | None = None # [L, n_features]
47 ttt_metrics: dict[str, Any] | None = None
48
49
50@dataclass
51class MolecularComplexMetadata:
52 """Metadata for MolecularComplex objects."""
53
54 entity_lookup: dict[int, str]
55 chain_lookup: dict[int, str]
56 assembly_composition: dict[str, list[str]] | None = None
57
58
59@dataclass
60class Molecule:
61 """Represents a single molecule/token within a MolecularComplex."""
62
63 token: str
64 token_idx: int
65 atom_positions: np.ndarray # [N_atoms, 3]
66 atom_elements: np.ndarray # [N_atoms] element strings
67 atom_names: np.ndarray | None = None # [N_atoms] atom names (optional)
68 atom_hetero: np.ndarray | None = None # [N_atoms] hetero flags (optional)
69 residue_type: int = 0
70 molecule_type: int = 0 # PROTEIN=0, RNA=1, DNA=2, LIGAND=3
71 confidence: float = 0.0
72
73
74@dataclass(frozen=True)
75class MolecularComplex:
76 """
77 Dataclass representing a molecular complex with support for proteins, nucleic acids, and ligands.
78
79 Uses a flat atom representation with token-based sequence indexing, supporting all atom types
80 beyond the traditional atom37 protein representation.
81 """
82
83 id: str
84 sequence: list[str] # Token sequence like ['MET', 'LYS', 'A', 'G', 'ATP']
85
86 # Flat atom arrays - simplified representation
87 atom_positions: np.ndarray # [N_atoms, 3] 3D coordinates
88 atom_elements: np.ndarray # [N_atoms] element strings
89
90 # Token-to-atom mapping for efficient access
91 token_to_atoms: np.ndarray # [N_tokens, 2] start/end indices into atoms array
92
93 # Chain information
94 chain_id: np.ndarray # [N_tokens] chain identifier for each token
95
96 # Confidence data
97 plddt: np.ndarray # Per-token confidence scores [N_tokens]
98
99 # Metadata
100 metadata: MolecularComplexMetadata
101
102 # Optional atom names and hetero flags (preserved from original structures)
103 atom_names: np.ndarray | None = None # [N_atoms] atom names (optional)
104 atom_hetero: np.ndarray | None = None # [N_atoms] hetero flags (optional)
105
106 def __post_init__(self):
107 """Validate array dimensions."""
108 n_tokens = len(self.sequence)
109 n_atoms = len(self.atom_positions)
110 assert (
111 self.token_to_atoms.shape[0] == n_tokens
112 ), f"token_to_atoms shape {self.token_to_atoms.shape} != {n_tokens} tokens"
113 assert (
114 self.chain_id.shape[0] == n_tokens
115 ), f"chain_id shape {self.chain_id.shape} != {n_tokens} tokens"
116 assert (
117 self.plddt.shape[0] == n_tokens
118 ), f"plddt shape {self.plddt.shape} != {n_tokens} tokens"
119 if self.atom_names is not None:
120 assert (
121 self.atom_names.shape[0] == n_atoms
122 ), f"atom_names shape {self.atom_names.shape} != {n_atoms} atoms"
123 if self.atom_hetero is not None:
124 assert (
125 self.atom_hetero.shape[0] == n_atoms
126 ), f"atom_hetero shape {self.atom_hetero.shape} != {n_atoms} atoms"
127
128 def __len__(self) -> int:
129 """Return number of tokens."""
130 return len(self.sequence)
131
132 def __getitem__(self, idx: int) -> Molecule:
133 """Access individual molecules/tokens by index."""
134 if idx >= len(self.sequence) or idx < 0:
135 raise IndexError(
136 f"Token index {idx} out of range for {len(self.sequence)} tokens"
137 )
138
139 token = self.sequence[idx]
140 start_atom, end_atom = self.token_to_atoms[idx]
141
142 # Extract atom data for this token
143 token_atom_positions = self.atom_positions[start_atom:end_atom]
144 token_atom_elements = self.atom_elements[start_atom:end_atom]
145 token_atom_names = None
146 if self.atom_names is not None:
147 token_atom_names = self.atom_names[start_atom:end_atom]
148 token_atom_hetero = None
149 if self.atom_hetero is not None:
150 token_atom_hetero = self.atom_hetero[start_atom:end_atom]
151
152 # Default values for residue/molecule type (would be extended based on actual implementation)
153 residue_type = 0 # Default to standard residue
154 molecule_type = 0 # Default to protein
155
156 return Molecule(
157 token=token,
158 token_idx=idx,
159 atom_positions=token_atom_positions,
160 atom_elements=token_atom_elements,
161 atom_names=token_atom_names,
162 atom_hetero=token_atom_hetero,
163 residue_type=residue_type,
164 molecule_type=molecule_type,
165 confidence=self.plddt[idx],
166 )
167
168 @property
169 def atom_coordinates(self) -> np.ndarray:
170 """Get flat array of all atom coordinates [N_atoms, 3]."""
171 return self.atom_positions
172
173 # Conversion methods
174 @classmethod
175 def from_protein_complex(cls, pc: ProteinComplex) -> "MolecularComplex":
176 """Convert a ProteinComplex to MolecularComplex.
177
178 Args:
179 pc: ProteinComplex object with atom37 representation
180
181 Returns:
182 MolecularComplex with flat atom arrays and token-based indexing
183 """
184 from . import esmfold2_residue_constants
185
186 # Extract sequence without chain breaks
187 sequence_no_breaks = pc.sequence.replace("|", "")
188 sequence_tokens = [
189 residue_constants.restype_1to3.get(aa, "UNK") for aa in sequence_no_breaks
190 ]
191
192 # Convert atom37 to flat arrays
193 flat_positions = []
194 flat_elements = []
195 flat_names = []
196 flat_hetero = []
197 token_to_atoms = []
198
199 atom_idx = 0
200
201 for i, aa in enumerate(pc.sequence):
202 if aa == "|":
203 # Skip chain break tokens
204 continue
205
206 # Get atom37 positions and mask for this residue.
207 # ProteinComplex arrays are indexed by sequence position (including |),
208 # so use `i` not a separate residue counter.
209 res_positions = pc.atom37_positions[i] # [37, 3]
210 res_mask = pc.atom37_mask[i] # [37]
211
212 # Track start position for this token
213 token_start = atom_idx
214
215 # Process each atom type in atom37 representation
216 for atom_type_idx, atom_name in enumerate(residue_constants.atom_types):
217 if res_mask[atom_type_idx]: # Atom is present
218 # Add position
219 flat_positions.append(res_positions[atom_type_idx])
220
221 # Determine element from atom name
222 element = (
223 atom_name[0] if atom_name else "C"
224 ) # First character is element
225 flat_elements.append(element)
226
227 # Add atom name
228 flat_names.append(atom_name)
229
230 # Add hetero flag (all proteins are non-hetero)
231 flat_hetero.append(False)
232
233 atom_idx += 1
234
235 # Record token-to-atom mapping [start_idx, end_idx)
236 token_to_atoms.append([token_start, atom_idx])
237
238 # Convert to numpy arrays
239 atom_positions = np.array(flat_positions, dtype=np.float32)
240 atom_elements = np.array(flat_elements, dtype=object)
241 atom_names = np.array(flat_names, dtype=object)
242 atom_hetero = np.array(flat_hetero, dtype=bool)
243 token_to_atoms_array = np.array(token_to_atoms, dtype=np.int32)
244
245 # Extract confidence scores and chain_ids (skip chain breaks)
246 confidence_scores = []
247 chain_ids = []
248 for seq_idx, aa in enumerate(pc.sequence):
249 if aa != "|":
250 confidence_scores.append(pc.confidence[seq_idx])
251 chain_ids.append(pc.chain_id[seq_idx])
252
253 confidence_array = np.array(confidence_scores, dtype=np.float32)
254 chain_id_array = np.array(chain_ids, dtype=np.int64)
255
256 # Create metadata - convert entity IDs to strings for MolecularComplexMetadata
257 entity_lookup_str = {k: str(v) for k, v in pc.metadata.entity_lookup.items()}
258 metadata = MolecularComplexMetadata(
259 entity_lookup=entity_lookup_str,
260 chain_lookup=pc.metadata.chain_lookup,
261 assembly_composition=pc.metadata.assembly_composition,
262 )
263
264 return cls(
265 id=pc.id,
266 sequence=sequence_tokens,
267 atom_positions=atom_positions,
268 atom_elements=atom_elements,
269 token_to_atoms=token_to_atoms_array,
270 chain_id=chain_id_array,
271 plddt=confidence_array,
272 metadata=metadata,
273 atom_names=atom_names,
274 atom_hetero=atom_hetero,
275 )
276
277 def to_protein_complex(self) -> ProteinComplex:
278 """Convert MolecularComplex back to ProteinComplex format.
279
280 Extracts only protein tokens and converts from flat atom representation
281 back to atom37 format used by ProteinComplex.
282
283 Returns:
284 ProteinComplex with protein residues only, excluding ligands/nucleic acids
285 """
286 from . import esmfold2_residue_constants
287
288 # No need for element mapping - already using element characters
289
290 # Filter for protein tokens only (skip ligands, nucleic acids)
291 protein_tokens = []
292 protein_indices = []
293
294 for i, token in enumerate(self.sequence):
295 # Check if token is a standard 3-letter amino acid code
296 if token in residue_constants.restype_3to1:
297 protein_tokens.append(token)
298 protein_indices.append(i)
299
300 if not protein_tokens:
301 raise ValueError("No protein tokens found in MolecularComplex")
302
303 n_residues = len(protein_tokens)
304
305 # Initialize atom37 arrays
306 atom37_positions = np.full((n_residues, 37, 3), np.nan, dtype=np.float32)
307 atom37_mask = np.zeros((n_residues, 37), dtype=bool)
308
309 # Extract confidence scores and chain_ids for protein residues only
310 protein_confidence = self.plddt[protein_indices]
311 protein_chain_ids = self.chain_id[protein_indices]
312
313 # Convert tokens back to single-letter sequence with chain breaks
314 single_letter_residues = []
315 prev_chain_id = None
316
317 for i, (token, chain_id_val) in enumerate(
318 zip(protein_tokens, protein_chain_ids)
319 ):
320 # Add chain break if we're switching to a new chain
321 if prev_chain_id is not None and chain_id_val != prev_chain_id:
322 single_letter_residues.append("|")
323 single_letter_residues.append(residue_constants.restype_3to1[token])
324 prev_chain_id = chain_id_val
325
326 single_letter_sequence = "".join(single_letter_residues)
327
328 # Calculate final sequence length (includes chain breaks)
329 sequence_length = len(single_letter_sequence)
330
331 # Convert flat atoms back to atom37 representation using atom names
332 for res_idx, token_idx in enumerate(protein_indices):
333 token = self.sequence[token_idx]
334 start_atom, end_atom = self.token_to_atoms[token_idx]
335
336 res_atom_positions = self.atom_positions[start_atom:end_atom]
337 res_atom_names = (
338 np.array(self.atom_names[start_atom:end_atom], dtype=str)
339 if self.atom_names is not None
340 else np.array([], dtype=str)
341 )
342
343 # Build a mapping from normalized atom name -> position for this residue
344 # Normalize to uppercase and strip whitespace for robust matching
345 name_to_pos: dict[str, np.ndarray] = {}
346 for i, nm in enumerate(res_atom_names):
347 key = nm.upper().strip()
348 # Prefer first occurrence; ignore duplicates/altlocs
349 if key not in name_to_pos:
350 name_to_pos[key] = res_atom_positions[i]
351
352 # Place atoms into atom37 by matching stored atom names to atom37 indices.
353 # This handles all atoms present in the flat representation, not just
354 # the canonical residue_atoms for this residue type. This preserves
355 # atoms that were in the original atom37_mask even if they're atypical
356 # for the residue (e.g., from alternate conformations or data quirks).
357 for atom_name_str, pos in name_to_pos.items():
358 idx37 = residue_constants.atom_order.get(atom_name_str)
359 if idx37 is not None:
360 atom37_positions[res_idx, idx37] = pos
361 atom37_mask[res_idx, idx37] = True
362
363 # Create arrays that match sequence length (including chain breaks)
364 # Initialize arrays with proper size
365 chain_id_expanded = np.full(sequence_length, -1, dtype=np.int64)
366 entity_id_expanded = np.full(sequence_length, -1, dtype=np.int64)
367 sym_id_expanded = np.zeros(sequence_length, dtype=np.int64)
368 residue_index_expanded = np.zeros(sequence_length, dtype=np.int64)
369 insertion_code_expanded = np.array([""] * sequence_length, dtype=object)
370 confidence_expanded = np.zeros(sequence_length, dtype=np.float32)
371 atom37_positions_expanded = np.full(
372 (sequence_length, 37, 3), np.nan, dtype=np.float32
373 )
374 atom37_mask_expanded = np.zeros((sequence_length, 37), dtype=bool)
375
376 # Map residue data to sequence positions (skipping chain breaks)
377 residue_idx = 0
378 residue_counter_per_chain = {}
379
380 for seq_pos, char in enumerate(single_letter_sequence):
381 if char != "|":
382 # This is a residue position
383 chain_id_val = protein_chain_ids[residue_idx]
384
385 chain_id_expanded[seq_pos] = chain_id_val
386 entity_id_expanded[seq_pos] = chain_id_val # Simplified mapping
387
388 # Track residue numbering per chain
389 if chain_id_val not in residue_counter_per_chain:
390 residue_counter_per_chain[chain_id_val] = 1
391 else:
392 residue_counter_per_chain[chain_id_val] += 1
393
394 residue_index_expanded[seq_pos] = residue_counter_per_chain[
395 chain_id_val
396 ]
397 confidence_expanded[seq_pos] = protein_confidence[residue_idx]
398 atom37_positions_expanded[seq_pos] = atom37_positions[residue_idx]
399 atom37_mask_expanded[seq_pos] = atom37_mask[residue_idx]
400
401 residue_idx += 1
402 # Chain break positions keep default values (-1, False, etc.)
403
404 # Use the expanded arrays
405 chain_id = chain_id_expanded
406 entity_id = entity_id_expanded
407 sym_id = sym_id_expanded
408 residue_index = residue_index_expanded
409 insertion_code = insertion_code_expanded
410 protein_confidence = confidence_expanded
411 atom37_positions = atom37_positions_expanded
412 atom37_mask = atom37_mask_expanded
413
414 # Create protein complex metadata preserving chain information
415 # Convert MolecularComplex metadata to ProteinComplex format
416 unique_chain_ids = np.unique(protein_chain_ids)
417 entity_lookup = {int(cid): int(cid) for cid in unique_chain_ids}
418 chain_lookup = {
419 int(cid): self.metadata.chain_lookup.get(int(cid), chr(65 + int(cid)))
420 for cid in unique_chain_ids
421 }
422
423 protein_metadata = ProteinComplexMetadata(
424 entity_lookup=entity_lookup,
425 chain_lookup=chain_lookup,
426 assembly_composition=self.metadata.assembly_composition,
427 )
428
429 return ProteinComplex(
430 id=self.id,
431 sequence=single_letter_sequence,
432 entity_id=entity_id,
433 chain_id=chain_id,
434 sym_id=sym_id,
435 residue_index=residue_index,
436 insertion_code=insertion_code,
437 atom37_positions=atom37_positions,
438 atom37_mask=atom37_mask,
439 confidence=protein_confidence,
440 metadata=protein_metadata,
441 )
442
443 @classmethod
444 def from_mmcif(cls, inp: str, id: str | None = None) -> "MolecularComplex":
445 """Read MolecularComplex from mmcif file or string.
446
447 Args:
448 inp: Path to mmCIF file or mmCIF content as string
449 id: Optional identifier to assign to the complex
450
451 Returns:
452 MolecularComplex with all molecules (proteins, ligands, nucleic acids)
453 """
454 from io import StringIO
455
456 # Check if input is a file path or mmCIF string content
457 if os.path.exists(inp):
458 # Input is a file path
459 mmcif_file = pdbx.CIFFile.read(inp)
460 else:
461 # Input is mmCIF string content
462 mmcif_file = pdbx.CIFFile.read(StringIO(inp))
463
464 # Get structure - handle missing model information gracefully
465 try:
466 structure = pdbx.get_structure(
467 mmcif_file, model=1, extra_fields=["b_factor"]
468 )
469 except (KeyError, ValueError):
470 # Fallback for mmCIF files without model information
471 try:
472 structure = pdbx.get_structure(mmcif_file)
473 except Exception:
474 # Last resort: use the first available model or all atoms
475 structure = pdbx.get_structure(mmcif_file, model=None)
476 # Type hint for pyright - structure is an AtomArray which is iterable
477 if TYPE_CHECKING:
478 structure: Any = structure
479
480 # Read label_asym_id from the raw CIF atom_site category.
481 # Biotite's atom.chain_id uses auth_asym_id, which collapses ligands
482 # onto their parent protein chain. label_asym_id gives each entity a
483 # distinct chain identifier.
484 block = mmcif_file.block
485 label_asym_ids: list[str] | None = None
486 if "atom_site" in block:
487 atom_site = block["atom_site"]
488 if "label_asym_id" in atom_site:
489 _col = atom_site["label_asym_id"]
490 _raw = (
491 _col.as_array(str)
492 if hasattr(_col, "as_array")
493 else np.array(list(_col), dtype=str) # type: ignore[arg-type]
494 )
495 # biotite's get_structure(model=1) filters to model 1 AND
496 # removes alternate conformations. We must apply the same
497 # filters to label_asym_id to keep arrays aligned.
498 keep = np.ones(len(_raw), dtype=bool)
499 if "pdbx_PDB_model_num" in atom_site:
500 _mc = atom_site["pdbx_PDB_model_num"]
501 _models = (
502 _mc.as_array(str)
503 if hasattr(_mc, "as_array")
504 else np.array(list(_mc), dtype=str) # type: ignore[arg-type]
505 )
506 keep &= _models == "1"
507 if "label_alt_id" in atom_site:
508 _ac = atom_site["label_alt_id"]
509 _alts = (
510 _ac.as_array(str)
511 if hasattr(_ac, "as_array")
512 else np.array(list(_ac), dtype=str) # type: ignore[arg-type]
513 )
514 keep &= np.isin(_alts, [".", "?", "", "A"])
515 filtered = _raw[keep]
516 if len(filtered) == len(structure):
517 label_asym_ids = filtered.tolist()
518 # If lengths still don't match, fall back to atom.chain_id
519
520 # Get entity information from mmCIF
521 entity_info = {}
522 try:
523 if "entity" in block:
524 entity_category = block["entity"]
525 if "id" in entity_category and "type" in entity_category:
526 entity_ids = entity_category["id"]
527 entity_types = entity_category["type"]
528 # Convert CIFColumn to list for iteration
529 if hasattr(entity_ids, "__iter__") and hasattr(
530 entity_types, "__iter__"
531 ):
532 # Type annotation to help pyright understand these are iterable
533 entity_ids_list = list(entity_ids) # type: ignore
534 entity_types_list = list(entity_types) # type: ignore
535 for eid, etype in zip(entity_ids_list, entity_types_list):
536 entity_info[eid] = etype
537 except Exception:
538 pass
539
540 # Initialize arrays for flat atom representation
541 sequence_tokens = []
542 flat_positions = []
543 flat_elements = []
544 flat_names = []
545 flat_hetero = []
546 token_to_atoms = []
547 confidence_scores = []
548 chain_ids = [] # Track chain IDs for each token
549
550 atom_idx = 0
551
552 # Group atoms by chain and residue.
553 # Use label_asym_id (distinct per entity) when available, otherwise
554 # fall back to biotite's chain_id (auth_asym_id).
555 chain_residue_groups: dict[str, dict[tuple[int, str], dict]] = {}
556 for atom_i, atom in enumerate(structure):
557 chain_id = (
558 label_asym_ids[atom_i] if label_asym_ids is not None else atom.chain_id
559 )
560 res_id = atom.res_id
561 res_name = atom.res_name
562
563 if chain_id not in chain_residue_groups:
564 chain_residue_groups[chain_id] = {}
565 # Key by (res_id, res_name) to distinguish residues that share
566 # the same res_id but have different res_name (e.g. a protein
567 # residue and a ligand that were on the same auth chain).
568 res_key = (res_id, res_name)
569 if res_key not in chain_residue_groups[chain_id]:
570 chain_residue_groups[chain_id][res_key] = {
571 "atoms": [],
572 "res_name": res_name,
573 "is_hetero": atom.hetero,
574 }
575 chain_residue_groups[chain_id][res_key]["atoms"].append(atom)
576
577 # Create a mapping from chain_id to numeric indices
578 chain_id_to_numeric = {
579 chain_id: idx
580 for idx, chain_id in enumerate(sorted(chain_residue_groups.keys()))
581 }
582
583 # Process each chain and residue
584 for chain_id in sorted(chain_residue_groups.keys()):
585 residues = chain_residue_groups[chain_id]
586 numeric_chain_id = chain_id_to_numeric[chain_id]
587
588 for res_key in sorted(residues.keys()):
589 residue_data = residues[res_key]
590 res_name = residue_data["res_name"]
591 atoms = residue_data["atoms"]
592 is_hetero = residue_data["is_hetero"]
593
594 # Skip water molecules
595 if res_name == "HOH":
596 continue
597
598 # Determine token name
599 if not is_hetero and res_name in residue_constants.restype_3to1:
600 # Standard amino acid
601 token_name = res_name
602 elif res_name in ["A", "T", "G", "C", "U", "DA", "DT", "DG", "DC"]:
603 # Nucleotide
604 token_name = res_name
605 else:
606 # Ligand or other molecule
607 token_name = res_name
608
609 sequence_tokens.append(token_name)
610 chain_ids.append(
611 numeric_chain_id
612 ) # Store the numeric chain ID for this token
613 token_start = atom_idx
614
615 # Add all atoms from this residue
616 for atom in atoms:
617 flat_positions.append(atom.coord)
618
619 # Get element character
620 element = atom.element
621 flat_elements.append(element)
622
623 # Get atom name
624 atom_name = atom.atom_name
625 flat_names.append(atom_name)
626
627 # Get hetero flag
628 hetero_flag = atom.hetero
629 flat_hetero.append(hetero_flag)
630
631 atom_idx += 1
632
633 # Record token-to-atom mapping
634 token_to_atoms.append([token_start, atom_idx])
635
636 # Add confidence score (B-factor if available, otherwise 1.0)
637 bfactor = getattr(atoms[0], "b_factor", 50.0) if atoms else 50.0
638 confidence_scores.append(min(bfactor / 100.0, 1.0))
639
640 # Convert to numpy arrays
641 if not flat_positions:
642 # Create minimal arrays if no atoms found
643 atom_positions = np.zeros((0, 3), dtype=np.float32)
644 atom_elements = np.zeros(0, dtype=object)
645 atom_names = np.zeros(0, dtype=object)
646 atom_hetero = np.zeros(0, dtype=bool)
647 token_to_atoms_array = np.zeros((len(sequence_tokens), 2), dtype=np.int32)
648 chain_id_array = (
649 np.array(chain_ids, dtype=np.int64)
650 if chain_ids
651 else np.zeros(len(sequence_tokens), dtype=np.int64)
652 )
653 else:
654 atom_positions = np.array(flat_positions, dtype=np.float32)
655 atom_elements = np.array(flat_elements, dtype=object)
656 atom_names = np.array(flat_names, dtype=object)
657 atom_hetero = np.array(flat_hetero, dtype=bool)
658 token_to_atoms_array = np.array(token_to_atoms, dtype=np.int32)
659 chain_id_array = np.array(chain_ids, dtype=np.int64)
660
661 confidence_array = np.array(confidence_scores, dtype=np.float32)
662
663 # Create metadata using the chain_id_to_numeric mapping
664 if chain_residue_groups:
665 chain_lookup = {
666 numeric_id: chain_id
667 for chain_id, numeric_id in chain_id_to_numeric.items()
668 }
669 else:
670 chain_lookup = {}
671
672 metadata = MolecularComplexMetadata(
673 entity_lookup=entity_info,
674 chain_lookup=chain_lookup,
675 assembly_composition=None,
676 )
677
678 # Set complex ID - if input was a path, use the stem; otherwise use default
679 if os.path.exists(inp):
680 complex_id = id or Path(inp).stem
681 else:
682 complex_id = id or "complex_from_string"
683
684 return cls(
685 id=complex_id,
686 sequence=sequence_tokens,
687 atom_positions=atom_positions,
688 atom_elements=atom_elements,
689 token_to_atoms=token_to_atoms_array,
690 chain_id=chain_id_array,
691 plddt=confidence_array,
692 metadata=metadata,
693 atom_names=atom_names,
694 atom_hetero=atom_hetero,
695 )
696
697 def _get_entity_mapping(
698 self,
699 ) -> tuple[dict[str, list[str]], dict[str, int], dict[int, tuple[str, ...]]]:
700 """Compute chain→sequence, chain→entity_id, and entity_id→sequence mappings.
701
702 Returns:
703 (chain_sequences, chain_to_entity, entity_sequences)
704 """
705 chain_sequences: dict[str, list[str]] = {}
706 for token_idx in range(len(self.token_to_atoms)):
707 chain_id_numeric = self.chain_id[token_idx]
708 chain_id_str = self.metadata.chain_lookup.get(
709 int(chain_id_numeric), chr(65 + int(chain_id_numeric))
710 )
711 if chain_id_str not in chain_sequences:
712 chain_sequences[chain_id_str] = []
713 chain_sequences[chain_id_str].append(self.sequence[token_idx])
714
715 sequence_to_entity: dict[tuple[str, ...], int] = {}
716 chain_to_entity: dict[str, int] = {}
717 entity_sequences: dict[int, tuple[str, ...]] = {}
718 entity_id_counter = 1
719 for chain_id_str, sequence in chain_sequences.items():
720 seq_tuple = tuple(sequence)
721 if seq_tuple not in sequence_to_entity:
722 sequence_to_entity[seq_tuple] = entity_id_counter
723 entity_sequences[entity_id_counter] = seq_tuple
724 entity_id_counter += 1
725 chain_to_entity[chain_id_str] = sequence_to_entity[seq_tuple]
726
727 return chain_sequences, chain_to_entity, entity_sequences
728
729 def _add_entity_information(
730 self, cif_file: CIFFile, entity_sequences: dict[int, tuple[str, ...]]
731 ) -> None:
732 """Add _entity category to CIF file so OST can identify ligands vs polymers."""
733
734 entity_ids: list[str] = []
735 entity_types: list[str] = []
736 entity_descriptions: list[str] = []
737 for eid in sorted(entity_sequences.keys()):
738 seq = entity_sequences[eid]
739 entity_ids.append(str(eid))
740 has_protein = any(t in residue_constants.restype_3to1 for t in seq)
741 has_na = any(
742 t in ("A", "T", "G", "C", "U", "DA", "DT", "DG", "DC") for t in seq
743 )
744 if has_protein or has_na:
745 entity_types.append("polymer")
746 if has_protein:
747 entity_descriptions.append(f"Polymer entity {eid} (protein)")
748 else:
749 entity_descriptions.append(f"Polymer entity {eid} (nucleic acid)")
750 else:
751 entity_types.append("non-polymer")
752 entity_descriptions.append(f"Non-polymer entity {eid}")
753
754 if entity_ids:
755 cif_file.block["entity"] = CIFCategory(
756 name="entity",
757 columns={
758 "id": CIFColumn(
759 data=CIFData(array=np.array(entity_ids), dtype=np.str_)
760 ),
761 "type": CIFColumn(
762 data=CIFData(array=np.array(entity_types), dtype=np.str_)
763 ),
764 "pdbx_description": CIFColumn(
765 data=CIFData(array=np.array(entity_descriptions), dtype=np.str_)
766 ),
767 },
768 )
769
770 # Add _struct_asym to map chain IDs to entity IDs
771 _, chain_to_entity, _ = self._get_entity_mapping()
772 if chain_to_entity:
773 asym_ids = sorted(chain_to_entity.keys())
774 asym_entity_ids = [str(chain_to_entity[c]) for c in asym_ids]
775 cif_file.block["struct_asym"] = CIFCategory(
776 name="struct_asym",
777 columns={
778 "id": CIFColumn(
779 data=CIFData(array=np.array(asym_ids), dtype=np.str_)
780 ),
781 "entity_id": CIFColumn(
782 data=CIFData(array=np.array(asym_entity_ids), dtype=np.str_)
783 ),
784 },
785 )
786
787 def to_mmcif(self) -> str:
788 """Write MolecularComplex to mmcif string using biotite.
789
790 Returns:
791 String representation of the complex in mmCIF format
792 """
793 # Pre-allocate AtomArray
794 n_atoms = len(self.atom_positions)
795 atom_array = bs.AtomArray(length=n_atoms)
796
797 # Set coordinates directly (already vectorized)
798 atom_array.coord = self.atom_positions
799
800 # Pre-allocate per-atom arrays
801 atom_res_ids = np.zeros(n_atoms, dtype=np.int32)
802 atom_chain_ids = np.empty(n_atoms, dtype=object)
803 atom_res_names = np.empty(n_atoms, dtype=object)
804 atom_hetero = np.zeros(n_atoms, dtype=bool)
805 atom_bfactors = np.zeros(n_atoms, dtype=np.float32)
806 atom_names = np.empty(n_atoms, dtype=object)
807
808 # Build entity mappings: chains with identical sequences share entity ID
809 _, chain_to_entity, entity_sequences = self._get_entity_mapping()
810
811 atom_entity_ids = np.zeros(n_atoms, dtype=np.int32)
812
813 # Track residue IDs per chain
814 chain_res_counters: dict[int, int] = {}
815
816 # Vectorized expansion of token-level to atom-level annotations
817 for token_idx, (start, end) in enumerate(self.token_to_atoms):
818 token = self.sequence[token_idx]
819 chain_id_numeric = self.chain_id[token_idx]
820 chain_id_str = self.metadata.chain_lookup.get(
821 int(chain_id_numeric), chr(65 + int(chain_id_numeric))
822 )
823
824 # Track residue numbering per chain
825 if chain_id_numeric not in chain_res_counters:
826 chain_res_counters[chain_id_numeric] = 1
827 res_id = chain_res_counters[chain_id_numeric]
828 chain_res_counters[chain_id_numeric] += 1
829
830 # Determine if protein
831 is_protein = token in residue_constants.restype_3to1
832
833 # Get atom names for this residue
834 if self.atom_names is not None:
835 # Use stored atom names (preserves original names from mmCIF)
836 names = list(self.atom_names[start:end])
837 elif is_protein:
838 # Fallback: use standard protein atom names
839 standard_names = residue_constants.residue_atoms.get(
840 token, ["N", "CA", "C", "O"]
841 )
842 names = standard_names[: end - start]
843 # Pad if needed
844 while len(names) < (end - start):
845 names.append(f"X{len(names)+1}")
846 else:
847 # Fallback: generate names for ligands/nucleic acids
848 names = [f"C{i+1}" for i in range(end - start)]
849
850 # Vectorized assignment for this token's atoms
851 atom_res_ids[start:end] = res_id
852 atom_chain_ids[start:end] = chain_id_str
853 atom_res_names[start:end] = token
854 # Use stored hetero flags if available, otherwise guess based on protein status
855 if self.atom_hetero is not None:
856 atom_hetero[start:end] = self.atom_hetero[start:end]
857 else:
858 atom_hetero[start:end] = not is_protein
859 atom_bfactors[start:end] = self.plddt[token_idx] * 100.0
860 atom_names[start:end] = names
861 atom_entity_ids[start:end] = chain_to_entity.get(chain_id_str, 1)
862
863 # Set all AtomArray attributes at once (convert object arrays to proper string arrays)
864 # res_name uses U8 to accommodate CCD codes up to 5 characters (e.g., A1AZ2);
865 # chain_id uses U16 because chain names like ``ligand_1`` / ``ligand_2`` /
866 # auth-asym ids of arbitrary length are possible.
867 atom_array.res_id = atom_res_ids
868 atom_array.chain_id = np.array(atom_chain_ids, dtype="U16")
869 atom_array.res_name = np.array(atom_res_names, dtype="U8")
870 atom_array.hetero = atom_hetero
871 atom_array.atom_name = np.array(atom_names, dtype="U4")
872 atom_array.add_annotation("b_factor", dtype=float)
873 atom_array.b_factor = atom_bfactors
874 atom_array.add_annotation("entity_id", dtype=int)
875 atom_array.entity_id = atom_entity_ids
876
877 # Use existing elements or infer them from atom names
878 if self.atom_elements is not None and len(self.atom_elements) == n_atoms:
879 # Convert object array to proper string array for biotite
880 atom_array.element = np.array(self.atom_elements, dtype="U4")
881 else:
882 # Use biotite's built-in element inference
883 atom_array.element = bs.infer_elements(atom_array)
884
885 # Create CIF file and set structure
886 cif_file = CIFFile()
887 set_structure(cif_file, atom_array, data_block=self.id)
888
889 # Manually fix label_entity_id (biotite doesn't use entity_id annotation correctly)
890 if "atom_site" in cif_file.block:
891 atom_site = cif_file.block["atom_site"]
892 if "label_asym_id" in atom_site and "label_entity_id" in atom_site:
893 label_asym_ids = atom_site["label_asym_id"]
894 if hasattr(label_asym_ids, "as_array"):
895 chain_ids_list = label_asym_ids.as_array(str).tolist()
896 elif hasattr(label_asym_ids, "__iter__"):
897 chain_ids_list = list(label_asym_ids) # type: ignore[arg-type]
898 else:
899 chain_ids_list = []
900 updated_entity_ids = [
901 str(chain_to_entity.get(cid, 1)) for cid in chain_ids_list
902 ]
903 if updated_entity_ids:
904 atom_site["label_entity_id"] = CIFColumn(
905 data=CIFData(array=np.array(updated_entity_ids), dtype=np.str_)
906 )
907
908 # Add _entity category for OST compatibility
909 self._add_entity_information(cif_file, entity_sequences)
910
911 # Convert to string
912 output = io.StringIO()
913 cif_file.write(output)
914 return output.getvalue()
915
916 def dockq(self, native: "MolecularComplex") -> Any:
917 """Compute DockQ score against native structure.
918
919 Args:
920 native: Native MolecularComplex to compute DockQ against
921
922 Returns:
923 DockQ result containing score and alignment information
924 """
925 # Imports moved to top of file
926
927 # Convert both complexes to ProteinComplex format for DockQ computation
928 # This extracts only the protein portion and converts to PDB format
929 try:
930 self_pc = self.to_protein_complex()
931 native_pc = native.to_protein_complex()
932 except ValueError as e:
933 raise ValueError(
934 f"Cannot convert MolecularComplex to ProteinComplex for DockQ: {e}"
935 )
936
937 # Normalize chain IDs for PDB compatibility
938 self_pc = self_pc.normalize_chain_ids_for_pdb()
939 native_pc = native_pc.normalize_chain_ids_for_pdb()
940
941 # Use the existing ProteinComplex.dockq() method
942 try:
943 dockq_result = self_pc.dockq(native_pc)
944 return dockq_result
945 except Exception:
946 # Fallback to manual DockQ computation if ProteinComplex.dockq() fails
947 return self._compute_dockq_manual(native)
948
949 def _compute_dockq_manual(self, native: "MolecularComplex") -> Any:
950 """Manual DockQ computation fallback."""
951 # Imports moved to top of file
952
953 # Convert both complexes to ProteinComplex format
954 try:
955 self_pc = self.to_protein_complex()
956 native_pc = native.to_protein_complex()
957 except ValueError as e:
958 raise ValueError(
959 f"Cannot convert MolecularComplex to ProteinComplex for DockQ: {e}"
960 )
961
962 # Normalize chain IDs for PDB compatibility
963 self_pc = self_pc.normalize_chain_ids_for_pdb()
964 native_pc = native_pc.normalize_chain_ids_for_pdb()
965
966 # Write temporary PDB files and run DockQ
967 with TemporaryDirectory() as tdir:
968 dir_path = Path(tdir)
969 self_pdb = dir_path / "self.pdb"
970 native_pdb = dir_path / "native.pdb"
971
972 # Write PDB files
973 self_pc.to_pdb(self_pdb)
974 native_pc.to_pdb(native_pdb)
975
976 # Run DockQ
977 try:
978 output = check_output(["DockQ", str(self_pdb), str(native_pdb)])
979 output_text = output.decode()
980
981 # Parse DockQ output
982 lines = output_text.split("\n")
983
984 # Find the total DockQ score
985 dockq_score = None
986 for line in lines:
987 if "Total DockQ" in line:
988 match = re.search(r"Total DockQ.*: ([\d.]+)", line)
989 if match:
990 dockq_score = float(match.group(1))
991 break
992
993 if dockq_score is None:
994 # Try to find individual DockQ scores
995 for line in lines:
996 if line.startswith("DockQ") and ":" in line:
997 try:
998 dockq_score = float(line.split(":")[1].strip())
999 break
1000 except (ValueError, IndexError):
1001 continue
1002
1003 if dockq_score is None:
1004 raise ValueError("Could not parse DockQ score from output")
1005
1006 # Return a simple result structure
1007 return {
1008 "total_dockq": dockq_score,
1009 "raw_output": output_text,
1010 "aligned": self, # Return self as aligned structure
1011 }
1012
1013 except FileNotFoundError:
1014 raise RuntimeError(
1015 "DockQ is not installed. Please install DockQ to use this method."
1016 )
1017 except Exception as e:
1018 raise RuntimeError(f"DockQ computation failed: {e}")
1019
1020 def rmsd(self, target: "MolecularComplex", **kwargs) -> float:
1021 """Compute RMSD against target structure.
1022
1023 Args:
1024 target: Target MolecularComplex to compute RMSD against
1025 **kwargs: Additional arguments passed to compute_rmsd
1026
1027 Returns:
1028 float: RMSD value between the two structures
1029 """
1030 # Imports moved to top of file
1031
1032 # Ensure both complexes have the same number of tokens
1033 if len(self) != len(target):
1034 raise ValueError(
1035 f"Complexes must have the same number of tokens: {len(self)} vs {len(target)}"
1036 )
1037
1038 # Extract center positions for each token (using centroid of atoms)
1039 mobile_coords = []
1040 target_coords = []
1041 atom_mask = []
1042
1043 for i in range(len(self)):
1044 # Get atom positions for this token
1045 mobile_start, mobile_end = self.token_to_atoms[i]
1046 target_start, target_end = target.token_to_atoms[i]
1047
1048 # Extract atom positions
1049 mobile_atoms = self.atom_positions[mobile_start:mobile_end]
1050 target_atoms = target.atom_positions[target_start:target_end]
1051
1052 # Check if both tokens have atoms
1053 if len(mobile_atoms) == 0 or len(target_atoms) == 0:
1054 # Skip tokens with no atoms
1055 continue
1056
1057 # For simplicity, use the centroid of atoms as the representative position
1058 mobile_center = mobile_atoms.mean(axis=0)
1059 target_center = target_atoms.mean(axis=0)
1060
1061 mobile_coords.append(mobile_center)
1062 target_coords.append(target_center)
1063 atom_mask.append(True)
1064
1065 if len(mobile_coords) == 0:
1066 raise ValueError("No valid atoms found for RMSD computation")
1067
1068 # Convert to tensors
1069 mobile_tensor = torch.from_numpy(np.stack(mobile_coords, axis=0)).unsqueeze(
1070 0
1071 ) # [1, N, 3]
1072 target_tensor = torch.from_numpy(np.stack(target_coords, axis=0)).unsqueeze(
1073 0
1074 ) # [1, N, 3]
1075 mask_tensor = torch.tensor(atom_mask, dtype=torch.bool).unsqueeze(0) # [1, N]
1076
1077 # Compute RMSD using existing infrastructure
1078 rmsd_value = compute_rmsd(
1079 mobile=mobile_tensor,
1080 target=target_tensor,
1081 atom_exists_mask=mask_tensor,
1082 reduction="batch",
1083 **kwargs,
1084 )
1085
1086 return float(rmsd_value)
1087
1088 def lddt_ca(self, target: "MolecularComplex", **kwargs) -> float:
1089 """Compute LDDT score against target structure.
1090
1091 Args:
1092 target: Target MolecularComplex to compute LDDT against
1093 **kwargs: Additional arguments passed to compute_lddt
1094
1095 Returns:
1096 float: LDDT value between the two structures
1097 """
1098 # Imports moved to top of file
1099
1100 # Ensure both complexes have the same number of tokens
1101 if len(self) != len(target):
1102 raise ValueError(
1103 f"Complexes must have the same number of tokens: {len(self)} vs {len(target)}"
1104 )
1105
1106 # Extract center positions for each token (using centroid of atoms)
1107 mobile_coords = []
1108 target_coords = []
1109 atom_mask = []
1110
1111 for i in range(len(self)):
1112 # Get atom positions for this token
1113 mobile_start, mobile_end = self.token_to_atoms[i]
1114 target_start, target_end = target.token_to_atoms[i]
1115
1116 # Extract atom positions
1117 mobile_atoms = self.atom_positions[mobile_start:mobile_end]
1118 target_atoms = target.atom_positions[target_start:target_end]
1119
1120 # Check if both tokens have atoms
1121 if len(mobile_atoms) == 0 or len(target_atoms) == 0:
1122 # Skip tokens with no atoms
1123 mobile_coords.append(np.full(3, np.nan))
1124 target_coords.append(np.full(3, np.nan))
1125 atom_mask.append(False)
1126 continue
1127
1128 # For simplicity, use the centroid of atoms as the representative position
1129 mobile_center = mobile_atoms.mean(axis=0)
1130 target_center = target_atoms.mean(axis=0)
1131
1132 mobile_coords.append(mobile_center)
1133 target_coords.append(target_center)
1134 atom_mask.append(True)
1135
1136 if not any(atom_mask):
1137 raise ValueError("No valid atoms found for LDDT computation")
1138
1139 # Convert to tensors
1140 mobile_tensor = torch.from_numpy(np.stack(mobile_coords, axis=0)).unsqueeze(
1141 0
1142 ) # [1, N, 3]
1143 target_tensor = torch.from_numpy(np.stack(target_coords, axis=0)).unsqueeze(
1144 0
1145 ) # [1, N, 3]
1146 mask_tensor = torch.tensor(atom_mask, dtype=torch.bool).unsqueeze(0) # [1, N]
1147
1148 # Compute LDDT using existing infrastructure
1149 lddt_value = compute_lddt(
1150 all_atom_pred_pos=mobile_tensor,
1151 all_atom_positions=target_tensor,
1152 all_atom_mask=mask_tensor,
1153 per_residue=False, # Return overall LDDT score
1154 **kwargs,
1155 )
1156
1157 return float(lddt_value)
1158
1159 def state_dict(self):
1160 """This state dict is optimized for storage, so it turns things to fp16 whenever
1161 possible and converts numpy arrays to lists for JSON serialization.
1162 """
1163 dct = {k: v for k, v in vars(self).items()}
1164 for k, v in dct.items():
1165 if isinstance(v, np.ndarray):
1166 match v.dtype:
1167 case np.int64:
1168 dct[k] = v.astype(np.int32).tolist()
1169 case np.float64 | np.float32:
1170 dct[k] = v.astype(np.float16).tolist()
1171 case _:
1172 dct[k] = v.tolist()
1173 elif isinstance(v, MolecularComplexMetadata):
1174 dct[k] = asdict(v)
1175
1176 return dct
1177
1178 def to_blob(self) -> bytes:
1179 return brotli.compress(msgpack.dumps(self.state_dict()), quality=5)
1180
1181 @classmethod
1182 def from_state_dict(cls, dct):
1183 for k, v in dct.items():
1184 if isinstance(v, list) and k in [
1185 "atom_positions",
1186 "atom_elements",
1187 "atom_names",
1188 "atom_hetero",
1189 "token_to_atoms",
1190 "chain_id",
1191 "plddt",
1192 ]:
1193 dct[k] = np.array(v)
1194
1195 for k, v in dct.items():
1196 if isinstance(v, np.ndarray):
1197 if k in ["atom_positions", "plddt"]:
1198 dct[k] = v.astype(np.float32)
1199 elif k in ["token_to_atoms", "chain_id"]:
1200 dct[k] = (