Synthyra/ESMFold2-Experimental-Cutoff2025
0101
1import random
2from contextlib import contextmanager, nullcontext
3from pathlib import Path
4from typing import Any
5
6import numpy as np
7import torch
8
9from .esmfold2_conformers import load_ccd
10from .esmfold2_output import build_molecular_complex_from_features
11from .esmfold2_prepare_input import ChainInfo, prepare_esmfold2_input
12from .esmfold2_types import (
13 MSA,
14 Modification,
15 ProteinInput,
16 StructurePredictionInput,
17)
18from .esmfold2_molecular_complex import MolecularComplexResult
19
20
21@contextmanager
22def _seed_context(seed: int | None):
23 if seed is None:
24 yield
25 return
26 py_state = random.getstate()
27 np_state = np.random.get_state()
28 torch_state = torch.random.get_rng_state()
29 cuda_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
30 random.seed(seed)
31 np.random.seed(seed)
32 torch.manual_seed(seed)
33 if torch.cuda.is_available():
34 torch.cuda.manual_seed_all(seed)
35 try:
36 yield
37 finally:
38 random.setstate(py_state)
39 np.random.set_state(np_state)
40 torch.random.set_rng_state(torch_state)
41 if cuda_state is not None:
42 torch.cuda.set_rng_state_all(cuda_state)
43
44
45def clean_esmfold2_input(input: StructurePredictionInput) -> StructurePredictionInput:
46 """Group identical protein sequences into the same ProteinInput with multiple ids.
47
48 Example: Passing a tetramer like [ProteinInput(id=["0"], seq="AAA|AAA|BBB|BBB")]
49 gets converted into [ProteinInput(id=["0_0", "0_1"], seq="AAA"),
50 ProteinInput(id=["0_2", "0_3"], seq="BBB")]
51
52 Preserves the original order of unique sequences. Also converts "|" chainbreak
53 tokens to ":" in the sequence.
54 """
55 cleaned_sequences: list = []
56 chain_to_ids: dict[str, list[str]] = {}
57 chain_to_modifications: dict[str, list] = {}
58 chain_to_msa: dict[str, MSA | None] = {}
59
60 for item in input.sequences:
61 if isinstance(item, ProteinInput):
62 sequence = ":".join(item.sequence.split("|"))
63 if ":" not in sequence:
64 cleaned_sequences.append(item)
65 continue
66
67 if ":" in sequence and input.covalent_bonds is not None:
68 raise ValueError(
69 "Covalent bonds are not supported when using chainbreaks. "
70 "Chains must be separated into multiple ProteinInput objects."
71 )
72
73 base_id = item.id[0] if isinstance(item.id, list) else item.id
74 chain_to_ids = {}
75 chain_to_modifications = {}
76 chain_to_msa = {}
77 chains = sequence.split(":")
78
79 chain_start_positions = []
80 pos = 0
81 for chain in chains:
82 chain_start_positions.append(pos)
83 pos += len(chain) + 1
84
85 if item.modifications is not None:
86 for chain_idx, chain in enumerate(chains):
87 chain_start = chain_start_positions[chain_idx]
88 chain_end = chain_start + len(chain)
89 chain_modifications = []
90 for mod in item.modifications:
91 if chain_start <= mod.position < chain_end:
92 adjusted_mod = Modification(
93 position=mod.position - chain_start, ccd=mod.ccd
94 )
95 chain_modifications.append(adjusted_mod)
96 if chain not in chain_to_modifications:
97 chain_to_modifications[chain] = chain_modifications
98 else:
99 chain_to_modifications[chain].extend(chain_modifications)
100
101 if item.msa is not None:
102 for chain_idx, chain in enumerate(chains):
103 if chain not in chain_to_msa:
104 chain_start = chain_start_positions[chain_idx]
105 chain_end = chain_start + len(chain)
106 chain_msa = item.msa.select_positions( # type: ignore
107 np.arange(chain_start, chain_end)
108 )
109 chain_to_msa[chain] = chain_msa
110
111 for i, chain in enumerate(chains):
112 chain_id = base_id + "_" + str(i)
113 if chain in chain_to_ids:
114 chain_to_ids[chain].append(chain_id)
115 else:
116 chain_to_ids[chain] = [chain_id]
117 cleaned_sequences.append((item, chain))
118 else:
119 cleaned_sequences.append(item)
120
121 for i in range(len(cleaned_sequences)):
122 if isinstance(cleaned_sequences[i], tuple):
123 item, chain = cleaned_sequences[i]
124 chain_ids = chain_to_ids[chain]
125 chain_modifications = (
126 chain_to_modifications.get(chain) if item.modifications else None
127 )
128 chain_msa = chain_to_msa.get(chain) if item.msa else None
129 cleaned_sequences[i] = ProteinInput(
130 id=chain_ids,
131 sequence=chain,
132 msa=chain_msa,
133 modifications=chain_modifications,
134 )
135
136 return StructurePredictionInput(
137 sequences=cleaned_sequences,
138 distogram_conditioning=input.distogram_conditioning,
139 covalent_bonds=input.covalent_bonds,
140 )
141
142
143class ESMFold2InputBuilder:
144 def __init__(self, ccd_cache: Path | None = None):
145 load_ccd(ccd_cache)
146
147 def prepare_input(
148 self,
149 input: StructurePredictionInput,
150 seed: int | None = None,
151 device: torch.device | str | None = None,
152 ) -> tuple[dict, list[ChainInfo]]:
153 """Prepare raw input for the folding model.
154
155 Converts user-provided StructurePredictionInput into batched tensors
156 ready for model inference.
157
158 Parameters
159 ----------
160 input : StructurePredictionInput
161 Input specification (sequences, structures, constraints, etc.).
162 seed : int, optional
163 Random seed for reproducibility.
164 device : torch.device or str, optional
165 Target device for the returned tensors. Defaults to CPU; pass
166 ``model.device`` to skip a separate ``.to(...)`` step. ``fold()``
167 forwards ``model.device`` automatically.
168
169 Returns
170 -------
171 tuple[dict, list[ChainInfo]]
172 Batched input tensors and chain metadata for output processing.
173 """
174 structure_prediction_input = clean_esmfold2_input(input)
175 with _seed_context(seed) if seed is not None else nullcontext():
176 features, chain_infos = prepare_esmfold2_input(
177 structure_prediction_input, seed=seed
178 )
179 features = {
180 k: (v[None].to(device) if device is not None else v[None])
181 if isinstance(v, torch.Tensor)
182 else v
183 for k, v in features.items()
184 }
185
186 return features, chain_infos
187
188 def __call__(
189 self,
190 input: StructurePredictionInput,
191 seed: int | None = None,
192 device: torch.device | str | None = None,
193 ) -> tuple[dict, list[ChainInfo]]:
194 return self.prepare_input(input, seed=seed, device=device)
195
196 def decode(
197 self,
198 output: dict[str, torch.Tensor],
199 features: dict[str, torch.Tensor],
200 chain_infos: list[ChainInfo],
201 *,
202 num_diffusion_samples: int = 1,
203 complex_id: str = "pred",
204 ) -> MolecularComplexResult | list[MolecularComplexResult]:
205 """Convert raw model outputs into one MolecularComplexResult per sample.
206
207 Parameters
208 ----------
209 output : dict[str, Tensor]
210 Output dict returned by ESMFold2Model.forward.
211 features : dict[str, Tensor]
212 Feature dict from :meth:`prepare_input` (batched, on the model device).
213 chain_infos : list[ChainInfo]
214 Chain metadata returned alongside `features`.
215 num_diffusion_samples : int
216 Number of diffusion samples present in the output (Bm = B * num_diffusion_samples).
217 complex_id : str
218 Identifier assigned to each MolecularComplex.
219
220 Returns
221 -------
222 MolecularComplexResult or list[MolecularComplexResult]
223 A single result when num_diffusion_samples == 1, otherwise a list of length Bm.
224 """
225 atom_mask = features["atom_attention_mask"][0]
226 ref_element = features["ref_element"][0]
227 ref_atom_name_chars = features["ref_atom_name_chars"][0]
228
229 sample_coords = output["sample_atom_coords"]
230 plddts = output["plddt"]
231 Bm = sample_coords.shape[0]
232
233 ptm_t = output.get("ptm")
234 iptm_t = output.get("iptm")
235 pae_t = output.get("pae")
236 distogram_t = output.get("distogram_logits")
237 pair_chains_t = output.get("pair_chains_iptm")
238 residue_index_t = output.get("residue_index")
239 entity_id_t = output.get("entity_id")
240
241 results: list[MolecularComplexResult] = []
242 for i in range(Bm):
243 mc = build_molecular_complex_from_features(
244 coords=sample_coords[i],
245 plddt=plddts[i],
246 atom_mask=atom_mask,
247 ref_element=ref_element,
248 ref_atom_name_chars=ref_atom_name_chars,
249 chain_infos=chain_infos,
250 complex_id=complex_id,
251 )
252 results.append(
253 MolecularComplexResult(
254 complex=mc,
255 plddt=plddts[i].detach().cpu(),
256 ptm=float(ptm_t[i].item()) if ptm_t is not None else None,
257 iptm=float(iptm_t[i].item()) if iptm_t is not None else None,
258 pae=pae_t[i].detach().cpu() if pae_t is not None else None,
259 distogram=(
260 distogram_t[0].detach().cpu()
261 if distogram_t is not None
262 else None
263 ),
264 pair_chains_iptm=(
265 pair_chains_t[i].detach().cpu()
266 if pair_chains_t is not None
267 else None
268 ),
269 residue_index=(
270 residue_index_t[0].detach().cpu()
271 if residue_index_t is not None
272 else None
273 ),
274 entity_id=(
275 entity_id_t[0].detach().cpu()
276 if entity_id_t is not None
277 else None
278 ),
279 )
280 )
281
282 if num_diffusion_samples == 1 and len(results) == 1:
283 return results[0]
284 return results
285
286 def fold(
287 self,
288 model: Any,
289 input: StructurePredictionInput,
290 *,
291 num_loops: int = 3,
292 num_sampling_steps: int = 200,
293 num_diffusion_samples: int = 1,
294 seed: int | None = None,
295 noise_scale: float | None = None,
296 step_scale: float | None = None,
297 max_inference_sigma: int | None = None,
298 early_exit: bool = False,
299 complex_id: str = "pred",
300 ) -> MolecularComplexResult | list[MolecularComplexResult]:
301 """Fold a structure end-to-end: encode → model → decode.
302
303 Parameters
304 ----------
305 model : ESMFold2Model
306 The folding model. Must already be on the target device and in eval mode.
307 input : StructurePredictionInput
308 User-facing input specification.
309 num_loops, num_sampling_steps, num_diffusion_samples : int
310 Inference knobs forwarded to the model.
311 seed : int, optional
312 Seeds both input prep (SMILES conformer generation) and diffusion sampling.
313 noise_scale, step_scale, max_inference_sigma, early_exit
314 Optional sampler overrides forwarded to the model when not None.
315 complex_id : str
316 Identifier assigned to the predicted MolecularComplex(es).
317
318 Returns
319 -------
320 MolecularComplexResult or list[MolecularComplexResult]
321 A single result when num_diffusion_samples == 1, otherwise a list.
322 """
323 features, chain_infos = self.prepare_input(
324 input, seed=seed, device=model.device
325 )
326
327 sampler_kwargs: dict[str, Any] = {}
328 if noise_scale is not None:
329 sampler_kwargs["noise_scale"] = noise_scale
330 if step_scale is not None:
331 sampler_kwargs["step_scale"] = step_scale
332 if max_inference_sigma is not None:
333 sampler_kwargs["max_inference_sigma"] = max_inference_sigma
334
335 with torch.no_grad():
336 with _seed_context(seed) if seed is not None else nullcontext():
337 output = model(
338 **features,
339 num_loops=num_loops,
340 num_sampling_steps=num_sampling_steps,
341 num_diffusion_samples=num_diffusion_samples,
342 early_exit=early_exit,
343 **sampler_kwargs,
344 )
345
346 return self.decode(
347 output,
348 features,
349 chain_infos,
350 num_diffusion_samples=num_diffusion_samples,
351 complex_id=complex_id,
352 )
353
354
355__all__ = ["ESMFold2InputBuilder", "clean_esmfold2_input"]
356