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