Upload folder using huggingface_hub

LICENSE

@@ -1,9 +1,9 @@
-**License (MIT)**
-
-Copyright 2026 Chan Zuckerberg Biohub, Inc.
-
-Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+**License (MIT)**
+
+Copyright 2026 Chan Zuckerberg Biohub, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

README.md

@@ -1,122 +1,247 @@
----
-library_name: transformers
-tags:
-  - biology
-  - protein-structure
-  - esmfold2
-  - multimodal-protein-model
----
-
-# FastPLMs ESMFold2
-
-FastPLMs ESMFold2 is a self-contained Hugging Face `AutoModel` wrapper for Biohub's ESMFold2 and ESMFold2-Fast structure predictors. It vendors the released Biohub ESMFold2 model code, ESMC backbone code, input builder, MSA helpers, and structure export utilities needed for remote-code loading.
-
-## Load With AutoModel
+---
+library_name: transformers
+tags:
+  - biology
+  - protein-structure
+  - esmfold2
+  - multimodal-protein-model
+---
+
+# FastPLMs ESMFold2
+
+FastPLMs ESMFold2 is a self-contained Hugging Face `AutoModel` wrapper for
+Biohub's ESMFold2, ESMFold2-Fast, and experimental ESMFold2 structure
+predictors. It vendors the released Biohub ESMFold2 model code, input builder,
+MSA helpers, and structure export utilities, while loading the PLM backbone
+through FastPLMs ESM++.
+
+## Load With AutoModel
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "Synthyra/ESMFold2-Fast",
+    trust_remote_code=True,
+    dtype=torch.float32,
+).eval().cuda()
+```
+
+Use `Synthyra/ESMFold2` for the full model, `Synthyra/ESMFold2-Fast` for the
+faster release variant, and the `Synthyra/ESMFold2-Experimental*` checkpoints
+for differentiable binder design and experimental critic ensembles.
+The folding trunk runs in fp32; the 6B FastPLMs ESM++ backbone is loaded in
+bf16 by default via `esmc_precision="bf16"` and uses the flex attention backend
+by default inside ESMFold2.
+
+## Fold One Protein
+
+```python
+sequence = "MKTLLILAVVAAALA"
+
+result = model.fold_protein(
+    sequence,
+    num_loops=3,
+    num_sampling_steps=50,
+    num_diffusion_samples=1,
+    seed=0,
+)
+
+print(float(result.plddt.mean()))
+print(float(result.ptm))
+```
+
+## Experimental Test-Time Training
+
+TTT is disabled by default. Standard `fold_protein(...)`, `fold(...)`, raw tensor
+inference, and `state_dict()` keys are unchanged unless you explicitly pass
+`ttt=True` or call `fold_protein_ttt(...)`.
+
+The ESMFold2 TTT path is experimental and protein-only in v1. It trains local
+LoRA adapters only on `_esmc` with a masked language modeling objective. The
+folding trunk, confidence head, diffusion head, and structure input pipeline are
+frozen. TTT can improve difficult low-confidence folds, but it adds substantial
+test-time compute and can degrade already confident predictions.
+
+```python
+result = model.fold_protein(
+    "MSTNPKPQRKTKRNT",
+    num_loops=1,
+    num_sampling_steps=10,
+    num_diffusion_samples=1,
+    seed=0,
+    ttt=True,
+    ttt_config={
+        "steps": 1,
+        "ags": 1,
+        "batch_size": 1,
+        "lora_rank": 8,
+        "lora_alpha": 32.0,
+    },
+)
+
+print(result.ttt_metrics["losses"])
+print(result.ttt_metrics["step_plddts"])
+print(result.ttt_metrics["best_step"])
+```
+
+`load_esmc=True` is required for TTT because the ESM++ MLM head is loaded lazily
+from `config.esmc_id`. If that pretrained MLM head cannot be loaded, TTT raises
+an assertion instead of silently using a random head.
+
+## Save mmCIF or PDB
+
+```python
+model.save_as_cif(result, "prediction.cif")
+model.save_as_pdb(result, "prediction.pdb")
+
+cif_text = model.result_to_cif(result)
+pdb_text = model.result_to_pdb(result)
+```
+
+`result_to_cif` preserves the full `MolecularComplex`. `result_to_pdb` converts through Biohub's protein-only `ProteinComplex` representation, so use mmCIF for complexes with ligands or nucleic acids.
+
+## Fold Complexes
+
+```python
+types = model.input_types
+
+complex_input = types.StructurePredictionInput(
+    sequences=[
+        types.ProteinInput(id="A", sequence="MKTLLILAVVAAALA"),
+        types.DNAInput(id="B", sequence="GATAGC"),
+        types.LigandInput(id="L", ccd=["SAH"]),
+    ]
+)
+
+result = model.fold(
+    complex_input,
+    num_loops=3,
+    num_sampling_steps=50,
+    num_diffusion_samples=1,
+    seed=0,
+)
+
+model.save_as_cif(result, "complex_prediction.cif")
+```
+
+## Binder Design With FastPLMs ESMFold2
+
+FastPLMs includes a FastPLMs-only port of the Biohub ESMFold2 binder design
+tutorial at `cookbook/tutorials/binder_design_fastplms.py`. The workflow uses
+ESMFold2 experimental checkpoints for differentiable folding losses, ESM++ for
+sequence regularization, and ESMFold2 hero critics for final confidence scoring.
+
+![FastPLMs EGFR minibinder design](https://raw.githubusercontent.com/Synthyra/FastPLMs/main/docs/assets/egfr_fastplms_binder_design.png)
+
+The optimizer follows the official strategy:
+
+1. Optimize mutable `#` residues as continuous amino acid logits.
+2. Suppress cysteine design by masking cysteine logits and gradients.
+3. Backpropagate through ESMFold2 `res_type_soft` using intra-contact,
+   inter-contact, and globularity losses from the distogram.
+4. Add an ESM++ masked-LM pseudoperplexity regularizer on mutable binder
+   residues.
+5. Keep the late-trajectory sequence with the best iPTM.
+6. Fold the selected sequence with the final critic ensemble and write
+   `results.parquet`, `selection.parquet`, `trajectory.jsonl`,
+   `best_sequences.fasta`, and per-critic PDB/CIF/logit files.
+
+Run the verified EGFR 128 amino acid de novo minibinder example:
+
+```bash
+cd /home/ubuntu/FastPLMs
+
+sudo -n docker run --gpus all --rm \
+  -v /home/ubuntu/FastPLMs:/app \
+  -v /home/ubuntu/FastPLMs:/workspace \
+  -v /home/ubuntu/.cache/huggingface:/workspace/.cache/huggingface \
+  -w /workspace fastplms-esmfold2 \
+  python /app/cookbook/tutorials/binder_design_fastplms.py \
+    --backend local \
+    --target-name egfr \
+    --binder-sequence '################################################################################################################################' \
+    --not-antibody \
+    --steps 150 \
+    --batch-size 1 \
+    --seed 103 \
+    --output-dir /workspace/campaign_egfr_len128_b1_s150_seed103_consensus_cli
+```
+
+Verified result:
+
+| Metric | Value |
+| :--- | :--- |
+| Binder length | `128` |
+| Seed | `103` |
+| Steps | `150` |
+| Hero mean iPTM | `0.913870` |
+| Hero min iPTM | `0.904600` |
+| All four hero critics above 0.9 | `True` |
+
+Binder sequence:
+
+```text
+SAVKHLLEIVKYLEEAIEKALEVDPVFLVPPAAEELLIAAKVIKELAKENPELIEVYELLMKAVKGLKKLVRSNDKEILREVIRLLRKAAKVIREILKNNPDLDPELRKALEELAKVLEEIAEVLEQQ
+```
+
+See the full guide in [`docs/binder_design.md`](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md)
+for Modal execution, official pI and selection scoring, per-critic metrics, and
+the tested cheaper step-count boundary.
+
+## Use MSAs
+
+```python
+types = model.input_types
+
+msa = types.MSA.from_a3m("query.a3m", max_sequences=128)
+input_with_msa = types.StructurePredictionInput(
+    sequences=[
+        types.ProteinInput(id="A", sequence=msa.query, msa=msa),
+    ]
+)
+
+result = model.fold(input_with_msa, num_sampling_steps=50, seed=0)
+```
+
+## Raw Tensor Inference
+
+```python
+features, chain_infos = model.prepare_structure_input(complex_input, seed=0)
+
+with torch.inference_mode():
+    output = model(
+        **features,
+        num_loops=3,
+        num_sampling_steps=50,
+        num_diffusion_samples=1,
+    )
+
+decoded = model.input_builder.decode(output, features, chain_infos)
+```
+
+Set `load_esmc=False` when loading if you want to provide precomputed `lm_hidden_states` manually or run folding-trunk tests without loading the 6B ESM++ backbone:
 
 ```python
-import torch
-from transformers import AutoModel
-
 model = AutoModel.from_pretrained(
     "Synthyra/ESMFold2-Fast",
-    trust_remote_code=True,
-    dtype=torch.float32,
-).eval().cuda()
-```
-
-Use `Synthyra/ESMFold2` for the full model and `Synthyra/ESMFold2-Fast` for the faster release variant.
-The folding trunk runs in fp32; the 6B ESMC backbone is loaded in bf16 by default via `esmc_precision="bf16"`.
-
-## Fold One Protein
-
-```python
-sequence = "MKTLLILAVVAAALA"
-
-result = model.fold_protein(
-    sequence,
-    num_loops=3,
-    num_sampling_steps=50,
-    num_diffusion_samples=1,
-    seed=0,
-)
-
-print(float(result.plddt.mean()))
-print(float(result.ptm))
-```
-
-## Save mmCIF or PDB
-
-```python
-model.save_as_cif(result, "prediction.cif")
-model.save_as_pdb(result, "prediction.pdb")
-
-cif_text = model.result_to_cif(result)
-pdb_text = model.result_to_pdb(result)
-```
-
-`result_to_cif` preserves the full `MolecularComplex`. `result_to_pdb` converts through Biohub's protein-only `ProteinComplex` representation, so use mmCIF for complexes with ligands or nucleic acids.
-
-## Fold Complexes
-
-```python
-types = model.input_types
-
-complex_input = types.StructurePredictionInput(
-    sequences=[
-        types.ProteinInput(id="A", sequence="MKTLLILAVVAAALA"),
-        types.DNAInput(id="B", sequence="GATAGC"),
-        types.LigandInput(id="L", ccd=["SAH"]),
-    ]
-)
-
-result = model.fold(
-    complex_input,
-    num_loops=3,
-    num_sampling_steps=50,
-    num_diffusion_samples=1,
-    seed=0,
-)
-
-model.save_as_cif(result, "complex_prediction.cif")
-```
-
-## Use MSAs
-
-```python
-types = model.input_types
-
-msa = types.MSA.from_a3m("query.a3m", max_sequences=128)
-input_with_msa = types.StructurePredictionInput(
-    sequences=[
-        types.ProteinInput(id="A", sequence=msa.query, msa=msa),
-    ]
-)
-
-result = model.fold(input_with_msa, num_sampling_steps=50, seed=0)
-```
-
-## Raw Tensor Inference
-
-```python
-features, chain_infos = model.prepare_structure_input(complex_input, seed=0)
-
-with torch.inference_mode():
-    output = model(
-        **features,
-        num_loops=3,
-        num_sampling_steps=50,
-        num_diffusion_samples=1,
-    )
-
-decoded = model.input_builder.decode(output, features, chain_infos)
+    trust_remote_code=True,
+    load_esmc=False,
+).cuda().eval()
 ```
 
-Set `load_esmc=False` when loading if you want to provide precomputed `lm_hidden_states` manually or run folding-trunk tests without loading the 6B ESMC backbone:
+For FP8 LM inference, install `transformer_engine.pytorch` in a CUDA
+environment with FP8-capable hardware and load the shared FastPLMs ESM++
+backbone with:
 
 ```python
 model = AutoModel.from_pretrained(
     "Synthyra/ESMFold2-Fast",
     trust_remote_code=True,
-    load_esmc=False,
+    esmc_precision="fp8",
 ).cuda().eval()
 ```
+
+FP8 is inference-only for the ESMFold2 LM backbone. TTT remains a bf16/fp32
+path.

__init__.py

@@ -1,6 +1,5 @@
-from .configuration_esmfold2 import ESMFold2Config
-from .modeling_esmfold2_experimental import ESMFold2ExperimentalModel
-from .modeling_esmfold2 import ESMFold2Model
-
-__all__ = ["ESMFold2Config", "ESMFold2ExperimentalModel", "ESMFold2Model"]
+from .configuration_esmfold2 import ESMFold2Config
+from .modeling_esmfold2_experimental import ESMFold2ExperimentalModel
+from .modeling_esmfold2 import ESMFold2Model
 
+__all__ = ["ESMFold2Config", "ESMFold2ExperimentalModel", "ESMFold2Model"]

configuration_esmfold2.py

@@ -1,194 +1,206 @@
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""ESMFold2 model configuration."""
-
-from __future__ import annotations
-
-from dataclasses import asdict, dataclass, field
-
-from transformers.configuration_utils import PretrainedConfig
-
-# ---------------------------------------------------------------------------
-# Nested dataclass configs
-# ---------------------------------------------------------------------------
-
-_DEFAULT_ESMC_HF_REPO = "biohub/ESMC-6B"
-
-
-@dataclass
-class MSAEncoderConfig:
-    """Config for the optional MSA encoder module (Large MSA models only)."""
-
-    enabled: bool = False
-    d_msa: int = 128
-    d_hidden: int = 32
-    n_layers: int = 4
-    n_heads_msa: int = 8
-    msa_head_width: int = 32
-
-
-@dataclass
-class ParcaeConfig:
-    """Release-only config for the parcae diffusion-loop scheduler."""
-
-    enabled: bool = True
-    poisson_mean: float = 3.0
-    min_steps: int = 1
-    max_steps: int | None = 6
-    coda_n_layers: int = 2
-
-
-@dataclass
-class LMEncoderConfig:
-    """Release-only config for the LM-side pair encoder."""
-
-    enabled: bool = True
-    n_layers: int = 4
-    lm_dropout: float = 0.25
-    per_loop_lm_dropout: bool = True
-
-
-@dataclass
-class AtomAttentionConfig:
-    """Config for SWA atom encoder/decoder with 3D RoPE."""
-
-    d_atom: int = 128
-    d_token: int = 768
-    n_blocks: int = 3
-    n_heads: int = 4
-    swa_window_size: int = 128
-    expansion_ratio: int = 2
-    # 3D RoPE config
-    spatial_rope_base_frequency: float = 20.0
-    n_spatial_rope_pairs_per_axis: int = 2
-    n_uid_rope_pairs: int = 10
-    uid_rope_base_frequency: float = 10000.0
-
-
-@dataclass
-class FoldingTrunkConfig:
-    n_layers: int = 24
-    n_heads: int = 8
-    dropout: float = 0.0
-
-
-@dataclass
-class InputsEmbedderConfig:
-    d_inputs: int = 451
-    atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig)
-
-    def __post_init__(self):
-        if isinstance(self.atom_encoder, dict):
-            self.atom_encoder = AtomAttentionConfig(**self.atom_encoder)
-
-
-@dataclass
-class DiffusionModuleConfig:
-    """Config for the DiffusionModule."""
-
-    sigma_data: float = 16.0
-    c_atom: int = 128
-    c_token: int = 768
-    c_z: int = 256
-    c_s_inputs: int = 451
-    fourier_dim: int = 256
-    relpos_r_max: int = 32
-    relpos_s_max: int = 2
-    atom_num_blocks: int = 3
-    atom_num_heads: int = 4
-    token_num_blocks: int = 12
-    token_num_heads: int = 16
-    transition_multiplier: int = 2
-
-
-@dataclass
-class DiffusionStructureHeadConfig:
-    """Config for the diffusion-based structure prediction head."""
-
-    diffusion_module: DiffusionModuleConfig = field(
-        default_factory=DiffusionModuleConfig
-    )
-    distogram_bins: int = 128
-
-    # Training noise: sigma ~ sigma_data * exp(mu + sigma * N(0,1))
-    train_noise_log_mean: float = -1.2
-    train_noise_log_std: float = 1.5
-
-    # Sampling defaults (ODE)
-    gamma_0: float = 0.605
-    gamma_min: float = 1.107
-    noise_scale: float = 0.0
-    step_scale: float = 1.0
-
-    # Inference schedule defaults
-    inference_s_max: float = 160.0
-    inference_s_min: float = 4e-4
-    inference_p: float = 8.0
-    inference_num_steps: int = 68
-
-    def __post_init__(self):
-        if isinstance(self.diffusion_module, dict):
-            self.diffusion_module = DiffusionModuleConfig(**self.diffusion_module)
-
-
-@dataclass
-class ConfidenceHeadConfig:
-    enabled: bool = True
-    num_plddt_bins: int = 50
-    num_pde_bins: int = 64
-    num_pae_bins: int = 64
-    min_dist: float = 2.0
-    max_dist: float = 52.0
-    distogram_bins: int = 128
-    folding_trunk: FoldingTrunkConfig = field(
-        default_factory=lambda: FoldingTrunkConfig(n_layers=4)
-    )
-
-    def __post_init__(self):
-        if isinstance(self.folding_trunk, dict):
-            self.folding_trunk = FoldingTrunkConfig(**self.folding_trunk)
-
-
-# ---------------------------------------------------------------------------
-# Top-level config
-# ---------------------------------------------------------------------------
-
-
-class ESMFold2Config(PretrainedConfig):
-    """
-    Configuration for the ESMFold2 structure prediction model.
-
-    Uses SWA atom encoders with 3D RoPE, a diffusion transformer,
-    a folding trunk, and an ESMC 6B PLM backbone.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control
-    the model outputs. Read the documentation from [`PretrainedConfig`] for more
-    information.
-
-    Args:
-        d_single (`int`, defaults to 384):
-            Dimensionality of single (per-residue) representations.
-        d_pair (`int`, defaults to 256):
-            Dimensionality of pair (residue-residue) representations.
-        n_relative_residx_bins (`int`, defaults to 32):
-            Number of bins for relative residue index encoding.
-        n_relative_chain_bins (`int`, defaults to 2):
-            Number of bins for relative chain encoding.
-        num_loops (`int`, defaults to 10):
-            Number of trunk loops for iterative refinement.
-        num_diffusion_samples (`int`, defaults to 8):
-            Number of parallel structure predictions to generate.
+# Copyright 2026 Biohub. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""ESMFold2 model configuration."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+
+from transformers.configuration_utils import PretrainedConfig
+
+# ---------------------------------------------------------------------------
+# Nested dataclass configs
+# ---------------------------------------------------------------------------
+
+_DEFAULT_ESMC_HF_REPO = "Synthyra/ESMplusplus_6B"
+_DEFAULT_ESMC_ATTN_BACKEND = "flex"
+_ESMC_ID_ALIASES = {
+    "biohub/ESMC-300M": "Synthyra/ESMplusplus_small",
+    "biohub/ESMC-600M": "Synthyra/ESMplusplus_large",
+    "biohub/ESMC-6B": "Synthyra/ESMplusplus_6B",
+}
+
+
+def normalize_esmc_id(esmc_id: str) -> str:
+    if esmc_id in _ESMC_ID_ALIASES:
+        return _ESMC_ID_ALIASES[esmc_id]
+    return esmc_id
+
+
+@dataclass
+class MSAEncoderConfig:
+    """Config for the optional MSA encoder module (Large MSA models only)."""
+
+    enabled: bool = False
+    d_msa: int = 128
+    d_hidden: int = 32
+    n_layers: int = 4
+    n_heads_msa: int = 8
+    msa_head_width: int = 32
+
+
+@dataclass
+class ParcaeConfig:
+    """Release-only config for the parcae diffusion-loop scheduler."""
+
+    enabled: bool = True
+    poisson_mean: float = 3.0
+    min_steps: int = 1
+    max_steps: int | None = 6
+    coda_n_layers: int = 2
+
+
+@dataclass
+class LMEncoderConfig:
+    """Release-only config for the LM-side pair encoder."""
+
+    enabled: bool = True
+    n_layers: int = 4
+    lm_dropout: float = 0.25
+    per_loop_lm_dropout: bool = True
+
+
+@dataclass
+class AtomAttentionConfig:
+    """Config for SWA atom encoder/decoder with 3D RoPE."""
+
+    d_atom: int = 128
+    d_token: int = 768
+    n_blocks: int = 3
+    n_heads: int = 4
+    swa_window_size: int = 128
+    expansion_ratio: int = 2
+    # 3D RoPE config
+    spatial_rope_base_frequency: float = 20.0
+    n_spatial_rope_pairs_per_axis: int = 2
+    n_uid_rope_pairs: int = 10
+    uid_rope_base_frequency: float = 10000.0
+
+
+@dataclass
+class FoldingTrunkConfig:
+    n_layers: int = 24
+    n_heads: int = 8
+    dropout: float = 0.0
+
+
+@dataclass
+class InputsEmbedderConfig:
+    d_inputs: int = 451
+    atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig)
+
+    def __post_init__(self):
+        if isinstance(self.atom_encoder, dict):
+            self.atom_encoder = AtomAttentionConfig(**self.atom_encoder)
+
+
+@dataclass
+class DiffusionModuleConfig:
+    """Config for the DiffusionModule."""
+
+    sigma_data: float = 16.0
+    c_atom: int = 128
+    c_token: int = 768
+    c_z: int = 256
+    c_s_inputs: int = 451
+    fourier_dim: int = 256
+    relpos_r_max: int = 32
+    relpos_s_max: int = 2
+    atom_num_blocks: int = 3
+    atom_num_heads: int = 4
+    token_num_blocks: int = 12
+    token_num_heads: int = 16
+    transition_multiplier: int = 2
+
+
+@dataclass
+class DiffusionStructureHeadConfig:
+    """Config for the diffusion-based structure prediction head."""
+
+    diffusion_module: DiffusionModuleConfig = field(
+        default_factory=DiffusionModuleConfig
+    )
+    distogram_bins: int = 128
+
+    # Training noise: sigma ~ sigma_data * exp(mu + sigma * N(0,1))
+    train_noise_log_mean: float = -1.2
+    train_noise_log_std: float = 1.5
+
+    # Sampling defaults (ODE)
+    gamma_0: float = 0.605
+    gamma_min: float = 1.107
+    noise_scale: float = 0.0
+    step_scale: float = 1.0
+
+    # Inference schedule defaults
+    inference_s_max: float = 160.0
+    inference_s_min: float = 4e-4
+    inference_p: float = 8.0
+    inference_num_steps: int = 68
+
+    def __post_init__(self):
+        if isinstance(self.diffusion_module, dict):
+            self.diffusion_module = DiffusionModuleConfig(**self.diffusion_module)
+
+
+@dataclass
+class ConfidenceHeadConfig:
+    enabled: bool = True
+    num_plddt_bins: int = 50
+    num_pde_bins: int = 64
+    num_pae_bins: int = 64
+    min_dist: float = 2.0
+    max_dist: float = 52.0
+    distogram_bins: int = 128
+    folding_trunk: FoldingTrunkConfig = field(
+        default_factory=lambda: FoldingTrunkConfig(n_layers=4)
+    )
+
+    def __post_init__(self):
+        if isinstance(self.folding_trunk, dict):
+            self.folding_trunk = FoldingTrunkConfig(**self.folding_trunk)
+
+
+# ---------------------------------------------------------------------------
+# Top-level config
+# ---------------------------------------------------------------------------
+
+
+class ESMFold2Config(PretrainedConfig):
+    """
+    Configuration for the ESMFold2 structure prediction model.
+
+    Uses SWA atom encoders with 3D RoPE, a diffusion transformer,
+    a folding trunk, and an ESMC 6B PLM backbone.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control
+    the model outputs. Read the documentation from [`PretrainedConfig`] for more
+    information.
+
+    Args:
+        d_single (`int`, defaults to 384):
+            Dimensionality of single (per-residue) representations.
+        d_pair (`int`, defaults to 256):
+            Dimensionality of pair (residue-residue) representations.
+        n_relative_residx_bins (`int`, defaults to 32):
+            Number of bins for relative residue index encoding.
+        n_relative_chain_bins (`int`, defaults to 2):
+            Number of bins for relative chain encoding.
+        num_loops (`int`, defaults to 10):
+            Number of trunk loops for iterative refinement.
+        num_diffusion_samples (`int`, defaults to 8):
+            Number of parallel structure predictions to generate.
         lm_dropout (`float`, defaults to 0.0):
             Dropout probability on LM pair embeddings. When > 0, dropout is
             applied with ``training=True`` (including at inference) to match
@@ -196,103 +208,121 @@ class ESMFold2Config(PretrainedConfig):
         force_lm_dropout_during_inference (`bool`, defaults to False):
             When True, apply ``lm_dropout`` even when ``model.eval()`` and
             ``lm_dropout`` > 0. Binder-design loads set this to True.
+        lm_mask_pct (`float`, defaults to 0.0):
+            Fraction of LM residue tokens randomly replaced with the LM mask
+            token before running the PLM backbone.
         disable_msa_features (`bool`, defaults to False):
             When True, zero out MSA-derived ``profile`` and ``deletion_mean``
             before the inputs embedder (experimental medium/large checkpoints).
-        inputs (`InputsEmbedderConfig`):
-            Configuration for the inputs embedder module.
-        folding_trunk (`FoldingTrunkConfig`):
-            Configuration for the folding trunk.
-        structure_head (`DiffusionStructureHeadConfig`):
-            Configuration for the diffusion-based structure prediction head.
-        confidence_head (`ConfidenceHeadConfig`):
-            Configuration for the confidence prediction head.
-
-    Examples:
-
-    ```python
-    >>> from transformers import ESMFold2Config, ESMFold2ExperimentalModel
-
-    >>> # Initializing an ESMFold2 configuration
-    >>> configuration = ESMFold2Config(type="experimental")
-
-    >>> # Initializing a model (with random weights) from the configuration
-    >>> model = ESMFold2ExperimentalModel(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```
-    """
-
-    model_type = "esmfold2"
-    has_no_defaults_at_init = True
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-        self.type: str = kwargs.get("type", "release")
-        if self.type not in ("release", "experimental"):
-            raise ValueError(
-                f"ESMFold2Config.type must be 'release' or 'experimental', "
-                f"got {self.type!r}"
-            )
-
-        # Top-level scalar fields
-        self.d_single: int = kwargs.get("d_single", 384)
-        self.d_pair: int = kwargs.get("d_pair", 256)
-        self.n_relative_residx_bins: int = kwargs.get("n_relative_residx_bins", 32)
-        self.n_relative_chain_bins: int = kwargs.get("n_relative_chain_bins", 2)
-        self.num_loops: int = kwargs.get("num_loops", 10)
-        self.num_diffusion_samples: int = kwargs.get("num_diffusion_samples", 8)
-        # If True, ``profile`` / ``deletion_mean`` are zeroed before the inputs
-        # embedder.
-        self.disable_msa_features: bool = kwargs.get("disable_msa_features", False)
-        self.lm_dropout: float = kwargs.get("lm_dropout", 0.0)
+        inputs (`InputsEmbedderConfig`):
+            Configuration for the inputs embedder module.
+        folding_trunk (`FoldingTrunkConfig`):
+            Configuration for the folding trunk.
+        structure_head (`DiffusionStructureHeadConfig`):
+            Configuration for the diffusion-based structure prediction head.
+        confidence_head (`ConfidenceHeadConfig`):
+            Configuration for the confidence prediction head.
+
+    Examples:
+
+    ```python
+    >>> from transformers import ESMFold2Config, ESMFold2ExperimentalModel
+
+    >>> # Initializing an ESMFold2 configuration
+    >>> configuration = ESMFold2Config(type="experimental")
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = ESMFold2ExperimentalModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "esmfold2"
+    has_no_defaults_at_init = True
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        self.type: str = kwargs.get("type", "release")
+        if self.type not in ("release", "experimental"):
+            raise ValueError(
+                f"ESMFold2Config.type must be 'release' or 'experimental', "
+                f"got {self.type!r}"
+            )
+
+        # Top-level scalar fields
+        self.d_single: int = kwargs.get("d_single", 384)
+        self.d_pair: int = kwargs.get("d_pair", 256)
+        self.n_relative_residx_bins: int = kwargs.get("n_relative_residx_bins", 32)
+        self.n_relative_chain_bins: int = kwargs.get("n_relative_chain_bins", 2)
+        self.num_loops: int = kwargs.get("num_loops", 10)
+        self.num_diffusion_samples: int = kwargs.get("num_diffusion_samples", 8)
+        # If True, ``profile`` / ``deletion_mean`` are zeroed before the inputs
+        # embedder.
+        self.disable_msa_features: bool = kwargs.get("disable_msa_features", False)
+        self.lm_dropout: float = kwargs.get("lm_dropout", 0.0)
         self.force_lm_dropout_during_inference: bool = kwargs.get(
             "force_lm_dropout_during_inference", False
         )
+        self.lm_mask_pct: float = kwargs.get("lm_mask_pct", 0.0)
 
         self.lm_d_model: int = kwargs.get("lm_d_model", 2560)
-        self.lm_num_layers: int = kwargs.get("lm_num_layers", 80)
-        # Required, no default — every shipped HF export must name its ESMC backbone.
-        self.esmc_id: str = kwargs.get("esmc_id", _DEFAULT_ESMC_HF_REPO)
-
-        def _init_nested(cls, val):
-            if isinstance(val, cls):
-                return val
-            if isinstance(val, dict):
-                return cls(**val)
-            return cls()
-
-        self.inputs = _init_nested(InputsEmbedderConfig, kwargs.get("inputs"))
-        self.folding_trunk = _init_nested(
-            FoldingTrunkConfig, kwargs.get("folding_trunk")
-        )
-        self.structure_head = _init_nested(
-            DiffusionStructureHeadConfig, kwargs.get("structure_head")
-        )
-        self.confidence_head = _init_nested(
-            ConfidenceHeadConfig, kwargs.get("confidence_head")
-        )
-        self.msa_encoder = _init_nested(MSAEncoderConfig, kwargs.get("msa_encoder"))
-        # Release-only modules — ignored when ``type == "experimental"``.
-        self.parcae = _init_nested(ParcaeConfig, kwargs.get("parcae"))
-        self.lm_encoder = _init_nested(LMEncoderConfig, kwargs.get("lm_encoder"))
-        # If True, MSA encoder output replaces the pair stream; if False, it is added.
-        self.msa_encoder_overwrite: bool = bool(
-            kwargs.get("msa_encoder_overwrite", True)
-        )
-
-    def to_dict(self):
-        output = super().to_dict()
-        output["inputs"] = asdict(self.inputs)
-        output["folding_trunk"] = asdict(self.folding_trunk)
-        output["structure_head"] = asdict(self.structure_head)
-        output["confidence_head"] = asdict(self.confidence_head)
-        output["msa_encoder"] = asdict(self.msa_encoder)
-        output["parcae"] = asdict(self.parcae)
-        output["lm_encoder"] = asdict(self.lm_encoder)
-        return output
-
-
-__all__ = ["ESMFold2Config", "MSAEncoderConfig", "ParcaeConfig", "LMEncoderConfig"]
+        self.lm_num_layers: int = kwargs.get("lm_num_layers", 80)
+        # Backward-compatible field name; values now point to FastPLMs ESM++.
+        raw_esmc_id = (
+            kwargs["esmc_id"] if "esmc_id" in kwargs else _DEFAULT_ESMC_HF_REPO
+        )
+        self.esmc_id: str = normalize_esmc_id(raw_esmc_id)
+        self.esmc_attn_backend: str = (
+            kwargs["esmc_attn_backend"]
+            if "esmc_attn_backend" in kwargs
+            else _DEFAULT_ESMC_ATTN_BACKEND
+        )
+
+        def _init_nested(cls, val):
+            if isinstance(val, cls):
+                return val
+            if isinstance(val, dict):
+                return cls(**val)
+            return cls()
+
+        self.inputs = _init_nested(InputsEmbedderConfig, kwargs.get("inputs"))
+        self.folding_trunk = _init_nested(
+            FoldingTrunkConfig, kwargs.get("folding_trunk")
+        )
+        self.structure_head = _init_nested(
+            DiffusionStructureHeadConfig, kwargs.get("structure_head")
+        )
+        self.confidence_head = _init_nested(
+            ConfidenceHeadConfig, kwargs.get("confidence_head")
+        )
+        self.msa_encoder = _init_nested(MSAEncoderConfig, kwargs.get("msa_encoder"))
+        # Release-only modules — ignored when ``type == "experimental"``.
+        self.parcae = _init_nested(ParcaeConfig, kwargs.get("parcae"))
+        self.lm_encoder = _init_nested(LMEncoderConfig, kwargs.get("lm_encoder"))
+        # If True, MSA encoder output replaces the pair stream; if False, it is added.
+        self.msa_encoder_overwrite: bool = bool(
+            kwargs.get("msa_encoder_overwrite", True)
+        )
+
+    def to_dict(self):
+        output = super().to_dict()
+        output["inputs"] = asdict(self.inputs)
+        output["folding_trunk"] = asdict(self.folding_trunk)
+        output["structure_head"] = asdict(self.structure_head)
+        output["confidence_head"] = asdict(self.confidence_head)
+        output["msa_encoder"] = asdict(self.msa_encoder)
+        output["parcae"] = asdict(self.parcae)
+        output["lm_encoder"] = asdict(self.lm_encoder)
+        return output
+
+
+__all__ = [
+    "ESMFold2Config",
+    "MSAEncoderConfig",
+    "ParcaeConfig",
+    "LMEncoderConfig",
+    "normalize_esmc_id",
+]

esmfold2_affine3d.py

@@ -1,560 +1,560 @@
-from __future__ import annotations
-
-import typing as T
-from abc import ABC
-from dataclasses import dataclass
-
-import torch
-from torch.nn import functional as F
-from typing_extensions import Self
-
-from .esmfold2_misc import fp32_autocast_context
-
-
-class Rotation(ABC):
-    @classmethod
-    def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ...
-
-    @classmethod
-    def random(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ...
-
-    def __getitem__(self, idx: T.Any) -> Self: ...
-
-    @property
-    def tensor(self) -> torch.Tensor:
-        # We claim that this should be zero-cost abstraction that returns the raw tensor backing this
-        # object. The raw tensor should always have exactly 1 more dim than self.shape, which should be
-        # implemented using reshaping
-        ...
-
-    @property
-    def shape(self) -> torch.Size:
-        # The "shape" of the rotation, as if it was a torch.tensor object
-        # This means that 1x4 quaternions are treated as size (1,) for example
-        ...
-
-    def as_matrix(self) -> RotationMatrix: ...
-
-    def as_quat(self, normalize: bool = False) -> RotationQuat: ...
-
-    def compose(self, other: Self) -> Self:
-        # To be safe, we force users to explicitly convert between rotation types.
-        ...
-
-    def convert_compose(self, other: Self) -> Self:
-        # This function will automatically convert between types of rotations
-        ...
-
-    def apply(self, p: torch.Tensor) -> torch.Tensor:
-        # rotates points by this rotation object
-        ...
-
-    def invert(self) -> Self: ...
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return self.tensor.dtype
-
-    @property
-    def device(self) -> torch.device:
-        return self.tensor.device
-
-    @property
-    def requires_grad(self) -> bool:
-        return self.tensor.requires_grad
-
-    @classmethod
-    def _from_tensor(cls, t: torch.Tensor) -> Self:
-        # This function exists to simplify the below functions, esp type signatures
-        # Its implementation is different from Affine3D.from_tensor and does not
-        # autodetect rotation types.
-        return cls(t)  # type: ignore
-
-    def to(self, **kwargs) -> Self:
-        return self._from_tensor(self.tensor.to(**kwargs))
-
-    def detach(self, *args, **kwargs) -> Self:
-        return self._from_tensor(self.tensor.detach(**kwargs))
-
-    def tensor_apply(self, func) -> Self:
-        # Applys a function to the underlying tensor
-        return self._from_tensor(
-            torch.stack([func(x) for x in self.tensor.unbind(dim=-1)], dim=-1)
-        )
-
-
-class RotationMatrix(Rotation):
-    def __init__(self, rots: torch.Tensor):
-        if rots.shape[-1] == 9:
-            rots = rots.unflatten(-1, (3, 3))
-        assert rots.shape[-1] == 3
-        assert rots.shape[-2] == 3
-        # Force full precision
-        rots = rots.to(torch.float32)
-        self._rots = rots
-
-    @classmethod
-    def identity(cls, shape, **tensor_kwargs):
-        rots = torch.eye(3, **tensor_kwargs)
-        rots = rots.view(*[1 for _ in range(len(shape))], 3, 3)
-        rots = rots.expand(*shape, -1, -1)
-        return cls(rots)
-
-    @classmethod
-    def random(cls, shape, **tensor_kwargs):
-        return RotationQuat.random(shape, **tensor_kwargs).as_matrix()
-
-    def __getitem__(self, idx: T.Any) -> RotationMatrix:
-        indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx)
-        return RotationMatrix(self._rots[indices + (slice(None), slice(None))])
-
-    @property
-    def shape(self) -> torch.Size:
-        return self._rots.shape[:-2]
-
-    def as_matrix(self) -> RotationMatrix:
-        return self
-
-    def as_quat(self, normalize: bool = False) -> RotationQuat:
-        m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
-            self._rots.flatten(-2), dim=-1
-        )
-        q_abs = _sqrt_subgradient(
-            torch.stack(
-                [
-                    1.0 + m00 + m11 + m22,
-                    1.0 + m00 - m11 - m22,
-                    1.0 - m00 + m11 - m22,
-                    1.0 - m00 - m11 + m22,
-                ],
-                dim=-1,
-            )
-        )
-        # we produce the desired quaternion multiplied by each of r, i, j, k
-        quat_by_rijk = torch.stack(
-            [
-                x
-                for lst in [
-                    [q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01],
-                    [m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20],
-                    [m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21],
-                    [m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2],
-                ]
-                for x in lst
-            ],
-            dim=-1,
-        ).unflatten(-1, (4, 4))
-
-        # We floor here at 0.1 but the exact level is not important; if q_abs is small,
-        # the candidate won't be picked.
-        flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
-        quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
-
-        # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
-        # forall i; we pick the best-conditioned one (with the largest denominator)
-        # We manually implement one_hot so torch.compile works
-        one_hot = torch.zeros_like(q_abs, dtype=torch.bool)
-        one_hot.scatter_(-1, q_abs.argmax(dim=-1, keepdim=True), True)
-        quat = quat_candidates[one_hot, :].reshape(q_abs.shape)
-        return RotationQuat(quat)
-
-    def compose(self, other: RotationMatrix) -> RotationMatrix:
-        with fp32_autocast_context(self._rots.device.type):
-            return RotationMatrix(self._rots @ other._rots)
-
-    def convert_compose(self, other: Rotation):
-        return self.compose(other.as_matrix())
-
-    def apply(self, p: torch.Tensor) -> torch.Tensor:
-        with fp32_autocast_context(self.device.type):
-            if self._rots.shape[-3] == 1:
-                # This is a slight speedup over einsum for batched rotations
-                return p @ self._rots.transpose(-1, -2).squeeze(-3)
-            else:
-                # einsum way faster than bmm!
-                return torch.einsum("...ij,...j", self._rots, p)
-
-    def invert(self) -> RotationMatrix:
-        return RotationMatrix(self._rots.transpose(-1, -2))
-
-    @property
-    def tensor(self) -> torch.Tensor:
-        return self._rots.flatten(-2)
-
-    def to_3x3(self) -> torch.Tensor:
-        return self._rots
-
-    @staticmethod
-    def from_graham_schmidt(
-        x_axis: torch.Tensor, xy_plane: torch.Tensor, eps: float = 1e-12
-    ) -> RotationMatrix:
-        # A low eps here is necessary for good stability!
-        return RotationMatrix(_graham_schmidt(x_axis, xy_plane, eps))
-
-
-class RotationQuat(Rotation):
-    def __init__(self, quats: torch.Tensor, normalized=False):
-        assert quats.shape[-1] == 4
-        self._normalized = normalized
-        # Force float32 as well
-        if normalized:
-            self._quats = F.normalize(quats.to(torch.float32), dim=-1)
-            self._quats = self._quats.where(self._quats[..., :1] >= 0, -self._quats)
-        else:
-            self._quats = quats.to(torch.float32)
-
-    @classmethod
-    def identity(cls, shape, **tensor_kwargs):
-        q = torch.ones((*shape, 4), **tensor_kwargs)
-        mult = torch.tensor([1, 0, 0, 0], device=q.device)
-        return RotationQuat(q * mult)
-
-    @classmethod
-    def random(cls, shape, **tensor_kwargs):
-        quat = torch.randn((*shape, 4), **tensor_kwargs)
-        return RotationQuat(quat, normalized=True)
-
-    def __getitem__(self, idx: T.Any) -> RotationQuat:
-        indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx)
-        return RotationQuat(self._quats[indices + (slice(None),)])
-
-    @property
-    def shape(self) -> torch.Size:
-        return self._quats.shape[:-1]
-
-    def compose(self, other: RotationQuat) -> RotationQuat:
-        with fp32_autocast_context(self._quats.device.type):
-            return RotationQuat(_quat_mult(self._quats, other._quats))
-
-    def convert_compose(self, other: Rotation):
-        return self.compose(other.as_quat())
-
-    def as_matrix(self) -> RotationMatrix:
-        q = self.normalized().tensor
-        r, i, j, k = torch.unbind(q, -1)
-        two_s = 2.0 / torch.linalg.norm(q, dim=-1)
-
-        o = torch.stack(
-            (
-                1 - two_s * (j * j + k * k),
-                two_s * (i * j - k * r),
-                two_s * (i * k + j * r),
-                two_s * (i * j + k * r),
-                1 - two_s * (i * i + k * k),
-                two_s * (j * k - i * r),
-                two_s * (i * k - j * r),
-                two_s * (j * k + i * r),
-                1 - two_s * (i * i + j * j),
-            ),
-            -1,
-        )
-        return RotationMatrix(o.reshape(q.shape[:-1] + (3, 3)))
-
-    def as_quat(self, normalize: bool = False) -> RotationQuat:
-        return self
-
-    def apply(self, p: torch.Tensor) -> torch.Tensor:
-        return _quat_rotation(self.normalized()._quats, p)
-
-    def invert(self) -> RotationQuat:
-        return RotationQuat(_quat_invert(self._quats))
-
-    @property
-    def tensor(self) -> torch.Tensor:
-        return self._quats
-
-    def normalized(self) -> RotationQuat:
-        return self if self._normalized else RotationQuat(self._quats, normalized=True)
-
-
-@dataclass(frozen=True)
-class Affine3D:
-    trans: torch.Tensor
-    rot: Rotation
-
-    def __post_init__(self):
-        assert self.trans.shape[:-1] == self.rot.shape
-
-    @staticmethod
-    def identity(
-        shape_or_affine: T.Union[tuple[int, ...], "Affine3D"],
-        rotation_type: T.Type[Rotation] = RotationMatrix,
-        **tensor_kwargs,
-    ):
-        # Creates a new identity Affine3D object with a specified shape
-        # or the same shape as another Affine3D object.
-        if isinstance(shape_or_affine, Affine3D):
-            kwargs = {"dtype": shape_or_affine.dtype, "device": shape_or_affine.device}
-            kwargs.update(tensor_kwargs)
-            shape = shape_or_affine.shape
-            rotation_type = type(shape_or_affine.rot)
-        else:
-            kwargs = tensor_kwargs
-            shape = shape_or_affine
-        return Affine3D(
-            torch.zeros((*shape, 3), **kwargs), rotation_type.identity(shape, **kwargs)
-        )
-
-    @staticmethod
-    def random(
-        shape: tuple[int, ...],
-        std: float = 1,
-        rotation_type: T.Type[Rotation] = RotationMatrix,
-        **tensor_kwargs,
-    ) -> "Affine3D":
-        return Affine3D(
-            trans=torch.randn((*shape, 3), **tensor_kwargs).mul(std),
-            rot=rotation_type.random(shape, **tensor_kwargs),
-        )
-
-    def __getitem__(self, idx: T.Any) -> "Affine3D":
-        indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx)
-        return Affine3D(trans=self.trans[indices + (slice(None),)], rot=self.rot[idx])
-
-    @property
-    def shape(self) -> torch.Size:
-        return self.trans.shape[:-1]
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return self.trans.dtype
-
-    @property
-    def device(self) -> torch.device:
-        return self.trans.device
-
-    @property
-    def requires_grad(self) -> bool:
-        return self.trans.requires_grad
-
-    def to(self, **kwargs) -> "Affine3D":
-        return Affine3D(self.trans.to(**kwargs), self.rot.to(**kwargs))
-
-    def detach(self, *args, **kwargs) -> "Affine3D":
-        return Affine3D(self.trans.detach(**kwargs), self.rot.detach(**kwargs))
-
-    def tensor_apply(self, func) -> "Affine3D":
-        # Applys a function to the underlying tensor
-        return self.from_tensor(
-            torch.stack([func(x) for x in self.tensor.unbind(dim=-1)], dim=-1)
-        )
-
-    def as_matrix(self):
-        return Affine3D(trans=self.trans, rot=self.rot.as_matrix())
-
-    def as_quat(self, normalize: bool = False):
-        return Affine3D(trans=self.trans, rot=self.rot.as_quat(normalize))
-
-    def compose(self, other: "Affine3D", autoconvert: bool = False):
-        rot = self.rot
-        new_rot = (rot.convert_compose if autoconvert else rot.compose)(other.rot)
-        new_trans = rot.apply(other.trans) + self.trans
-        return Affine3D(trans=new_trans, rot=new_rot)
-
-    def compose_rotation(self, other: Rotation, autoconvert: bool = False):
-        return Affine3D(
-            trans=self.trans,
-            rot=(self.rot.convert_compose if autoconvert else self.rot.compose)(other),
-        )
-
-    def scale(self, v: torch.Tensor | float):
-        return Affine3D(self.trans * v, self.rot)
-
-    def mask(self, mask: torch.Tensor, with_zero=False):
-        # Returns a transform where True positions in mask is identity
-        if with_zero:
-            tensor = self.tensor
-            return Affine3D.from_tensor(
-                torch.zeros_like(tensor).where(mask[..., None], tensor)
-            )
-        else:
-            identity = self.identity(
-                self.shape,
-                rotation_type=type(self.rot),
-                device=self.device,
-                dtype=self.dtype,
-            ).tensor
-            return Affine3D.from_tensor(identity.where(mask[..., None], self.tensor))
-
-    def apply(self, p: torch.Tensor) -> torch.Tensor:
-        return self.rot.apply(p) + self.trans
-
-    def invert(self):
-        inv_rot = self.rot.invert()
-        return Affine3D(trans=-inv_rot.apply(self.trans), rot=inv_rot)
-
-    @property
-    def tensor(self) -> torch.Tensor:
-        return torch.cat([self.rot.tensor, self.trans], dim=-1)
-
-    @staticmethod
-    def from_tensor(t: torch.Tensor) -> "Affine3D":
-        match t.shape[-1]:
-            case 4:
-                # Assume tensor 4x4 for backward compat with alphafold
-                trans = t[..., :3, 3]
-                rot = RotationMatrix(t[..., :3, :3])
-            case 6:
-                # Assume quaternion representation with real part = 1
-                trans = t[..., -3:]
-                rot = RotationQuat(F.pad(t[..., :3], (1, 0), value=1))
-            case 7:
-                trans = t[..., -3:]
-                rot = RotationQuat(t[..., :4])
-            case 12:
-                trans = t[..., -3:]
-                rot = RotationMatrix(t[..., :-3].unflatten(-1, (3, 3)))
-            case _:
-                raise RuntimeError(
-                    f"Cannot detect rotation fromat from {t.shape[-1] -3}-d flat vector"
-                )
-        return Affine3D(trans, rot)
-
-    @staticmethod
-    def from_tensor_pair(t: torch.Tensor, r: torch.Tensor) -> "Affine3D":
-        return Affine3D(t, RotationMatrix(r))
-
-    @staticmethod
-    def from_graham_schmidt(
-        neg_x_axis: torch.Tensor,
-        origin: torch.Tensor,
-        xy_plane: torch.Tensor,
-        eps: float = 1e-10,
-    ):
-        # The arguments of this function is for parity with AlphaFold
-        x_axis = origin - neg_x_axis
-        xy_plane = xy_plane - origin
-        return Affine3D(
-            trans=origin, rot=RotationMatrix.from_graham_schmidt(x_axis, xy_plane, eps)
-        )
-
-    @staticmethod
-    def cat(affines: list["Affine3D"], dim: int = 0):
-        if dim < 0:
-            dim = len(affines[0].shape) + dim
-        return Affine3D.from_tensor(torch.cat([x.tensor for x in affines], dim=dim))
-
-
-def _quat_mult(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    """
-    Multiply two quaternions.
-    Usual torch rules for broadcasting apply.
-
-    Args:
-        a: Quaternions as tensor of shape (..., 4), real part first.
-        b: Quaternions as tensor of shape (..., 4), real part first.
-
-    Returns:
-        The product of a and b, a tensor of quaternions shape (..., 4).
-    """
-    aw, ax, ay, az = torch.unbind(a, -1)
-    bw, bx, by, bz = torch.unbind(b, -1)
-    ow = aw * bw - ax * bx - ay * by - az * bz
-    ox = aw * bx + ax * bw + ay * bz - az * by
-    oy = aw * by - ax * bz + ay * bw + az * bx
-    oz = aw * bz + ax * by - ay * bx + az * bw
-    return torch.stack((ow, ox, oy, oz), -1)
-
-
-def _quat_rotation(q: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
-    """
-    Rotates p by quaternion q. Usual torch rules for broadcasting apply.
-
-    Args:
-        q: Quaternions as tensor of shape (..., 4), real part first.
-        p: Points as tensor of shape (..., 3)
-
-    Returns:
-        The rotated version of p, of shape (..., 3)
-    """
-    aw, ax, ay, az = torch.unbind(q, -1)
-    bx, by, bz = torch.unbind(p, -1)
-    # fmt: off
-    ow =         - ax * bx - ay * by - az * bz
-    ox = aw * bx           + ay * bz - az * by
-    oy = aw * by - ax * bz           + az * bx
-    oz = aw * bz + ax * by - ay * bx
-    # fmt: on
-    q_mul_pts = torch.stack((ow, ox, oy, oz), -1)
-    return _quat_mult(q_mul_pts, _quat_invert(q))[..., 1:]
-
-
-def _quat_invert(q: torch.Tensor):
-    return q * torch.tensor([1, -1, -1, -1], device=q.device)
-
-
-def _sqrt_subgradient(x: torch.Tensor) -> torch.Tensor:
-    # Returns torch.sqrt(torch.max(0, x)) but with a zero subgradient where x is 0.
-    ret = torch.zeros_like(x)
-    positive_mask = x > 0
-    ret[positive_mask] = torch.sqrt(x[positive_mask])
-    return ret
-
-
-def _graham_schmidt(x_axis: torch.Tensor, xy_plane: torch.Tensor, eps: float = 1e-12):
-    # A low eps here is necessary for good stability!
-    with fp32_autocast_context(x_axis.device.type):
-        e1 = xy_plane
-
-        denom = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps)
-        x_axis = x_axis / denom
-        dot = (x_axis * e1).sum(dim=-1, keepdim=True)
-        e1 = e1 - x_axis * dot
-        denom = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps)
-        e1 = e1 / denom
-        e2 = torch.cross(x_axis, e1, dim=-1)
-
-        rots = torch.stack([x_axis, e1, e2], dim=-1)
-
-        return rots
-
-
-def build_affine3d_from_coordinates(
-    coords: torch.Tensor,  # (N, CA, C).
-) -> tuple[Affine3D, torch.Tensor]:
-    _MAX_SUPPORTED_DISTANCE = 1e6
-    coord_mask = torch.all(
-        torch.all(torch.isfinite(coords) & (coords < _MAX_SUPPORTED_DISTANCE), dim=-1),
-        dim=-1,
-    )
-
-    def atom3_to_backbone_affine(bb_positions: torch.Tensor) -> Affine3D:
-        N, CA, C = bb_positions.unbind(dim=-2)
-        return Affine3D.from_graham_schmidt(C, CA, N)
-
-    coords = coords.clone().float()
-    coords[~coord_mask] = 0
-
-    # NOTE(thayes): If you have already normalized the coordinates, then
-    # the black hole affine translations will be zeros and the rotations will be
-    # the identity.
-    average_per_n_ca_c = coords.masked_fill(~coord_mask[..., None, None], 0).sum(1) / (
-        coord_mask.sum(-1)[..., None, None] + 1e-8
-    )
-    affine_from_average = atom3_to_backbone_affine(
-        average_per_n_ca_c.float()
-    ).as_matrix()
-
-    B, S, _, _ = coords.shape
-    assert isinstance(B, int)
-    assert isinstance(S, int)
-    affine_rot_mats = affine_from_average.rot.tensor[..., None, :].expand(B, S, 9)
-    affine_trans = affine_from_average.trans[..., None, :].expand(B, S, 3)
-
-    # We use the identity rotation whereever we have no coordinates. This is
-    # important because otherwise the rotation matrices will be all zeros, which
-    # will cause collapse in the distance/direction attention mechanism.
-    identity_rot = RotationMatrix.identity(
-        (B, S), dtype=torch.float32, device=coords.device, requires_grad=False
-    )
-    affine_rot_mats = affine_rot_mats.where(
-        coord_mask.any(-1)[..., None, None], identity_rot.tensor
-    )
-    black_hole_affine = Affine3D(affine_trans, RotationMatrix(affine_rot_mats))
-
-    affine = atom3_to_backbone_affine(coords.float())
-    affine = Affine3D.from_tensor(
-        affine.tensor.where(coord_mask[..., None], black_hole_affine.tensor)
-    )
-
-    return affine, coord_mask
+from __future__ import annotations
+
+import typing as T
+from abc import ABC
+from dataclasses import dataclass
+
+import torch
+from torch.nn import functional as F
+from typing_extensions import Self
+
+from .esmfold2_misc import fp32_autocast_context
+
+
+class Rotation(ABC):
+    @classmethod
+    def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ...
+
+    @classmethod
+    def random(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ...
+
+    def __getitem__(self, idx: T.Any) -> Self: ...
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        # We claim that this should be zero-cost abstraction that returns the raw tensor backing this
+        # object. The raw tensor should always have exactly 1 more dim than self.shape, which should be
+        # implemented using reshaping
+        ...
+
+    @property
+    def shape(self) -> torch.Size:
+        # The "shape" of the rotation, as if it was a torch.tensor object
+        # This means that 1x4 quaternions are treated as size (1,) for example
+        ...
+
+    def as_matrix(self) -> RotationMatrix: ...
+
+    def as_quat(self, normalize: bool = False) -> RotationQuat: ...
+
+    def compose(self, other: Self) -> Self:
+        # To be safe, we force users to explicitly convert between rotation types.
+        ...
+
+    def convert_compose(self, other: Self) -> Self:
+        # This function will automatically convert between types of rotations
+        ...
+
+    def apply(self, p: torch.Tensor) -> torch.Tensor:
+        # rotates points by this rotation object
+        ...
+
+    def invert(self) -> Self: ...
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.tensor.dtype
+
+    @property
+    def device(self) -> torch.device:
+        return self.tensor.device
+
+    @property
+    def requires_grad(self) -> bool:
+        return self.tensor.requires_grad
+
+    @classmethod
+    def _from_tensor(cls, t: torch.Tensor) -> Self:
+        # This function exists to simplify the below functions, esp type signatures
+        # Its implementation is different from Affine3D.from_tensor and does not
+        # autodetect rotation types.
+        return cls(t)  # type: ignore
+
+    def to(self, **kwargs) -> Self:
+        return self._from_tensor(self.tensor.to(**kwargs))
+
+    def detach(self, *args, **kwargs) -> Self:
+        return self._from_tensor(self.tensor.detach(**kwargs))
+
+    def tensor_apply(self, func) -> Self:
+        # Applys a function to the underlying tensor
+        return self._from_tensor(
+            torch.stack([func(x) for x in self.tensor.unbind(dim=-1)], dim=-1)
+        )
+
+
+class RotationMatrix(Rotation):
+    def __init__(self, rots: torch.Tensor):
+        if rots.shape[-1] == 9:
+            rots = rots.unflatten(-1, (3, 3))
+        assert rots.shape[-1] == 3
+        assert rots.shape[-2] == 3
+        # Force full precision
+        rots = rots.to(torch.float32)
+        self._rots = rots
+
+    @classmethod
+    def identity(cls, shape, **tensor_kwargs):
+        rots = torch.eye(3, **tensor_kwargs)
+        rots = rots.view(*[1 for _ in range(len(shape))], 3, 3)
+        rots = rots.expand(*shape, -1, -1)
+        return cls(rots)
+
+    @classmethod
+    def random(cls, shape, **tensor_kwargs):
+        return RotationQuat.random(shape, **tensor_kwargs).as_matrix()
+
+    def __getitem__(self, idx: T.Any) -> RotationMatrix:
+        indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx)
+        return RotationMatrix(self._rots[indices + (slice(None), slice(None))])
+
+    @property
+    def shape(self) -> torch.Size:
+        return self._rots.shape[:-2]
+
+    def as_matrix(self) -> RotationMatrix:
+        return self
+
+    def as_quat(self, normalize: bool = False) -> RotationQuat:
+        m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
+            self._rots.flatten(-2), dim=-1
+        )
+        q_abs = _sqrt_subgradient(
+            torch.stack(
+                [
+                    1.0 + m00 + m11 + m22,
+                    1.0 + m00 - m11 - m22,
+                    1.0 - m00 + m11 - m22,
+                    1.0 - m00 - m11 + m22,
+                ],
+                dim=-1,
+            )
+        )
+        # we produce the desired quaternion multiplied by each of r, i, j, k
+        quat_by_rijk = torch.stack(
+            [
+                x
+                for lst in [
+                    [q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01],
+                    [m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20],
+                    [m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21],
+                    [m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2],
+                ]
+                for x in lst
+            ],
+            dim=-1,
+        ).unflatten(-1, (4, 4))
+
+        # We floor here at 0.1 but the exact level is not important; if q_abs is small,
+        # the candidate won't be picked.
+        flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
+        quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
+
+        # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
+        # forall i; we pick the best-conditioned one (with the largest denominator)
+        # We manually implement one_hot so torch.compile works
+        one_hot = torch.zeros_like(q_abs, dtype=torch.bool)
+        one_hot.scatter_(-1, q_abs.argmax(dim=-1, keepdim=True), True)
+        quat = quat_candidates[one_hot, :].reshape(q_abs.shape)
+        return RotationQuat(quat)
+
+    def compose(self, other: RotationMatrix) -> RotationMatrix:
+        with fp32_autocast_context(self._rots.device.type):
+            return RotationMatrix(self._rots @ other._rots)
+
+    def convert_compose(self, other: Rotation):
+        return self.compose(other.as_matrix())
+
+    def apply(self, p: torch.Tensor) -> torch.Tensor:
+        with fp32_autocast_context(self.device.type):
+            if self._rots.shape[-3] == 1:
+                # This is a slight speedup over einsum for batched rotations
+                return p @ self._rots.transpose(-1, -2).squeeze(-3)
+            else:
+                # einsum way faster than bmm!
+                return torch.einsum("...ij,...j", self._rots, p)
+
+    def invert(self) -> RotationMatrix:
+        return RotationMatrix(self._rots.transpose(-1, -2))
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        return self._rots.flatten(-2)
+
+    def to_3x3(self) -> torch.Tensor:
+        return self._rots
+
+    @staticmethod
+    def from_graham_schmidt(
+        x_axis: torch.Tensor, xy_plane: torch.Tensor, eps: float = 1e-12
+    ) -> RotationMatrix:
+        # A low eps here is necessary for good stability!
+        return RotationMatrix(_graham_schmidt(x_axis, xy_plane, eps))
+
+
+class RotationQuat(Rotation):
+    def __init__(self, quats: torch.Tensor, normalized=False):
+        assert quats.shape[-1] == 4
+        self._normalized = normalized
+        # Force float32 as well
+        if normalized:
+            self._quats = F.normalize(quats.to(torch.float32), dim=-1)
+            self._quats = self._quats.where(self._quats[..., :1] >= 0, -self._quats)
+        else:
+            self._quats = quats.to(torch.float32)
+
+    @classmethod
+    def identity(cls, shape, **tensor_kwargs):
+        q = torch.ones((*shape, 4), **tensor_kwargs)
+        mult = torch.tensor([1, 0, 0, 0], device=q.device)
+        return RotationQuat(q * mult)
+
+    @classmethod
+    def random(cls, shape, **tensor_kwargs):
+        quat = torch.randn((*shape, 4), **tensor_kwargs)
+        return RotationQuat(quat, normalized=True)
+
+    def __getitem__(self, idx: T.Any) -> RotationQuat:
+        indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx)
+        return RotationQuat(self._quats[indices + (slice(None),)])
+
+    @property
+    def shape(self) -> torch.Size:
+        return self._quats.shape[:-1]
+
+    def compose(self, other: RotationQuat) -> RotationQuat:
+        with fp32_autocast_context(self._quats.device.type):
+            return RotationQuat(_quat_mult(self._quats, other._quats))
+
+    def convert_compose(self, other: Rotation):
+        return self.compose(other.as_quat())
+
+    def as_matrix(self) -> RotationMatrix:
+        q = self.normalized().tensor
+        r, i, j, k = torch.unbind(q, -1)
+        two_s = 2.0 / torch.linalg.norm(q, dim=-1)
+
+        o = torch.stack(
+            (
+                1 - two_s * (j * j + k * k),
+                two_s * (i * j - k * r),
+                two_s * (i * k + j * r),
+                two_s * (i * j + k * r),
+                1 - two_s * (i * i + k * k),
+                two_s * (j * k - i * r),
+                two_s * (i * k - j * r),
+                two_s * (j * k + i * r),
+                1 - two_s * (i * i + j * j),
+            ),
+            -1,
+        )
+        return RotationMatrix(o.reshape(q.shape[:-1] + (3, 3)))
+
+    def as_quat(self, normalize: bool = False) -> RotationQuat:
+        return self
+
+    def apply(self, p: torch.Tensor) -> torch.Tensor:
+        return _quat_rotation(self.normalized()._quats, p)
+
+    def invert(self) -> RotationQuat:
+        return RotationQuat(_quat_invert(self._quats))
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        return self._quats
+
+    def normalized(self) -> RotationQuat:
+        return self if self._normalized else RotationQuat(self._quats, normalized=True)
+
+
+@dataclass(frozen=True)
+class Affine3D:
+    trans: torch.Tensor
+    rot: Rotation
+
+    def __post_init__(self):
+        assert self.trans.shape[:-1] == self.rot.shape
+
+    @staticmethod
+    def identity(
+        shape_or_affine: T.Union[tuple[int, ...], "Affine3D"],
+        rotation_type: T.Type[Rotation] = RotationMatrix,
+        **tensor_kwargs,
+    ):
+        # Creates a new identity Affine3D object with a specified shape
+        # or the same shape as another Affine3D object.
+        if isinstance(shape_or_affine, Affine3D):
+            kwargs = {"dtype": shape_or_affine.dtype, "device": shape_or_affine.device}
+            kwargs.update(tensor_kwargs)
+            shape = shape_or_affine.shape
+            rotation_type = type(shape_or_affine.rot)
+        else:
+            kwargs = tensor_kwargs
+            shape = shape_or_affine
+        return Affine3D(
+            torch.zeros((*shape, 3), **kwargs), rotation_type.identity(shape, **kwargs)
+        )
+
+    @staticmethod
+    def random(
+        shape: tuple[int, ...],
+        std: float = 1,
+        rotation_type: T.Type[Rotation] = RotationMatrix,
+        **tensor_kwargs,
+    ) -> "Affine3D":
+        return Affine3D(
+            trans=torch.randn((*shape, 3), **tensor_kwargs).mul(std),
+            rot=rotation_type.random(shape, **tensor_kwargs),
+        )
+
+    def __getitem__(self, idx: T.Any) -> "Affine3D":
+        indices = (idx,) if isinstance(idx, int) or idx is None else tuple(idx)
+        return Affine3D(trans=self.trans[indices + (slice(None),)], rot=self.rot[idx])
+
+    @property
+    def shape(self) -> torch.Size:
+        return self.trans.shape[:-1]
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.trans.dtype
+
+    @property
+    def device(self) -> torch.device:
+        return self.trans.device
+
+    @property
+    def requires_grad(self) -> bool:
+        return self.trans.requires_grad
+
+    def to(self, **kwargs) -> "Affine3D":
+        return Affine3D(self.trans.to(**kwargs), self.rot.to(**kwargs))
+
+    def detach(self, *args, **kwargs) -> "Affine3D":
+        return Affine3D(self.trans.detach(**kwargs), self.rot.detach(**kwargs))
+
+    def tensor_apply(self, func) -> "Affine3D":
+        # Applys a function to the underlying tensor
+        return self.from_tensor(
+            torch.stack([func(x) for x in self.tensor.unbind(dim=-1)], dim=-1)
+        )
+
+    def as_matrix(self):
+        return Affine3D(trans=self.trans, rot=self.rot.as_matrix())
+
+    def as_quat(self, normalize: bool = False):
+        return Affine3D(trans=self.trans, rot=self.rot.as_quat(normalize))
+
+    def compose(self, other: "Affine3D", autoconvert: bool = False):
+        rot = self.rot
+        new_rot = (rot.convert_compose if autoconvert else rot.compose)(other.rot)
+        new_trans = rot.apply(other.trans) + self.trans
+        return Affine3D(trans=new_trans, rot=new_rot)
+
+    def compose_rotation(self, other: Rotation, autoconvert: bool = False):
+        return Affine3D(
+            trans=self.trans,
+            rot=(self.rot.convert_compose if autoconvert else self.rot.compose)(other),
+        )
+
+    def scale(self, v: torch.Tensor | float):
+        return Affine3D(self.trans * v, self.rot)
+
+    def mask(self, mask: torch.Tensor, with_zero=False):
+        # Returns a transform where True positions in mask is identity
+        if with_zero:
+            tensor = self.tensor
+            return Affine3D.from_tensor(
+                torch.zeros_like(tensor).where(mask[..., None], tensor)
+            )
+        else:
+            identity = self.identity(
+                self.shape,
+                rotation_type=type(self.rot),
+                device=self.device,
+                dtype=self.dtype,
+            ).tensor
+            return Affine3D.from_tensor(identity.where(mask[..., None], self.tensor))
+
+    def apply(self, p: torch.Tensor) -> torch.Tensor:
+        return self.rot.apply(p) + self.trans
+
+    def invert(self):
+        inv_rot = self.rot.invert()
+        return Affine3D(trans=-inv_rot.apply(self.trans), rot=inv_rot)
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        return torch.cat([self.rot.tensor, self.trans], dim=-1)
+
+    @staticmethod
+    def from_tensor(t: torch.Tensor) -> "Affine3D":
+        match t.shape[-1]:
+            case 4:
+                # Assume tensor 4x4 for backward compat with alphafold
+                trans = t[..., :3, 3]
+                rot = RotationMatrix(t[..., :3, :3])
+            case 6:
+                # Assume quaternion representation with real part = 1
+                trans = t[..., -3:]
+                rot = RotationQuat(F.pad(t[..., :3], (1, 0), value=1))
+            case 7:
+                trans = t[..., -3:]
+                rot = RotationQuat(t[..., :4])
+            case 12:
+                trans = t[..., -3:]
+                rot = RotationMatrix(t[..., :-3].unflatten(-1, (3, 3)))
+            case _:
+                raise RuntimeError(
+                    f"Cannot detect rotation fromat from {t.shape[-1] -3}-d flat vector"
+                )
+        return Affine3D(trans, rot)
+
+    @staticmethod
+    def from_tensor_pair(t: torch.Tensor, r: torch.Tensor) -> "Affine3D":
+        return Affine3D(t, RotationMatrix(r))
+
+    @staticmethod
+    def from_graham_schmidt(
+        neg_x_axis: torch.Tensor,
+        origin: torch.Tensor,
+        xy_plane: torch.Tensor,
+        eps: float = 1e-10,
+    ):
+        # The arguments of this function is for parity with AlphaFold
+        x_axis = origin - neg_x_axis
+        xy_plane = xy_plane - origin
+        return Affine3D(
+            trans=origin, rot=RotationMatrix.from_graham_schmidt(x_axis, xy_plane, eps)
+        )
+
+    @staticmethod
+    def cat(affines: list["Affine3D"], dim: int = 0):
+        if dim < 0:
+            dim = len(affines[0].shape) + dim
+        return Affine3D.from_tensor(torch.cat([x.tensor for x in affines], dim=dim))
+
+
+def _quat_mult(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """
+    Multiply two quaternions.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions shape (..., 4).
+    """
+    aw, ax, ay, az = torch.unbind(a, -1)
+    bw, bx, by, bz = torch.unbind(b, -1)
+    ow = aw * bw - ax * bx - ay * by - az * bz
+    ox = aw * bx + ax * bw + ay * bz - az * by
+    oy = aw * by - ax * bz + ay * bw + az * bx
+    oz = aw * bz + ax * by - ay * bx + az * bw
+    return torch.stack((ow, ox, oy, oz), -1)
+
+
+def _quat_rotation(q: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
+    """
+    Rotates p by quaternion q. Usual torch rules for broadcasting apply.
+
+    Args:
+        q: Quaternions as tensor of shape (..., 4), real part first.
+        p: Points as tensor of shape (..., 3)
+
+    Returns:
+        The rotated version of p, of shape (..., 3)
+    """
+    aw, ax, ay, az = torch.unbind(q, -1)
+    bx, by, bz = torch.unbind(p, -1)
+    # fmt: off
+    ow =         - ax * bx - ay * by - az * bz
+    ox = aw * bx           + ay * bz - az * by
+    oy = aw * by - ax * bz           + az * bx
+    oz = aw * bz + ax * by - ay * bx
+    # fmt: on
+    q_mul_pts = torch.stack((ow, ox, oy, oz), -1)
+    return _quat_mult(q_mul_pts, _quat_invert(q))[..., 1:]
+
+
+def _quat_invert(q: torch.Tensor):
+    return q * torch.tensor([1, -1, -1, -1], device=q.device)
+
+
+def _sqrt_subgradient(x: torch.Tensor) -> torch.Tensor:
+    # Returns torch.sqrt(torch.max(0, x)) but with a zero subgradient where x is 0.
+    ret = torch.zeros_like(x)
+    positive_mask = x > 0
+    ret[positive_mask] = torch.sqrt(x[positive_mask])
+    return ret
+
+
+def _graham_schmidt(x_axis: torch.Tensor, xy_plane: torch.Tensor, eps: float = 1e-12):
+    # A low eps here is necessary for good stability!
+    with fp32_autocast_context(x_axis.device.type):
+        e1 = xy_plane
+
+        denom = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps)
+        x_axis = x_axis / denom
+        dot = (x_axis * e1).sum(dim=-1, keepdim=True)
+        e1 = e1 - x_axis * dot
+        denom = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps)
+        e1 = e1 / denom
+        e2 = torch.cross(x_axis, e1, dim=-1)
+
+        rots = torch.stack([x_axis, e1, e2], dim=-1)
+
+        return rots
+
+
+def build_affine3d_from_coordinates(
+    coords: torch.Tensor,  # (N, CA, C).
+) -> tuple[Affine3D, torch.Tensor]:
+    _MAX_SUPPORTED_DISTANCE = 1e6
+    coord_mask = torch.all(
+        torch.all(torch.isfinite(coords) & (coords < _MAX_SUPPORTED_DISTANCE), dim=-1),
+        dim=-1,
+    )
+
+    def atom3_to_backbone_affine(bb_positions: torch.Tensor) -> Affine3D:
+        N, CA, C = bb_positions.unbind(dim=-2)
+        return Affine3D.from_graham_schmidt(C, CA, N)
+
+    coords = coords.clone().float()
+    coords[~coord_mask] = 0
+
+    # NOTE(thayes): If you have already normalized the coordinates, then
+    # the black hole affine translations will be zeros and the rotations will be
+    # the identity.
+    average_per_n_ca_c = coords.masked_fill(~coord_mask[..., None, None], 0).sum(1) / (
+        coord_mask.sum(-1)[..., None, None] + 1e-8
+    )
+    affine_from_average = atom3_to_backbone_affine(
+        average_per_n_ca_c.float()
+    ).as_matrix()
+
+    B, S, _, _ = coords.shape
+    assert isinstance(B, int)
+    assert isinstance(S, int)
+    affine_rot_mats = affine_from_average.rot.tensor[..., None, :].expand(B, S, 9)
+    affine_trans = affine_from_average.trans[..., None, :].expand(B, S, 3)
+
+    # We use the identity rotation whereever we have no coordinates. This is
+    # important because otherwise the rotation matrices will be all zeros, which
+    # will cause collapse in the distance/direction attention mechanism.
+    identity_rot = RotationMatrix.identity(
+        (B, S), dtype=torch.float32, device=coords.device, requires_grad=False
+    )
+    affine_rot_mats = affine_rot_mats.where(
+        coord_mask.any(-1)[..., None, None], identity_rot.tensor
+    )
+    black_hole_affine = Affine3D(affine_trans, RotationMatrix(affine_rot_mats))
+
+    affine = atom3_to_backbone_affine(coords.float())
+    affine = Affine3D.from_tensor(
+        affine.tensor.where(coord_mask[..., None], black_hole_affine.tensor)
+    )
+
+    return affine, coord_mask

esmfold2_aligner.py

@@ -1,101 +1,101 @@
-from __future__ import annotations
-
-from dataclasses import Field, replace
-from typing import Any, ClassVar, Protocol, TypeVar
-
-import numpy as np
-import torch
-
-from .esmfold2_protein_structure import compute_affine_and_rmsd
-
-
-class Alignable(Protocol):
-    # Trick to detect whether an object is a dataclass
-    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
-
-    @property
-    def atom37_positions(self) -> np.ndarray:  # type: ignore
-        pass
-
-    @property
-    def atom37_mask(self) -> np.ndarray:  # type: ignore
-        pass
-
-    def __len__(self) -> int: ...
-
-
-T = TypeVar("T", bound=Alignable)
-
-
-class Aligner:
-    def __init__(
-        self,
-        mobile: Alignable,
-        target: Alignable,
-        only_use_backbone: bool = False,
-        use_reflection: bool = False,
-    ):
-        """
-        Aligns a mobile protein chain against a target protein chain.
-
-        Args:
-            mobile (ProteinChain): Protein chain to be aligned.
-            target (ProteinChain): Protein chain target.
-            only_use_backbone (bool): Whether to only use backbone atoms.
-            use_reflection (bool): Whether to align to target reflection.
-        """
-        # Check proteins must have same number of residues
-        assert len(mobile) == len(target)
-
-        # Determine overlapping atoms
-        joint_atom37_mask = mobile.atom37_mask.astype(bool) & target.atom37_mask.astype(
-            bool
-        )
-
-        # Backbone atoms are first sites in atom37 representation
-        if only_use_backbone:
-            joint_atom37_mask[:, 3:] = False
-
-        # Extract matching atom positions and convert to batched tensors
-        mobile_atom_tensor = (
-            torch.from_numpy(mobile.atom37_positions).type(torch.double).unsqueeze(0)
-        )
-        target_atom_tensor = (
-            torch.from_numpy(target.atom37_positions).type(torch.double).unsqueeze(0)
-        )
-        joint_atom37_mask = (
-            torch.from_numpy(joint_atom37_mask).type(torch.bool).unsqueeze(0)
-        )
-
-        # If using reflection flip target
-        if use_reflection:
-            target_atom_tensor = -target_atom_tensor
-
-        # Compute alignment and rmsd
-        affine3D, rmsd = compute_affine_and_rmsd(
-            mobile_atom_tensor, target_atom_tensor, atom_exists_mask=joint_atom37_mask
-        )
-        self._affine3D = affine3D
-        self._rmsd = rmsd.item()
-
-    @property
-    def rmsd(self):
-        return self._rmsd
-
-    def apply(self, mobile: T) -> T:
-        """Apply alignment to a protein chain"""
-        # Extract atom positions and convert to batched tensors
-        mobile_atom_tensor = (
-            torch.from_numpy(mobile.atom37_positions[mobile.atom37_mask])
-            .type(torch.float32)
-            .unsqueeze(0)
-        )
-
-        # Transform atom arrays
-        aligned_atom_tensor = self._affine3D.apply(mobile_atom_tensor).squeeze(0)
-
-        # Rebuild atom37 positions
-        aligned_atom37_positions = np.full_like(mobile.atom37_positions, np.nan)
-        aligned_atom37_positions[mobile.atom37_mask] = aligned_atom_tensor
-
-        return replace(mobile, atom37_positions=aligned_atom37_positions)
+from __future__ import annotations
+
+from dataclasses import Field, replace
+from typing import Any, ClassVar, Protocol, TypeVar
+
+import numpy as np
+import torch
+
+from .esmfold2_protein_structure import compute_affine_and_rmsd
+
+
+class Alignable(Protocol):
+    # Trick to detect whether an object is a dataclass
+    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
+
+    @property
+    def atom37_positions(self) -> np.ndarray:  # type: ignore
+        pass
+
+    @property
+    def atom37_mask(self) -> np.ndarray:  # type: ignore
+        pass
+
+    def __len__(self) -> int: ...
+
+
+T = TypeVar("T", bound=Alignable)
+
+
+class Aligner:
+    def __init__(
+        self,
+        mobile: Alignable,
+        target: Alignable,
+        only_use_backbone: bool = False,
+        use_reflection: bool = False,
+    ):
+        """
+        Aligns a mobile protein chain against a target protein chain.
+
+        Args:
+            mobile (ProteinChain): Protein chain to be aligned.
+            target (ProteinChain): Protein chain target.
+            only_use_backbone (bool): Whether to only use backbone atoms.
+            use_reflection (bool): Whether to align to target reflection.
+        """
+        # Check proteins must have same number of residues
+        assert len(mobile) == len(target)
+
+        # Determine overlapping atoms
+        joint_atom37_mask = mobile.atom37_mask.astype(bool) & target.atom37_mask.astype(
+            bool
+        )
+
+        # Backbone atoms are first sites in atom37 representation
+        if only_use_backbone:
+            joint_atom37_mask[:, 3:] = False
+
+        # Extract matching atom positions and convert to batched tensors
+        mobile_atom_tensor = (
+            torch.from_numpy(mobile.atom37_positions).type(torch.double).unsqueeze(0)
+        )
+        target_atom_tensor = (
+            torch.from_numpy(target.atom37_positions).type(torch.double).unsqueeze(0)
+        )
+        joint_atom37_mask = (
+            torch.from_numpy(joint_atom37_mask).type(torch.bool).unsqueeze(0)
+        )
+
+        # If using reflection flip target
+        if use_reflection:
+            target_atom_tensor = -target_atom_tensor
+
+        # Compute alignment and rmsd
+        affine3D, rmsd = compute_affine_and_rmsd(
+            mobile_atom_tensor, target_atom_tensor, atom_exists_mask=joint_atom37_mask
+        )
+        self._affine3D = affine3D
+        self._rmsd = rmsd.item()
+
+    @property
+    def rmsd(self):
+        return self._rmsd
+
+    def apply(self, mobile: T) -> T:
+        """Apply alignment to a protein chain"""
+        # Extract atom positions and convert to batched tensors
+        mobile_atom_tensor = (
+            torch.from_numpy(mobile.atom37_positions[mobile.atom37_mask])
+            .type(torch.float32)
+            .unsqueeze(0)
+        )
+
+        # Transform atom arrays
+        aligned_atom_tensor = self._affine3D.apply(mobile_atom_tensor).squeeze(0)
+
+        # Rebuild atom37 positions
+        aligned_atom37_positions = np.full_like(mobile.atom37_positions, np.nan)
+        aligned_atom37_positions[mobile.atom37_mask] = aligned_atom_tensor
+
+        return replace(mobile, atom37_positions=aligned_atom37_positions)

esmfold2_atom_indexer.py

@@ -1,15 +1,15 @@
-import numpy as np
-
-from .esmfold2_protein_structure import index_by_atom_name
-
-
-class AtomIndexer:
-    def __init__(self, structure, property: str, dim: int):
-        self.structure = structure
-        self.property = property
-        self.dim = dim
-
-    def __getitem__(self, atom_names: str | list[str]) -> np.ndarray:
-        return index_by_atom_name(
-            getattr(self.structure, self.property), atom_names, self.dim
-        )
+import numpy as np
+
+from .esmfold2_protein_structure import index_by_atom_name
+
+
+class AtomIndexer:
+    def __init__(self, structure, property: str, dim: int):
+        self.structure = structure
+        self.property = property
+        self.dim = dim
+
+    def __getitem__(self, atom_names: str | list[str]) -> np.ndarray:
+        return index_by_atom_name(
+            getattr(self.structure, self.property), atom_names, self.dim
+        )

esmfold2_conformers.py

@@ -1,291 +1,291 @@
-"""CCD conformer loading utilities.
-
-Loads idealized conformer coordinates from a CCD pickle file containing RDKit molecules.
-Conformer priority follows AF3 Section 2.8: Computed > Ideal > first available.
-"""
-
-from __future__ import annotations
-
-import os
-import pickle
-from pathlib import Path
-
-import numpy as np
-from huggingface_hub import hf_hub_download
-
-from .esmfold2_constants import RES_TYPE_TO_CCD
-
-if os.environ.get("ESMCFOLD_CCD_PATH"):
-    CCD_PICKLE_PATH = Path(os.environ["ESMCFOLD_CCD_PATH"])
-else:
-    CCD_PICKLE_PATH = None
-
-
-# Lazily loaded CCD dictionary
-_CCD_MOLECULES: dict | None = None
-
-# Caches
-_CCD_CONFORMERS: dict[str, dict[str, np.ndarray]] = {}
-_CCD_ATOM_CACHE: dict[str, list[tuple[str, str, int]]] = {}
-_CCD_BONDS_CACHE: dict[str, list[tuple[str, str]]] = {}
-_CCD_LEAVING_ATOMS_CACHE: dict[str, set[str]] = {}
-_IDEALIZED_POS_CACHE: dict[tuple[int, str], np.ndarray | None] = {}
-_LIGAND_IDEALIZED_POS_CACHE: dict[tuple[str, str], np.ndarray | None] = {}
-
-
-def load_ccd(cache_dir: Path | str | None = None) -> dict:
-    """Load CCD molecules from pickle file, downloading if needed.
-
-    Args:
-        cache_dir: Directory to cache the downloaded CCD pickle.
-            If None, uses CCD_PICKLE_PATH env var or downloads to ~/.cache/esmcfold/.
-    """
-    global _CCD_MOLECULES
-    if _CCD_MOLECULES is not None:
-        return _CCD_MOLECULES
-
-    # Determine pickle path
-    if CCD_PICKLE_PATH is not None and CCD_PICKLE_PATH.exists():
-        pkl_path = CCD_PICKLE_PATH
-    elif cache_dir is not None:
-        cache_dir = Path(cache_dir)
-        cache_dir.mkdir(parents=True, exist_ok=True)
-        pkl_path = cache_dir / "ccd.pkl"
-    else:
-        try:
-            pkl_path = Path(
-                hf_hub_download(repo_id="biohub/ESMFold2", filename="ccd.pkl")
-            )
-        except Exception as e:
-            raise FileNotFoundError(
-                f"Failed to download CCD pickle file from Hugging Face repository: {e}"
-            )
-
-    if not pkl_path.exists():
-        raise FileNotFoundError(
-            f"CCD pickle file not found: {pkl_path}. Please set the ESMCFOLD_CCD_PATH environment variable to the path of a valid CCD pickle file or download the file from the Hugging Face repository."
-        )
-
-    print(f"Loading CCD dictionary from {pkl_path}")
-    with open(pkl_path, "rb") as f:
-        _CCD_MOLECULES = pickle.load(f)
-
-    if _CCD_MOLECULES is None:
-        _CCD_MOLECULES = {}
-
-    return _CCD_MOLECULES
-
-
-def _get_ccd_molecules() -> dict:
-    """Get CCD molecules, loading lazily on first call."""
-    global _CCD_MOLECULES
-    if _CCD_MOLECULES is None:
-        return load_ccd()
-    return _CCD_MOLECULES
-
-
-def _get_ccd_mol_with_significant_h(comp_id: str):
-    """Get CCD molecule with only chemically significant hydrogens.
-
-    Returns (mol, conformer) tuple or (None, None) if not available.
-    """
-    ccd = _get_ccd_molecules()
-    if comp_id not in ccd:
-        return None, None
-
-    mol = ccd[comp_id]
-    if mol.GetNumConformers() == 0:
-        return None, None
-
-    # Find the "Computed" conformer (RDKit ETKDGv3), fall back to "Ideal"
-    conf_idx = 0
-    for i, c in enumerate(mol.GetConformers()):
-        props = c.GetPropsAsDict()
-        if props.get("name") == "Computed":
-            conf_idx = i
-            break
-    else:
-        for i, c in enumerate(mol.GetConformers()):
-            props = c.GetPropsAsDict()
-            if props.get("name") == "Ideal":
-                conf_idx = i
-                break
-
-    from rdkit import Chem
-
-    mol_no_h = Chem.RemoveHs(mol, sanitize=False)
-
-    if mol_no_h.GetNumConformers() == 0:
-        return None, None
-
-    return mol_no_h, mol_no_h.GetConformer(
-        min(conf_idx, mol_no_h.GetNumConformers() - 1)
-    )
-
-
-def get_ccd_conformer(comp_id: str) -> dict[str, np.ndarray] | None:
-    """Get idealized conformer as dict of atom_name -> position [3].
-
-    Conformer priority: Computed > Ideal > first available.
-    """
-    if comp_id in _CCD_CONFORMERS:
-        cached = _CCD_CONFORMERS[comp_id]
-        return cached if cached else None
-
-    mol, conf = _get_ccd_mol_with_significant_h(comp_id)
-    if mol is None or conf is None:
-        _CCD_CONFORMERS[comp_id] = {}
-        return None
-
-    conformer: dict[str, np.ndarray] = {}
-    for atom in mol.GetAtoms():
-        props = atom.GetPropsAsDict()
-        atom_name = props.get("name")
-        if not isinstance(atom_name, str) or not atom_name:
-            continue
-        idx = atom.GetIdx()
-        pos = conf.GetAtomPosition(idx)
-        conformer[atom_name] = np.array([pos.x, pos.y, pos.z], dtype=np.float32)
-
-    _CCD_CONFORMERS[comp_id] = conformer
-    return conformer if conformer else None
-
-
-def get_idealized_atom_pos(res_type: int, atom_name: str) -> np.ndarray | None:
-    """Get idealized position for a standard residue atom.
-
-    Uses res_type index to look up CCD component, then returns position.
-    Returns None if not found.
-    """
-    cache_key = (res_type, atom_name)
-    if cache_key in _IDEALIZED_POS_CACHE:
-        return _IDEALIZED_POS_CACHE[cache_key]
-
-    comp_id = RES_TYPE_TO_CCD.get(res_type)
-    if comp_id:
-        ccd_conformer = get_ccd_conformer(comp_id)
-        if ccd_conformer and atom_name in ccd_conformer:
-            pos = ccd_conformer[atom_name]
-            _IDEALIZED_POS_CACHE[cache_key] = pos
-            return pos
-
-    _IDEALIZED_POS_CACHE[cache_key] = None
-    return None
-
-
-def get_ligand_idealized_atom_pos(res_name: str, atom_name: str) -> np.ndarray | None:
-    """Get idealized position for a ligand/modified residue atom.
-
-    Returns None if not found.
-    """
-    cache_key = (res_name, atom_name)
-    if cache_key in _LIGAND_IDEALIZED_POS_CACHE:
-        return _LIGAND_IDEALIZED_POS_CACHE[cache_key]
-
-    ccd_conformer = get_ccd_conformer(res_name)
-    if ccd_conformer and atom_name in ccd_conformer:
-        pos = ccd_conformer[atom_name]
-        _LIGAND_IDEALIZED_POS_CACHE[cache_key] = pos
-        return pos
-
-    _LIGAND_IDEALIZED_POS_CACHE[cache_key] = None
-    return None
-
-
-def get_ligand_ccd_atoms_with_charges(
-    comp_id: str,
-) -> list[tuple[str, str, int]] | None:
-    """Get list of (atom_name, element, charge) for a CCD component.
-
-    Uses RDKit RemoveHs(sanitize=False) to keep chemically significant hydrogens.
-    Returns None if CCD data not available.
-    """
-    if comp_id in _CCD_ATOM_CACHE:
-        cached = _CCD_ATOM_CACHE[comp_id]
-        return cached if cached else None
-
-    mol, _ = _get_ccd_mol_with_significant_h(comp_id)
-    if mol is None:
-        _CCD_ATOM_CACHE[comp_id] = []
-        return None
-
-    atoms: list[tuple[str, str, int]] = []
-    for atom in mol.GetAtoms():
-        props = atom.GetPropsAsDict()
-        atom_name = props.get("name")
-        if not isinstance(atom_name, str) or not atom_name:
-            continue
-        element = atom.GetSymbol()
-        charge = atom.GetFormalCharge()
-        atoms.append((atom_name, element, charge))
-
-    _CCD_ATOM_CACHE[comp_id] = atoms
-    return atoms if atoms else None
-
-
-def get_ligand_ccd_bonds(comp_id: str) -> list[tuple[str, str]] | None:
-    """Get list of (atom1_name, atom2_name) bonds for a CCD component.
-
-    Returns None if CCD data not available.
-    """
-    if comp_id in _CCD_BONDS_CACHE:
-        cached = _CCD_BONDS_CACHE[comp_id]
-        return cached if cached else None
-
-    mol, _ = _get_ccd_mol_with_significant_h(comp_id)
-    if mol is None:
-        _CCD_BONDS_CACHE[comp_id] = []
-        return None
-
-    # Get included atom names
-    included_atoms = set()
-    for atom in mol.GetAtoms():
-        props = atom.GetPropsAsDict()
-        atom_name = props.get("name")
-        if isinstance(atom_name, str) and atom_name:
-            included_atoms.add(atom_name)
-
-    bonds: list[tuple[str, str]] = []
-    for bond in mol.GetBonds():
-        a1 = bond.GetBeginAtom()
-        a2 = bond.GetEndAtom()
-        n1 = a1.GetPropsAsDict().get("name")
-        n2 = a2.GetPropsAsDict().get("name")
-        if (
-            isinstance(n1, str)
-            and isinstance(n2, str)
-            and n1
-            and n2
-            and n1 in included_atoms
-            and n2 in included_atoms
-        ):
-            bonds.append((n1, n2))
-
-    _CCD_BONDS_CACHE[comp_id] = bonds
-    return bonds if bonds else None
-
-
-def get_ccd_leaving_atoms(comp_id: str) -> set[str]:
-    """Get set of atom names marked as leaving atoms in CCD.
-
-    Leaving atoms are removed during polymerization (e.g., OP3 in nucleotides).
-    """
-    if comp_id in _CCD_LEAVING_ATOMS_CACHE:
-        return _CCD_LEAVING_ATOMS_CACHE[comp_id]
-
-    ccd = _get_ccd_molecules()
-    if comp_id not in ccd:
-        _CCD_LEAVING_ATOMS_CACHE[comp_id] = set()
-        return set()
-
-    mol = ccd[comp_id]
-    leaving_atoms = set()
-    for atom in mol.GetAtoms():
-        if atom.HasProp("leaving_atom"):
-            if atom.GetProp("leaving_atom") == "1":
-                name = atom.GetProp("name") if atom.HasProp("name") else ""
-                if name:
-                    leaving_atoms.add(name)
-
-    _CCD_LEAVING_ATOMS_CACHE[comp_id] = leaving_atoms
-    return leaving_atoms
+"""CCD conformer loading utilities.
+
+Loads idealized conformer coordinates from a CCD pickle file containing RDKit molecules.
+Conformer priority follows AF3 Section 2.8: Computed > Ideal > first available.
+"""
+
+from __future__ import annotations
+
+import os
+import pickle
+from pathlib import Path
+
+import numpy as np
+from huggingface_hub import hf_hub_download
+
+from .esmfold2_constants import RES_TYPE_TO_CCD
+
+if os.environ.get("ESMCFOLD_CCD_PATH"):
+    CCD_PICKLE_PATH = Path(os.environ["ESMCFOLD_CCD_PATH"])
+else:
+    CCD_PICKLE_PATH = None
+
+
+# Lazily loaded CCD dictionary
+_CCD_MOLECULES: dict | None = None
+
+# Caches
+_CCD_CONFORMERS: dict[str, dict[str, np.ndarray]] = {}
+_CCD_ATOM_CACHE: dict[str, list[tuple[str, str, int]]] = {}
+_CCD_BONDS_CACHE: dict[str, list[tuple[str, str]]] = {}
+_CCD_LEAVING_ATOMS_CACHE: dict[str, set[str]] = {}
+_IDEALIZED_POS_CACHE: dict[tuple[int, str], np.ndarray | None] = {}
+_LIGAND_IDEALIZED_POS_CACHE: dict[tuple[str, str], np.ndarray | None] = {}
+
+
+def load_ccd(cache_dir: Path | str | None = None) -> dict:
+    """Load CCD molecules from pickle file, downloading if needed.
+
+    Args:
+        cache_dir: Directory to cache the downloaded CCD pickle.
+            If None, uses CCD_PICKLE_PATH env var or downloads to ~/.cache/esmcfold/.
+    """
+    global _CCD_MOLECULES
+    if _CCD_MOLECULES is not None:
+        return _CCD_MOLECULES
+
+    # Determine pickle path
+    if CCD_PICKLE_PATH is not None and CCD_PICKLE_PATH.exists():
+        pkl_path = CCD_PICKLE_PATH
+    elif cache_dir is not None:
+        cache_dir = Path(cache_dir)
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        pkl_path = cache_dir / "ccd.pkl"
+    else:
+        try:
+            pkl_path = Path(
+                hf_hub_download(repo_id="biohub/ESMFold2", filename="ccd.pkl")
+            )
+        except Exception as e:
+            raise FileNotFoundError(
+                f"Failed to download CCD pickle file from Hugging Face repository: {e}"
+            )
+
+    if not pkl_path.exists():
+        raise FileNotFoundError(
+            f"CCD pickle file not found: {pkl_path}. Please set the ESMCFOLD_CCD_PATH environment variable to the path of a valid CCD pickle file or download the file from the Hugging Face repository."
+        )
+
+    print(f"Loading CCD dictionary from {pkl_path}")
+    with open(pkl_path, "rb") as f:
+        _CCD_MOLECULES = pickle.load(f)
+
+    if _CCD_MOLECULES is None:
+        _CCD_MOLECULES = {}
+
+    return _CCD_MOLECULES
+
+
+def _get_ccd_molecules() -> dict:
+    """Get CCD molecules, loading lazily on first call."""
+    global _CCD_MOLECULES
+    if _CCD_MOLECULES is None:
+        return load_ccd()
+    return _CCD_MOLECULES
+
+
+def _get_ccd_mol_with_significant_h(comp_id: str):
+    """Get CCD molecule with only chemically significant hydrogens.
+
+    Returns (mol, conformer) tuple or (None, None) if not available.
+    """
+    ccd = _get_ccd_molecules()
+    if comp_id not in ccd:
+        return None, None
+
+    mol = ccd[comp_id]
+    if mol.GetNumConformers() == 0:
+        return None, None
+
+    # Find the "Computed" conformer (RDKit ETKDGv3), fall back to "Ideal"
+    conf_idx = 0
+    for i, c in enumerate(mol.GetConformers()):
+        props = c.GetPropsAsDict()
+        if props.get("name") == "Computed":
+            conf_idx = i
+            break
+    else:
+        for i, c in enumerate(mol.GetConformers()):
+            props = c.GetPropsAsDict()
+            if props.get("name") == "Ideal":
+                conf_idx = i
+                break
+
+    from rdkit import Chem
+
+    mol_no_h = Chem.RemoveHs(mol, sanitize=False)
+
+    if mol_no_h.GetNumConformers() == 0:
+        return None, None
+
+    return mol_no_h, mol_no_h.GetConformer(
+        min(conf_idx, mol_no_h.GetNumConformers() - 1)
+    )
+
+
+def get_ccd_conformer(comp_id: str) -> dict[str, np.ndarray] | None:
+    """Get idealized conformer as dict of atom_name -> position [3].
+
+    Conformer priority: Computed > Ideal > first available.
+    """
+    if comp_id in _CCD_CONFORMERS:
+        cached = _CCD_CONFORMERS[comp_id]
+        return cached if cached else None
+
+    mol, conf = _get_ccd_mol_with_significant_h(comp_id)
+    if mol is None or conf is None:
+        _CCD_CONFORMERS[comp_id] = {}
+        return None
+
+    conformer: dict[str, np.ndarray] = {}
+    for atom in mol.GetAtoms():
+        props = atom.GetPropsAsDict()
+        atom_name = props.get("name")
+        if not isinstance(atom_name, str) or not atom_name:
+            continue
+        idx = atom.GetIdx()
+        pos = conf.GetAtomPosition(idx)
+        conformer[atom_name] = np.array([pos.x, pos.y, pos.z], dtype=np.float32)
+
+    _CCD_CONFORMERS[comp_id] = conformer
+    return conformer if conformer else None
+
+
+def get_idealized_atom_pos(res_type: int, atom_name: str) -> np.ndarray | None:
+    """Get idealized position for a standard residue atom.
+
+    Uses res_type index to look up CCD component, then returns position.
+    Returns None if not found.
+    """
+    cache_key = (res_type, atom_name)
+    if cache_key in _IDEALIZED_POS_CACHE:
+        return _IDEALIZED_POS_CACHE[cache_key]
+
+    comp_id = RES_TYPE_TO_CCD.get(res_type)
+    if comp_id:
+        ccd_conformer = get_ccd_conformer(comp_id)
+        if ccd_conformer and atom_name in ccd_conformer:
+            pos = ccd_conformer[atom_name]
+            _IDEALIZED_POS_CACHE[cache_key] = pos
+            return pos
+
+    _IDEALIZED_POS_CACHE[cache_key] = None
+    return None
+
+
+def get_ligand_idealized_atom_pos(res_name: str, atom_name: str) -> np.ndarray | None:
+    """Get idealized position for a ligand/modified residue atom.
+
+    Returns None if not found.
+    """
+    cache_key = (res_name, atom_name)
+    if cache_key in _LIGAND_IDEALIZED_POS_CACHE:
+        return _LIGAND_IDEALIZED_POS_CACHE[cache_key]
+
+    ccd_conformer = get_ccd_conformer(res_name)
+    if ccd_conformer and atom_name in ccd_conformer:
+        pos = ccd_conformer[atom_name]
+        _LIGAND_IDEALIZED_POS_CACHE[cache_key] = pos
+        return pos
+
+    _LIGAND_IDEALIZED_POS_CACHE[cache_key] = None
+    return None
+
+
+def get_ligand_ccd_atoms_with_charges(
+    comp_id: str,
+) -> list[tuple[str, str, int]] | None:
+    """Get list of (atom_name, element, charge) for a CCD component.
+
+    Uses RDKit RemoveHs(sanitize=False) to keep chemically significant hydrogens.
+    Returns None if CCD data not available.
+    """
+    if comp_id in _CCD_ATOM_CACHE:
+        cached = _CCD_ATOM_CACHE[comp_id]
+        return cached if cached else None
+
+    mol, _ = _get_ccd_mol_with_significant_h(comp_id)
+    if mol is None:
+        _CCD_ATOM_CACHE[comp_id] = []
+        return None
+
+    atoms: list[tuple[str, str, int]] = []
+    for atom in mol.GetAtoms():
+        props = atom.GetPropsAsDict()
+        atom_name = props.get("name")
+        if not isinstance(atom_name, str) or not atom_name:
+            continue
+        element = atom.GetSymbol()
+        charge = atom.GetFormalCharge()
+        atoms.append((atom_name, element, charge))
+
+    _CCD_ATOM_CACHE[comp_id] = atoms
+    return atoms if atoms else None
+
+
+def get_ligand_ccd_bonds(comp_id: str) -> list[tuple[str, str]] | None:
+    """Get list of (atom1_name, atom2_name) bonds for a CCD component.
+
+    Returns None if CCD data not available.
+    """
+    if comp_id in _CCD_BONDS_CACHE:
+        cached = _CCD_BONDS_CACHE[comp_id]
+        return cached if cached else None
+
+    mol, _ = _get_ccd_mol_with_significant_h(comp_id)
+    if mol is None:
+        _CCD_BONDS_CACHE[comp_id] = []
+        return None
+
+    # Get included atom names
+    included_atoms = set()
+    for atom in mol.GetAtoms():
+        props = atom.GetPropsAsDict()
+        atom_name = props.get("name")
+        if isinstance(atom_name, str) and atom_name:
+            included_atoms.add(atom_name)
+
+    bonds: list[tuple[str, str]] = []
+    for bond in mol.GetBonds():
+        a1 = bond.GetBeginAtom()
+        a2 = bond.GetEndAtom()
+        n1 = a1.GetPropsAsDict().get("name")
+        n2 = a2.GetPropsAsDict().get("name")
+        if (
+            isinstance(n1, str)
+            and isinstance(n2, str)
+            and n1
+            and n2
+            and n1 in included_atoms
+            and n2 in included_atoms
+        ):
+            bonds.append((n1, n2))
+
+    _CCD_BONDS_CACHE[comp_id] = bonds
+    return bonds if bonds else None
+
+
+def get_ccd_leaving_atoms(comp_id: str) -> set[str]:
+    """Get set of atom names marked as leaving atoms in CCD.
+
+    Leaving atoms are removed during polymerization (e.g., OP3 in nucleotides).
+    """
+    if comp_id in _CCD_LEAVING_ATOMS_CACHE:
+        return _CCD_LEAVING_ATOMS_CACHE[comp_id]
+
+    ccd = _get_ccd_molecules()
+    if comp_id not in ccd:
+        _CCD_LEAVING_ATOMS_CACHE[comp_id] = set()
+        return set()
+
+    mol = ccd[comp_id]
+    leaving_atoms = set()
+    for atom in mol.GetAtoms():
+        if atom.HasProp("leaving_atom"):
+            if atom.GetProp("leaving_atom") == "1":
+                name = atom.GetProp("name") if atom.HasProp("name") else ""
+                if name:
+                    leaving_atoms.add(name)
+
+    _CCD_LEAVING_ATOMS_CACHE[comp_id] = leaving_atoms
+    return leaving_atoms

esmfold2_constants.py

@@ -1,562 +1,562 @@
-"""Constants for the ESMFold2 input pipeline.
-
-Includes molecule types, residue types, vocabularies, atom lists, and element data.
-"""
-
-# =============================================================================
-# Molecule types
-# =============================================================================
-
-MOL_TYPE_PROTEIN = 0
-MOL_TYPE_DNA = 1
-MOL_TYPE_RNA = 2
-MOL_TYPE_NONPOLYMER = 3
-
-# =============================================================================
-# Residue type indices
-# =============================================================================
-
-# Standard amino acids (indices 2-21), MSE mapped to MET
-PROTEIN_RESIDUE_TO_RES_TYPE = {
-    "ALA": 2,
-    "ARG": 3,
-    "ASN": 4,
-    "ASP": 5,
-    "CYS": 6,
-    "GLN": 7,
-    "GLU": 8,
-    "GLY": 9,
-    "HIS": 10,
-    "ILE": 11,
-    "LEU": 12,
-    "LYS": 13,
-    "MET": 14,
-    "PHE": 15,
-    "PRO": 16,
-    "SER": 17,
-    "THR": 18,
-    "TRP": 19,
-    "TYR": 20,
-    "VAL": 21,
-    "MSE": 14,  # Selenomethionine -> MET
-}
-PROTEIN_UNK_RES_TYPE = 22
-
-# RNA nucleotides (indices 23-26, unknown=27)
-RNA_RESIDUE_TO_RES_TYPE = {"A": 23, "G": 24, "C": 25, "U": 26}
-RNA_UNK_RES_TYPE = 27
-
-# DNA nucleotides (indices 28-31, unknown=32)
-DNA_RESIDUE_TO_RES_TYPE = {"DA": 28, "DG": 29, "DC": 30, "DT": 31}
-DNA_UNK_RES_TYPE = 32
-
-GAP_RES_TYPE = 32
-
-# =============================================================================
-# Vocabularies
-# =============================================================================
-
-# 3-letter to 1-letter codes for proteins
-PROTEIN_3TO1 = {
-    "ALA": "A",
-    "ARG": "R",
-    "ASN": "N",
-    "ASP": "D",
-    "CYS": "C",
-    "GLN": "Q",
-    "GLU": "E",
-    "GLY": "G",
-    "HIS": "H",
-    "ILE": "I",
-    "LEU": "L",
-    "LYS": "K",
-    "MET": "M",
-    "PHE": "F",
-    "PRO": "P",
-    "SER": "S",
-    "THR": "T",
-    "TRP": "W",
-    "TYR": "Y",
-    "VAL": "V",
-    "MSE": "M",
-}
-
-# 1-letter to 3-letter codes
-PROTEIN_1TO3 = {v: k for k, v in PROTEIN_3TO1.items() if k != "MSE"}
-PROTEIN_1TO3["X"] = "UNK"
-
-# DNA 1-letter to CCD code
-DNA_1TO3 = {"A": "DA", "T": "DT", "C": "DC", "G": "DG"}
-
-# RNA 1-letter to CCD code
-RNA_1TO3 = {"A": "A", "U": "U", "C": "C", "G": "G"}
-
-# ESM-2 input_ids vocabulary for proteins
-ESM_PROTEIN_VOCAB = {
-    "L": 4,
-    "A": 5,
-    "G": 6,
-    "V": 7,
-    "S": 8,
-    "E": 9,
-    "R": 10,
-    "T": 11,
-    "I": 12,
-    "D": 13,
-    "P": 14,
-    "K": 15,
-    "Q": 16,
-    "N": 17,
-    "F": 18,
-    "Y": 19,
-    "M": 20,
-    "H": 21,
-    "W": 22,
-    "C": 23,
-    "X": 3,  # Unknown
-}
-
-# For DNA/RNA/ligands
-DNA_RNA_LIGAND_INPUT_ID = 24
-
-# MSA tokens
-MSA_PAD_TOKEN_ID = 0
-MSA_GAP_TOKEN_ID = 1  # Gap/insertion token for MSA
-
-# res_type int -> CCD component ID (for conformer lookup)
-RES_TYPE_TO_CCD = {
-    # Proteins (2-22)
-    2: "ALA",
-    3: "ARG",
-    4: "ASN",
-    5: "ASP",
-    6: "CYS",
-    7: "GLN",
-    8: "GLU",
-    9: "GLY",
-    10: "HIS",
-    11: "ILE",
-    12: "LEU",
-    13: "LYS",
-    14: "MET",
-    15: "PHE",
-    16: "PRO",
-    17: "SER",
-    18: "THR",
-    19: "TRP",
-    20: "TYR",
-    21: "VAL",
-    22: "UNK",
-    # RNA (23-27)
-    23: "A",
-    24: "G",
-    25: "C",
-    26: "U",
-    27: "N",
-    # DNA (28-32)
-    28: "DA",
-    29: "DG",
-    30: "DC",
-    31: "DT",
-    32: "DN",
-}
-
-# =============================================================================
-# Charged atoms at physiological pH
-# =============================================================================
-
-CHARGED_ATOMS: dict[tuple[str, str], int] = {
-    ("LYS", "NZ"): 1,
-    ("ARG", "NH2"): 1,
-    ("HIS", "ND1"): 1,
-    ("PO4", "O2"): -1,
-    ("PO4", "O3"): -1,
-    ("PO4", "O4"): -1,
-    ("SO4", "O3"): -1,
-    ("SO4", "O4"): -1,
-    ("MG", "MG"): 2,
-    ("ZN", "ZN"): 2,
-    ("CA", "CA"): 2,
-    ("FE2", "FE"): 2,
-    ("MN", "MN"): 2,
-    ("CO", "CO"): 2,
-    ("NCO", "CO"): 3,
-    ("CU", "CU"): 2,
-    ("NI", "NI"): 2,
-    ("K", "K"): 1,
-    ("NA", "NA"): 1,
-    ("CD", "CD"): 2,
-    ("CL", "CL"): -1,
-    ("ACT", "OXT"): -1,
-    ("NAD", "O2N"): -1,
-    ("NAD", "N1N"): 1,
-    ("NAP", "O2N"): -1,
-    ("NAP", "N1N"): 1,
-    ("IMD", "N3"): 1,
-    ("SAM", "SD"): 1,
-    ("FE", "FE"): 3,
-    ("A1BH3", "N3"): 1,
-}
-
-# =============================================================================
-# Element atomic numbers (Z=1 to 92)
-# =============================================================================
-
-ELEMENT_TO_ATOMIC_NUM = {
-    "H": 1,
-    "LI": 3,
-    "BE": 4,
-    "B": 5,
-    "C": 6,
-    "N": 7,
-    "O": 8,
-    "F": 9,
-    "NE": 10,
-    "NA": 11,
-    "MG": 12,
-    "AL": 13,
-    "SI": 14,
-    "P": 15,
-    "S": 16,
-    "CL": 17,
-    "AR": 18,
-    "K": 19,
-    "CA": 20,
-    "SC": 21,
-    "TI": 22,
-    "V": 23,
-    "CR": 24,
-    "MN": 25,
-    "FE": 26,
-    "CO": 27,
-    "NI": 28,
-    "CU": 29,
-    "ZN": 30,
-    "GA": 31,
-    "GE": 32,
-    "AS": 33,
-    "SE": 34,
-    "BR": 35,
-    "KR": 36,
-    "RB": 37,
-    "SR": 38,
-    "Y": 39,
-    "ZR": 40,
-    "NB": 41,
-    "MO": 42,
-    "TC": 43,
-    "RU": 44,
-    "RH": 45,
-    "PD": 46,
-    "AG": 47,
-    "CD": 48,
-    "IN": 49,
-    "SN": 50,
-    "SB": 51,
-    "TE": 52,
-    "I": 53,
-    "XE": 54,
-    "CS": 55,
-    "BA": 56,
-    "LA": 57,
-    "CE": 58,
-    "PR": 59,
-    "ND": 60,
-    "PM": 61,
-    "SM": 62,
-    "EU": 63,
-    "GD": 64,
-    "TB": 65,
-    "DY": 66,
-    "HO": 67,
-    "ER": 68,
-    "TM": 69,
-    "YB": 70,
-    "LU": 71,
-    "HF": 72,
-    "TA": 73,
-    "W": 74,
-    "RE": 75,
-    "OS": 76,
-    "IR": 77,
-    "PT": 78,
-    "AU": 79,
-    "HG": 80,
-    "TL": 81,
-    "PB": 82,
-    "BI": 83,
-    "PO": 84,
-    "AT": 85,
-    "RN": 86,
-    "FR": 87,
-    "RA": 88,
-    "AC": 89,
-    "TH": 90,
-    "PA": 91,
-    "U": 92,
-}
-
-# Inverse mapping: atomic number → element symbol
-ELEMENT_NUMBER_TO_SYMBOL = {v: k for k, v in ELEMENT_TO_ATOMIC_NUM.items()}
-
-# =============================================================================
-# Standard heavy atoms per residue type
-# =============================================================================
-
-PROTEIN_HEAVY_ATOMS = {
-    "ALA": ["N", "CA", "C", "O", "CB"],
-    "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"],
-    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"],
-    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"],
-    "CYS": ["N", "CA", "C", "O", "CB", "SG"],
-    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"],
-    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"],
-    "GLY": ["N", "CA", "C", "O"],
-    "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"],
-    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"],
-    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"],
-    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"],
-    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
-    "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"],
-    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"],
-    "SER": ["N", "CA", "C", "O", "CB", "OG"],
-    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"],
-    "TRP": [
-        "N",
-        "CA",
-        "C",
-        "O",
-        "CB",
-        "CG",
-        "CD1",
-        "CD2",
-        "NE1",
-        "CE2",
-        "CE3",
-        "CZ2",
-        "CZ3",
-        "CH2",
-    ],
-    "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"],
-    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"],
-    "MSE": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
-    "UNK": ["N", "CA", "C", "O"],
-}
-
-DNA_HEAVY_ATOMS = {
-    "DA": [
-        "P",
-        "OP1",
-        "OP2",
-        "O5'",
-        "C5'",
-        "C4'",
-        "O4'",
-        "C3'",
-        "O3'",
-        "C2'",
-        "C1'",
-        "N9",
-        "C8",
-        "N7",
-        "C5",
-        "C6",
-        "N6",
-        "N1",
-        "C2",
-        "N3",
-        "C4",
-    ],
-    "DG": [
-        "P",
-        "OP1",
-        "OP2",
-        "O5'",
-        "C5'",
-        "C4'",
-        "O4'",
-        "C3'",
-        "O3'",
-        "C2'",
-        "C1'",
-        "N9",
-        "C8",
-        "N7",
-        "C5",
-        "C6",
-        "O6",
-        "N1",
-        "C2",
-        "N2",
-        "N3",
-        "C4",
-    ],
-    "DC": [
-        "P",
-        "OP1",
-        "OP2",
-        "O5'",
-        "C5'",
-        "C4'",
-        "O4'",
-        "C3'",
-        "O3'",
-        "C2'",
-        "C1'",
-        "N1",
-        "C2",
-        "O2",
-        "N3",
-        "C4",
-        "N4",
-        "C5",
-        "C6",
-    ],
-    "DT": [
-        "P",
-        "OP1",
-        "OP2",
-        "O5'",
-        "C5'",
-        "C4'",
-        "O4'",
-        "C3'",
-        "O3'",
-        "C2'",
-        "C1'",
-        "N1",
-        "C2",
-        "O2",
-        "N3",
-        "C4",
-        "O4",
-        "C5",
-        "C7",
-        "C6",
-    ],
-}
-
-RNA_HEAVY_ATOMS = {
-    "A": [
-        "P",
-        "OP1",
-        "OP2",
-        "O5'",
-        "C5'",
-        "C4'",
-        "O4'",
-        "C3'",
-        "O3'",
-        "C2'",
-        "O2'",
-        "C1'",
-        "N9",
-        "C8",
-        "N7",
-        "C5",
-        "C6",
-        "N6",
-        "N1",
-        "C2",
-        "N3",
-        "C4",
-    ],
-    "G": [
-        "P",
-        "OP1",
-        "OP2",
-        "O5'",
-        "C5'",
-        "C4'",
-        "O4'",
-        "C3'",
-        "O3'",
-        "C2'",
-        "O2'",
-        "C1'",
-        "N9",
-        "C8",
-        "N7",
-        "C5",
-        "C6",
-        "O6",
-        "N1",
-        "C2",
-        "N2",
-        "N3",
-        "C4",
-    ],
-    "C": [
-        "P",
-        "OP1",
-        "OP2",
-        "O5'",
-        "C5'",
-        "C4'",
-        "O4'",
-        "C3'",
-        "O3'",
-        "C2'",
-        "O2'",
-        "C1'",
-        "N1",
-        "C2",
-        "O2",
-        "N3",
-        "C4",
-        "N4",
-        "C5",
-        "C6",
-    ],
-    "U": [
-        "P",
-        "OP1",
-        "OP2",
-        "O5'",
-        "C5'",
-        "C4'",
-        "O4'",
-        "C3'",
-        "O3'",
-        "C2'",
-        "O2'",
-        "C1'",
-        "N1",
-        "C2",
-        "O2",
-        "N3",
-        "C4",
-        "O4",
-        "C5",
-        "C6",
-    ],
-}
-
-# Unknown nucleotide backbone atoms
-DNA_BACKBONE_ATOMS = [
-    "P",
-    "OP1",
-    "OP2",
-    "O5'",
-    "C5'",
-    "C4'",
-    "O4'",
-    "C3'",
-    "O3'",
-    "C2'",
-    "C1'",
-]
-RNA_BACKBONE_ATOMS = [
-    "P",
-    "OP1",
-    "OP2",
-    "O5'",
-    "C5'",
-    "C4'",
-    "O4'",
-    "C3'",
-    "O3'",
-    "C2'",
-    "O2'",
-    "C1'",
-]
+"""Constants for the ESMFold2 input pipeline.
+
+Includes molecule types, residue types, vocabularies, atom lists, and element data.
+"""
+
+# =============================================================================
+# Molecule types
+# =============================================================================
+
+MOL_TYPE_PROTEIN = 0
+MOL_TYPE_DNA = 1
+MOL_TYPE_RNA = 2
+MOL_TYPE_NONPOLYMER = 3
+
+# =============================================================================
+# Residue type indices
+# =============================================================================
+
+# Standard amino acids (indices 2-21), MSE mapped to MET
+PROTEIN_RESIDUE_TO_RES_TYPE = {
+    "ALA": 2,
+    "ARG": 3,
+    "ASN": 4,
+    "ASP": 5,
+    "CYS": 6,
+    "GLN": 7,
+    "GLU": 8,
+    "GLY": 9,
+    "HIS": 10,
+    "ILE": 11,
+    "LEU": 12,
+    "LYS": 13,
+    "MET": 14,
+    "PHE": 15,
+    "PRO": 16,
+    "SER": 17,
+    "THR": 18,
+    "TRP": 19,
+    "TYR": 20,
+    "VAL": 21,
+    "MSE": 14,  # Selenomethionine -> MET
+}
+PROTEIN_UNK_RES_TYPE = 22
+
+# RNA nucleotides (indices 23-26, unknown=27)
+RNA_RESIDUE_TO_RES_TYPE = {"A": 23, "G": 24, "C": 25, "U": 26}
+RNA_UNK_RES_TYPE = 27
+
+# DNA nucleotides (indices 28-31, unknown=32)
+DNA_RESIDUE_TO_RES_TYPE = {"DA": 28, "DG": 29, "DC": 30, "DT": 31}
+DNA_UNK_RES_TYPE = 32
+
+GAP_RES_TYPE = 32
+
+# =============================================================================
+# Vocabularies
+# =============================================================================
+
+# 3-letter to 1-letter codes for proteins
+PROTEIN_3TO1 = {
+    "ALA": "A",
+    "ARG": "R",
+    "ASN": "N",
+    "ASP": "D",
+    "CYS": "C",
+    "GLN": "Q",
+    "GLU": "E",
+    "GLY": "G",
+    "HIS": "H",
+    "ILE": "I",
+    "LEU": "L",
+    "LYS": "K",
+    "MET": "M",
+    "PHE": "F",
+    "PRO": "P",
+    "SER": "S",
+    "THR": "T",
+    "TRP": "W",
+    "TYR": "Y",
+    "VAL": "V",
+    "MSE": "M",
+}
+
+# 1-letter to 3-letter codes
+PROTEIN_1TO3 = {v: k for k, v in PROTEIN_3TO1.items() if k != "MSE"}
+PROTEIN_1TO3["X"] = "UNK"
+
+# DNA 1-letter to CCD code
+DNA_1TO3 = {"A": "DA", "T": "DT", "C": "DC", "G": "DG"}
+
+# RNA 1-letter to CCD code
+RNA_1TO3 = {"A": "A", "U": "U", "C": "C", "G": "G"}
+
+# ESM-2 input_ids vocabulary for proteins
+ESM_PROTEIN_VOCAB = {
+    "L": 4,
+    "A": 5,
+    "G": 6,
+    "V": 7,
+    "S": 8,
+    "E": 9,
+    "R": 10,
+    "T": 11,
+    "I": 12,
+    "D": 13,
+    "P": 14,
+    "K": 15,
+    "Q": 16,
+    "N": 17,
+    "F": 18,
+    "Y": 19,
+    "M": 20,
+    "H": 21,
+    "W": 22,
+    "C": 23,
+    "X": 3,  # Unknown
+}
+
+# For DNA/RNA/ligands
+DNA_RNA_LIGAND_INPUT_ID = 24
+
+# MSA tokens
+MSA_PAD_TOKEN_ID = 0
+MSA_GAP_TOKEN_ID = 1  # Gap/insertion token for MSA
+
+# res_type int -> CCD component ID (for conformer lookup)
+RES_TYPE_TO_CCD = {
+    # Proteins (2-22)
+    2: "ALA",
+    3: "ARG",
+    4: "ASN",
+    5: "ASP",
+    6: "CYS",
+    7: "GLN",
+    8: "GLU",
+    9: "GLY",
+    10: "HIS",
+    11: "ILE",
+    12: "LEU",
+    13: "LYS",
+    14: "MET",
+    15: "PHE",
+    16: "PRO",
+    17: "SER",
+    18: "THR",
+    19: "TRP",
+    20: "TYR",
+    21: "VAL",
+    22: "UNK",
+    # RNA (23-27)
+    23: "A",
+    24: "G",
+    25: "C",
+    26: "U",
+    27: "N",
+    # DNA (28-32)
+    28: "DA",
+    29: "DG",
+    30: "DC",
+    31: "DT",
+    32: "DN",
+}
+
+# =============================================================================
+# Charged atoms at physiological pH
+# =============================================================================
+
+CHARGED_ATOMS: dict[tuple[str, str], int] = {
+    ("LYS", "NZ"): 1,
+    ("ARG", "NH2"): 1,
+    ("HIS", "ND1"): 1,
+    ("PO4", "O2"): -1,
+    ("PO4", "O3"): -1,
+    ("PO4", "O4"): -1,
+    ("SO4", "O3"): -1,
+    ("SO4", "O4"): -1,
+    ("MG", "MG"): 2,
+    ("ZN", "ZN"): 2,
+    ("CA", "CA"): 2,
+    ("FE2", "FE"): 2,
+    ("MN", "MN"): 2,
+    ("CO", "CO"): 2,
+    ("NCO", "CO"): 3,
+    ("CU", "CU"): 2,
+    ("NI", "NI"): 2,
+    ("K", "K"): 1,
+    ("NA", "NA"): 1,
+    ("CD", "CD"): 2,
+    ("CL", "CL"): -1,
+    ("ACT", "OXT"): -1,
+    ("NAD", "O2N"): -1,
+    ("NAD", "N1N"): 1,
+    ("NAP", "O2N"): -1,
+    ("NAP", "N1N"): 1,
+    ("IMD", "N3"): 1,
+    ("SAM", "SD"): 1,
+    ("FE", "FE"): 3,
+    ("A1BH3", "N3"): 1,
+}
+
+# =============================================================================
+# Element atomic numbers (Z=1 to 92)
+# =============================================================================
+
+ELEMENT_TO_ATOMIC_NUM = {
+    "H": 1,
+    "LI": 3,
+    "BE": 4,
+    "B": 5,
+    "C": 6,
+    "N": 7,
+    "O": 8,
+    "F": 9,
+    "NE": 10,
+    "NA": 11,
+    "MG": 12,
+    "AL": 13,
+    "SI": 14,
+    "P": 15,
+    "S": 16,
+    "CL": 17,
+    "AR": 18,
+    "K": 19,
+    "CA": 20,
+    "SC": 21,
+    "TI": 22,
+    "V": 23,
+    "CR": 24,
+    "MN": 25,
+    "FE": 26,
+    "CO": 27,
+    "NI": 28,
+    "CU": 29,
+    "ZN": 30,
+    "GA": 31,
+    "GE": 32,
+    "AS": 33,
+    "SE": 34,
+    "BR": 35,
+    "KR": 36,
+    "RB": 37,
+    "SR": 38,
+    "Y": 39,
+    "ZR": 40,
+    "NB": 41,
+    "MO": 42,
+    "TC": 43,
+    "RU": 44,
+    "RH": 45,
+    "PD": 46,
+    "AG": 47,
+    "CD": 48,
+    "IN": 49,
+    "SN": 50,
+    "SB": 51,
+    "TE": 52,
+    "I": 53,
+    "XE": 54,
+    "CS": 55,
+    "BA": 56,
+    "LA": 57,
+    "CE": 58,
+    "PR": 59,
+    "ND": 60,
+    "PM": 61,
+    "SM": 62,
+    "EU": 63,
+    "GD": 64,
+    "TB": 65,
+    "DY": 66,
+    "HO": 67,
+    "ER": 68,
+    "TM": 69,
+    "YB": 70,
+    "LU": 71,
+    "HF": 72,
+    "TA": 73,
+    "W": 74,
+    "RE": 75,
+    "OS": 76,
+    "IR": 77,
+    "PT": 78,
+    "AU": 79,
+    "HG": 80,
+    "TL": 81,
+    "PB": 82,
+    "BI": 83,
+    "PO": 84,
+    "AT": 85,
+    "RN": 86,
+    "FR": 87,
+    "RA": 88,
+    "AC": 89,
+    "TH": 90,
+    "PA": 91,
+    "U": 92,
+}
+
+# Inverse mapping: atomic number → element symbol
+ELEMENT_NUMBER_TO_SYMBOL = {v: k for k, v in ELEMENT_TO_ATOMIC_NUM.items()}
+
+# =============================================================================
+# Standard heavy atoms per residue type
+# =============================================================================
+
+PROTEIN_HEAVY_ATOMS = {
+    "ALA": ["N", "CA", "C", "O", "CB"],
+    "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"],
+    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"],
+    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"],
+    "CYS": ["N", "CA", "C", "O", "CB", "SG"],
+    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"],
+    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"],
+    "GLY": ["N", "CA", "C", "O"],
+    "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"],
+    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"],
+    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"],
+    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"],
+    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
+    "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"],
+    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"],
+    "SER": ["N", "CA", "C", "O", "CB", "OG"],
+    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"],
+    "TRP": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "NE1",
+        "CE2",
+        "CE3",
+        "CZ2",
+        "CZ3",
+        "CH2",
+    ],
+    "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"],
+    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"],
+    "MSE": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
+    "UNK": ["N", "CA", "C", "O"],
+}
+
+DNA_HEAVY_ATOMS = {
+    "DA": [
+        "P",
+        "OP1",
+        "OP2",
+        "O5'",
+        "C5'",
+        "C4'",
+        "O4'",
+        "C3'",
+        "O3'",
+        "C2'",
+        "C1'",
+        "N9",
+        "C8",
+        "N7",
+        "C5",
+        "C6",
+        "N6",
+        "N1",
+        "C2",
+        "N3",
+        "C4",
+    ],
+    "DG": [
+        "P",
+        "OP1",
+        "OP2",
+        "O5'",
+        "C5'",
+        "C4'",
+        "O4'",
+        "C3'",
+        "O3'",
+        "C2'",
+        "C1'",
+        "N9",
+        "C8",
+        "N7",
+        "C5",
+        "C6",
+        "O6",
+        "N1",
+        "C2",
+        "N2",
+        "N3",
+        "C4",
+    ],
+    "DC": [
+        "P",
+        "OP1",
+        "OP2",
+        "O5'",
+        "C5'",
+        "C4'",
+        "O4'",
+        "C3'",
+        "O3'",
+        "C2'",
+        "C1'",
+        "N1",
+        "C2",
+        "O2",
+        "N3",
+        "C4",
+        "N4",
+        "C5",
+        "C6",
+    ],
+    "DT": [
+        "P",
+        "OP1",
+        "OP2",
+        "O5'",
+        "C5'",
+        "C4'",
+        "O4'",
+        "C3'",
+        "O3'",
+        "C2'",
+        "C1'",
+        "N1",
+        "C2",
+        "O2",
+        "N3",
+        "C4",
+        "O4",
+        "C5",
+        "C7",
+        "C6",
+    ],
+}
+
+RNA_HEAVY_ATOMS = {
+    "A": [
+        "P",
+        "OP1",
+        "OP2",
+        "O5'",
+        "C5'",
+        "C4'",
+        "O4'",
+        "C3'",
+        "O3'",
+        "C2'",
+        "O2'",
+        "C1'",
+        "N9",
+        "C8",
+        "N7",
+        "C5",
+        "C6",
+        "N6",
+        "N1",
+        "C2",
+        "N3",
+        "C4",
+    ],
+    "G": [
+        "P",
+        "OP1",
+        "OP2",
+        "O5'",
+        "C5'",
+        "C4'",
+        "O4'",
+        "C3'",
+        "O3'",
+        "C2'",
+        "O2'",
+        "C1'",
+        "N9",
+        "C8",
+        "N7",
+        "C5",
+        "C6",
+        "O6",
+        "N1",
+        "C2",
+        "N2",
+        "N3",
+        "C4",
+    ],
+    "C": [
+        "P",
+        "OP1",
+        "OP2",
+        "O5'",
+        "C5'",
+        "C4'",
+        "O4'",
+        "C3'",
+        "O3'",
+        "C2'",
+        "O2'",
+        "C1'",
+        "N1",
+        "C2",
+        "O2",
+        "N3",
+        "C4",
+        "N4",
+        "C5",
+        "C6",
+    ],
+    "U": [
+        "P",
+        "OP1",
+        "OP2",
+        "O5'",
+        "C5'",
+        "C4'",
+        "O4'",
+        "C3'",
+        "O3'",
+        "C2'",
+        "O2'",
+        "C1'",
+        "N1",
+        "C2",
+        "O2",
+        "N3",
+        "C4",
+        "O4",
+        "C5",
+        "C6",
+    ],
+}
+
+# Unknown nucleotide backbone atoms
+DNA_BACKBONE_ATOMS = [
+    "P",
+    "OP1",
+    "OP2",
+    "O5'",
+    "C5'",
+    "C4'",
+    "O4'",
+    "C3'",
+    "O3'",
+    "C2'",
+    "C1'",
+]
+RNA_BACKBONE_ATOMS = [
+    "P",
+    "OP1",
+    "OP2",
+    "O5'",
+    "C5'",
+    "C4'",
+    "O4'",
+    "C3'",
+    "O3'",
+    "C2'",
+    "O2'",
+    "C1'",
+]

esmfold2_constants_esm3.py

@@ -1,137 +1,137 @@
-import os
-from functools import cache
-from pathlib import Path
-
-from huggingface_hub import snapshot_download
-
-SEQUENCE_BOS_TOKEN = 0
-SEQUENCE_PAD_TOKEN = 1
-SEQUENCE_EOS_TOKEN = 2
-SEQUENCE_CHAINBREAK_TOKEN = 31
-SEQUENCE_MASK_TOKEN = 32
-
-VQVAE_CODEBOOK_SIZE = 4096
-VQVAE_SPECIAL_TOKENS = {
-    "MASK": VQVAE_CODEBOOK_SIZE,
-    "EOS": VQVAE_CODEBOOK_SIZE + 1,
-    "BOS": VQVAE_CODEBOOK_SIZE + 2,
-    "PAD": VQVAE_CODEBOOK_SIZE + 3,
-    "CHAINBREAK": VQVAE_CODEBOOK_SIZE + 4,
-}
-VQVAE_DIRECTION_LOSS_BINS = 16
-VQVAE_PAE_BINS = 64
-VQVAE_MAX_PAE_BIN = 31.0
-VQVAE_PLDDT_BINS = 50
-
-STRUCTURE_MASK_TOKEN = VQVAE_SPECIAL_TOKENS["MASK"]
-STRUCTURE_BOS_TOKEN = VQVAE_SPECIAL_TOKENS["BOS"]
-STRUCTURE_EOS_TOKEN = VQVAE_SPECIAL_TOKENS["EOS"]
-STRUCTURE_PAD_TOKEN = VQVAE_SPECIAL_TOKENS["PAD"]
-STRUCTURE_CHAINBREAK_TOKEN = VQVAE_SPECIAL_TOKENS["CHAINBREAK"]
-STRUCTURE_UNDEFINED_TOKEN = 955
-
-SASA_PAD_TOKEN = 0
-
-SS8_PAD_TOKEN = 0
-
-INTERPRO_PAD_TOKEN = 0
-
-RESIDUE_PAD_TOKEN = 0
-
-CHAIN_BREAK_STR = "|"
-
-SEQUENCE_BOS_STR = "<cls>"
-SEQUENCE_EOS_STR = "<eos>"
-
-MASK_STR_SHORT = "_"
-SEQUENCE_MASK_STR = "<mask>"
-SASA_MASK_STR = "<unk>"
-SS8_MASK_STR = "<unk>"
-
-# fmt: off
-SEQUENCE_VOCAB = [
-    "<cls>", "<pad>", "<eos>", "<unk>",
-    "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K",
-    "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z",
-    "O", ".", "-", "|",
-    "<mask>",
-]
-# fmt: on
-
-SEQUENCE_STANDARD_AA_MIN_TOKEN = 4  # L
-SEQUENCE_STANDARD_AA_MAX_TOKEN = 24  # X (exclusive)
-
-SSE_8CLASS_VOCAB = "GHITEBSC"
-SSE_3CLASS_VOCAB = "HEC"
-SSE_8CLASS_TO_3CLASS_MAP = {
-    "G": "H",
-    "H": "H",
-    "I": "H",
-    "T": "C",
-    "E": "E",
-    "B": "E",
-    "S": "C",
-    "C": "C",
-}
-
-SASA_DISCRETIZATION_BOUNDARIES = [
-    0.8,
-    4.0,
-    9.6,
-    16.4,
-    24.5,
-    32.9,
-    42.0,
-    51.5,
-    61.2,
-    70.9,
-    81.6,
-    93.3,
-    107.2,
-    125.4,
-    151.4,
-]
-
-MAX_RESIDUE_ANNOTATIONS = 16
-
-
-TFIDF_VECTOR_SIZE = 58641
-
-FUNCTION_TOKENS_DEPTH = 8
-
-
-@staticmethod
-@cache
-def data_root(model: str):
-    if "INFRA_PROVIDER" in os.environ:
-        return Path("")
-    # Try to download from huggingface if it doesn't exist
-    if model.startswith("esm3"):
-        path = Path(snapshot_download(repo_id="biohub/esm3-sm-open-v1"))
-    elif model.startswith("esmc-300"):
-        path = Path(snapshot_download(repo_id="biohub/esmc-300m-2024-12"))
-    elif model.startswith("esmc-600"):
-        path = Path(snapshot_download(repo_id="biohub/esmc-600m-2024-12"))
-    elif model.startswith("esmc-6b"):
-        path = Path(snapshot_download(repo_id="biohub/esmc-6b-2024-12"))
-    else:
-        raise ValueError(f"{model=} is an invalid model name.")
-    return path
-
-
-IN_REPO_DATA_FOLDER = Path(__file__).parents[2] / "data"
-
-INTERPRO_ENTRY = IN_REPO_DATA_FOLDER / "entry_list_safety_29026.list"
-INTERPRO_HIERARCHY = IN_REPO_DATA_FOLDER / "ParentChildTreeFile.txt"
-INTERPRO2GO = IN_REPO_DATA_FOLDER / "ParentChildTreeFile.txt"
-INTERPRO_2ID = "data/tag_dict_4_safety_filtered.json"
-
-LSH_TABLE_PATHS = {"8bit": "data/hyperplanes_8bit_58641.npz"}
-
-KEYWORDS_VOCABULARY = (
-    IN_REPO_DATA_FOLDER / "keyword_vocabulary_safety_filtered_58641.txt"
-)
-KEYWORDS_IDF = IN_REPO_DATA_FOLDER / "keyword_idf_safety_filtered_58641.npy"
-
-RESID_CSV = "data/uniref90_and_mgnify90_residue_annotations_gt_1k_proteins.csv"
-INTERPRO2KEYWORDS = IN_REPO_DATA_FOLDER / "interpro_29026_to_keywords_58641.csv"
+import os
+from functools import cache
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+SEQUENCE_BOS_TOKEN = 0
+SEQUENCE_PAD_TOKEN = 1
+SEQUENCE_EOS_TOKEN = 2
+SEQUENCE_CHAINBREAK_TOKEN = 31
+SEQUENCE_MASK_TOKEN = 32
+
+VQVAE_CODEBOOK_SIZE = 4096
+VQVAE_SPECIAL_TOKENS = {
+    "MASK": VQVAE_CODEBOOK_SIZE,
+    "EOS": VQVAE_CODEBOOK_SIZE + 1,
+    "BOS": VQVAE_CODEBOOK_SIZE + 2,
+    "PAD": VQVAE_CODEBOOK_SIZE + 3,
+    "CHAINBREAK": VQVAE_CODEBOOK_SIZE + 4,
+}
+VQVAE_DIRECTION_LOSS_BINS = 16
+VQVAE_PAE_BINS = 64
+VQVAE_MAX_PAE_BIN = 31.0
+VQVAE_PLDDT_BINS = 50
+
+STRUCTURE_MASK_TOKEN = VQVAE_SPECIAL_TOKENS["MASK"]
+STRUCTURE_BOS_TOKEN = VQVAE_SPECIAL_TOKENS["BOS"]
+STRUCTURE_EOS_TOKEN = VQVAE_SPECIAL_TOKENS["EOS"]
+STRUCTURE_PAD_TOKEN = VQVAE_SPECIAL_TOKENS["PAD"]
+STRUCTURE_CHAINBREAK_TOKEN = VQVAE_SPECIAL_TOKENS["CHAINBREAK"]
+STRUCTURE_UNDEFINED_TOKEN = 955
+
+SASA_PAD_TOKEN = 0
+
+SS8_PAD_TOKEN = 0
+
+INTERPRO_PAD_TOKEN = 0
+
+RESIDUE_PAD_TOKEN = 0
+
+CHAIN_BREAK_STR = "|"
+
+SEQUENCE_BOS_STR = "<cls>"
+SEQUENCE_EOS_STR = "<eos>"
+
+MASK_STR_SHORT = "_"
+SEQUENCE_MASK_STR = "<mask>"
+SASA_MASK_STR = "<unk>"
+SS8_MASK_STR = "<unk>"
+
+# fmt: off
+SEQUENCE_VOCAB = [
+    "<cls>", "<pad>", "<eos>", "<unk>",
+    "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K",
+    "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z",
+    "O", ".", "-", "|",
+    "<mask>",
+]
+# fmt: on
+
+SEQUENCE_STANDARD_AA_MIN_TOKEN = 4  # L
+SEQUENCE_STANDARD_AA_MAX_TOKEN = 24  # X (exclusive)
+
+SSE_8CLASS_VOCAB = "GHITEBSC"
+SSE_3CLASS_VOCAB = "HEC"
+SSE_8CLASS_TO_3CLASS_MAP = {
+    "G": "H",
+    "H": "H",
+    "I": "H",
+    "T": "C",
+    "E": "E",
+    "B": "E",
+    "S": "C",
+    "C": "C",
+}
+
+SASA_DISCRETIZATION_BOUNDARIES = [
+    0.8,
+    4.0,
+    9.6,
+    16.4,
+    24.5,
+    32.9,
+    42.0,
+    51.5,
+    61.2,
+    70.9,
+    81.6,
+    93.3,
+    107.2,
+    125.4,
+    151.4,
+]
+
+MAX_RESIDUE_ANNOTATIONS = 16
+
+
+TFIDF_VECTOR_SIZE = 58641
+
+FUNCTION_TOKENS_DEPTH = 8
+
+
+@staticmethod
+@cache
+def data_root(model: str):
+    if "INFRA_PROVIDER" in os.environ:
+        return Path("")
+    # Try to download from huggingface if it doesn't exist
+    if model.startswith("esm3"):
+        path = Path(snapshot_download(repo_id="biohub/esm3-sm-open-v1"))
+    elif model.startswith("esmc-300"):
+        path = Path(snapshot_download(repo_id="biohub/esmc-300m-2024-12"))
+    elif model.startswith("esmc-600"):
+        path = Path(snapshot_download(repo_id="biohub/esmc-600m-2024-12"))
+    elif model.startswith("esmc-6b"):
+        path = Path(snapshot_download(repo_id="biohub/esmc-6b-2024-12"))
+    else:
+        raise ValueError(f"{model=} is an invalid model name.")
+    return path
+
+
+IN_REPO_DATA_FOLDER = Path(__file__).parents[2] / "data"
+
+INTERPRO_ENTRY = IN_REPO_DATA_FOLDER / "entry_list_safety_29026.list"
+INTERPRO_HIERARCHY = IN_REPO_DATA_FOLDER / "ParentChildTreeFile.txt"
+INTERPRO2GO = IN_REPO_DATA_FOLDER / "ParentChildTreeFile.txt"
+INTERPRO_2ID = "data/tag_dict_4_safety_filtered.json"
+
+LSH_TABLE_PATHS = {"8bit": "data/hyperplanes_8bit_58641.npz"}
+
+KEYWORDS_VOCABULARY = (
+    IN_REPO_DATA_FOLDER / "keyword_vocabulary_safety_filtered_58641.txt"
+)
+KEYWORDS_IDF = IN_REPO_DATA_FOLDER / "keyword_idf_safety_filtered_58641.npy"
+
+RESID_CSV = "data/uniref90_and_mgnify90_residue_annotations_gt_1k_proteins.csv"
+INTERPRO2KEYWORDS = IN_REPO_DATA_FOLDER / "interpro_29026_to_keywords_58641.csv"

esmfold2_input_builder.py

@@ -1,254 +1,254 @@
-from dataclasses import dataclass
-from typing import Any, Sequence, TypeAlias, Union
-
-import numpy as np
-
-from .esmfold2_msa import MSA
-
-# fmt: off
-MSAInput: TypeAlias = Union[
-    MSA,
-    None,
-]
-# fmt: on
-
-
-@dataclass
-class Modification:
-    position: int  # zero-indexed
-    ccd: str
-    smiles: str | None = None  # TODO(mlee): add smiles support
-
-
-@dataclass
-class ProteinInput:
-    id: str | list[str]
-    sequence: str
-    modifications: list[Modification] | None = None
-    msa: MSAInput = None
-
-
-@dataclass
-class RNAInput:
-    id: str | list[str]
-    sequence: str
-    modifications: list[Modification] | None = None
-
-
-@dataclass
-class DNAInput:
-    id: str | list[str]
-    sequence: str
-    modifications: list[Modification] | None = None
-
-
-@dataclass
-class LigandInput:
-    id: str | list[str]
-    smiles: str | None = None
-    ccd: list[str] | None = None
-
-
-@dataclass
-class DistogramConditioning:
-    chain_id: str
-    distogram: np.ndarray
-
-
-@dataclass
-class PocketConditioning:
-    binder_chain_id: str
-    contacts: list[tuple[str, int]]
-
-
-@dataclass
-class CovalentBond:
-    chain_id1: str
-    res_idx1: int
-    atom_idx1: int
-    chain_id2: str
-    res_idx2: int
-    atom_idx2: int
-
-
-@dataclass
-class StructurePredictionInput:
-    sequences: Sequence[ProteinInput | RNAInput | DNAInput | LigandInput]
-    pocket: PocketConditioning | None = None
-    distogram_conditioning: list[DistogramConditioning] | None = None
-    covalent_bonds: list[CovalentBond] | None = None
-
-
-def serialize_structure_prediction_input(all_atom_input: StructurePredictionInput):
-    def create_chain_data(seq_input, chain_type: str) -> dict[str, Any]:
-        chain_data: dict[str, Any] = {
-            "sequence": seq_input.sequence,
-            "id": seq_input.id,
-            "type": chain_type,
-        }
-        if hasattr(seq_input, "modifications") and seq_input.modifications:
-            mods = [
-                {"position": mod.position, "ccd": mod.ccd}
-                for mod in seq_input.modifications
-            ]
-            chain_data["modifications"] = mods
-        if not hasattr(seq_input, "msa"):
-            pass
-        elif seq_input.msa is None:
-            chain_data["msa"] = None
-        elif isinstance(seq_input.msa, MSA):
-            chain_data["msa"] = {"sequences": seq_input.msa.sequences}
-        else:
-            error_msg = f"MSA must be None or MSA. Got {seq_input.msa} instead."
-            raise AttributeError(error_msg)
-        return chain_data
-
-    sequences = []
-    for seq_input in all_atom_input.sequences:
-        if isinstance(seq_input, ProteinInput):
-            sequences.append(create_chain_data(seq_input, "protein"))
-        elif isinstance(seq_input, RNAInput):
-            sequences.append(create_chain_data(seq_input, "rna"))
-        elif isinstance(seq_input, DNAInput):
-            sequences.append(create_chain_data(seq_input, "dna"))
-        elif isinstance(seq_input, LigandInput):
-            sequences.append(
-                {
-                    "smiles": seq_input.smiles,
-                    "id": seq_input.id,
-                    "ccd": seq_input.ccd,
-                    "type": "ligand",
-                }
-            )
-        else:
-            raise ValueError(f"Unsupported sequence input type: {type(seq_input)}")
-
-    result: dict[str, Any] = {"sequences": sequences}
-
-    if all_atom_input.covalent_bonds is not None:
-        result["covalent_bonds"] = [
-            {
-                "chain_id1": bond.chain_id1,
-                "res_idx1": bond.res_idx1,
-                "atom_idx1": bond.atom_idx1,
-                "chain_id2": bond.chain_id2,
-                "res_idx2": bond.res_idx2,
-                "atom_idx2": bond.atom_idx2,
-            }
-            for bond in all_atom_input.covalent_bonds
-        ]
-
-    if all_atom_input.pocket is not None:
-        result["pocket"] = {
-            "binder_chain_id": all_atom_input.pocket.binder_chain_id,
-            "contacts": all_atom_input.pocket.contacts,
-        }
-
-    if all_atom_input.distogram_conditioning is not None:
-        result["distogram_conditioning"] = [
-            {"chain_id": disto.chain_id, "distogram": disto.distogram.tolist()}
-            for disto in all_atom_input.distogram_conditioning
-        ]
-
-    return result
-
-
-def deserialize_structure_prediction_input(
-    data: dict[str, Any],
-) -> StructurePredictionInput:
-    """Inverse of :func:`serialize_structure_prediction_input`.
-
-    Reconstructs a :class:`StructurePredictionInput` from the JSON-safe dict
-    produced by ``serialize_structure_prediction_input``. Values round-trip;
-    ``DistogramConditioning.distogram`` dtype follows from JSON (``int64``
-    for integer entries, ``float64`` for floats) — cast back to the original
-    dtype if downstream code requires a specific one.
-    """
-
-    def _mods(chain: dict[str, Any]) -> list[Modification] | None:
-        raw = chain.get("modifications")
-        if not raw:
-            return None
-        return [Modification(position=m["position"], ccd=m["ccd"]) for m in raw]
-
-    def _msa(chain: dict[str, Any]) -> MSAInput:
-        if "msa" not in chain or chain["msa"] is None:
-            return None
-        msa_blk = chain["msa"]
-        if isinstance(msa_blk, str):
-            raise ValueError(f"Unexpected MSA string value: {msa_blk!r}")
-        return MSA.from_sequences(msa_blk["sequences"])
-
-    sequences: list[ProteinInput | RNAInput | DNAInput | LigandInput] = []
-    for chain in data["sequences"]:
-        t = chain["type"]
-        if t == "protein":
-            sequences.append(
-                ProteinInput(
-                    id=chain["id"],
-                    sequence=chain["sequence"],
-                    modifications=_mods(chain),
-                    msa=_msa(chain),
-                )
-            )
-        elif t == "rna":
-            sequences.append(
-                RNAInput(
-                    id=chain["id"],
-                    sequence=chain["sequence"],
-                    modifications=_mods(chain),
-                )
-            )
-        elif t == "dna":
-            sequences.append(
-                DNAInput(
-                    id=chain["id"],
-                    sequence=chain["sequence"],
-                    modifications=_mods(chain),
-                )
-            )
-        elif t == "ligand":
-            sequences.append(
-                LigandInput(
-                    id=chain["id"], smiles=chain.get("smiles"), ccd=chain.get("ccd")
-                )
-            )
-        else:
-            raise ValueError(f"Unsupported sequence type: {t!r}")
-
-    pocket: PocketConditioning | None = None
-    if (pocket_blk := data.get("pocket")) is not None:
-        pocket = PocketConditioning(
-            binder_chain_id=pocket_blk["binder_chain_id"],
-            contacts=[tuple(c) for c in pocket_blk["contacts"]],
-        )
-
-    distogram_conditioning: list[DistogramConditioning] | None = None
-    if (disto_blk := data.get("distogram_conditioning")) is not None:
-        distogram_conditioning = [
-            DistogramConditioning(
-                chain_id=d["chain_id"], distogram=np.asarray(d["distogram"])
-            )
-            for d in disto_blk
-        ]
-
-    covalent_bonds: list[CovalentBond] | None = None
-    if (bonds_blk := data.get("covalent_bonds")) is not None:
-        covalent_bonds = [
-            CovalentBond(
-                chain_id1=b["chain_id1"],
-                res_idx1=b["res_idx1"],
-                atom_idx1=b["atom_idx1"],
-                chain_id2=b["chain_id2"],
-                res_idx2=b["res_idx2"],
-                atom_idx2=b["atom_idx2"],
-            )
-            for b in bonds_blk
-        ]
-
-    return StructurePredictionInput(
-        sequences=sequences,
-        pocket=pocket,
-        distogram_conditioning=distogram_conditioning,
-        covalent_bonds=covalent_bonds,
-    )
+from dataclasses import dataclass
+from typing import Any, Sequence, TypeAlias, Union
+
+import numpy as np
+
+from .esmfold2_msa import MSA
+
+# fmt: off
+MSAInput: TypeAlias = Union[
+    MSA,
+    None,
+]
+# fmt: on
+
+
+@dataclass
+class Modification:
+    position: int  # zero-indexed
+    ccd: str
+    smiles: str | None = None  # TODO(mlee): add smiles support
+
+
+@dataclass
+class ProteinInput:
+    id: str | list[str]
+    sequence: str
+    modifications: list[Modification] | None = None
+    msa: MSAInput = None
+
+
+@dataclass
+class RNAInput:
+    id: str | list[str]
+    sequence: str
+    modifications: list[Modification] | None = None
+
+
+@dataclass
+class DNAInput:
+    id: str | list[str]
+    sequence: str
+    modifications: list[Modification] | None = None
+
+
+@dataclass
+class LigandInput:
+    id: str | list[str]
+    smiles: str | None = None
+    ccd: list[str] | None = None
+
+
+@dataclass
+class DistogramConditioning:
+    chain_id: str
+    distogram: np.ndarray
+
+
+@dataclass
+class PocketConditioning:
+    binder_chain_id: str
+    contacts: list[tuple[str, int]]
+
+
+@dataclass
+class CovalentBond:
+    chain_id1: str
+    res_idx1: int
+    atom_idx1: int
+    chain_id2: str
+    res_idx2: int
+    atom_idx2: int
+
+
+@dataclass
+class StructurePredictionInput:
+    sequences: Sequence[ProteinInput | RNAInput | DNAInput | LigandInput]
+    pocket: PocketConditioning | None = None
+    distogram_conditioning: list[DistogramConditioning] | None = None
+    covalent_bonds: list[CovalentBond] | None = None
+
+
+def serialize_structure_prediction_input(all_atom_input: StructurePredictionInput):
+    def create_chain_data(seq_input, chain_type: str) -> dict[str, Any]:
+        chain_data: dict[str, Any] = {
+            "sequence": seq_input.sequence,
+            "id": seq_input.id,
+            "type": chain_type,
+        }
+        if hasattr(seq_input, "modifications") and seq_input.modifications:
+            mods = [
+                {"position": mod.position, "ccd": mod.ccd}
+                for mod in seq_input.modifications
+            ]
+            chain_data["modifications"] = mods
+        if not hasattr(seq_input, "msa"):
+            pass
+        elif seq_input.msa is None:
+            chain_data["msa"] = None
+        elif isinstance(seq_input.msa, MSA):
+            chain_data["msa"] = {"sequences": seq_input.msa.sequences}
+        else:
+            error_msg = f"MSA must be None or MSA. Got {seq_input.msa} instead."
+            raise AttributeError(error_msg)
+        return chain_data
+
+    sequences = []
+    for seq_input in all_atom_input.sequences:
+        if isinstance(seq_input, ProteinInput):
+            sequences.append(create_chain_data(seq_input, "protein"))
+        elif isinstance(seq_input, RNAInput):
+            sequences.append(create_chain_data(seq_input, "rna"))
+        elif isinstance(seq_input, DNAInput):
+            sequences.append(create_chain_data(seq_input, "dna"))
+        elif isinstance(seq_input, LigandInput):
+            sequences.append(
+                {
+                    "smiles": seq_input.smiles,
+                    "id": seq_input.id,
+                    "ccd": seq_input.ccd,
+                    "type": "ligand",
+                }
+            )
+        else:
+            raise ValueError(f"Unsupported sequence input type: {type(seq_input)}")
+
+    result: dict[str, Any] = {"sequences": sequences}
+
+    if all_atom_input.covalent_bonds is not None:
+        result["covalent_bonds"] = [
+            {
+                "chain_id1": bond.chain_id1,
+                "res_idx1": bond.res_idx1,
+                "atom_idx1": bond.atom_idx1,
+                "chain_id2": bond.chain_id2,
+                "res_idx2": bond.res_idx2,
+                "atom_idx2": bond.atom_idx2,
+            }
+            for bond in all_atom_input.covalent_bonds
+        ]
+
+    if all_atom_input.pocket is not None:
+        result["pocket"] = {
+            "binder_chain_id": all_atom_input.pocket.binder_chain_id,
+            "contacts": all_atom_input.pocket.contacts,
+        }
+
+    if all_atom_input.distogram_conditioning is not None:
+        result["distogram_conditioning"] = [
+            {"chain_id": disto.chain_id, "distogram": disto.distogram.tolist()}
+            for disto in all_atom_input.distogram_conditioning
+        ]
+
+    return result
+
+
+def deserialize_structure_prediction_input(
+    data: dict[str, Any],
+) -> StructurePredictionInput:
+    """Inverse of :func:`serialize_structure_prediction_input`.
+
+    Reconstructs a :class:`StructurePredictionInput` from the JSON-safe dict
+    produced by ``serialize_structure_prediction_input``. Values round-trip;
+    ``DistogramConditioning.distogram`` dtype follows from JSON (``int64``
+    for integer entries, ``float64`` for floats) — cast back to the original
+    dtype if downstream code requires a specific one.
+    """
+
+    def _mods(chain: dict[str, Any]) -> list[Modification] | None:
+        raw = chain.get("modifications")
+        if not raw:
+            return None
+        return [Modification(position=m["position"], ccd=m["ccd"]) for m in raw]
+
+    def _msa(chain: dict[str, Any]) -> MSAInput:
+        if "msa" not in chain or chain["msa"] is None:
+            return None
+        msa_blk = chain["msa"]
+        if isinstance(msa_blk, str):
+            raise ValueError(f"Unexpected MSA string value: {msa_blk!r}")
+        return MSA.from_sequences(msa_blk["sequences"])
+
+    sequences: list[ProteinInput | RNAInput | DNAInput | LigandInput] = []
+    for chain in data["sequences"]:
+        t = chain["type"]
+        if t == "protein":
+            sequences.append(
+                ProteinInput(
+                    id=chain["id"],
+                    sequence=chain["sequence"],
+                    modifications=_mods(chain),
+                    msa=_msa(chain),
+                )
+            )
+        elif t == "rna":
+            sequences.append(
+                RNAInput(
+                    id=chain["id"],
+                    sequence=chain["sequence"],
+                    modifications=_mods(chain),
+                )
+            )
+        elif t == "dna":
+            sequences.append(
+                DNAInput(
+                    id=chain["id"],
+                    sequence=chain["sequence"],
+                    modifications=_mods(chain),
+                )
+            )
+        elif t == "ligand":
+            sequences.append(
+                LigandInput(
+                    id=chain["id"], smiles=chain.get("smiles"), ccd=chain.get("ccd")
+                )
+            )
+        else:
+            raise ValueError(f"Unsupported sequence type: {t!r}")
+
+    pocket: PocketConditioning | None = None
+    if (pocket_blk := data.get("pocket")) is not None:
+        pocket = PocketConditioning(
+            binder_chain_id=pocket_blk["binder_chain_id"],
+            contacts=[tuple(c) for c in pocket_blk["contacts"]],
+        )
+
+    distogram_conditioning: list[DistogramConditioning] | None = None
+    if (disto_blk := data.get("distogram_conditioning")) is not None:
+        distogram_conditioning = [
+            DistogramConditioning(
+                chain_id=d["chain_id"], distogram=np.asarray(d["distogram"])
+            )
+            for d in disto_blk
+        ]
+
+    covalent_bonds: list[CovalentBond] | None = None
+    if (bonds_blk := data.get("covalent_bonds")) is not None:
+        covalent_bonds = [
+            CovalentBond(
+                chain_id1=b["chain_id1"],
+                res_idx1=b["res_idx1"],
+                atom_idx1=b["atom_idx1"],
+                chain_id2=b["chain_id2"],
+                res_idx2=b["res_idx2"],
+                atom_idx2=b["atom_idx2"],
+            )
+            for b in bonds_blk
+        ]
+
+    return StructurePredictionInput(
+        sequences=sequences,
+        pocket=pocket,
+        distogram_conditioning=distogram_conditioning,
+        covalent_bonds=covalent_bonds,
+    )

esmfold2_metrics.py

@@ -1,373 +1,373 @@
-import numpy as np
-import torch
-import torch.nn.functional as F
-from einops import rearrange
-from torch import Tensor
-from torch.amp import autocast  # type: ignore
-
-from . import esmfold2_residue_constants as residue_constants
-from .esmfold2_misc import binpack, unbinpack
-from .esmfold2_protein_structure import (
-    compute_alignment_tensors,
-    compute_gdt_ts_no_alignment,
-    compute_rmsd_no_alignment,
-)
-
-
-def contact_precision(
-    predictions: Tensor,
-    targets: Tensor,
-    src_lengths: Tensor | None = None,
-    minsep: int = 6,
-    maxsep: int | None = None,
-    override_length: int | None = None,  # for casp
-):
-    """Computes contact precisions.
-
-    For protein contact prediction, precision is measured for the top (L/K) highest confidence predictions,
-    with L being the length of the protein sequence and K generally being equal to 1 or 5.
-
-    K = 5 measures the predictions of the very highest confidence contacts, while K = 1 is a more general measure
-    over all relatively high confidence predictions.
-
-    Since there are roughly ~L true contacts in a protein, this is a reasonable cutoff.
-
-
-    Args:
-        predictions (Tensor): Tensor of probabilities of size (B, L, L)
-        targets (Tensor): Tensor of true contacts of size (B, L, L)
-        src_lengths (Tensor, optional): Lengths of each sample in the batch, if using variable lengths.
-            If not provided, inferred from the size of the predictions.
-        minsep (int): Minimum separation distance to consider. We often want to measure contacts at a
-            certain range. Typical ranges are short [6, 12), medium [12, 24), and long [24, inf).
-        maxsep (int, optional): Used in conjunction with minsep to specify a contact range. If not provided uses
-            assumes no maximum range
-        override_length (int, optional): Used for casp evaluation where sometimes the "true" length is not
-            the same as the length of the input. Kept for posterity, we probably don't need this argument.
-    """
-    if predictions.dim() == 2:
-        predictions = predictions.unsqueeze(0)
-    if targets.dim() == 2:
-        targets = targets.unsqueeze(0)
-
-    # Check sizes
-    if predictions.size() != targets.size():
-        raise ValueError(
-            f"Size mismatch. Received predictions of size {predictions.size()}, "
-            f"targets of size {targets.size()}"
-        )
-    device = predictions.device
-
-    batch_size, seqlen, _ = predictions.size()
-
-    # Step 1) Construct a mask of size [B, L, L] to mask invalid contacts
-    seqlen_range = torch.arange(seqlen, device=device)
-    sep = seqlen_range.unsqueeze(0) - seqlen_range.unsqueeze(1)
-    sep = sep.unsqueeze(0)
-    # Mask contacts that are closer than minsep
-    valid_mask = sep >= minsep
-    # Mask contacts where target is negative (padding or unknown)
-    valid_mask = valid_mask & (targets >= 0)  # negative targets are invalid
-
-    # Mask contacts that are farther than maxsep, if provided
-    if maxsep is not None:
-        valid_mask &= sep < maxsep
-
-    if src_lengths is not None:
-        # If the lengths of the individual sequences are provided, mask positions
-        # that are farther than the end of the sequence.
-        valid = seqlen_range.unsqueeze(0) < src_lengths.unsqueeze(1)
-        valid_mask &= valid.unsqueeze(1) & valid.unsqueeze(2)
-    else:
-        src_lengths = torch.full([batch_size], seqlen, device=device, dtype=torch.long)
-
-    # Fill in the logit tensor with -inf for all invalid positions
-    predictions = predictions.masked_fill(~valid_mask, float("-inf"))
-
-    # Step 2) Select the top half of the prediction (should be symmetric)
-    x_ind, y_ind = np.triu_indices(seqlen, minsep)
-    predictions_upper = predictions[:, x_ind, y_ind]
-    targets_upper = targets[:, x_ind, y_ind]
-
-    # Step 3) Select the topk values in each batch where k = L (length of sequence)
-    topk = seqlen if override_length is None else max(seqlen, override_length)
-    # Indices are the indices into the predictions corresponding to the most confident predictions
-    indices = predictions_upper.argsort(dim=-1, descending=True)[:, :topk]
-    # topk_targets are the target values corresponding to the above indices
-    topk_targets = targets_upper[torch.arange(batch_size).unsqueeze(1), indices]
-    if topk_targets.size(1) < topk:
-        # If there aren't enough targets, pad to the output.
-        topk_targets = F.pad(topk_targets, [0, topk - topk_targets.size(1)])
-
-    # Step 4) Sum the accuracy at of the top-i predictions for i in 1, L
-    # topk_targets => 1/0 true vs. false contact, sorted by confidence of prediction
-    # cmumulative sum => Number of correct answers for the top-i predictions.
-    cumulative_dist = topk_targets.type_as(predictions).cumsum(-1)
-
-    # Step 5) Find the gather indices. This should be P@(L / K) for varous values of K
-    # The values will differ for each batch.
-    gather_lengths = src_lengths.unsqueeze(1)
-    if override_length is not None:
-        gather_lengths = override_length * torch.ones_like(
-            gather_lengths, device=device
-        )
-
-    # This gets you (0.1 * L, 0.2 * L, 0.3 * L, etc.)
-    gather_indices = (
-        (torch.arange(0.1, 1.1, 0.1, device=device).unsqueeze(0) * gather_lengths).type(
-            torch.long
-        )
-        - 1
-    ).clamp_min(0)
-
-    # Step 6) Gather the results and divide by the number of guesses to get the precision.
-    binned_cumulative_dist = cumulative_dist.gather(1, gather_indices)
-    binned_precisions = binned_cumulative_dist / (gather_indices + 1).type_as(
-        binned_cumulative_dist
-    )
-
-    # Select specific P@L/k. pl5 is index 1 b/c that corresponds to L * 0.2 in
-    # gather_indices above
-    pl5 = binned_precisions[:, 1]
-    # pl2 = binned_precisions[:, 4]
-    pl = binned_precisions[:, 9]
-    # AUC is the integral wrt K of P@L/K for K in range(1, L)
-    auc = binned_precisions.mean(-1)
-
-    return {"AUC": auc, "P@L": pl, "P@L5": pl5}
-
-
-def compute_lddt(
-    all_atom_pred_pos: torch.Tensor,
-    all_atom_positions: torch.Tensor,
-    all_atom_mask: torch.Tensor,
-    pairwise_all_atom_mask: torch.Tensor | None = None,
-    cutoff: float | torch.Tensor = 15.0,
-    eps: float = 1e-10,
-    per_residue: bool = True,
-    sequence_id: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """
-    Computes LDDT for a protein. Tensor sizes below include some optional dimensions. Specifically:
-        Nstates:
-            all_atom_pred_pos can contain multiple states in the first dimension which corresponds to outputs from different layers of a model (e.g. each IPA block). The return size will be [Nstates x Batch size] if this is included.
-        Natoms:
-            LDDT can be computed for all atoms or some atoms. The second to last dimension should contain the *FLATTENED* representation of L x Natoms. If you want to calculate for atom37, e.g., this will be of size (L * 37). If you are only calculating CA LDDT, it will be of size L.
-
-    Args:
-        all_atom_pred_pos (Tensor[float], [(Nstates x) B x (L * Natoms x) 3]): Tensor of predicted positions
-        all_atom_positions (Tensor[float], [B x (L * Natoms x) 3]): Tensor of true positions
-        all_atom_mask (Tensor[float], [B x (L * Natoms)]): Tensor of masks, indicating whether an atom exists.
-        pairwise_all_atom_mask (Tensor[float], [B x (L * Natoms x L * Natoms)], optional): Tensor of masks, indicating whether a pair of atoms should be considered in the LDDT calculation.
-        cutoff (float): Max distance to score lddt over. This can either be a float, or a tensor of shape [B, L, L] to allow for per-residue cutoffs, e.g. if you want to use a different cutoff for nucleic acids.
-        per_residue (bool): Whether to return per-residue or full-protein lddt.
-        sequence_id (Tensor, optional): Sequence id tensor for binpacking. NOTE: only supported for lddt_ca calculations, not when Natoms is passed!
-
-    Returns:
-        LDDT Tensor:
-            if per_residue:
-                Tensor[float], [(Nstates x) B x (L * Natoms)]
-            else:
-                Tensor[float], [(Nstates x) B]
-    """
-    all_atom_mask = all_atom_mask[..., None]  # add a dimension for broadcasting
-    dmat_true = torch.sqrt(
-        eps
-        + torch.sum(
-            (all_atom_positions[..., None, :] - all_atom_positions[..., None, :, :])
-            ** 2,
-            dim=-1,
-        )
-    )
-
-    dmat_pred = torch.sqrt(
-        eps
-        + torch.sum(
-            (all_atom_pred_pos[..., None, :] - all_atom_pred_pos[..., None, :, :]) ** 2,
-            dim=-1,
-        )
-    )
-    mask = all_atom_mask * rearrange(all_atom_mask, "... a b -> ... b a")
-    if pairwise_all_atom_mask is not None:
-        mask = mask * pairwise_all_atom_mask
-
-    if sequence_id is not None:
-        # TODO: This will work for lddt_ca, but not for regular lddt
-        # Problem is that regular lddt has natoms * nres scores, so would need to repeat this mask by natoms
-        # Leaving for now because it won't fail silently so should be ook.
-        seqid_mask = sequence_id[..., None] == sequence_id[..., None, :]
-        mask = mask * seqid_mask.type_as(mask)
-
-    return compute_lddt_from_dmat(
-        dmat_pred, dmat_true, mask, cutoff=cutoff, eps=eps, per_residue=per_residue
-    )
-
-
-def compute_lddt_from_dmat(
-    dmat_pred: torch.Tensor,
-    dmat_true: torch.Tensor,
-    pairwise_mask: torch.Tensor,
-    cutoff: float | torch.Tensor = 15.0,
-    eps: float = 1e-10,
-    per_residue: bool = True,
-):
-    """
-    Compute LDDT from pre-computed distance matrices.
-    This is useful when you want to compute LDDT with multiple different masks or cutoffs, e.g. for different molecule types (protein, nucleic acid, etc.).
-
-    Args:
-        dmat_pred (Tensor[float], [B x L x L]): Predicted distance matrix
-        dmat_true (Tensor[float], [B x L x L]): True distance matrix
-        pairwise_mask (Tensor[float], [B x L x L]): Pairwise mask indicating which pairs of atoms to consider
-        cutoff (float): Max distance to score lddt over. This can either be a float, or a tensor of shape [B, L, L] to allow for per-residue cutoffs, e.g. if you want to use a different cutoff for nucleic acids.
-        per_residue (bool): Whether to return per-residue or full-protein lddt.
-
-    Returns:
-        LDDT Tensor:
-            if per_residue:
-                Tensor[float], [B x L]
-            else:
-                Tensor[float], [B]
-    """
-    n = dmat_true.size(-1)
-    dists_to_score = (
-        (dmat_true < cutoff)
-        * pairwise_mask
-        * (1.0 - torch.eye(n, device=dmat_true.device))
-    )
-
-    dist_l1 = torch.abs(dmat_true - dmat_pred)
-    score = (
-        (dist_l1 < 0.5).type(dist_l1.dtype)
-        + (dist_l1 < 1.0).type(dist_l1.dtype)
-        + (dist_l1 < 2.0).type(dist_l1.dtype)
-        + (dist_l1 < 4.0).type(dist_l1.dtype)
-    )
-    score = score * 0.25
-
-    dims = (-1,) if per_residue else (-2, -1)
-    norm = 1.0 / (eps + torch.sum(dists_to_score, dim=dims))
-    score = norm * (eps + torch.sum(dists_to_score * score, dim=dims))
-    return score
-
-
-def compute_lddt_ca(
-    all_atom_pred_pos: torch.Tensor,
-    all_atom_positions: torch.Tensor,
-    all_atom_mask: torch.Tensor,
-    cutoff: float = 15.0,
-    eps: float = 1e-10,
-    per_residue: bool = True,
-    sequence_id: torch.Tensor | None = None,
-) -> torch.Tensor:
-    ca_pos = residue_constants.atom_order["CA"]
-    if all_atom_pred_pos.dim() != 3:
-        all_atom_pred_pos = all_atom_pred_pos[..., ca_pos, :]
-    all_atom_positions = all_atom_positions[..., ca_pos, :]
-    all_atom_mask = all_atom_mask[..., ca_pos]
-
-    return compute_lddt(
-        all_atom_pred_pos,
-        all_atom_positions,
-        all_atom_mask,
-        cutoff=cutoff,
-        eps=eps,
-        per_residue=per_residue,
-        sequence_id=sequence_id,
-    )
-
-
-# NOTE(roshan): no_grad required for stack_variable_length_tensors apparently... let's revisit if we want to backprop
-@torch.no_grad()
-@autocast("cuda", enabled=False)
-def compute_rmsd(
-    mobile: torch.Tensor,
-    target: torch.Tensor,
-    atom_exists_mask: torch.Tensor | None = None,
-    sequence_id: torch.Tensor | None = None,
-    reduction: str = "batch",
-):
-    """
-    Compute RMSD between two batches of structures with support for masking invalid atoms using PyTorch.
-
-    Args:
-    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
-    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
-    - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N)
-    - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking.
-    - reduction (str): One of "batch", "per_sample", "per_residue".
-
-    Returns:
-    If reduction == "batch":
-        (torch.Tensor): 0-dim, Average Root Mean Square Deviation between the structures for each batch
-    If reduction == "per_sample":
-        (torch.Tensor): (B,)-dim, Root Mean Square Deviation between the structures for each batch
-    If reduction == "per_residue":
-        (torch.Tensor): (B, N)-dim, Root Mean Square Deviation between the structures for residue in the batch
-    """
-
-    (centered_mobile, _, centered_target, _, rotation_matrix, num_valid_atoms) = (
-        compute_alignment_tensors(
-            mobile=mobile,
-            target=target,
-            atom_exists_mask=atom_exists_mask,
-            sequence_id=sequence_id,
-        )
-    )
-
-    # Apply transformation to centered structure
-    rotated_mobile = torch.matmul(centered_mobile, rotation_matrix)
-
-    # Compute rmsd for centered structures
-    rmsd = compute_rmsd_no_alignment(
-        rotated_mobile, centered_target, num_valid_atoms, reduction=reduction
-    )
-    if reduction == "per_residue" and sequence_id is not None:
-        rmsd = binpack(rmsd, sequence_id, pad_value=0)
-    return rmsd
-
-
-def compute_gdt_ts(
-    mobile: torch.Tensor,
-    target: torch.Tensor,
-    atom_exists_mask: torch.Tensor | None = None,
-    sequence_id: torch.Tensor | None = None,
-    reduction: str = "per_sample",
-):
-    """
-    Compute GDT_TS between two batches of structures with support for masking invalid atoms using PyTorch.
-
-    Args:
-    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
-    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
-    - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N)
-    - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking.
-    - reduction (str): One of "batch", "per_sample", "per_residue".
-
-    Returns:
-    If reduction == "batch":
-        (torch.Tensor): 0-dim, GDT_TS between the structures for each batch
-    If reduction == "per_sample":
-        (torch.Tensor): (B,)-dim, GDT_TS between the structures for each sample in the batch
-    """
-    if atom_exists_mask is None:
-        atom_exists_mask = torch.isfinite(target).all(dim=-1)
-    (centered_mobile, _, centered_target, _, rotation_matrix, _) = (
-        compute_alignment_tensors(
-            mobile=mobile,
-            target=target,
-            atom_exists_mask=atom_exists_mask,
-            sequence_id=sequence_id,
-        )
-    )
-
-    # Apply transformation to centered structure
-    rotated_mobile = torch.matmul(centered_mobile, rotation_matrix)
-
-    # the coordinate tensors returned by `compute_alignment_tensors` are unbinpacked and contain zeros for invalid positions
-    # so `compute_gdt_ts_no_alignment` requires `atom_exists_mask` to be passed and be unbinpacked
-    if sequence_id is not None:
-        atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=False)
-    return compute_gdt_ts_no_alignment(
-        rotated_mobile, centered_target, atom_exists_mask, reduction
-    )
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor
+from torch.amp import autocast  # type: ignore
+
+from . import esmfold2_residue_constants as residue_constants
+from .esmfold2_misc import binpack, unbinpack
+from .esmfold2_protein_structure import (
+    compute_alignment_tensors,
+    compute_gdt_ts_no_alignment,
+    compute_rmsd_no_alignment,
+)
+
+
+def contact_precision(
+    predictions: Tensor,
+    targets: Tensor,
+    src_lengths: Tensor | None = None,
+    minsep: int = 6,
+    maxsep: int | None = None,
+    override_length: int | None = None,  # for casp
+):
+    """Computes contact precisions.
+
+    For protein contact prediction, precision is measured for the top (L/K) highest confidence predictions,
+    with L being the length of the protein sequence and K generally being equal to 1 or 5.
+
+    K = 5 measures the predictions of the very highest confidence contacts, while K = 1 is a more general measure
+    over all relatively high confidence predictions.
+
+    Since there are roughly ~L true contacts in a protein, this is a reasonable cutoff.
+
+
+    Args:
+        predictions (Tensor): Tensor of probabilities of size (B, L, L)
+        targets (Tensor): Tensor of true contacts of size (B, L, L)
+        src_lengths (Tensor, optional): Lengths of each sample in the batch, if using variable lengths.
+            If not provided, inferred from the size of the predictions.
+        minsep (int): Minimum separation distance to consider. We often want to measure contacts at a
+            certain range. Typical ranges are short [6, 12), medium [12, 24), and long [24, inf).
+        maxsep (int, optional): Used in conjunction with minsep to specify a contact range. If not provided uses
+            assumes no maximum range
+        override_length (int, optional): Used for casp evaluation where sometimes the "true" length is not
+            the same as the length of the input. Kept for posterity, we probably don't need this argument.
+    """
+    if predictions.dim() == 2:
+        predictions = predictions.unsqueeze(0)
+    if targets.dim() == 2:
+        targets = targets.unsqueeze(0)
+
+    # Check sizes
+    if predictions.size() != targets.size():
+        raise ValueError(
+            f"Size mismatch. Received predictions of size {predictions.size()}, "
+            f"targets of size {targets.size()}"
+        )
+    device = predictions.device
+
+    batch_size, seqlen, _ = predictions.size()
+
+    # Step 1) Construct a mask of size [B, L, L] to mask invalid contacts
+    seqlen_range = torch.arange(seqlen, device=device)
+    sep = seqlen_range.unsqueeze(0) - seqlen_range.unsqueeze(1)
+    sep = sep.unsqueeze(0)
+    # Mask contacts that are closer than minsep
+    valid_mask = sep >= minsep
+    # Mask contacts where target is negative (padding or unknown)
+    valid_mask = valid_mask & (targets >= 0)  # negative targets are invalid
+
+    # Mask contacts that are farther than maxsep, if provided
+    if maxsep is not None:
+        valid_mask &= sep < maxsep
+
+    if src_lengths is not None:
+        # If the lengths of the individual sequences are provided, mask positions
+        # that are farther than the end of the sequence.
+        valid = seqlen_range.unsqueeze(0) < src_lengths.unsqueeze(1)
+        valid_mask &= valid.unsqueeze(1) & valid.unsqueeze(2)
+    else:
+        src_lengths = torch.full([batch_size], seqlen, device=device, dtype=torch.long)
+
+    # Fill in the logit tensor with -inf for all invalid positions
+    predictions = predictions.masked_fill(~valid_mask, float("-inf"))
+
+    # Step 2) Select the top half of the prediction (should be symmetric)
+    x_ind, y_ind = np.triu_indices(seqlen, minsep)
+    predictions_upper = predictions[:, x_ind, y_ind]
+    targets_upper = targets[:, x_ind, y_ind]
+
+    # Step 3) Select the topk values in each batch where k = L (length of sequence)
+    topk = seqlen if override_length is None else max(seqlen, override_length)
+    # Indices are the indices into the predictions corresponding to the most confident predictions
+    indices = predictions_upper.argsort(dim=-1, descending=True)[:, :topk]
+    # topk_targets are the target values corresponding to the above indices
+    topk_targets = targets_upper[torch.arange(batch_size).unsqueeze(1), indices]
+    if topk_targets.size(1) < topk:
+        # If there aren't enough targets, pad to the output.
+        topk_targets = F.pad(topk_targets, [0, topk - topk_targets.size(1)])
+
+    # Step 4) Sum the accuracy at of the top-i predictions for i in 1, L
+    # topk_targets => 1/0 true vs. false contact, sorted by confidence of prediction
+    # cmumulative sum => Number of correct answers for the top-i predictions.
+    cumulative_dist = topk_targets.type_as(predictions).cumsum(-1)
+
+    # Step 5) Find the gather indices. This should be P@(L / K) for varous values of K
+    # The values will differ for each batch.
+    gather_lengths = src_lengths.unsqueeze(1)
+    if override_length is not None:
+        gather_lengths = override_length * torch.ones_like(
+            gather_lengths, device=device
+        )
+
+    # This gets you (0.1 * L, 0.2 * L, 0.3 * L, etc.)
+    gather_indices = (
+        (torch.arange(0.1, 1.1, 0.1, device=device).unsqueeze(0) * gather_lengths).type(
+            torch.long
+        )
+        - 1
+    ).clamp_min(0)
+
+    # Step 6) Gather the results and divide by the number of guesses to get the precision.
+    binned_cumulative_dist = cumulative_dist.gather(1, gather_indices)
+    binned_precisions = binned_cumulative_dist / (gather_indices + 1).type_as(
+        binned_cumulative_dist
+    )
+
+    # Select specific P@L/k. pl5 is index 1 b/c that corresponds to L * 0.2 in
+    # gather_indices above
+    pl5 = binned_precisions[:, 1]
+    # pl2 = binned_precisions[:, 4]
+    pl = binned_precisions[:, 9]
+    # AUC is the integral wrt K of P@L/K for K in range(1, L)
+    auc = binned_precisions.mean(-1)
+
+    return {"AUC": auc, "P@L": pl, "P@L5": pl5}
+
+
+def compute_lddt(
+    all_atom_pred_pos: torch.Tensor,
+    all_atom_positions: torch.Tensor,
+    all_atom_mask: torch.Tensor,
+    pairwise_all_atom_mask: torch.Tensor | None = None,
+    cutoff: float | torch.Tensor = 15.0,
+    eps: float = 1e-10,
+    per_residue: bool = True,
+    sequence_id: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """
+    Computes LDDT for a protein. Tensor sizes below include some optional dimensions. Specifically:
+        Nstates:
+            all_atom_pred_pos can contain multiple states in the first dimension which corresponds to outputs from different layers of a model (e.g. each IPA block). The return size will be [Nstates x Batch size] if this is included.
+        Natoms:
+            LDDT can be computed for all atoms or some atoms. The second to last dimension should contain the *FLATTENED* representation of L x Natoms. If you want to calculate for atom37, e.g., this will be of size (L * 37). If you are only calculating CA LDDT, it will be of size L.
+
+    Args:
+        all_atom_pred_pos (Tensor[float], [(Nstates x) B x (L * Natoms x) 3]): Tensor of predicted positions
+        all_atom_positions (Tensor[float], [B x (L * Natoms x) 3]): Tensor of true positions
+        all_atom_mask (Tensor[float], [B x (L * Natoms)]): Tensor of masks, indicating whether an atom exists.
+        pairwise_all_atom_mask (Tensor[float], [B x (L * Natoms x L * Natoms)], optional): Tensor of masks, indicating whether a pair of atoms should be considered in the LDDT calculation.
+        cutoff (float): Max distance to score lddt over. This can either be a float, or a tensor of shape [B, L, L] to allow for per-residue cutoffs, e.g. if you want to use a different cutoff for nucleic acids.
+        per_residue (bool): Whether to return per-residue or full-protein lddt.
+        sequence_id (Tensor, optional): Sequence id tensor for binpacking. NOTE: only supported for lddt_ca calculations, not when Natoms is passed!
+
+    Returns:
+        LDDT Tensor:
+            if per_residue:
+                Tensor[float], [(Nstates x) B x (L * Natoms)]
+            else:
+                Tensor[float], [(Nstates x) B]
+    """
+    all_atom_mask = all_atom_mask[..., None]  # add a dimension for broadcasting
+    dmat_true = torch.sqrt(
+        eps
+        + torch.sum(
+            (all_atom_positions[..., None, :] - all_atom_positions[..., None, :, :])
+            ** 2,
+            dim=-1,
+        )
+    )
+
+    dmat_pred = torch.sqrt(
+        eps
+        + torch.sum(
+            (all_atom_pred_pos[..., None, :] - all_atom_pred_pos[..., None, :, :]) ** 2,
+            dim=-1,
+        )
+    )
+    mask = all_atom_mask * rearrange(all_atom_mask, "... a b -> ... b a")
+    if pairwise_all_atom_mask is not None:
+        mask = mask * pairwise_all_atom_mask
+
+    if sequence_id is not None:
+        # TODO: This will work for lddt_ca, but not for regular lddt
+        # Problem is that regular lddt has natoms * nres scores, so would need to repeat this mask by natoms
+        # Leaving for now because it won't fail silently so should be ook.
+        seqid_mask = sequence_id[..., None] == sequence_id[..., None, :]
+        mask = mask * seqid_mask.type_as(mask)
+
+    return compute_lddt_from_dmat(
+        dmat_pred, dmat_true, mask, cutoff=cutoff, eps=eps, per_residue=per_residue
+    )
+
+
+def compute_lddt_from_dmat(
+    dmat_pred: torch.Tensor,
+    dmat_true: torch.Tensor,
+    pairwise_mask: torch.Tensor,
+    cutoff: float | torch.Tensor = 15.0,
+    eps: float = 1e-10,
+    per_residue: bool = True,
+):
+    """
+    Compute LDDT from pre-computed distance matrices.
+    This is useful when you want to compute LDDT with multiple different masks or cutoffs, e.g. for different molecule types (protein, nucleic acid, etc.).
+
+    Args:
+        dmat_pred (Tensor[float], [B x L x L]): Predicted distance matrix
+        dmat_true (Tensor[float], [B x L x L]): True distance matrix
+        pairwise_mask (Tensor[float], [B x L x L]): Pairwise mask indicating which pairs of atoms to consider
+        cutoff (float): Max distance to score lddt over. This can either be a float, or a tensor of shape [B, L, L] to allow for per-residue cutoffs, e.g. if you want to use a different cutoff for nucleic acids.
+        per_residue (bool): Whether to return per-residue or full-protein lddt.
+
+    Returns:
+        LDDT Tensor:
+            if per_residue:
+                Tensor[float], [B x L]
+            else:
+                Tensor[float], [B]
+    """
+    n = dmat_true.size(-1)
+    dists_to_score = (
+        (dmat_true < cutoff)
+        * pairwise_mask
+        * (1.0 - torch.eye(n, device=dmat_true.device))
+    )
+
+    dist_l1 = torch.abs(dmat_true - dmat_pred)
+    score = (
+        (dist_l1 < 0.5).type(dist_l1.dtype)
+        + (dist_l1 < 1.0).type(dist_l1.dtype)
+        + (dist_l1 < 2.0).type(dist_l1.dtype)
+        + (dist_l1 < 4.0).type(dist_l1.dtype)
+    )
+    score = score * 0.25
+
+    dims = (-1,) if per_residue else (-2, -1)
+    norm = 1.0 / (eps + torch.sum(dists_to_score, dim=dims))
+    score = norm * (eps + torch.sum(dists_to_score * score, dim=dims))
+    return score
+
+
+def compute_lddt_ca(
+    all_atom_pred_pos: torch.Tensor,
+    all_atom_positions: torch.Tensor,
+    all_atom_mask: torch.Tensor,
+    cutoff: float = 15.0,
+    eps: float = 1e-10,
+    per_residue: bool = True,
+    sequence_id: torch.Tensor | None = None,
+) -> torch.Tensor:
+    ca_pos = residue_constants.atom_order["CA"]
+    if all_atom_pred_pos.dim() != 3:
+        all_atom_pred_pos = all_atom_pred_pos[..., ca_pos, :]
+    all_atom_positions = all_atom_positions[..., ca_pos, :]
+    all_atom_mask = all_atom_mask[..., ca_pos]
+
+    return compute_lddt(
+        all_atom_pred_pos,
+        all_atom_positions,
+        all_atom_mask,
+        cutoff=cutoff,
+        eps=eps,
+        per_residue=per_residue,
+        sequence_id=sequence_id,
+    )
+
+
+# NOTE(roshan): no_grad required for stack_variable_length_tensors apparently... let's revisit if we want to backprop
+@torch.no_grad()
+@autocast("cuda", enabled=False)
+def compute_rmsd(
+    mobile: torch.Tensor,
+    target: torch.Tensor,
+    atom_exists_mask: torch.Tensor | None = None,
+    sequence_id: torch.Tensor | None = None,
+    reduction: str = "batch",
+):
+    """
+    Compute RMSD between two batches of structures with support for masking invalid atoms using PyTorch.
+
+    Args:
+    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
+    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
+    - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N)
+    - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking.
+    - reduction (str): One of "batch", "per_sample", "per_residue".
+
+    Returns:
+    If reduction == "batch":
+        (torch.Tensor): 0-dim, Average Root Mean Square Deviation between the structures for each batch
+    If reduction == "per_sample":
+        (torch.Tensor): (B,)-dim, Root Mean Square Deviation between the structures for each batch
+    If reduction == "per_residue":
+        (torch.Tensor): (B, N)-dim, Root Mean Square Deviation between the structures for residue in the batch
+    """
+
+    (centered_mobile, _, centered_target, _, rotation_matrix, num_valid_atoms) = (
+        compute_alignment_tensors(
+            mobile=mobile,
+            target=target,
+            atom_exists_mask=atom_exists_mask,
+            sequence_id=sequence_id,
+        )
+    )
+
+    # Apply transformation to centered structure
+    rotated_mobile = torch.matmul(centered_mobile, rotation_matrix)
+
+    # Compute rmsd for centered structures
+    rmsd = compute_rmsd_no_alignment(
+        rotated_mobile, centered_target, num_valid_atoms, reduction=reduction
+    )
+    if reduction == "per_residue" and sequence_id is not None:
+        rmsd = binpack(rmsd, sequence_id, pad_value=0)
+    return rmsd
+
+
+def compute_gdt_ts(
+    mobile: torch.Tensor,
+    target: torch.Tensor,
+    atom_exists_mask: torch.Tensor | None = None,
+    sequence_id: torch.Tensor | None = None,
+    reduction: str = "per_sample",
+):
+    """
+    Compute GDT_TS between two batches of structures with support for masking invalid atoms using PyTorch.
+
+    Args:
+    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
+    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
+    - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N)
+    - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking.
+    - reduction (str): One of "batch", "per_sample", "per_residue".
+
+    Returns:
+    If reduction == "batch":
+        (torch.Tensor): 0-dim, GDT_TS between the structures for each batch
+    If reduction == "per_sample":
+        (torch.Tensor): (B,)-dim, GDT_TS between the structures for each sample in the batch
+    """
+    if atom_exists_mask is None:
+        atom_exists_mask = torch.isfinite(target).all(dim=-1)
+    (centered_mobile, _, centered_target, _, rotation_matrix, _) = (
+        compute_alignment_tensors(
+            mobile=mobile,
+            target=target,
+            atom_exists_mask=atom_exists_mask,
+            sequence_id=sequence_id,
+        )
+    )
+
+    # Apply transformation to centered structure
+    rotated_mobile = torch.matmul(centered_mobile, rotation_matrix)
+
+    # the coordinate tensors returned by `compute_alignment_tensors` are unbinpacked and contain zeros for invalid positions
+    # so `compute_gdt_ts_no_alignment` requires `atom_exists_mask` to be passed and be unbinpacked
+    if sequence_id is not None:
+        atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=False)
+    return compute_gdt_ts_no_alignment(
+        rotated_mobile, centered_target, atom_exists_mask, reduction
+    )

esmfold2_misc.py

@@ -1,504 +1,504 @@
-from __future__ import annotations
-
-import os
-from collections import defaultdict
-from contextlib import nullcontext
-from dataclasses import is_dataclass
-from io import BytesIO
-from typing import (
-    Any,
-    ContextManager,
-    Generator,
-    Iterable,
-    Protocol,
-    Sequence,
-    TypeVar,
-    runtime_checkable,
-)
-from warnings import warn
-
-import huggingface_hub
-import numpy as np
-import torch
-import zstd
-
-from .esmfold2_constants_esm3 import CHAIN_BREAK_STR
-from .esmfold2_utils_types import FunctionAnnotation
-
-MAX_SUPPORTED_DISTANCE = 1e6
-
-
-TSequence = TypeVar("TSequence", bound=Sequence)
-
-
-@runtime_checkable
-class Concatable(Protocol):
-    @classmethod
-    def concat(cls, objs: list[Concatable]) -> Concatable: ...
-
-
-def slice_python_object_as_numpy(
-    obj: TSequence, idx: int | list[int] | slice | np.ndarray
-) -> TSequence:
-    """
-    Slice a python object (like a list, string, or tuple) as if it was a numpy object.
-
-    Example:
-        >>> obj = "ABCDE"
-        >>> slice_python_object_as_numpy(obj, [1, 3, 4])
-        "BDE"
-
-        >>> obj = [1, 2, 3, 4, 5]
-        >>> slice_python_object_as_numpy(obj, np.arange(5) < 3)
-        [1, 2, 3]
-    """
-    if np.isscalar(idx):
-        idx = [int(idx)]  # type: ignore
-
-    if isinstance(idx, np.ndarray) and idx.dtype == bool:
-        sliced_obj = [obj[i] for i in np.where(idx)[0]]
-    elif isinstance(idx, slice):
-        sliced_obj = obj[idx]
-    else:
-        sliced_obj = [obj[i] for i in idx]  # type: ignore
-
-    match obj, sliced_obj:
-        case str(), list():
-            sliced_obj = "".join(sliced_obj)
-        case _:
-            sliced_obj = obj.__class__(sliced_obj)  # type: ignore
-
-    return sliced_obj  # type: ignore
-
-
-def slice_any_object(
-    obj: TSequence, idx: int | list[int] | slice | np.ndarray
-) -> TSequence:
-    """
-    Slice a arbitrary object (like a list, string, or tuple) as if it was a numpy object. Similar to `slice_python_object_as_numpy`, but detects if it's a numpy array or Tensor and uses the existing slice method if so.
-
-    If the object is a dataclass, it will simply apply the index to the object, under the assumption that the object has correcty implemented numpy indexing.
-
-    Example:
-        >>> obj = "ABCDE"
-        >>> slice_any_object(obj, [1, 3, 4])
-        "BDE"
-
-        >>> obj = np.array([1, 2, 3, 4, 5])
-        >>> slice_any_object(obj, np.arange(5) < 3)
-        np.array([1, 2, 3])
-
-        >>> obj = ProteinChain.from_rcsb("1a3a", "A")
-        >>> slice_any_object(obj, np.arange(len(obj)) < 10)
-        # ProteinChain w/ length 10
-
-    """
-    if isinstance(obj, (np.ndarray, torch.Tensor)):
-        return obj[idx]  # type: ignore
-    elif is_dataclass(obj):
-        # if passing a dataclass, assume it implements a custom slice
-        return obj[idx]  # type: ignore
-    else:
-        return slice_python_object_as_numpy(obj, idx)
-
-
-def rbf(values, v_min, v_max, n_bins=16):
-    """
-    Returns RBF encodings in a new dimension at the end.
-    """
-    rbf_centers = torch.linspace(
-        v_min, v_max, n_bins, device=values.device, dtype=values.dtype
-    )
-    rbf_centers = rbf_centers.view([1] * len(values.shape) + [-1])
-    rbf_std = (v_max - v_min) / n_bins
-    z = (values.unsqueeze(-1) - rbf_centers) / rbf_std
-    return torch.exp(-(z**2))
-
-
-def batched_gather(data, inds, dim=0, no_batch_dims=0):
-    ranges = []
-    for i, s in enumerate(data.shape[:no_batch_dims]):
-        r = torch.arange(s)
-        r = r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape) - i - 1))))
-        ranges.append(r)
-
-    remaining_dims = [slice(None) for _ in range(len(data.shape) - no_batch_dims)]
-    remaining_dims[dim - no_batch_dims if dim >= 0 else dim] = inds
-    ranges.extend(remaining_dims)
-    return data[ranges]
-
-
-def node_gather(s: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
-    return batched_gather(s.unsqueeze(-3), edges, -2, no_batch_dims=len(s.shape) - 1)
-
-
-def knn_graph(
-    coords: torch.Tensor,
-    coord_mask: torch.Tensor,
-    padding_mask: torch.Tensor,
-    sequence_id: torch.Tensor,
-    *,
-    no_knn: int,
-):
-    L = coords.shape[-2]
-    num_by_dist = min(no_knn, L)
-    device = coords.device
-
-    coords = coords.nan_to_num()
-    coord_mask = ~(coord_mask[..., None, :] & coord_mask[..., :, None])
-    padding_pairwise_mask = padding_mask[..., None, :] | padding_mask[..., :, None]
-    if sequence_id is not None:
-        padding_pairwise_mask |= torch.unsqueeze(sequence_id, 1) != torch.unsqueeze(
-            sequence_id, 2
-        )
-    dists = (coords.unsqueeze(-2) - coords.unsqueeze(-3)).norm(dim=-1)
-    arange = torch.arange(L, device=device)
-    seq_dists = (arange.unsqueeze(-1) - arange.unsqueeze(-2)).abs()
-    # We only support up to a certain distance, above that, we use sequence distance
-    # instead. This is so that when a large portion of the structure is masked out,
-    # the edges are built according to sequence distance.
-    max_dist = MAX_SUPPORTED_DISTANCE
-    if not (dists[~coord_mask] < max_dist).all():
-        raise ValueError(
-            f"Coordinate pairwise distances exceed max supported distance ({max_dist}). "
-        )
-    struct_then_seq_dist = (
-        seq_dists.to(dists.dtype)
-        .mul(1e2)
-        .add(max_dist)
-        .where(coord_mask, dists)
-        .masked_fill(padding_pairwise_mask, torch.inf)
-    )
-    dists, edges = struct_then_seq_dist.sort(dim=-1, descending=False)
-    # This is a L x L tensor, where we index by rows first,
-    # and columns are the edges we should pick.
-    chosen_edges = edges[..., :num_by_dist]
-    chosen_mask = dists[..., :num_by_dist].isfinite()
-    return chosen_edges, chosen_mask
-
-
-def stack_variable_length_tensors(
-    sequences: Sequence[torch.Tensor],
-    constant_value: int | float = 0,
-    dtype: torch.dtype | None = None,
-) -> torch.Tensor:
-    """Automatically stack tensors together, padding variable lengths with the
-    value in constant_value. Handles an arbitrary number of dimensions.
-
-    Examples:
-        >>> tensor1, tensor2 = torch.ones([2]), torch.ones([5])
-        >>> stack_variable_length_tensors(tensor1, tensor2)
-        tensor of shape [2, 5]. First row is [1, 1, 0, 0, 0]. Second row is all ones.
-
-        >>> tensor1, tensor2 = torch.ones([2, 4]), torch.ones([5, 3])
-        >>> stack_variable_length_tensors(tensor1, tensor2)
-        tensor of shape [2, 5, 4]
-    """
-    batch_size = len(sequences)
-    shape = [batch_size] + np.max([seq.shape for seq in sequences], 0).tolist()
-
-    if dtype is None:
-        dtype = sequences[0].dtype
-    device = sequences[0].device
-
-    array = torch.full(shape, constant_value, dtype=dtype, device=device)
-    for arr, seq in zip(array, sequences):
-        arrslice = tuple(slice(dim) for dim in seq.shape)
-        arr[arrslice] = seq
-
-    return array
-
-
-def binpack(
-    tensor: torch.Tensor, sequence_id: torch.Tensor | None, pad_value: int | float
-):
-    """
-    Args:
-        tensor (Tensor): [B, L, ...]
-
-    Returns:
-        Tensor: [B_binpacked, L_binpacked, ...]
-    """
-    if sequence_id is None:
-        return tensor
-
-    num_sequences = sequence_id.max(dim=-1).values + 1
-
-    dims = sequence_id.shape + tensor.shape[2:]
-    output_tensor = torch.full(
-        dims, fill_value=pad_value, dtype=tensor.dtype, device=tensor.device
-    )
-
-    idx = 0
-    for batch_idx, (batch_seqid, batch_num_sequences) in enumerate(
-        zip(sequence_id, num_sequences)
-    ):
-        for seqid in range(batch_num_sequences):
-            mask = batch_seqid == seqid
-            output_tensor[batch_idx, mask] = tensor[idx, : mask.sum()]
-            idx += 1
-    return output_tensor
-
-
-def unbinpack(
-    tensor: torch.Tensor, sequence_id: torch.Tensor | None, pad_value: int | float
-):
-    """
-    Args:
-        tensor (Tensor): [B, L, ...]
-
-    Returns:
-        Tensor: [B_unbinpacked, L_unbinpack, ...]
-    """
-    if sequence_id is None:
-        return tensor
-
-    unpacked_tensors = []
-    num_sequences = sequence_id.max(dim=-1).values + 1
-    for batch_idx, (batch_seqid, batch_num_sequences) in enumerate(
-        zip(sequence_id, num_sequences)
-    ):
-        for seqid in range(batch_num_sequences):
-            mask = batch_seqid == seqid
-            unpacked = tensor[batch_idx, mask]
-            unpacked_tensors.append(unpacked)
-    return stack_variable_length_tensors(unpacked_tensors, pad_value)
-
-
-def fp32_autocast_context(device_type: str) -> ContextManager[Any]:  # type: ignore
-    """
-    Returns an autocast context manager that disables downcasting by AMP.
-
-    Args:
-        device_type: The device type ('cpu' or 'cuda')
-
-    Returns:
-        An autocast context manager with the specified behavior.
-    """
-    if device_type == "cpu":
-        return torch.amp.autocast(device_type, enabled=False)  # type: ignore
-    elif device_type == "mps":
-        # For MPS, just return a no-op context manager (nullcontext) since MPS does not support autocast.
-        return nullcontext()
-    elif device_type == "cuda":
-        return torch.amp.autocast(device_type, dtype=torch.float32)  # type: ignore
-    else:
-        raise ValueError(f"Unsupported device type: {device_type}")
-
-
-def merge_ranges(ranges: list[range], merge_gap_max: int | None = None) -> list[range]:
-    """Merge overlapping ranges into sorted, non-overlapping segments.
-
-    Args:
-        ranges: collection of ranges to merge.
-        merge_gap_max: optionally merge neighboring ranges that are separated by a gap
-          no larger than this size.
-    Returns:
-        non-overlapping ranges merged from the inputs, sorted by position.
-    """
-    ranges = sorted(ranges, key=lambda r: r.start)
-    merge_gap_max = merge_gap_max if merge_gap_max is not None else 0
-    assert merge_gap_max >= 0, f"Invalid merge_gap_max: {merge_gap_max}"
-
-    merged = []
-    for r in ranges:
-        if not merged:
-            merged.append(r)
-        else:
-            last = merged[-1]
-            if last.stop + merge_gap_max >= r.start:
-                merged[-1] = range(last.start, max(last.stop, r.stop))
-            else:
-                merged.append(r)
-    return merged
-
-
-def merge_annotations(
-    annotations: list[FunctionAnnotation], merge_gap_max: int | None = None
-) -> list[FunctionAnnotation]:
-    """Merges annotations into non-overlapping segments.
-
-    Args:
-        annotations: annotations to merge.
-        merge_gap_max: optionally merge neighboring ranges that are separated by a gap
-          no larger than this size.
-    Returns:
-        non-overlapping annotations with gaps merged.
-    """
-    grouped: dict[str, list[range]] = defaultdict(list)
-    for a in annotations:
-        # +1 since FunctionAnnotation.end is inlcusive.
-        grouped[a.label].append(range(a.start, a.end + 1))
-
-    merged = []
-    for label, ranges in grouped.items():
-        merged_ranges = merge_ranges(ranges, merge_gap_max=merge_gap_max)
-        for range_ in merged_ranges:
-            annotation = FunctionAnnotation(
-                label=label,
-                start=range_.start,
-                end=range_.stop - 1,  # convert range.stop exclusive -> inclusive.
-            )
-            merged.append(annotation)
-    return merged
-
-
-def replace_inf(data):
-    if data is None:
-        return None
-    array = np.asarray(data, dtype=np.float32)
-    array = np.where(np.isinf(array), 1000, array)
-    return array.tolist()
-
-
-def maybe_tensor(x, convert_none_to_nan: bool = False) -> torch.Tensor | None:
-    if x is None:
-        return None
-    if isinstance(x, torch.Tensor):
-        return x
-    if isinstance(x, list) and all(isinstance(t, torch.Tensor) for t in x):
-        return torch.stack(x)
-    if convert_none_to_nan:
-        x = np.asarray(x, dtype=np.float32)
-        x = np.where(x is None, np.nan, x)
-    return torch.tensor(x)
-
-
-def maybe_list(x, convert_nan_to_none: bool = False) -> list | None:
-    if x is None:
-        return None
-    if not convert_nan_to_none:
-        return x.tolist()
-
-    # Handle both torch.tensor and np.ndarray input.
-    if isinstance(x, torch.Tensor):
-        nan_mask = torch.isnan(x).cpu().numpy()
-        np_arr = x.cpu().numpy().astype(object)
-    elif isinstance(x, np.ndarray):
-        nan_mask = np.isnan(x)
-        np_arr = x.astype(object)
-    else:
-        raise TypeError("maybe_list can only work with torch.tensor or np.ndarray.")
-
-    np_arr[nan_mask] = None
-    return np_arr.tolist()
-
-
-def huggingfacehub_login():
-    """Authenticates with the Hugging Face Hub using the HF_TOKEN environment
-    variable, else by prompting the user"""
-    token = os.environ.get("HF_TOKEN")
-    huggingface_hub.login(token=token)
-
-
-def get_chainbreak_boundaries_from_sequence(sequence: Sequence[str]) -> np.ndarray:
-    chain_boundaries = [0]
-    for i, aa in enumerate(sequence):
-        if aa == CHAIN_BREAK_STR:
-            if i == (len(sequence) - 1):
-                raise ValueError(
-                    "Encountered chain break token at end of sequence, this is unexpected."
-                )
-            if i == (len(sequence) - 2):
-                warn(
-                    "Encountered chain break token at penultimate position, this is unexpected."
-                )
-            chain_boundaries.append(i)
-            chain_boundaries.append(i + 1)
-    chain_boundaries.append(len(sequence))
-    assert len(chain_boundaries) % 2 == 0
-    chain_boundaries = np.array(chain_boundaries).reshape(-1, 2)
-    return chain_boundaries
-
-
-def deserialize_tensors(b: bytes) -> Any:
-    buf = BytesIO(zstd.ZSTD_uncompress(b))
-    d = torch.load(buf, map_location="cpu", weights_only=False)
-    return d
-
-
-def join_lists(
-    lists: Sequence[Sequence[Any]], separator: Sequence[Any] | None = None
-) -> list[Any]:
-    """Joins multiple lists with separator element. Like str.join but for lists.
-
-    Example: [[1, 2], [3], [4]], separator=[0] -> [1, 2, 0, 3, 0, 4]
-
-    Args:
-        lists: Lists of elements to chain
-        separator: separators to intsert between chained output.
-    Returns:
-        Joined lists.
-    """
-    if not lists:
-        return []
-    joined = []
-    joined.extend(lists[0])
-    for l in lists[1:]:
-        if separator:
-            joined.extend(separator)
-        joined.extend(l)
-    return joined
-
-
-def iterate_with_intermediate(
-    lists: Iterable, intermediate
-) -> Generator[Any, None, None]:
-    """
-    Iterate over the iterable, yielding the intermediate value between
-    every element of the intermediate. Useful for joining objects with
-    separator tokens.
-    """
-    it = iter(lists)
-    yield next(it)
-    for l in it:
-        yield intermediate
-        yield l
-
-
-def concat_objects(objs: Sequence[Any], separator: Any | None = None):
-    """
-    Concat objects with each other using a separator token.
-
-    Supports:
-        - Concatable (objects that implement `concat` classmethod)
-        - strings
-        - lists
-        - numpy arrays
-        - torch Tensors
-
-    Example:
-        >>> foo = "abc"
-        >>> bar = "def"
-        >>> concat_objects([foo, bar], "|")
-        "abc|def"
-    """
-    match objs[0]:
-        case Concatable():
-            return objs[0].__class__.concat(objs)  # type: ignore
-        case str():
-            assert isinstance(
-                separator, str
-            ), "Trying to join strings but separator is not a string"
-            return separator.join(objs)
-        case list():
-            if separator is not None:
-                return join_lists(objs, [separator])
-            else:
-                return join_lists(objs)
-        case np.ndarray():
-            if separator is not None:
-                return np.concatenate(
-                    list(iterate_with_intermediate(objs, np.array([separator])))
-                )
-            else:
-                return np.concatenate(objs)
-        case torch.Tensor():
-            if separator is not None:
-                return torch.cat(
-                    list(iterate_with_intermediate(objs, torch.tensor([separator])))
-                )
-            else:
-                return torch.cat(objs)  # type: ignore
-        case _:
-            raise TypeError(type(objs[0]))
+from __future__ import annotations
+
+import os
+from collections import defaultdict
+from contextlib import nullcontext
+from dataclasses import is_dataclass
+from io import BytesIO
+from typing import (
+    Any,
+    ContextManager,
+    Generator,
+    Iterable,
+    Protocol,
+    Sequence,
+    TypeVar,
+    runtime_checkable,
+)
+from warnings import warn
+
+import huggingface_hub
+import numpy as np
+import torch
+import zstd
+
+from .esmfold2_constants_esm3 import CHAIN_BREAK_STR
+from .esmfold2_utils_types import FunctionAnnotation
+
+MAX_SUPPORTED_DISTANCE = 1e6
+
+
+TSequence = TypeVar("TSequence", bound=Sequence)
+
+
+@runtime_checkable
+class Concatable(Protocol):
+    @classmethod
+    def concat(cls, objs: list[Concatable]) -> Concatable: ...
+
+
+def slice_python_object_as_numpy(
+    obj: TSequence, idx: int | list[int] | slice | np.ndarray
+) -> TSequence:
+    """
+    Slice a python object (like a list, string, or tuple) as if it was a numpy object.
+
+    Example:
+        >>> obj = "ABCDE"
+        >>> slice_python_object_as_numpy(obj, [1, 3, 4])
+        "BDE"
+
+        >>> obj = [1, 2, 3, 4, 5]
+        >>> slice_python_object_as_numpy(obj, np.arange(5) < 3)
+        [1, 2, 3]
+    """
+    if np.isscalar(idx):
+        idx = [int(idx)]  # type: ignore
+
+    if isinstance(idx, np.ndarray) and idx.dtype == bool:
+        sliced_obj = [obj[i] for i in np.where(idx)[0]]
+    elif isinstance(idx, slice):
+        sliced_obj = obj[idx]
+    else:
+        sliced_obj = [obj[i] for i in idx]  # type: ignore
+
+    match obj, sliced_obj:
+        case str(), list():
+            sliced_obj = "".join(sliced_obj)
+        case _:
+            sliced_obj = obj.__class__(sliced_obj)  # type: ignore
+
+    return sliced_obj  # type: ignore
+
+
+def slice_any_object(
+    obj: TSequence, idx: int | list[int] | slice | np.ndarray
+) -> TSequence:
+    """
+    Slice a arbitrary object (like a list, string, or tuple) as if it was a numpy object. Similar to `slice_python_object_as_numpy`, but detects if it's a numpy array or Tensor and uses the existing slice method if so.
+
+    If the object is a dataclass, it will simply apply the index to the object, under the assumption that the object has correcty implemented numpy indexing.
+
+    Example:
+        >>> obj = "ABCDE"
+        >>> slice_any_object(obj, [1, 3, 4])
+        "BDE"
+
+        >>> obj = np.array([1, 2, 3, 4, 5])
+        >>> slice_any_object(obj, np.arange(5) < 3)
+        np.array([1, 2, 3])
+
+        >>> obj = ProteinChain.from_rcsb("1a3a", "A")
+        >>> slice_any_object(obj, np.arange(len(obj)) < 10)
+        # ProteinChain w/ length 10
+
+    """
+    if isinstance(obj, (np.ndarray, torch.Tensor)):
+        return obj[idx]  # type: ignore
+    elif is_dataclass(obj):
+        # if passing a dataclass, assume it implements a custom slice
+        return obj[idx]  # type: ignore
+    else:
+        return slice_python_object_as_numpy(obj, idx)
+
+
+def rbf(values, v_min, v_max, n_bins=16):
+    """
+    Returns RBF encodings in a new dimension at the end.
+    """
+    rbf_centers = torch.linspace(
+        v_min, v_max, n_bins, device=values.device, dtype=values.dtype
+    )
+    rbf_centers = rbf_centers.view([1] * len(values.shape) + [-1])
+    rbf_std = (v_max - v_min) / n_bins
+    z = (values.unsqueeze(-1) - rbf_centers) / rbf_std
+    return torch.exp(-(z**2))
+
+
+def batched_gather(data, inds, dim=0, no_batch_dims=0):
+    ranges = []
+    for i, s in enumerate(data.shape[:no_batch_dims]):
+        r = torch.arange(s)
+        r = r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape) - i - 1))))
+        ranges.append(r)
+
+    remaining_dims = [slice(None) for _ in range(len(data.shape) - no_batch_dims)]
+    remaining_dims[dim - no_batch_dims if dim >= 0 else dim] = inds
+    ranges.extend(remaining_dims)
+    return data[ranges]
+
+
+def node_gather(s: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
+    return batched_gather(s.unsqueeze(-3), edges, -2, no_batch_dims=len(s.shape) - 1)
+
+
+def knn_graph(
+    coords: torch.Tensor,
+    coord_mask: torch.Tensor,
+    padding_mask: torch.Tensor,
+    sequence_id: torch.Tensor,
+    *,
+    no_knn: int,
+):
+    L = coords.shape[-2]
+    num_by_dist = min(no_knn, L)
+    device = coords.device
+
+    coords = coords.nan_to_num()
+    coord_mask = ~(coord_mask[..., None, :] & coord_mask[..., :, None])
+    padding_pairwise_mask = padding_mask[..., None, :] | padding_mask[..., :, None]
+    if sequence_id is not None:
+        padding_pairwise_mask |= torch.unsqueeze(sequence_id, 1) != torch.unsqueeze(
+            sequence_id, 2
+        )
+    dists = (coords.unsqueeze(-2) - coords.unsqueeze(-3)).norm(dim=-1)
+    arange = torch.arange(L, device=device)
+    seq_dists = (arange.unsqueeze(-1) - arange.unsqueeze(-2)).abs()
+    # We only support up to a certain distance, above that, we use sequence distance
+    # instead. This is so that when a large portion of the structure is masked out,
+    # the edges are built according to sequence distance.
+    max_dist = MAX_SUPPORTED_DISTANCE
+    if not (dists[~coord_mask] < max_dist).all():
+        raise ValueError(
+            f"Coordinate pairwise distances exceed max supported distance ({max_dist}). "
+        )
+    struct_then_seq_dist = (
+        seq_dists.to(dists.dtype)
+        .mul(1e2)
+        .add(max_dist)
+        .where(coord_mask, dists)
+        .masked_fill(padding_pairwise_mask, torch.inf)
+    )
+    dists, edges = struct_then_seq_dist.sort(dim=-1, descending=False)
+    # This is a L x L tensor, where we index by rows first,
+    # and columns are the edges we should pick.
+    chosen_edges = edges[..., :num_by_dist]
+    chosen_mask = dists[..., :num_by_dist].isfinite()
+    return chosen_edges, chosen_mask
+
+
+def stack_variable_length_tensors(
+    sequences: Sequence[torch.Tensor],
+    constant_value: int | float = 0,
+    dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Automatically stack tensors together, padding variable lengths with the
+    value in constant_value. Handles an arbitrary number of dimensions.
+
+    Examples:
+        >>> tensor1, tensor2 = torch.ones([2]), torch.ones([5])
+        >>> stack_variable_length_tensors(tensor1, tensor2)
+        tensor of shape [2, 5]. First row is [1, 1, 0, 0, 0]. Second row is all ones.
+
+        >>> tensor1, tensor2 = torch.ones([2, 4]), torch.ones([5, 3])
+        >>> stack_variable_length_tensors(tensor1, tensor2)
+        tensor of shape [2, 5, 4]
+    """
+    batch_size = len(sequences)
+    shape = [batch_size] + np.max([seq.shape for seq in sequences], 0).tolist()
+
+    if dtype is None:
+        dtype = sequences[0].dtype
+    device = sequences[0].device
+
+    array = torch.full(shape, constant_value, dtype=dtype, device=device)
+    for arr, seq in zip(array, sequences):
+        arrslice = tuple(slice(dim) for dim in seq.shape)
+        arr[arrslice] = seq
+
+    return array
+
+
+def binpack(
+    tensor: torch.Tensor, sequence_id: torch.Tensor | None, pad_value: int | float
+):
+    """
+    Args:
+        tensor (Tensor): [B, L, ...]
+
+    Returns:
+        Tensor: [B_binpacked, L_binpacked, ...]
+    """
+    if sequence_id is None:
+        return tensor
+
+    num_sequences = sequence_id.max(dim=-1).values + 1
+
+    dims = sequence_id.shape + tensor.shape[2:]
+    output_tensor = torch.full(
+        dims, fill_value=pad_value, dtype=tensor.dtype, device=tensor.device
+    )
+
+    idx = 0
+    for batch_idx, (batch_seqid, batch_num_sequences) in enumerate(
+        zip(sequence_id, num_sequences)
+    ):
+        for seqid in range(batch_num_sequences):
+            mask = batch_seqid == seqid
+            output_tensor[batch_idx, mask] = tensor[idx, : mask.sum()]
+            idx += 1
+    return output_tensor
+
+
+def unbinpack(
+    tensor: torch.Tensor, sequence_id: torch.Tensor | None, pad_value: int | float
+):
+    """
+    Args:
+        tensor (Tensor): [B, L, ...]
+
+    Returns:
+        Tensor: [B_unbinpacked, L_unbinpack, ...]
+    """
+    if sequence_id is None:
+        return tensor
+
+    unpacked_tensors = []
+    num_sequences = sequence_id.max(dim=-1).values + 1
+    for batch_idx, (batch_seqid, batch_num_sequences) in enumerate(
+        zip(sequence_id, num_sequences)
+    ):
+        for seqid in range(batch_num_sequences):
+            mask = batch_seqid == seqid
+            unpacked = tensor[batch_idx, mask]
+            unpacked_tensors.append(unpacked)
+    return stack_variable_length_tensors(unpacked_tensors, pad_value)
+
+
+def fp32_autocast_context(device_type: str) -> ContextManager[Any]:  # type: ignore
+    """
+    Returns an autocast context manager that disables downcasting by AMP.
+
+    Args:
+        device_type: The device type ('cpu' or 'cuda')
+
+    Returns:
+        An autocast context manager with the specified behavior.
+    """
+    if device_type == "cpu":
+        return torch.amp.autocast(device_type, enabled=False)  # type: ignore
+    elif device_type == "mps":
+        # For MPS, just return a no-op context manager (nullcontext) since MPS does not support autocast.
+        return nullcontext()
+    elif device_type == "cuda":
+        return torch.amp.autocast(device_type, dtype=torch.float32)  # type: ignore
+    else:
+        raise ValueError(f"Unsupported device type: {device_type}")
+
+
+def merge_ranges(ranges: list[range], merge_gap_max: int | None = None) -> list[range]:
+    """Merge overlapping ranges into sorted, non-overlapping segments.
+
+    Args:
+        ranges: collection of ranges to merge.
+        merge_gap_max: optionally merge neighboring ranges that are separated by a gap
+          no larger than this size.
+    Returns:
+        non-overlapping ranges merged from the inputs, sorted by position.
+    """
+    ranges = sorted(ranges, key=lambda r: r.start)
+    merge_gap_max = merge_gap_max if merge_gap_max is not None else 0
+    assert merge_gap_max >= 0, f"Invalid merge_gap_max: {merge_gap_max}"
+
+    merged = []
+    for r in ranges:
+        if not merged:
+            merged.append(r)
+        else:
+            last = merged[-1]
+            if last.stop + merge_gap_max >= r.start:
+                merged[-1] = range(last.start, max(last.stop, r.stop))
+            else:
+                merged.append(r)
+    return merged
+
+
+def merge_annotations(
+    annotations: list[FunctionAnnotation], merge_gap_max: int | None = None
+) -> list[FunctionAnnotation]:
+    """Merges annotations into non-overlapping segments.
+
+    Args:
+        annotations: annotations to merge.
+        merge_gap_max: optionally merge neighboring ranges that are separated by a gap
+          no larger than this size.
+    Returns:
+        non-overlapping annotations with gaps merged.
+    """
+    grouped: dict[str, list[range]] = defaultdict(list)
+    for a in annotations:
+        # +1 since FunctionAnnotation.end is inlcusive.
+        grouped[a.label].append(range(a.start, a.end + 1))
+
+    merged = []
+    for label, ranges in grouped.items():
+        merged_ranges = merge_ranges(ranges, merge_gap_max=merge_gap_max)
+        for range_ in merged_ranges:
+            annotation = FunctionAnnotation(
+                label=label,
+                start=range_.start,
+                end=range_.stop - 1,  # convert range.stop exclusive -> inclusive.
+            )
+            merged.append(annotation)
+    return merged
+
+
+def replace_inf(data):
+    if data is None:
+        return None
+    array = np.asarray(data, dtype=np.float32)
+    array = np.where(np.isinf(array), 1000, array)
+    return array.tolist()
+
+
+def maybe_tensor(x, convert_none_to_nan: bool = False) -> torch.Tensor | None:
+    if x is None:
+        return None
+    if isinstance(x, torch.Tensor):
+        return x
+    if isinstance(x, list) and all(isinstance(t, torch.Tensor) for t in x):
+        return torch.stack(x)
+    if convert_none_to_nan:
+        x = np.asarray(x, dtype=np.float32)
+        x = np.where(x is None, np.nan, x)
+    return torch.tensor(x)
+
+
+def maybe_list(x, convert_nan_to_none: bool = False) -> list | None:
+    if x is None:
+        return None
+    if not convert_nan_to_none:
+        return x.tolist()
+
+    # Handle both torch.tensor and np.ndarray input.
+    if isinstance(x, torch.Tensor):
+        nan_mask = torch.isnan(x).cpu().numpy()
+        np_arr = x.cpu().numpy().astype(object)
+    elif isinstance(x, np.ndarray):
+        nan_mask = np.isnan(x)
+        np_arr = x.astype(object)
+    else:
+        raise TypeError("maybe_list can only work with torch.tensor or np.ndarray.")
+
+    np_arr[nan_mask] = None
+    return np_arr.tolist()
+
+
+def huggingfacehub_login():
+    """Authenticates with the Hugging Face Hub using the HF_TOKEN environment
+    variable, else by prompting the user"""
+    token = os.environ.get("HF_TOKEN")
+    huggingface_hub.login(token=token)
+
+
+def get_chainbreak_boundaries_from_sequence(sequence: Sequence[str]) -> np.ndarray:
+    chain_boundaries = [0]
+    for i, aa in enumerate(sequence):
+        if aa == CHAIN_BREAK_STR:
+            if i == (len(sequence) - 1):
+                raise ValueError(
+                    "Encountered chain break token at end of sequence, this is unexpected."
+                )
+            if i == (len(sequence) - 2):
+                warn(
+                    "Encountered chain break token at penultimate position, this is unexpected."
+                )
+            chain_boundaries.append(i)
+            chain_boundaries.append(i + 1)
+    chain_boundaries.append(len(sequence))
+    assert len(chain_boundaries) % 2 == 0
+    chain_boundaries = np.array(chain_boundaries).reshape(-1, 2)
+    return chain_boundaries
+
+
+def deserialize_tensors(b: bytes) -> Any:
+    buf = BytesIO(zstd.ZSTD_uncompress(b))
+    d = torch.load(buf, map_location="cpu", weights_only=False)
+    return d
+
+
+def join_lists(
+    lists: Sequence[Sequence[Any]], separator: Sequence[Any] | None = None
+) -> list[Any]:
+    """Joins multiple lists with separator element. Like str.join but for lists.
+
+    Example: [[1, 2], [3], [4]], separator=[0] -> [1, 2, 0, 3, 0, 4]
+
+    Args:
+        lists: Lists of elements to chain
+        separator: separators to intsert between chained output.
+    Returns:
+        Joined lists.
+    """
+    if not lists:
+        return []
+    joined = []
+    joined.extend(lists[0])
+    for l in lists[1:]:
+        if separator:
+            joined.extend(separator)
+        joined.extend(l)
+    return joined
+
+
+def iterate_with_intermediate(
+    lists: Iterable, intermediate
+) -> Generator[Any, None, None]:
+    """
+    Iterate over the iterable, yielding the intermediate value between
+    every element of the intermediate. Useful for joining objects with
+    separator tokens.
+    """
+    it = iter(lists)
+    yield next(it)
+    for l in it:
+        yield intermediate
+        yield l
+
+
+def concat_objects(objs: Sequence[Any], separator: Any | None = None):
+    """
+    Concat objects with each other using a separator token.
+
+    Supports:
+        - Concatable (objects that implement `concat` classmethod)
+        - strings
+        - lists
+        - numpy arrays
+        - torch Tensors
+
+    Example:
+        >>> foo = "abc"
+        >>> bar = "def"
+        >>> concat_objects([foo, bar], "|")
+        "abc|def"
+    """
+    match objs[0]:
+        case Concatable():
+            return objs[0].__class__.concat(objs)  # type: ignore
+        case str():
+            assert isinstance(
+                separator, str
+            ), "Trying to join strings but separator is not a string"
+            return separator.join(objs)
+        case list():
+            if separator is not None:
+                return join_lists(objs, [separator])
+            else:
+                return join_lists(objs)
+        case np.ndarray():
+            if separator is not None:
+                return np.concatenate(
+                    list(iterate_with_intermediate(objs, np.array([separator])))
+                )
+            else:
+                return np.concatenate(objs)
+        case torch.Tensor():
+            if separator is not None:
+                return torch.cat(
+                    list(iterate_with_intermediate(objs, torch.tensor([separator])))
+                )
+            else:
+                return torch.cat(objs)  # type: ignore
+        case _:
+            raise TypeError(type(objs[0]))

esmfold2_mmcif_parsing.py

@@ -1,469 +1,469 @@
-from __future__ import annotations
-
-import functools
-import io
-import os
-from dataclasses import dataclass
-from datetime import datetime
-from typing import Union
-
-import biotite.structure as bs
-import biotite.structure.io.pdbx as pdbx
-
-from . import esmfold2_residue_constants as residue_constants
-
-# Define PathOrBuffer for the opensource version
-PathOrBuffer = Union[str, os.PathLike, io.StringIO]
-
-
-class NoProteinError(Exception):
-    pass
-
-
-@dataclass
-class Residue:
-    residue_number: int | None = None
-    insertion_code: str = ""
-    hetflag: bool = False
-
-
-@dataclass
-class MmcifHeader:
-    release_date: datetime | None = None
-    resolution: float | None = None
-    structure_method: str = "UNKNOWN"
-
-
-class MmcifWrapper:
-    def __init__(self, id: str | None = None):
-        self.id: str = id or ""
-        self.raw: pdbx.CIFFile | None = None
-        self.structure: bs.AtomArray
-        self.header: MmcifHeader = MmcifHeader()
-        self.entities: dict[int, list[str]] = {}
-        self.chain_to_seqres: dict[str, str] = {}
-        self.seqres_to_structure: dict[str, dict[int, Residue]] = {}
-
-    @classmethod
-    def read(cls, path: PathOrBuffer, id: str | None = None) -> MmcifWrapper:
-        obj = cls(id=id)
-        obj._load(path)
-        return obj
-
-    def _load(self, path: PathOrBuffer, fileid: str | None = None):
-        """Load mmCIF data from file."""
-        self.raw = pdbx.CIFFile.read(path)
-
-        self._parse_structure()
-        self._parse_header()
-        self._parse_entities()
-        self._parse_sequences()
-
-    def _parse_structure(self):
-        """Parse the atomic structure from mmCIF."""
-        try:
-            structure = pdbx.get_structure(self.raw, model=1)
-            if structure is None or not isinstance(structure, bs.AtomArray):
-                raise NoProteinError("No structure found in mmCIF file")
-            if len(structure) == 0:
-                raise NoProteinError("Empty structure in mmCIF file")
-            self.structure = structure
-        except Exception as e:
-            raise ValueError(f"Failed to parse structure: {e}")
-
-    def _parse_header(self):
-        """Parse header information from mmCIF."""
-        if not self.raw:
-            return
-
-        try:
-            # Get the first (and usually only) block
-            block = self.raw.block
-
-            # Parse release date
-            if "pdbx_database_status" in block:
-                status_cat = block["pdbx_database_status"]
-                if "recvd_initial_deposition_date" in status_cat:
-                    date_str = status_cat["recvd_initial_deposition_date"].as_item()
-                    if date_str and date_str != "?":
-                        try:
-                            self.header.release_date = datetime.strptime(
-                                date_str, "%Y-%m-%d"
-                            )
-                        except ValueError:
-                            pass
-
-            # Parse resolution
-            if "refine" in block:
-                refine_cat = block["refine"]
-                if "ls_d_res_high" in refine_cat:
-                    res_str = refine_cat["ls_d_res_high"].as_item()
-                    if res_str and res_str != "?":
-                        try:
-                            self.header.resolution = float(res_str)
-                        except ValueError:
-                            pass
-
-            # Parse structure method
-            if "exptl" in block:
-                exptl_cat = block["exptl"]
-                if "method" in exptl_cat:
-                    method = exptl_cat["method"].as_item()
-                    if method and method != "?":
-                        self.header.structure_method = method.upper()
-
-        except Exception:
-            # If parsing fails, keep default values
-            pass
-
-    def _parse_entities(self):
-        """Parse entity information and map to chains."""
-        if not self.raw:
-            return
-
-        try:
-            block = self.raw.block
-
-            # Parse entity information
-            if "entity" in block:
-                entity_cat = block["entity"]
-                entity_ids = entity_cat["id"].as_array(str)
-                entity_types = entity_cat["type"].as_array(str)
-
-                # Initialize entities dict with all entities (not just polymers)
-                for i, (entity_id, entity_type) in enumerate(
-                    zip(entity_ids, entity_types)
-                ):
-                    self.entities[int(entity_id)] = []
-
-            # Map polymer chains to entities using entity_poly
-            if "entity_poly" in block:
-                poly_cat = block["entity_poly"]
-                entity_ids = poly_cat["entity_id"].as_array(str)
-                chain_lists = poly_cat["pdbx_strand_id"].as_array(str)
-
-                for entity_id, chain_list in zip(entity_ids, chain_lists):
-                    entity_id = int(entity_id)
-                    # Chain list is comma-separated
-                    chains = [c.strip() for c in chain_list.split(",") if c.strip()]
-                    if entity_id in self.entities:
-                        self.entities[entity_id] = chains
-
-            # Map non-polymer chains using struct_asym for entities not covered by entity_poly
-            if "struct_asym" in block:
-                asym_cat = block["struct_asym"]
-                asym_ids = asym_cat["id"].as_array(str)
-                entity_ids = asym_cat["entity_id"].as_array(str)
-
-                for asym_id, entity_id in zip(asym_ids, entity_ids):
-                    entity_id = int(entity_id)
-                    # Only add if entity exists but has no chains yet (non-polymer entities)
-                    if entity_id in self.entities and not self.entities[entity_id]:
-                        self.entities[entity_id].append(asym_id)
-
-        except Exception:
-            # If parsing fails, try to infer from structure
-            if (
-                self.structure
-                and hasattr(self.structure, "chain_id")
-                and self.structure.chain_id is not None
-                and hasattr(self.structure.chain_id, "__iter__")
-            ):
-                chain_ids = list(set(self.structure.chain_id))
-                self.entities = {1: chain_ids}
-
-    def _parse_sequences(self):
-        """Parse sequence information from mmCIF."""
-        if not self.raw:
-            return
-
-        block = self.raw.block
-
-        # Parse polymer sequences
-        if "entity_poly" in block:
-            poly_cat = block["entity_poly"]
-            entity_ids = poly_cat["entity_id"].as_array(str)
-            sequences = poly_cat["pdbx_seq_one_letter_code_can"].as_array(str)
-            chain_lists = poly_cat["pdbx_strand_id"].as_array(str)
-
-            for entity_id, sequence, chain_list in zip(
-                entity_ids, sequences, chain_lists
-            ):
-                # Clean up sequence (remove whitespace and newlines)
-                clean_seq = "".join(sequence.split())
-                chains = [c.strip() for c in chain_list.split(",") if c.strip()]
-
-                for chain_id in chains:
-                    self.chain_to_seqres[chain_id] = clean_seq
-
-        # Parse sequence to structure mapping
-        if "pdbx_poly_seq_scheme" in block:
-            seq_cat = block["pdbx_poly_seq_scheme"]
-            asym_ids = seq_cat["asym_id"].as_array(str)  # Internal chain IDs
-            seq_positions = seq_cat["seq_id"].as_array(str)
-            auth_seq_nums = seq_cat["auth_seq_num"].as_array(str)
-            ins_codes = (
-                seq_cat["pdb_ins_code"].as_array(str)
-                if "pdb_ins_code" in seq_cat
-                else [""] * len(asym_ids)
-            )
-            hetflags = (
-                seq_cat["hetflag"].as_array(str)
-                if "hetflag" in seq_cat
-                else ["N"] * len(asym_ids)
-            )
-
-            # Get author chain IDs if available
-            auth_chain_ids = (
-                seq_cat["pdb_strand_id"].as_array(str)
-                if "pdb_strand_id" in seq_cat
-                else asym_ids  # Fallback to internal IDs
-            )
-
-            # Build mapping from internal chain ID to author chain ID
-            asym_to_auth_mapping = {}
-            for asym_id, auth_id in zip(asym_ids, auth_chain_ids):
-                asym_to_auth_mapping[asym_id] = auth_id
-
-            # Group by internal chain ID first, then map to author chain ID
-            chain_data = {}
-            for asym_id, seq_pos, auth_seq, ins_code, hetflag in zip(
-                asym_ids, seq_positions, auth_seq_nums, ins_codes, hetflags
-            ):
-                if asym_id not in chain_data:
-                    chain_data[asym_id] = {}
-
-                try:
-                    seq_index = int(seq_pos) - 1  # Convert to 0-based indexing
-                    res_num = int(auth_seq) if auth_seq != "?" else None
-                except ValueError:
-                    continue
-
-                if res_num is not None:
-                    # Convert mmCIF "." and "?" to empty string
-                    clean_ins_code = "" if ins_code in [".", "?"] else ins_code
-                else:
-                    clean_ins_code = ""
-                    res_num = None
-
-                is_het = hetflag.upper() == "Y"  # type: ignore
-                chain_data[asym_id][seq_index] = Residue(
-                    residue_number=res_num,
-                    insertion_code=clean_ins_code,  # type: ignore
-                    hetflag=is_het,
-                )
-
-            # Handle cases where multiple residues have the same auth_seq_num
-            # by adjusting residue numbers to be unique within each chain
-            for asym_id, residue_data in chain_data.items():
-                # Check if there are duplicate residue numbers in this chain
-                positions_with_same_num = {}
-                for seq_idx, res_at_pos in residue_data.items():
-                    if res_at_pos.residue_number is not None:
-                        res_num = res_at_pos.residue_number
-                        if res_num not in positions_with_same_num:
-                            positions_with_same_num[res_num] = []
-                        positions_with_same_num[res_num].append(seq_idx)
-
-                # Fix duplicate residue numbers by making them sequential
-                for res_num, seq_indices in positions_with_same_num.items():
-                    if len(seq_indices) > 1:
-                        # Multiple residues have the same residue number
-                        # Make them sequential starting from the original number
-                        seq_indices.sort()  # Ensure consistent ordering
-                        for i, seq_idx in enumerate(seq_indices):
-                            original_pos = residue_data[seq_idx]
-                            new_pos = Residue(
-                                residue_number=res_num + i,
-                                insertion_code=original_pos.insertion_code,
-                                hetflag=original_pos.hetflag,
-                            )
-                            residue_data[seq_idx] = new_pos
-
-            # Create ordered mappings using author chain IDs
-            for asym_id in chain_data:
-                auth_chain_id = asym_to_auth_mapping.get(asym_id, asym_id)
-                if auth_chain_id in self.chain_to_seqres:
-                    seq_len = len(self.chain_to_seqres[auth_chain_id])
-                    ordered_mapping = {}
-
-                    for i in range(seq_len):
-                        if i in chain_data[asym_id]:
-                            ordered_mapping[i] = chain_data[asym_id][i]
-                        else:
-                            # Missing residue - no structure coordinates
-                            ordered_mapping[i] = Residue(
-                                residue_number=None, insertion_code="", hetflag=False
-                            )
-
-                    self.seqres_to_structure[auth_chain_id] = ordered_mapping
-                else:
-                    # Handle case where auth_chain_id is not in chain_to_seqres
-                    # This can happen if the chain is not a polymer or if there's a parsing issue
-                    # Create a basic mapping based on the chain_data
-                    if chain_data[asym_id]:
-                        # Sort by sequence index to create ordered mapping
-                        sorted_indices = sorted(chain_data[asym_id].keys())
-                        ordered_mapping = {}
-                        for i, seq_idx in enumerate(sorted_indices):
-                            ordered_mapping[i] = chain_data[asym_id][seq_idx]
-                        self.seqres_to_structure[auth_chain_id] = ordered_mapping
-
-        # Ensure all chains have complete mappings
-        for chain_id in self.chain_to_seqres:
-            if chain_id not in self.seqres_to_structure:
-                seq_len = len(self.chain_to_seqres[chain_id])
-                self.seqres_to_structure[chain_id] = {
-                    i: Residue(residue_number=None, insertion_code="", hetflag=False)
-                    for i in range(seq_len)
-                }
-            else:
-                # Fill in any missing indices
-                seq_len = len(self.chain_to_seqres[chain_id])
-                mapping = self.seqres_to_structure[chain_id]
-                for i in range(seq_len):
-                    if i not in mapping:
-                        mapping[i] = Residue(
-                            residue_number=None, insertion_code="", hetflag=False
-                        )
-
-        # Fallback: create basic mappings from structure for missing chains
-        if (
-            self.structure
-            and hasattr(self.structure, "chain_id")
-            and self.structure.chain_id is not None
-            and hasattr(self.structure.chain_id, "__iter__")
-        ):
-            for chain_id in set(self.structure.chain_id):
-                if chain_id not in self.seqres_to_structure:
-                    chain_structure = self.structure[
-                        self.structure.chain_id == chain_id
-                    ]
-                    if (
-                        hasattr(chain_structure, "res_id")
-                        and chain_structure.res_id is not None
-                        and hasattr(chain_structure.res_id, "__iter__")
-                    ):
-                        residue_ids = list(set(chain_structure.res_id))
-                        residue_ids.sort()
-
-                        self.seqres_to_structure[chain_id] = {
-                            i: Residue(
-                                residue_number=res_id, insertion_code="", hetflag=False
-                            )
-                            for i, res_id in enumerate(residue_ids)
-                        }
-
-    def _parse_nonpoly_from_mmcif(self) -> dict[tuple, bs.AtomArray]:
-        """Parse non-polymer coordinates from mmCIF block data."""
-        nonpoly_coords = {}
-
-        # Get non-polymer entities from the mmCIF block
-        assert self.raw is not None
-        block = self.raw.block
-        nonpoly_entities = set()
-
-        # Find non-polymer entities
-        if "entity" in block:
-            entity_cat = block["entity"]
-            entity_ids = entity_cat["id"].as_array(str)
-            entity_types = entity_cat["type"].as_array(str)
-
-            for entity_id, entity_type in zip(entity_ids, entity_types):
-                if entity_type.upper() in ["NON-POLYMER", "WATER", "BRANCHED"]:
-                    nonpoly_entities.add(entity_id)
-
-        # Map entities to chains for non-polymers
-        entity_to_chains = {}
-        if "pdbx_entity_nonpoly" in block:
-            nonpoly_cat = block["pdbx_entity_nonpoly"]
-            entity_ids = nonpoly_cat["entity_id"].as_array(str)
-            comp_ids = nonpoly_cat["comp_id"].as_array(str)
-
-            for entity_id, comp_id in zip(entity_ids, comp_ids):
-                if entity_id in nonpoly_entities:
-                    entity_to_chains[entity_id] = comp_id
-
-        # Get atom site information for non-polymers
-        if "atom_site" in block:
-            atom_cat = block["atom_site"]
-            atom_chain_ids = atom_cat["label_asym_id"].as_array(str)
-            atom_entity_ids = atom_cat["label_entity_id"].as_array(str)
-            atom_comp_ids = atom_cat["label_comp_id"].as_array(str)
-
-            # Group non-polymer atoms by entity and chain
-            nonpoly_atom_groups = {}
-            for i, (chain_id, entity_id, comp_id) in enumerate(
-                zip(atom_chain_ids, atom_entity_ids, atom_comp_ids)
-            ):
-                if entity_id in nonpoly_entities:
-                    key = (comp_id, chain_id)
-                    if key not in nonpoly_atom_groups:
-                        nonpoly_atom_groups[key] = []
-                    nonpoly_atom_groups[key].append(i)
-
-            # Extract coordinates for each non-polymer group
-            for (comp_id, chain_id), atom_indices in nonpoly_atom_groups.items():
-                # Match atoms by comparing chain_id and residue name
-                structure_mask = (self.structure.chain_id == chain_id) & (
-                    self.structure.res_name == comp_id
-                )
-
-                if structure_mask.any():
-                    nonpoly_array = self.structure[structure_mask]
-                    if (
-                        isinstance(nonpoly_array, (bs.AtomArray, bs.AtomArrayStack))
-                        and len(nonpoly_array) > 0
-                    ):
-                        nonpoly_coords[(comp_id, chain_id)] = nonpoly_array
-
-        return nonpoly_coords
-
-    def _parse_nonpoly_fallback(self) -> dict[tuple, bs.AtomArray]:
-        """Fallback method to extract heteroatoms directly from structure."""
-        nonpoly_coords = {}
-
-        if not (self.structure and hasattr(self.structure, "chain_id")):
-            return nonpoly_coords
-
-        # Create set of standard residues from residue_constants
-        standard_residues = set(residue_constants.resnames[:-1])  # Exclude 'UNK'
-        standard_residues.update({"A", "C", "G", "T", "U"})  # Add nucleic acids
-
-        if hasattr(self.structure, "chain_id") and self.structure.chain_id is not None:
-            for chain_id in set(self.structure.chain_id):
-                chain_structure = self.structure[self.structure.chain_id == chain_id]
-
-                # Find non-standard residues
-                if (
-                    hasattr(chain_structure, "res_name")
-                    and chain_structure.res_name is not None
-                    and hasattr(chain_structure.res_name, "__iter__")
-                ):
-                    for res_name in set(chain_structure.res_name):
-                        if res_name not in standard_residues:
-                            res_mask = (chain_structure.chain_id == chain_id) & (
-                                chain_structure.res_name == res_name
-                            )
-                            if res_mask.any() and isinstance(
-                                chain_structure, (bs.AtomArray, bs.AtomArrayStack)
-                            ):
-                                nonpoly_array = chain_structure[res_mask]
-                                nonpoly_coords[(res_name, chain_id)] = nonpoly_array
-
-        return nonpoly_coords
-
-    @functools.cached_property
-    def non_polymer_coords(self) -> dict[tuple, bs.AtomArray]:
-        """
-        Extract non-polymer coordinates (ligands, cofactors, etc.) from mmCIF structure.
-
-        Returns a dictionary mapping (nonpolymer_info, chain_id) tuples to AtomArrays.
-        """
-        if not self.structure or not self.raw:
-            return {}
-
-        try:
-            return self._parse_nonpoly_from_mmcif()
-        except Exception:
-            return self._parse_nonpoly_fallback()
+from __future__ import annotations
+
+import functools
+import io
+import os
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Union
+
+import biotite.structure as bs
+import biotite.structure.io.pdbx as pdbx
+
+from . import esmfold2_residue_constants as residue_constants
+
+# Define PathOrBuffer for the opensource version
+PathOrBuffer = Union[str, os.PathLike, io.StringIO]
+
+
+class NoProteinError(Exception):
+    pass
+
+
+@dataclass
+class Residue:
+    residue_number: int | None = None
+    insertion_code: str = ""
+    hetflag: bool = False
+
+
+@dataclass
+class MmcifHeader:
+    release_date: datetime | None = None
+    resolution: float | None = None
+    structure_method: str = "UNKNOWN"
+
+
+class MmcifWrapper:
+    def __init__(self, id: str | None = None):
+        self.id: str = id or ""
+        self.raw: pdbx.CIFFile | None = None
+        self.structure: bs.AtomArray
+        self.header: MmcifHeader = MmcifHeader()
+        self.entities: dict[int, list[str]] = {}
+        self.chain_to_seqres: dict[str, str] = {}
+        self.seqres_to_structure: dict[str, dict[int, Residue]] = {}
+
+    @classmethod
+    def read(cls, path: PathOrBuffer, id: str | None = None) -> MmcifWrapper:
+        obj = cls(id=id)
+        obj._load(path)
+        return obj
+
+    def _load(self, path: PathOrBuffer, fileid: str | None = None):
+        """Load mmCIF data from file."""
+        self.raw = pdbx.CIFFile.read(path)
+
+        self._parse_structure()
+        self._parse_header()
+        self._parse_entities()
+        self._parse_sequences()
+
+    def _parse_structure(self):
+        """Parse the atomic structure from mmCIF."""
+        try:
+            structure = pdbx.get_structure(self.raw, model=1)
+            if structure is None or not isinstance(structure, bs.AtomArray):
+                raise NoProteinError("No structure found in mmCIF file")
+            if len(structure) == 0:
+                raise NoProteinError("Empty structure in mmCIF file")
+            self.structure = structure
+        except Exception as e:
+            raise ValueError(f"Failed to parse structure: {e}")
+
+    def _parse_header(self):
+        """Parse header information from mmCIF."""
+        if not self.raw:
+            return
+
+        try:
+            # Get the first (and usually only) block
+            block = self.raw.block
+
+            # Parse release date
+            if "pdbx_database_status" in block:
+                status_cat = block["pdbx_database_status"]
+                if "recvd_initial_deposition_date" in status_cat:
+                    date_str = status_cat["recvd_initial_deposition_date"].as_item()
+                    if date_str and date_str != "?":
+                        try:
+                            self.header.release_date = datetime.strptime(
+                                date_str, "%Y-%m-%d"
+                            )
+                        except ValueError:
+                            pass
+
+            # Parse resolution
+            if "refine" in block:
+                refine_cat = block["refine"]
+                if "ls_d_res_high" in refine_cat:
+                    res_str = refine_cat["ls_d_res_high"].as_item()
+                    if res_str and res_str != "?":
+                        try:
+                            self.header.resolution = float(res_str)
+                        except ValueError:
+                            pass
+
+            # Parse structure method
+            if "exptl" in block:
+                exptl_cat = block["exptl"]
+                if "method" in exptl_cat:
+                    method = exptl_cat["method"].as_item()
+                    if method and method != "?":
+                        self.header.structure_method = method.upper()
+
+        except Exception:
+            # If parsing fails, keep default values
+            pass
+
+    def _parse_entities(self):
+        """Parse entity information and map to chains."""
+        if not self.raw:
+            return
+
+        try:
+            block = self.raw.block
+
+            # Parse entity information
+            if "entity" in block:
+                entity_cat = block["entity"]
+                entity_ids = entity_cat["id"].as_array(str)
+                entity_types = entity_cat["type"].as_array(str)
+
+                # Initialize entities dict with all entities (not just polymers)
+                for i, (entity_id, entity_type) in enumerate(
+                    zip(entity_ids, entity_types)
+                ):
+                    self.entities[int(entity_id)] = []
+
+            # Map polymer chains to entities using entity_poly
+            if "entity_poly" in block:
+                poly_cat = block["entity_poly"]
+                entity_ids = poly_cat["entity_id"].as_array(str)
+                chain_lists = poly_cat["pdbx_strand_id"].as_array(str)
+
+                for entity_id, chain_list in zip(entity_ids, chain_lists):
+                    entity_id = int(entity_id)
+                    # Chain list is comma-separated
+                    chains = [c.strip() for c in chain_list.split(",") if c.strip()]
+                    if entity_id in self.entities:
+                        self.entities[entity_id] = chains
+
+            # Map non-polymer chains using struct_asym for entities not covered by entity_poly
+            if "struct_asym" in block:
+                asym_cat = block["struct_asym"]
+                asym_ids = asym_cat["id"].as_array(str)
+                entity_ids = asym_cat["entity_id"].as_array(str)
+
+                for asym_id, entity_id in zip(asym_ids, entity_ids):
+                    entity_id = int(entity_id)
+                    # Only add if entity exists but has no chains yet (non-polymer entities)
+                    if entity_id in self.entities and not self.entities[entity_id]:
+                        self.entities[entity_id].append(asym_id)
+
+        except Exception:
+            # If parsing fails, try to infer from structure
+            if (
+                self.structure
+                and hasattr(self.structure, "chain_id")
+                and self.structure.chain_id is not None
+                and hasattr(self.structure.chain_id, "__iter__")
+            ):
+                chain_ids = list(set(self.structure.chain_id))
+                self.entities = {1: chain_ids}
+
+    def _parse_sequences(self):
+        """Parse sequence information from mmCIF."""
+        if not self.raw:
+            return
+
+        block = self.raw.block
+
+        # Parse polymer sequences
+        if "entity_poly" in block:
+            poly_cat = block["entity_poly"]
+            entity_ids = poly_cat["entity_id"].as_array(str)
+            sequences = poly_cat["pdbx_seq_one_letter_code_can"].as_array(str)
+            chain_lists = poly_cat["pdbx_strand_id"].as_array(str)
+
+            for entity_id, sequence, chain_list in zip(
+                entity_ids, sequences, chain_lists
+            ):
+                # Clean up sequence (remove whitespace and newlines)
+                clean_seq = "".join(sequence.split())
+                chains = [c.strip() for c in chain_list.split(",") if c.strip()]
+
+                for chain_id in chains:
+                    self.chain_to_seqres[chain_id] = clean_seq
+
+        # Parse sequence to structure mapping
+        if "pdbx_poly_seq_scheme" in block:
+            seq_cat = block["pdbx_poly_seq_scheme"]
+            asym_ids = seq_cat["asym_id"].as_array(str)  # Internal chain IDs
+            seq_positions = seq_cat["seq_id"].as_array(str)
+            auth_seq_nums = seq_cat["auth_seq_num"].as_array(str)
+            ins_codes = (
+                seq_cat["pdb_ins_code"].as_array(str)
+                if "pdb_ins_code" in seq_cat
+                else [""] * len(asym_ids)
+            )
+            hetflags = (
+                seq_cat["hetflag"].as_array(str)
+                if "hetflag" in seq_cat
+                else ["N"] * len(asym_ids)
+            )
+
+            # Get author chain IDs if available
+            auth_chain_ids = (
+                seq_cat["pdb_strand_id"].as_array(str)
+                if "pdb_strand_id" in seq_cat
+                else asym_ids  # Fallback to internal IDs
+            )
+
+            # Build mapping from internal chain ID to author chain ID
+            asym_to_auth_mapping = {}
+            for asym_id, auth_id in zip(asym_ids, auth_chain_ids):
+                asym_to_auth_mapping[asym_id] = auth_id
+
+            # Group by internal chain ID first, then map to author chain ID
+            chain_data = {}
+            for asym_id, seq_pos, auth_seq, ins_code, hetflag in zip(
+                asym_ids, seq_positions, auth_seq_nums, ins_codes, hetflags
+            ):
+                if asym_id not in chain_data:
+                    chain_data[asym_id] = {}
+
+                try:
+                    seq_index = int(seq_pos) - 1  # Convert to 0-based indexing
+                    res_num = int(auth_seq) if auth_seq != "?" else None
+                except ValueError:
+                    continue
+
+                if res_num is not None:
+                    # Convert mmCIF "." and "?" to empty string
+                    clean_ins_code = "" if ins_code in [".", "?"] else ins_code
+                else:
+                    clean_ins_code = ""
+                    res_num = None
+
+                is_het = hetflag.upper() == "Y"  # type: ignore
+                chain_data[asym_id][seq_index] = Residue(
+                    residue_number=res_num,
+                    insertion_code=clean_ins_code,  # type: ignore
+                    hetflag=is_het,
+                )
+
+            # Handle cases where multiple residues have the same auth_seq_num
+            # by adjusting residue numbers to be unique within each chain
+            for asym_id, residue_data in chain_data.items():
+                # Check if there are duplicate residue numbers in this chain
+                positions_with_same_num = {}
+                for seq_idx, res_at_pos in residue_data.items():
+                    if res_at_pos.residue_number is not None:
+                        res_num = res_at_pos.residue_number
+                        if res_num not in positions_with_same_num:
+                            positions_with_same_num[res_num] = []
+                        positions_with_same_num[res_num].append(seq_idx)
+
+                # Fix duplicate residue numbers by making them sequential
+                for res_num, seq_indices in positions_with_same_num.items():
+                    if len(seq_indices) > 1:
+                        # Multiple residues have the same residue number
+                        # Make them sequential starting from the original number
+                        seq_indices.sort()  # Ensure consistent ordering
+                        for i, seq_idx in enumerate(seq_indices):
+                            original_pos = residue_data[seq_idx]
+                            new_pos = Residue(
+                                residue_number=res_num + i,
+                                insertion_code=original_pos.insertion_code,
+                                hetflag=original_pos.hetflag,
+                            )
+                            residue_data[seq_idx] = new_pos
+
+            # Create ordered mappings using author chain IDs
+            for asym_id in chain_data:
+                auth_chain_id = asym_to_auth_mapping.get(asym_id, asym_id)
+                if auth_chain_id in self.chain_to_seqres:
+                    seq_len = len(self.chain_to_seqres[auth_chain_id])
+                    ordered_mapping = {}
+
+                    for i in range(seq_len):
+                        if i in chain_data[asym_id]:
+                            ordered_mapping[i] = chain_data[asym_id][i]
+                        else:
+                            # Missing residue - no structure coordinates
+                            ordered_mapping[i] = Residue(
+                                residue_number=None, insertion_code="", hetflag=False
+                            )
+
+                    self.seqres_to_structure[auth_chain_id] = ordered_mapping
+                else:
+                    # Handle case where auth_chain_id is not in chain_to_seqres
+                    # This can happen if the chain is not a polymer or if there's a parsing issue
+                    # Create a basic mapping based on the chain_data
+                    if chain_data[asym_id]:
+                        # Sort by sequence index to create ordered mapping
+                        sorted_indices = sorted(chain_data[asym_id].keys())
+                        ordered_mapping = {}
+                        for i, seq_idx in enumerate(sorted_indices):
+                            ordered_mapping[i] = chain_data[asym_id][seq_idx]
+                        self.seqres_to_structure[auth_chain_id] = ordered_mapping
+
+        # Ensure all chains have complete mappings
+        for chain_id in self.chain_to_seqres:
+            if chain_id not in self.seqres_to_structure:
+                seq_len = len(self.chain_to_seqres[chain_id])
+                self.seqres_to_structure[chain_id] = {
+                    i: Residue(residue_number=None, insertion_code="", hetflag=False)
+                    for i in range(seq_len)
+                }
+            else:
+                # Fill in any missing indices
+                seq_len = len(self.chain_to_seqres[chain_id])
+                mapping = self.seqres_to_structure[chain_id]
+                for i in range(seq_len):
+                    if i not in mapping:
+                        mapping[i] = Residue(
+                            residue_number=None, insertion_code="", hetflag=False
+                        )
+
+        # Fallback: create basic mappings from structure for missing chains
+        if (
+            self.structure
+            and hasattr(self.structure, "chain_id")
+            and self.structure.chain_id is not None
+            and hasattr(self.structure.chain_id, "__iter__")
+        ):
+            for chain_id in set(self.structure.chain_id):
+                if chain_id not in self.seqres_to_structure:
+                    chain_structure = self.structure[
+                        self.structure.chain_id == chain_id
+                    ]
+                    if (
+                        hasattr(chain_structure, "res_id")
+                        and chain_structure.res_id is not None
+                        and hasattr(chain_structure.res_id, "__iter__")
+                    ):
+                        residue_ids = list(set(chain_structure.res_id))
+                        residue_ids.sort()
+
+                        self.seqres_to_structure[chain_id] = {
+                            i: Residue(
+                                residue_number=res_id, insertion_code="", hetflag=False
+                            )
+                            for i, res_id in enumerate(residue_ids)
+                        }
+
+    def _parse_nonpoly_from_mmcif(self) -> dict[tuple, bs.AtomArray]:
+        """Parse non-polymer coordinates from mmCIF block data."""
+        nonpoly_coords = {}
+
+        # Get non-polymer entities from the mmCIF block
+        assert self.raw is not None
+        block = self.raw.block
+        nonpoly_entities = set()
+
+        # Find non-polymer entities
+        if "entity" in block:
+            entity_cat = block["entity"]
+            entity_ids = entity_cat["id"].as_array(str)
+            entity_types = entity_cat["type"].as_array(str)
+
+            for entity_id, entity_type in zip(entity_ids, entity_types):
+                if entity_type.upper() in ["NON-POLYMER", "WATER", "BRANCHED"]:
+                    nonpoly_entities.add(entity_id)
+
+        # Map entities to chains for non-polymers
+        entity_to_chains = {}
+        if "pdbx_entity_nonpoly" in block:
+            nonpoly_cat = block["pdbx_entity_nonpoly"]
+            entity_ids = nonpoly_cat["entity_id"].as_array(str)
+            comp_ids = nonpoly_cat["comp_id"].as_array(str)
+
+            for entity_id, comp_id in zip(entity_ids, comp_ids):
+                if entity_id in nonpoly_entities:
+                    entity_to_chains[entity_id] = comp_id
+
+        # Get atom site information for non-polymers
+        if "atom_site" in block:
+            atom_cat = block["atom_site"]
+            atom_chain_ids = atom_cat["label_asym_id"].as_array(str)
+            atom_entity_ids = atom_cat["label_entity_id"].as_array(str)
+            atom_comp_ids = atom_cat["label_comp_id"].as_array(str)
+
+            # Group non-polymer atoms by entity and chain
+            nonpoly_atom_groups = {}
+            for i, (chain_id, entity_id, comp_id) in enumerate(
+                zip(atom_chain_ids, atom_entity_ids, atom_comp_ids)
+            ):
+                if entity_id in nonpoly_entities:
+                    key = (comp_id, chain_id)
+                    if key not in nonpoly_atom_groups:
+                        nonpoly_atom_groups[key] = []
+                    nonpoly_atom_groups[key].append(i)
+
+            # Extract coordinates for each non-polymer group
+            for (comp_id, chain_id), atom_indices in nonpoly_atom_groups.items():
+                # Match atoms by comparing chain_id and residue name
+                structure_mask = (self.structure.chain_id == chain_id) & (
+                    self.structure.res_name == comp_id
+                )
+
+                if structure_mask.any():
+                    nonpoly_array = self.structure[structure_mask]
+                    if (
+                        isinstance(nonpoly_array, (bs.AtomArray, bs.AtomArrayStack))
+                        and len(nonpoly_array) > 0
+                    ):
+                        nonpoly_coords[(comp_id, chain_id)] = nonpoly_array
+
+        return nonpoly_coords
+
+    def _parse_nonpoly_fallback(self) -> dict[tuple, bs.AtomArray]:
+        """Fallback method to extract heteroatoms directly from structure."""
+        nonpoly_coords = {}
+
+        if not (self.structure and hasattr(self.structure, "chain_id")):
+            return nonpoly_coords
+
+        # Create set of standard residues from residue_constants
+        standard_residues = set(residue_constants.resnames[:-1])  # Exclude 'UNK'
+        standard_residues.update({"A", "C", "G", "T", "U"})  # Add nucleic acids
+
+        if hasattr(self.structure, "chain_id") and self.structure.chain_id is not None:
+            for chain_id in set(self.structure.chain_id):
+                chain_structure = self.structure[self.structure.chain_id == chain_id]
+
+                # Find non-standard residues
+                if (
+                    hasattr(chain_structure, "res_name")
+                    and chain_structure.res_name is not None
+                    and hasattr(chain_structure.res_name, "__iter__")
+                ):
+                    for res_name in set(chain_structure.res_name):
+                        if res_name not in standard_residues:
+                            res_mask = (chain_structure.chain_id == chain_id) & (
+                                chain_structure.res_name == res_name
+                            )
+                            if res_mask.any() and isinstance(
+                                chain_structure, (bs.AtomArray, bs.AtomArrayStack)
+                            ):
+                                nonpoly_array = chain_structure[res_mask]
+                                nonpoly_coords[(res_name, chain_id)] = nonpoly_array
+
+        return nonpoly_coords
+
+    @functools.cached_property
+    def non_polymer_coords(self) -> dict[tuple, bs.AtomArray]:
+        """
+        Extract non-polymer coordinates (ligands, cofactors, etc.) from mmCIF structure.
+
+        Returns a dictionary mapping (nonpolymer_info, chain_id) tuples to AtomArrays.
+        """
+        if not self.structure or not self.raw:
+            return {}
+
+        try:
+            return self._parse_nonpoly_from_mmcif()
+        except Exception:
+            return self._parse_nonpoly_fallback()

esmfold2_msa.py

@@ -1,506 +1,506 @@
-from __future__ import annotations
-
-import dataclasses
-import string
-from dataclasses import dataclass
-from functools import cached_property
-from itertools import islice
-from typing import Sequence
-
-import numpy as np
-from Bio import SeqIO
-from scipy.spatial.distance import cdist
-
-from .esmfold2_misc import slice_any_object
-from .esmfold2_msa_filter_sequences import greedy_select_indices, hhfilter
-from .esmfold2_parsing import FastaEntry, read_sequences, write_sequences
-from .esmfold2_sequential_dataclass import SequentialDataclass
-from .esmfold2_system import PathOrBuffer
-
-REMOVE_LOWERCASE_TRANSLATION = str.maketrans(dict.fromkeys(string.ascii_lowercase))
-
-
-def remove_insertions_from_sequence(seq: str) -> str:
-    return seq.translate(REMOVE_LOWERCASE_TRANSLATION)
-
-
-@dataclass(frozen=True)
-class MSA(SequentialDataclass):
-    """Object-oriented interface to an MSA.
-
-    Args:
-        sequences (list[str]): List of protein sequences
-        headers (list[str]): List of headers describing the sequences
-
-    """
-
-    entries: list[FastaEntry]
-
-    @cached_property
-    def sequences(self) -> list[str]:
-        return [entry.sequence for entry in self.entries]
-
-    @cached_property
-    def headers(self) -> list[str]:
-        return [entry.header for entry in self.entries]
-
-    def __repr__(self):
-        return (
-            f"MSA({self.entries[0].header}: Depth={self.depth}, Length={self.seqlen})"
-        )
-
-    def to_fast_msa(self) -> FastMSA:
-        return FastMSA(self.array, self.headers)
-
-    @classmethod
-    def from_a3m(
-        cls,
-        path: PathOrBuffer,
-        remove_insertions: bool = True,
-        max_sequences: int | None = None,
-    ) -> MSA:
-        entries = []
-        for header, seq in islice(read_sequences(path), max_sequences):
-            if remove_insertions:
-                seq = remove_insertions_from_sequence(seq)
-            if entries:
-                assert (
-                    len(seq) == len(entries[0].sequence)
-                ), f"Sequence length mismatch. Expected: {len(entries[0].sequence)}, Received: {len(seq)}"
-            entries.append(FastaEntry(header, seq))
-        return cls(entries)
-
-    def to_a3m(self, path: PathOrBuffer) -> None:
-        write_sequences(self.entries, path)
-
-    @classmethod
-    def from_stockholm(
-        cls,
-        path: PathOrBuffer,
-        remove_insertions: bool = True,
-        max_sequences: int | None = None,
-    ) -> MSA:
-        entries = []
-        for record in islice(SeqIO.parse(path, "stockholm"), max_sequences):
-            header = f"{record.id} {record.description}"
-            seq = str(record.seq)
-            if entries:
-                assert (
-                    len(seq) == len(entries[0].sequence)
-                ), f"Sequence length mismatch. Expected: {len(entries[0].sequence)}, Received: {len(seq)}"
-            entries.append(FastaEntry(header, seq))
-        msa = cls(entries)
-        if remove_insertions:
-            keep_inds = [i for i, aa in enumerate(msa.query) if aa != "-"]
-            msa = msa.select_positions(keep_inds)
-        return msa
-
-    def to_bytes(self) -> bytes:
-        version = 1
-        version_bytes = version.to_bytes(1, "little")
-        seqlen_bytes = self.seqlen.to_bytes(4, "little")
-        depth_bytes = self.depth.to_bytes(4, "little")
-        array_bytes = self.array.tobytes()
-        header_bytes = "\n".join(entry.header for entry in self.entries).encode()
-        all_bytes = (
-            version_bytes + seqlen_bytes + depth_bytes + array_bytes + header_bytes
-        )
-        return all_bytes
-
-    @classmethod
-    def from_bytes(cls, data: bytes) -> MSA:
-        version_bytes, seqlen_bytes, depth_bytes, data = (
-            data[:1],
-            data[1:5],
-            data[5:9],
-            data[9:],
-        )
-        version = int.from_bytes(version_bytes, "little")
-        if version != 1:
-            raise ValueError(f"Unsupported version: {version}")
-        seqlen = int.from_bytes(seqlen_bytes, "little")
-        depth = int.from_bytes(depth_bytes, "little")
-        array_bytes, header_bytes = data[: seqlen * depth], data[seqlen * depth :]
-        array = np.frombuffer(array_bytes, dtype="|S1")
-        array = array.reshape(depth, seqlen)
-        headers = header_bytes.decode().split("\n")
-        # Sometimes the separation is two newlines, which results in an empty header.
-        headers = [header for header in headers if header]
-        # If all headers were empty (e.g., saved from from_sequences), use empty headers
-        if len(headers) == 0 and depth > 0:
-            headers = [""] * depth
-        entries = [
-            FastaEntry(header, b"".join(row).decode())
-            for header, row in zip(headers, array)
-        ]
-        return cls(entries)
-
-    # TODO(jmaccarl): set remove_insertions to True by default here to match other utils
-    @classmethod
-    def from_sequences(
-        cls, sequences: list[str], remove_insertions: bool = False
-    ) -> MSA:
-        if remove_insertions:
-            entries = [
-                FastaEntry("", remove_insertions_from_sequence(seq))
-                for seq in sequences
-            ]
-        else:
-            entries = [FastaEntry("", seq) for seq in sequences]
-        return cls(entries)
-
-    def to_sequence_bytes(self) -> bytes:
-        """Stores ONLY SEQUENCES in array format as bytes. Header information will be lost."""
-        seqlen_bytes = self.seqlen.to_bytes(4, "little")
-        array_bytes = self.array.tobytes()
-        all_bytes = seqlen_bytes + array_bytes
-        return all_bytes
-
-    @classmethod
-    def from_sequence_bytes(cls, data: bytes) -> MSA:
-        seqlen_bytes, array_bytes = data[:4], data[4:]
-        seqlen = int.from_bytes(seqlen_bytes, "little")
-        array = np.frombuffer(array_bytes, dtype="|S1")
-        array = array.reshape(-1, seqlen)
-        entries = [FastaEntry("", b"".join(row).decode()) for row in array]
-        return cls(entries)
-
-    @property
-    def depth(self) -> int:
-        return len(self.entries)
-
-    @property
-    def seqlen(self) -> int:
-        return len(self.entries[0].sequence)
-
-    @cached_property
-    def array(self) -> np.ndarray:
-        return np.array([list(seq) for seq in self.sequences], dtype="|S1")
-
-    @property
-    def query(self) -> str:
-        return self.entries[0].sequence
-
-    def select_sequences(self, indices: Sequence[int] | np.ndarray) -> MSA:
-        """Subselect rows of the MSA."""
-        entries = [self.entries[idx] for idx in indices]
-        return dataclasses.replace(self, entries=entries)
-
-    def select_positions(self, indices: Sequence[int] | np.ndarray) -> MSA:
-        """Subselect columns of the MSA."""
-        entries = [
-            FastaEntry(header, "".join(seq[idx] for idx in indices))
-            for header, seq in self.entries
-        ]
-        return dataclasses.replace(self, entries=entries)
-
-    def __getitem__(self, indices: int | list[int] | slice | np.ndarray):
-        if isinstance(indices, int):
-            indices = [indices]
-
-        entries = [
-            FastaEntry(header, slice_any_object(seq, indices))
-            for header, seq in self.entries
-        ]
-        return dataclasses.replace(self, entries=entries)
-
-    def __len__(self):
-        return self.seqlen
-
-    def greedy_select(self, num_seqs: int, mode: str = "max") -> MSA:
-        """Greedily select sequences that either maximize or minimize hamming distance.
-
-        Algorithm proposed in the MSA Transformer paper. Starting from the query sequence,
-        iteratively add sequences to the list with the maximum (minimum) average Hamming
-        distance to the existing set of sequences.
-
-        Args:
-            num_seqs (int): Number of sequences to select.
-            mode (str): Whether to maximize or minimize diversity. DO NOT pick 'min' unless
-                you're doing it to prove a point for a paper.
-
-        Returns:
-            MSA object w/ subselected sequences.
-        """
-        assert mode in ("max", "min")
-        if self.depth <= num_seqs:
-            return self
-
-        indices = greedy_select_indices(self.array, num_seqs, mode)
-        return self.select_sequences(indices)
-
-    def hhfilter(
-        self,
-        seqid: int = 90,
-        diff: int = 0,
-        cov: int = 0,
-        qid: int = 0,
-        qsc: float = -20.0,
-        binary: str = "hhfilter",
-    ) -> MSA:
-        """Apply hhfilter to the sequences in the MSA and return a filtered MSA."""
-
-        indices = hhfilter(
-            self.sequences,
-            seqid=seqid,
-            diff=diff,
-            cov=cov,
-            qid=qid,
-            qsc=qsc,
-            binary=binary,
-        )
-        return self.select_sequences(indices)
-
-    def select_random_sequences(self, num_seqs: int) -> MSA:
-        """Uses random sampling to subselect sequences from the MSA. Always
-        keeps the query sequence.
-        """
-        if num_seqs >= self.depth:
-            return self
-
-        # Subselect random, always keeping the query sequence.
-        indices = np.sort(
-            np.append(
-                0, np.random.choice(self.depth - 1, num_seqs - 1, replace=False) + 1
-            )
-        )
-        msa = self.select_sequences(indices)  # type: ignore
-        return msa
-
-    def select_diverse_sequences(self, num_seqs: int) -> MSA:
-        """Applies hhfilter to select ~num_seqs sequences, then uses random sampling
-        to subselect if necessary.
-        """
-        if num_seqs >= self.depth:
-            return self
-
-        msa = self.hhfilter(diff=num_seqs)
-        if num_seqs < msa.depth:
-            msa = msa.select_random_sequences(num_seqs)
-        return msa
-
-    def pad_to_depth(self, depth: int) -> MSA:
-        if depth < self.depth:
-            raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}")
-        elif depth == self.depth:
-            return self
-
-        num_to_add = depth - self.depth
-        extra_entries = [FastaEntry("", "-" * self.seqlen) for _ in range(num_to_add)]
-        return dataclasses.replace(self, entries=self.entries + extra_entries)
-
-    @classmethod
-    def stack(
-        cls, msas: Sequence[MSA], remove_query_from_later_msas: bool = True
-    ) -> MSA:
-        """Stack a series of MSAs. Optionally remove the query from msas after the first."""
-        all_entries = []
-        for i, msa in enumerate(msas):
-            entries = msa.entries
-            if i > 0 and remove_query_from_later_msas:
-                entries = entries[1:]
-            all_entries.extend(entries)
-        return cls(entries=all_entries)
-
-    @cached_property
-    def seqid(self) -> np.ndarray:
-        array = self.array.view(np.uint8)
-        seqid = 1 - cdist(array[0][None], array, "hamming")
-        return seqid[0]
-
-    @classmethod
-    def concat(
-        cls,
-        msas: Sequence[MSA],
-        join_token: str | None = "|",
-        allow_depth_mismatch: bool = False,
-    ) -> MSA:
-        """Concatenate a series of MSAs horizontally, along the sequence dimension."""
-        if not msas:
-            raise ValueError("Cannot concatenate an empty list of MSAs")
-        msa_depths = [msa.depth for msa in msas]
-        if len(set(msa_depths)) != 1:
-            if not allow_depth_mismatch:
-                raise ValueError("Depth mismatch in concatenating MSAs")
-            else:
-                max_depth = max(msa_depths)
-                msas = [msa.pad_to_depth(max_depth) for msa in msas]
-        headers = [
-            "|".join([str(h) for h in headers])
-            for headers in zip(*(msa.headers for msa in msas))
-        ]
-
-        if join_token is None:
-            join_token = ""
-
-        seqs = [join_token.join(vals) for vals in zip(*(msa.sequences for msa in msas))]
-        entries = [FastaEntry(header, seq) for header, seq in zip(headers, seqs)]
-        return cls(entries)
-
-
-@dataclass(frozen=True)
-class FastMSA(SequentialDataclass):
-    """Object-oriented interface to an MSA stored as a numpy uint8 array."""
-
-    array: np.ndarray
-    headers: list[str] | None = None
-
-    def __post_init__(self):
-        if self.headers is not None:
-            assert (
-                len(self.headers) == self.depth
-            ), "Number of headers must match depth."
-
-    @classmethod
-    def from_bytes(cls, data: bytes) -> FastMSA:
-        version_bytes, seqlen_bytes, depth_bytes, data = (
-            data[:1],
-            data[1:5],
-            data[5:9],
-            data[9:],
-        )
-        version = int.from_bytes(version_bytes, "little")
-        if version != 1:
-            raise ValueError(f"Unsupported version: {version}")
-        seqlen = int.from_bytes(seqlen_bytes, "little")
-        depth = int.from_bytes(depth_bytes, "little")
-        array_bytes, header_bytes = data[: seqlen * depth], data[seqlen * depth :]
-        array = np.frombuffer(array_bytes, dtype="|S1")
-        array = array.reshape(depth, seqlen)
-        headers = header_bytes.decode().split("\n")
-        # Sometimes the separation is two newlines, which results in an empty header.
-        headers = [header for header in headers if header]
-        # If all headers were empty (e.g., saved from from_sequences), use empty headers
-        if len(headers) == 0 and depth > 0:
-            headers = [""] * depth
-        return cls(array, headers)
-
-    @classmethod
-    def from_sequence_bytes(cls, data: bytes) -> FastMSA:
-        seqlen_bytes, array_bytes = data[:4], data[4:]
-        seqlen = int.from_bytes(seqlen_bytes, "little")
-        array = np.frombuffer(array_bytes, dtype="|S1")
-        array = array.reshape(-1, seqlen)
-        return cls(array)
-
-    @property
-    def depth(self) -> int:
-        return self.array.shape[0]
-
-    @property
-    def seqlen(self) -> int:
-        return self.array.shape[1]
-
-    def __len__(self):
-        return self.seqlen
-
-    def __getitem__(self, indices: int | list[int] | slice | np.ndarray):
-        if isinstance(indices, int):
-            indices = [indices]
-
-        return dataclasses.replace(self, array=self.array[:, indices])
-
-    def select_sequences(self, indices: Sequence[int] | np.ndarray) -> FastMSA:
-        """Subselect rows of the MSA."""
-        array = self.array[indices]
-        headers = (
-            [self.headers[idx] for idx in indices] if self.headers is not None else None
-        )
-        return dataclasses.replace(self, array=array, headers=headers)
-
-    def select_random_sequences(self, num_seqs: int) -> FastMSA:
-        """Uses random sampling to subselect sequences from the MSA. Always
-        keeps the query sequence.
-        """
-        if num_seqs >= self.depth:
-            return self
-
-        # Subselect random, always keeping the query sequence.
-        indices = np.sort(
-            np.append(
-                0, np.random.choice(self.depth - 1, num_seqs - 1, replace=False) + 1
-            )
-        )
-        msa = self.select_sequences(indices)  # type: ignore
-        return msa
-
-    def pad_to_depth(self, depth: int) -> FastMSA:
-        if depth < self.depth:
-            raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}")
-        elif depth == self.depth:
-            return self
-
-        num_to_add = depth - self.depth
-        array = np.pad(
-            self.array,
-            [(0, num_to_add), (0, 0)],
-            constant_values=ord("-") if self.array.dtype == np.uint8 else b"-",
-        )
-        headers = self.headers
-        if headers is not None:
-            headers = headers + [""] * num_to_add
-        return dataclasses.replace(self, array=array, headers=headers)
-
-    @classmethod
-    def concat(
-        cls,
-        msas: Sequence[FastMSA],
-        join_token: str | None = None,
-        allow_depth_mismatch: bool = False,
-    ) -> FastMSA:
-        """Concatenate a series of MSAs horizontally, along the sequence dimension."""
-        if not msas:
-            raise ValueError("Cannot concatenate an empty list of MSAs")
-        if join_token is not None and join_token != "":
-            raise NotImplementedError("join_token is not supported for FastMSA")
-
-        msa_depths = [msa.depth for msa in msas]
-        if len(set(msa_depths)) != 1:
-            if not allow_depth_mismatch:
-                raise ValueError("Depth mismatch in concatenating MSAs")
-            else:
-                max_depth = max(msa_depths)
-                msas = [msa.pad_to_depth(max_depth) for msa in msas]
-        headers = [
-            "|".join([str(h) for h in headers])
-            for headers in zip(
-                *(
-                    msa.headers if msa.headers is not None else [""] * msa.depth
-                    for msa in msas
-                )
-            )
-        ]
-
-        array = np.concatenate([msa.array for msa in msas], axis=1)
-        return cls(array, headers)
-
-    def to_msa(self) -> MSA:
-        headers = (
-            self.headers
-            if self.headers is not None
-            else [f"seq{i}" for i in range(self.depth)]
-        )
-        entries = [
-            FastaEntry(header, b"".join(row).decode())
-            for header, row in zip(headers, self.array)
-        ]
-        return MSA(entries)
-
-    @classmethod
-    def stack(
-        cls, msas: Sequence[FastMSA], remove_query_from_later_msas: bool = True
-    ) -> FastMSA:
-        """Stack a series of MSAs. Optionally remove the query from msas after the first."""
-        arrays = []
-        all_headers = []
-        for i, msa in enumerate(msas):
-            array = msa.array
-            headers = msa.headers
-            if i > 0 and remove_query_from_later_msas:
-                array = array[1:]
-                if headers is not None:
-                    headers = headers[1:]
-            arrays.append(array)
-            if headers is not None:
-                all_headers.extend(headers)
-        return cls(np.concatenate(arrays, axis=0), all_headers)
+from __future__ import annotations
+
+import dataclasses
+import string
+from dataclasses import dataclass
+from functools import cached_property
+from itertools import islice
+from typing import Sequence
+
+import numpy as np
+from Bio import SeqIO
+from scipy.spatial.distance import cdist
+
+from .esmfold2_misc import slice_any_object
+from .esmfold2_msa_filter_sequences import greedy_select_indices, hhfilter
+from .esmfold2_parsing import FastaEntry, read_sequences, write_sequences
+from .esmfold2_sequential_dataclass import SequentialDataclass
+from .esmfold2_system import PathOrBuffer
+
+REMOVE_LOWERCASE_TRANSLATION = str.maketrans(dict.fromkeys(string.ascii_lowercase))
+
+
+def remove_insertions_from_sequence(seq: str) -> str:
+    return seq.translate(REMOVE_LOWERCASE_TRANSLATION)
+
+
+@dataclass(frozen=True)
+class MSA(SequentialDataclass):
+    """Object-oriented interface to an MSA.
+
+    Args:
+        sequences (list[str]): List of protein sequences
+        headers (list[str]): List of headers describing the sequences
+
+    """
+
+    entries: list[FastaEntry]
+
+    @cached_property
+    def sequences(self) -> list[str]:
+        return [entry.sequence for entry in self.entries]
+
+    @cached_property
+    def headers(self) -> list[str]:
+        return [entry.header for entry in self.entries]
+
+    def __repr__(self):
+        return (
+            f"MSA({self.entries[0].header}: Depth={self.depth}, Length={self.seqlen})"
+        )
+
+    def to_fast_msa(self) -> FastMSA:
+        return FastMSA(self.array, self.headers)
+
+    @classmethod
+    def from_a3m(
+        cls,
+        path: PathOrBuffer,
+        remove_insertions: bool = True,
+        max_sequences: int | None = None,
+    ) -> MSA:
+        entries = []
+        for header, seq in islice(read_sequences(path), max_sequences):
+            if remove_insertions:
+                seq = remove_insertions_from_sequence(seq)
+            if entries:
+                assert (
+                    len(seq) == len(entries[0].sequence)
+                ), f"Sequence length mismatch. Expected: {len(entries[0].sequence)}, Received: {len(seq)}"
+            entries.append(FastaEntry(header, seq))
+        return cls(entries)
+
+    def to_a3m(self, path: PathOrBuffer) -> None:
+        write_sequences(self.entries, path)
+
+    @classmethod
+    def from_stockholm(
+        cls,
+        path: PathOrBuffer,
+        remove_insertions: bool = True,
+        max_sequences: int | None = None,
+    ) -> MSA:
+        entries = []
+        for record in islice(SeqIO.parse(path, "stockholm"), max_sequences):
+            header = f"{record.id} {record.description}"
+            seq = str(record.seq)
+            if entries:
+                assert (
+                    len(seq) == len(entries[0].sequence)
+                ), f"Sequence length mismatch. Expected: {len(entries[0].sequence)}, Received: {len(seq)}"
+            entries.append(FastaEntry(header, seq))
+        msa = cls(entries)
+        if remove_insertions:
+            keep_inds = [i for i, aa in enumerate(msa.query) if aa != "-"]
+            msa = msa.select_positions(keep_inds)
+        return msa
+
+    def to_bytes(self) -> bytes:
+        version = 1
+        version_bytes = version.to_bytes(1, "little")
+        seqlen_bytes = self.seqlen.to_bytes(4, "little")
+        depth_bytes = self.depth.to_bytes(4, "little")
+        array_bytes = self.array.tobytes()
+        header_bytes = "\n".join(entry.header for entry in self.entries).encode()
+        all_bytes = (
+            version_bytes + seqlen_bytes + depth_bytes + array_bytes + header_bytes
+        )
+        return all_bytes
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> MSA:
+        version_bytes, seqlen_bytes, depth_bytes, data = (
+            data[:1],
+            data[1:5],
+            data[5:9],
+            data[9:],
+        )
+        version = int.from_bytes(version_bytes, "little")
+        if version != 1:
+            raise ValueError(f"Unsupported version: {version}")
+        seqlen = int.from_bytes(seqlen_bytes, "little")
+        depth = int.from_bytes(depth_bytes, "little")
+        array_bytes, header_bytes = data[: seqlen * depth], data[seqlen * depth :]
+        array = np.frombuffer(array_bytes, dtype="|S1")
+        array = array.reshape(depth, seqlen)
+        headers = header_bytes.decode().split("\n")
+        # Sometimes the separation is two newlines, which results in an empty header.
+        headers = [header for header in headers if header]
+        # If all headers were empty (e.g., saved from from_sequences), use empty headers
+        if len(headers) == 0 and depth > 0:
+            headers = [""] * depth
+        entries = [
+            FastaEntry(header, b"".join(row).decode())
+            for header, row in zip(headers, array)
+        ]
+        return cls(entries)
+
+    # TODO(jmaccarl): set remove_insertions to True by default here to match other utils
+    @classmethod
+    def from_sequences(
+        cls, sequences: list[str], remove_insertions: bool = False
+    ) -> MSA:
+        if remove_insertions:
+            entries = [
+                FastaEntry("", remove_insertions_from_sequence(seq))
+                for seq in sequences
+            ]
+        else:
+            entries = [FastaEntry("", seq) for seq in sequences]
+        return cls(entries)
+
+    def to_sequence_bytes(self) -> bytes:
+        """Stores ONLY SEQUENCES in array format as bytes. Header information will be lost."""
+        seqlen_bytes = self.seqlen.to_bytes(4, "little")
+        array_bytes = self.array.tobytes()
+        all_bytes = seqlen_bytes + array_bytes
+        return all_bytes
+
+    @classmethod
+    def from_sequence_bytes(cls, data: bytes) -> MSA:
+        seqlen_bytes, array_bytes = data[:4], data[4:]
+        seqlen = int.from_bytes(seqlen_bytes, "little")
+        array = np.frombuffer(array_bytes, dtype="|S1")
+        array = array.reshape(-1, seqlen)
+        entries = [FastaEntry("", b"".join(row).decode()) for row in array]
+        return cls(entries)
+
+    @property
+    def depth(self) -> int:
+        return len(self.entries)
+
+    @property
+    def seqlen(self) -> int:
+        return len(self.entries[0].sequence)
+
+    @cached_property
+    def array(self) -> np.ndarray:
+        return np.array([list(seq) for seq in self.sequences], dtype="|S1")
+
+    @property
+    def query(self) -> str:
+        return self.entries[0].sequence
+
+    def select_sequences(self, indices: Sequence[int] | np.ndarray) -> MSA:
+        """Subselect rows of the MSA."""
+        entries = [self.entries[idx] for idx in indices]
+        return dataclasses.replace(self, entries=entries)
+
+    def select_positions(self, indices: Sequence[int] | np.ndarray) -> MSA:
+        """Subselect columns of the MSA."""
+        entries = [
+            FastaEntry(header, "".join(seq[idx] for idx in indices))
+            for header, seq in self.entries
+        ]
+        return dataclasses.replace(self, entries=entries)
+
+    def __getitem__(self, indices: int | list[int] | slice | np.ndarray):
+        if isinstance(indices, int):
+            indices = [indices]
+
+        entries = [
+            FastaEntry(header, slice_any_object(seq, indices))
+            for header, seq in self.entries
+        ]
+        return dataclasses.replace(self, entries=entries)
+
+    def __len__(self):
+        return self.seqlen
+
+    def greedy_select(self, num_seqs: int, mode: str = "max") -> MSA:
+        """Greedily select sequences that either maximize or minimize hamming distance.
+
+        Algorithm proposed in the MSA Transformer paper. Starting from the query sequence,
+        iteratively add sequences to the list with the maximum (minimum) average Hamming
+        distance to the existing set of sequences.
+
+        Args:
+            num_seqs (int): Number of sequences to select.
+            mode (str): Whether to maximize or minimize diversity. DO NOT pick 'min' unless
+                you're doing it to prove a point for a paper.
+
+        Returns:
+            MSA object w/ subselected sequences.
+        """
+        assert mode in ("max", "min")
+        if self.depth <= num_seqs:
+            return self
+
+        indices = greedy_select_indices(self.array, num_seqs, mode)
+        return self.select_sequences(indices)
+
+    def hhfilter(
+        self,
+        seqid: int = 90,
+        diff: int = 0,
+        cov: int = 0,
+        qid: int = 0,
+        qsc: float = -20.0,
+        binary: str = "hhfilter",
+    ) -> MSA:
+        """Apply hhfilter to the sequences in the MSA and return a filtered MSA."""
+
+        indices = hhfilter(
+            self.sequences,
+            seqid=seqid,
+            diff=diff,
+            cov=cov,
+            qid=qid,
+            qsc=qsc,
+            binary=binary,
+        )
+        return self.select_sequences(indices)
+
+    def select_random_sequences(self, num_seqs: int) -> MSA:
+        """Uses random sampling to subselect sequences from the MSA. Always
+        keeps the query sequence.
+        """
+        if num_seqs >= self.depth:
+            return self
+
+        # Subselect random, always keeping the query sequence.
+        indices = np.sort(
+            np.append(
+                0, np.random.choice(self.depth - 1, num_seqs - 1, replace=False) + 1
+            )
+        )
+        msa = self.select_sequences(indices)  # type: ignore
+        return msa
+
+    def select_diverse_sequences(self, num_seqs: int) -> MSA:
+        """Applies hhfilter to select ~num_seqs sequences, then uses random sampling
+        to subselect if necessary.
+        """
+        if num_seqs >= self.depth:
+            return self
+
+        msa = self.hhfilter(diff=num_seqs)
+        if num_seqs < msa.depth:
+            msa = msa.select_random_sequences(num_seqs)
+        return msa
+
+    def pad_to_depth(self, depth: int) -> MSA:
+        if depth < self.depth:
+            raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}")
+        elif depth == self.depth:
+            return self
+
+        num_to_add = depth - self.depth
+        extra_entries = [FastaEntry("", "-" * self.seqlen) for _ in range(num_to_add)]
+        return dataclasses.replace(self, entries=self.entries + extra_entries)
+
+    @classmethod
+    def stack(
+        cls, msas: Sequence[MSA], remove_query_from_later_msas: bool = True
+    ) -> MSA:
+        """Stack a series of MSAs. Optionally remove the query from msas after the first."""
+        all_entries = []
+        for i, msa in enumerate(msas):
+            entries = msa.entries
+            if i > 0 and remove_query_from_later_msas:
+                entries = entries[1:]
+            all_entries.extend(entries)
+        return cls(entries=all_entries)
+
+    @cached_property
+    def seqid(self) -> np.ndarray:
+        array = self.array.view(np.uint8)
+        seqid = 1 - cdist(array[0][None], array, "hamming")
+        return seqid[0]
+
+    @classmethod
+    def concat(
+        cls,
+        msas: Sequence[MSA],
+        join_token: str | None = "|",
+        allow_depth_mismatch: bool = False,
+    ) -> MSA:
+        """Concatenate a series of MSAs horizontally, along the sequence dimension."""
+        if not msas:
+            raise ValueError("Cannot concatenate an empty list of MSAs")
+        msa_depths = [msa.depth for msa in msas]
+        if len(set(msa_depths)) != 1:
+            if not allow_depth_mismatch:
+                raise ValueError("Depth mismatch in concatenating MSAs")
+            else:
+                max_depth = max(msa_depths)
+                msas = [msa.pad_to_depth(max_depth) for msa in msas]
+        headers = [
+            "|".join([str(h) for h in headers])
+            for headers in zip(*(msa.headers for msa in msas))
+        ]
+
+        if join_token is None:
+            join_token = ""
+
+        seqs = [join_token.join(vals) for vals in zip(*(msa.sequences for msa in msas))]
+        entries = [FastaEntry(header, seq) for header, seq in zip(headers, seqs)]
+        return cls(entries)
+
+
+@dataclass(frozen=True)
+class FastMSA(SequentialDataclass):
+    """Object-oriented interface to an MSA stored as a numpy uint8 array."""
+
+    array: np.ndarray
+    headers: list[str] | None = None
+
+    def __post_init__(self):
+        if self.headers is not None:
+            assert (
+                len(self.headers) == self.depth
+            ), "Number of headers must match depth."
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> FastMSA:
+        version_bytes, seqlen_bytes, depth_bytes, data = (
+            data[:1],
+            data[1:5],
+            data[5:9],
+            data[9:],
+        )
+        version = int.from_bytes(version_bytes, "little")
+        if version != 1:
+            raise ValueError(f"Unsupported version: {version}")
+        seqlen = int.from_bytes(seqlen_bytes, "little")
+        depth = int.from_bytes(depth_bytes, "little")
+        array_bytes, header_bytes = data[: seqlen * depth], data[seqlen * depth :]
+        array = np.frombuffer(array_bytes, dtype="|S1")
+        array = array.reshape(depth, seqlen)
+        headers = header_bytes.decode().split("\n")
+        # Sometimes the separation is two newlines, which results in an empty header.
+        headers = [header for header in headers if header]
+        # If all headers were empty (e.g., saved from from_sequences), use empty headers
+        if len(headers) == 0 and depth > 0:
+            headers = [""] * depth
+        return cls(array, headers)
+
+    @classmethod
+    def from_sequence_bytes(cls, data: bytes) -> FastMSA:
+        seqlen_bytes, array_bytes = data[:4], data[4:]
+        seqlen = int.from_bytes(seqlen_bytes, "little")
+        array = np.frombuffer(array_bytes, dtype="|S1")
+        array = array.reshape(-1, seqlen)
+        return cls(array)
+
+    @property
+    def depth(self) -> int:
+        return self.array.shape[0]
+
+    @property
+    def seqlen(self) -> int:
+        return self.array.shape[1]
+
+    def __len__(self):
+        return self.seqlen
+
+    def __getitem__(self, indices: int | list[int] | slice | np.ndarray):
+        if isinstance(indices, int):
+            indices = [indices]
+
+        return dataclasses.replace(self, array=self.array[:, indices])
+
+    def select_sequences(self, indices: Sequence[int] | np.ndarray) -> FastMSA:
+        """Subselect rows of the MSA."""
+        array = self.array[indices]
+        headers = (
+            [self.headers[idx] for idx in indices] if self.headers is not None else None
+        )
+        return dataclasses.replace(self, array=array, headers=headers)
+
+    def select_random_sequences(self, num_seqs: int) -> FastMSA:
+        """Uses random sampling to subselect sequences from the MSA. Always
+        keeps the query sequence.
+        """
+        if num_seqs >= self.depth:
+            return self
+
+        # Subselect random, always keeping the query sequence.
+        indices = np.sort(
+            np.append(
+                0, np.random.choice(self.depth - 1, num_seqs - 1, replace=False) + 1
+            )
+        )
+        msa = self.select_sequences(indices)  # type: ignore
+        return msa
+
+    def pad_to_depth(self, depth: int) -> FastMSA:
+        if depth < self.depth:
+            raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}")
+        elif depth == self.depth:
+            return self
+
+        num_to_add = depth - self.depth
+        array = np.pad(
+            self.array,
+            [(0, num_to_add), (0, 0)],
+            constant_values=ord("-") if self.array.dtype == np.uint8 else b"-",
+        )
+        headers = self.headers
+        if headers is not None:
+            headers = headers + [""] * num_to_add
+        return dataclasses.replace(self, array=array, headers=headers)
+
+    @classmethod
+    def concat(
+        cls,
+        msas: Sequence[FastMSA],
+        join_token: str | None = None,
+        allow_depth_mismatch: bool = False,
+    ) -> FastMSA:
+        """Concatenate a series of MSAs horizontally, along the sequence dimension."""
+        if not msas:
+            raise ValueError("Cannot concatenate an empty list of MSAs")
+        if join_token is not None and join_token != "":
+            raise NotImplementedError("join_token is not supported for FastMSA")
+
+        msa_depths = [msa.depth for msa in msas]
+        if len(set(msa_depths)) != 1:
+            if not allow_depth_mismatch:
+                raise ValueError("Depth mismatch in concatenating MSAs")
+            else:
+                max_depth = max(msa_depths)
+                msas = [msa.pad_to_depth(max_depth) for msa in msas]
+        headers = [
+            "|".join([str(h) for h in headers])
+            for headers in zip(
+                *(
+                    msa.headers if msa.headers is not None else [""] * msa.depth
+                    for msa in msas
+                )
+            )
+        ]
+
+        array = np.concatenate([msa.array for msa in msas], axis=1)
+        return cls(array, headers)
+
+    def to_msa(self) -> MSA:
+        headers = (
+            self.headers
+            if self.headers is not None
+            else [f"seq{i}" for i in range(self.depth)]
+        )
+        entries = [
+            FastaEntry(header, b"".join(row).decode())
+            for header, row in zip(headers, self.array)
+        ]
+        return MSA(entries)
+
+    @classmethod
+    def stack(
+        cls, msas: Sequence[FastMSA], remove_query_from_later_msas: bool = True
+    ) -> FastMSA:
+        """Stack a series of MSAs. Optionally remove the query from msas after the first."""
+        arrays = []
+        all_headers = []
+        for i, msa in enumerate(msas):
+            array = msa.array
+            headers = msa.headers
+            if i > 0 and remove_query_from_later_msas:
+                array = array[1:]
+                if headers is not None:
+                    headers = headers[1:]
+            arrays.append(array)
+            if headers is not None:
+                all_headers.extend(headers)
+        return cls(np.concatenate(arrays, axis=0), all_headers)

esmfold2_msa_filter_sequences.py

@@ -1,82 +1,82 @@
-import os
-import tempfile
-from pathlib import Path
-
-import numpy as np
-from scipy.spatial.distance import cdist
-
-from .esmfold2_system import run_subprocess_with_errorcheck
-
-
-def greedy_select_indices(array, num_seqs: int, mode: str = "max") -> list[int]:
-    """Greedily select sequences that either maximize or minimize hamming distance.
-
-    Algorithm proposed in the MSA Transformer paper. Starting from the query sequence,
-    iteratively add sequences to the list with the maximum (minimum) average Hamming
-    distance to the existing set of sequences.
-
-    Args:
-        array (np.ndarray): Character array representing the sequences in the MSA
-        num_seqs (int): Number of sequences to select.
-        mode (str): Whether to maximize or minimize diversity. DO NOT pick 'min' unless
-            you're doing it to prove a point for a paper.
-
-    Returns:
-        list[int]: List of indices to select from the array
-    """
-    assert mode in ("max", "min")
-    depth = array.shape[0]
-    if depth <= num_seqs:
-        return list(range(depth))
-    array = array.view(np.uint8)
-
-    optfunc = np.argmax if mode == "max" else np.argmin
-    all_indices = np.arange(depth)
-    indices = [0]
-    pairwise_distances = np.zeros((0, depth))
-    for _ in range(num_seqs - 1):
-        dist = cdist(array[indices[-1:]], array, "hamming")
-        pairwise_distances = np.concatenate([pairwise_distances, dist])
-        shifted_distance = np.delete(pairwise_distances, indices, axis=1).mean(0)
-        shifted_index = optfunc(shifted_distance)
-        index = np.delete(all_indices, indices)[shifted_index]
-        indices.append(index)
-    indices = sorted(indices)
-    return indices
-
-
-def hhfilter(
-    sequences: list[str],
-    seqid: int = 90,
-    diff: int = 0,
-    cov: int = 0,
-    qid: int = 0,
-    qsc: float = -20.0,
-    binary: str = "hhfilter",
-) -> list[int]:
-    with tempfile.TemporaryDirectory(
-        dir="/dev/shm" if os.path.exists("/dev/shm") else None
-    ) as tempdirname:
-        tempdir = Path(tempdirname)
-        fasta_file = tempdir / "input.fasta"
-        fasta_file.write_text(
-            "\n".join(f">{i}\n{seq}" for i, seq in enumerate(sequences))
-        )
-        output_file = tempdir / "output.fasta"
-        command = " ".join(
-            [
-                f"{binary}",
-                f"-i {fasta_file}",
-                "-M a3m",
-                f"-o {output_file}",
-                f"-id {seqid}",
-                f"-diff {diff}",
-                f"-cov {cov}",
-                f"-qid {qid}",
-                f"-qsc {qsc}",
-            ]
-        ).split(" ")
-        run_subprocess_with_errorcheck(command, capture_output=True)
-        with output_file.open() as f:
-            indices = [int(line[1:].strip()) for line in f if line.startswith(">")]
-        return indices
+import os
+import tempfile
+from pathlib import Path
+
+import numpy as np
+from scipy.spatial.distance import cdist
+
+from .esmfold2_system import run_subprocess_with_errorcheck
+
+
+def greedy_select_indices(array, num_seqs: int, mode: str = "max") -> list[int]:
+    """Greedily select sequences that either maximize or minimize hamming distance.
+
+    Algorithm proposed in the MSA Transformer paper. Starting from the query sequence,
+    iteratively add sequences to the list with the maximum (minimum) average Hamming
+    distance to the existing set of sequences.
+
+    Args:
+        array (np.ndarray): Character array representing the sequences in the MSA
+        num_seqs (int): Number of sequences to select.
+        mode (str): Whether to maximize or minimize diversity. DO NOT pick 'min' unless
+            you're doing it to prove a point for a paper.
+
+    Returns:
+        list[int]: List of indices to select from the array
+    """
+    assert mode in ("max", "min")
+    depth = array.shape[0]
+    if depth <= num_seqs:
+        return list(range(depth))
+    array = array.view(np.uint8)
+
+    optfunc = np.argmax if mode == "max" else np.argmin
+    all_indices = np.arange(depth)
+    indices = [0]
+    pairwise_distances = np.zeros((0, depth))
+    for _ in range(num_seqs - 1):
+        dist = cdist(array[indices[-1:]], array, "hamming")
+        pairwise_distances = np.concatenate([pairwise_distances, dist])
+        shifted_distance = np.delete(pairwise_distances, indices, axis=1).mean(0)
+        shifted_index = optfunc(shifted_distance)
+        index = np.delete(all_indices, indices)[shifted_index]
+        indices.append(index)
+    indices = sorted(indices)
+    return indices
+
+
+def hhfilter(
+    sequences: list[str],
+    seqid: int = 90,
+    diff: int = 0,
+    cov: int = 0,
+    qid: int = 0,
+    qsc: float = -20.0,
+    binary: str = "hhfilter",
+) -> list[int]:
+    with tempfile.TemporaryDirectory(
+        dir="/dev/shm" if os.path.exists("/dev/shm") else None
+    ) as tempdirname:
+        tempdir = Path(tempdirname)
+        fasta_file = tempdir / "input.fasta"
+        fasta_file.write_text(
+            "\n".join(f">{i}\n{seq}" for i, seq in enumerate(sequences))
+        )
+        output_file = tempdir / "output.fasta"
+        command = " ".join(
+            [
+                f"{binary}",
+                f"-i {fasta_file}",
+                "-M a3m",
+                f"-o {output_file}",
+                f"-id {seqid}",
+                f"-diff {diff}",
+                f"-cov {cov}",
+                f"-qid {qid}",
+                f"-qsc {qsc}",
+            ]
+        ).split(" ")
+        run_subprocess_with_errorcheck(command, capture_output=True)
+        with output_file.open() as f:
+            indices = [int(line[1:].strip()) for line in f if line.startswith(">")]
+        return indices

esmfold2_normalize_coordinates.py

@@ -1,79 +1,79 @@
-from typing import TypeVar
-
-import numpy as np
-import torch
-from torch import Tensor
-
-from . import esmfold2_residue_constants as RC
-from .esmfold2_affine3d import Affine3D
-
-ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor)
-
-
-def atom3_to_backbone_frames(bb_positions: torch.Tensor) -> Affine3D:
-    N, CA, C = bb_positions.unbind(dim=-2)
-    return Affine3D.from_graham_schmidt(C, CA, N)
-
-
-def index_by_atom_name(
-    atom37: ArrayOrTensor, atom_names: str | list[str], dim: int = -2
-) -> ArrayOrTensor:
-    squeeze = False
-    if isinstance(atom_names, str):
-        atom_names = [atom_names]
-        squeeze = True
-    indices = [RC.atom_order[atom_name] for atom_name in atom_names]
-    dim = dim % atom37.ndim
-    index = tuple(slice(None) if dim != i else indices for i in range(atom37.ndim))
-    result = atom37[index]  # type: ignore
-    if squeeze:
-        result = result.squeeze(dim)
-    return result
-
-
-def get_protein_normalization_frame(coords: Tensor) -> Affine3D:
-    """Given a set of coordinates for a protein, compute a single frame that can be used to normalize the coordinates.
-    Specifically, we compute the average position of the N, CA, and C atoms use those 3 points to construct a frame
-    using the Gram-Schmidt algorithm. The average CA position is used as the origin of the frame.
-
-    Args:
-        coords (torch.FloatTensor): [L, 37, 3] tensor of coordinates
-
-    Returns:
-        Affine3D: tensor of Affine3D frame
-    """
-    bb_coords = index_by_atom_name(coords, ["N", "CA", "C"], dim=-2)
-    coord_mask = torch.all(torch.all(torch.isfinite(bb_coords), dim=-1), dim=-1)
-
-    average_position_per_n_ca_c = bb_coords.masked_fill(
-        ~coord_mask[..., None, None], 0
-    ).sum(-3) / (coord_mask.sum(-1)[..., None, None] + 1e-8)
-    frame = atom3_to_backbone_frames(average_position_per_n_ca_c.float())
-
-    return frame
-
-
-def apply_frame_to_coords(coords: Tensor, frame: Affine3D) -> Tensor:
-    """Given a set of coordinates and a single frame, apply the frame to the coordinates.
-
-    Args:
-        coords (torch.FloatTensor): [L, 37, 3] tensor of coordinates
-        frame (Affine3D): Affine3D frame
-
-    Returns:
-        torch.FloatTensor: [L, 37, 3] tensor of transformed coordinates
-    """
-    coords_trans_rot = frame[..., None, None].invert().apply(coords)
-
-    # only transform coordinates with frame that have a valid rotation
-    valid_frame = frame.trans.norm(dim=-1) > 0
-
-    is_inf = torch.isinf(coords)
-    coords = coords_trans_rot.where(valid_frame[..., None, None, None], coords)
-    coords.masked_fill_(is_inf, torch.inf)
-
-    return coords
-
-
-def normalize_coordinates(coords: Tensor) -> Tensor:
-    return apply_frame_to_coords(coords, get_protein_normalization_frame(coords))
+from typing import TypeVar
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from . import esmfold2_residue_constants as RC
+from .esmfold2_affine3d import Affine3D
+
+ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor)
+
+
+def atom3_to_backbone_frames(bb_positions: torch.Tensor) -> Affine3D:
+    N, CA, C = bb_positions.unbind(dim=-2)
+    return Affine3D.from_graham_schmidt(C, CA, N)
+
+
+def index_by_atom_name(
+    atom37: ArrayOrTensor, atom_names: str | list[str], dim: int = -2
+) -> ArrayOrTensor:
+    squeeze = False
+    if isinstance(atom_names, str):
+        atom_names = [atom_names]
+        squeeze = True
+    indices = [RC.atom_order[atom_name] for atom_name in atom_names]
+    dim = dim % atom37.ndim
+    index = tuple(slice(None) if dim != i else indices for i in range(atom37.ndim))
+    result = atom37[index]  # type: ignore
+    if squeeze:
+        result = result.squeeze(dim)
+    return result
+
+
+def get_protein_normalization_frame(coords: Tensor) -> Affine3D:
+    """Given a set of coordinates for a protein, compute a single frame that can be used to normalize the coordinates.
+    Specifically, we compute the average position of the N, CA, and C atoms use those 3 points to construct a frame
+    using the Gram-Schmidt algorithm. The average CA position is used as the origin of the frame.
+
+    Args:
+        coords (torch.FloatTensor): [L, 37, 3] tensor of coordinates
+
+    Returns:
+        Affine3D: tensor of Affine3D frame
+    """
+    bb_coords = index_by_atom_name(coords, ["N", "CA", "C"], dim=-2)
+    coord_mask = torch.all(torch.all(torch.isfinite(bb_coords), dim=-1), dim=-1)
+
+    average_position_per_n_ca_c = bb_coords.masked_fill(
+        ~coord_mask[..., None, None], 0
+    ).sum(-3) / (coord_mask.sum(-1)[..., None, None] + 1e-8)
+    frame = atom3_to_backbone_frames(average_position_per_n_ca_c.float())
+
+    return frame
+
+
+def apply_frame_to_coords(coords: Tensor, frame: Affine3D) -> Tensor:
+    """Given a set of coordinates and a single frame, apply the frame to the coordinates.
+
+    Args:
+        coords (torch.FloatTensor): [L, 37, 3] tensor of coordinates
+        frame (Affine3D): Affine3D frame
+
+    Returns:
+        torch.FloatTensor: [L, 37, 3] tensor of transformed coordinates
+    """
+    coords_trans_rot = frame[..., None, None].invert().apply(coords)
+
+    # only transform coordinates with frame that have a valid rotation
+    valid_frame = frame.trans.norm(dim=-1) > 0
+
+    is_inf = torch.isinf(coords)
+    coords = coords_trans_rot.where(valid_frame[..., None, None, None], coords)
+    coords.masked_fill_(is_inf, torch.inf)
+
+    return coords
+
+
+def normalize_coordinates(coords: Tensor) -> Tensor:
+    return apply_frame_to_coords(coords, get_protein_normalization_frame(coords))

esmfold2_output.py

@@ -1,224 +1,224 @@
-from itertools import groupby
-from typing import Any
-
-import numpy as np
-import torch
-
-from .esmfold2_constants import ELEMENT_NUMBER_TO_SYMBOL, MOL_TYPE_NONPOLYMER
-from .esmfold2_molecular_complex import (
-    MolecularComplex,
-    MolecularComplexMetadata,
-)
-
-
-def get_element_symbol(atomic_num: int) -> str:
-    return ELEMENT_NUMBER_TO_SYMBOL.get(atomic_num, "X")
-
-
-def build_molecular_complex_from_features(
-    coords: torch.Tensor,
-    plddt: torch.Tensor,
-    atom_mask: torch.Tensor,
-    ref_element: torch.Tensor,
-    ref_atom_name_chars: torch.Tensor,
-    chain_infos: list,
-    complex_id: str,
-) -> MolecularComplex:
-    """Construct a MolecularComplex from feature-dict tensors and chain metadata.
-
-    Non-polymer chains (ligands) collapse all per-atom tokens into a single
-    residue token whose pLDDT is the per-token average and whose hetero flag
-    is True.
-    """
-    mask_np = atom_mask.bool().cpu().numpy()
-    coords_np = coords.float().cpu().numpy()
-    name_chars_np = ref_atom_name_chars.cpu().numpy()
-    elements_np = ref_element.cpu().numpy()
-    plddt_np = plddt.float().cpu().numpy()
-
-    sequence_tokens: list[str] = []
-    chain_ids_per_token: list[int] = []
-    token_to_atoms: list[list[int]] = []
-    confidence: list[float] = []
-    flat_positions: list[list[float]] = []
-    flat_elements: list[str] = []
-    flat_names: list[str] = []
-    flat_hetero: list[bool] = []
-
-    chain_lookup: dict[int, str] = {}
-    entity_info: dict[int, str] = {}
-    out_atom_cursor = 0
-
-    for ci in chain_infos:
-        chain_lookup[ci.asym_id] = ci.chain_id
-        is_nonpolymer = ci.mol_type == MOL_TYPE_NONPOLYMER
-        entity_info[ci.entity_id] = "non-polymer" if is_nonpolymer else "polymer"
-
-        if is_nonpolymer:
-            residue_name = ci.tokens[0].residue_name if ci.tokens else "LIG"
-            sequence_tokens.append(residue_name)
-            chain_ids_per_token.append(ci.asym_id)
-            avg_plddt = (
-                float(np.mean([plddt_np[ti.token_index] for ti in ci.tokens]))
-                if ci.tokens
-                else 0.0
-            )
-            confidence.append(avg_plddt)
-            token_atom_start = out_atom_cursor
-            for ti in ci.tokens:
-                for atom_idx in range(ti.atom_start, ti.atom_start + ti.atom_count):
-                    if not mask_np[atom_idx]:
-                        continue
-                    flat_positions.append(coords_np[atom_idx].tolist())
-                    flat_elements.append(get_element_symbol(int(elements_np[atom_idx])))
-                    chars = name_chars_np[atom_idx]
-                    name = "".join(
-                        chr(int(c) + 32) for c in chars if int(c) != 0
-                    ).strip()
-                    flat_names.append(name)
-                    flat_hetero.append(True)
-                    out_atom_cursor += 1
-            token_to_atoms.append([token_atom_start, out_atom_cursor])
-            continue
-
-        # Atom-tokenized modified residues (HYP, MSE, ...) span multiple
-        # tokens per residue; collapse them back to one mmCIF residue.
-        for _residue_index, ti_iter in groupby(
-            ci.tokens, key=lambda t: t.residue_index
-        ):
-            ti_group = list(ti_iter)
-            sequence_tokens.append(ti_group[0].residue_name)
-            chain_ids_per_token.append(ci.asym_id)
-            confidence.append(
-                float(np.mean([plddt_np[ti.token_index] for ti in ti_group]))
-            )
-            token_atom_start = out_atom_cursor
-            for ti in ti_group:
-                for atom_idx in range(ti.atom_start, ti.atom_start + ti.atom_count):
-                    if not mask_np[atom_idx]:
-                        continue
-                    flat_positions.append(coords_np[atom_idx].tolist())
-                    flat_elements.append(get_element_symbol(int(elements_np[atom_idx])))
-                    chars = name_chars_np[atom_idx]
-                    name = "".join(
-                        chr(int(c) + 32) for c in chars if int(c) != 0
-                    ).strip()
-                    flat_names.append(name)
-                    flat_hetero.append(False)
-                    out_atom_cursor += 1
-            token_to_atoms.append([token_atom_start, out_atom_cursor])
-
-    return MolecularComplex(
-        id=complex_id,
-        sequence=sequence_tokens,
-        atom_positions=np.array(flat_positions, dtype=np.float32).reshape(-1, 3),
-        atom_elements=np.array(flat_elements, dtype=object),
-        token_to_atoms=np.array(token_to_atoms, dtype=np.int32).reshape(-1, 2),
-        chain_id=np.array(chain_ids_per_token, dtype=np.int64),
-        plddt=np.array(confidence, dtype=np.float32),
-        atom_names=np.array(flat_names, dtype=object),
-        atom_hetero=np.array(flat_hetero, dtype=bool),
-        metadata=MolecularComplexMetadata(
-            entity_lookup=entity_info,
-            chain_lookup=chain_lookup,
-            assembly_composition=None,
-        ),
-    )
-
-
-def build_molecular_complex(
-    structure: Any, coords: torch.Tensor, plddt: torch.Tensor, complex_id: str
-) -> MolecularComplex:
-    """Directly constructs a MolecularComplex from model outputs without intermediate files.
-
-    Args:
-        structure: Object with .chains, .residues, .atoms numpy structured arrays.
-        coords: [N_atoms, 3] predicted atom coordinates.
-        plddt: [N_residues] per-residue confidence scores.
-        complex_id: Identifier string for the resulting complex.
-    """
-    flat_positions = []
-    flat_elements = []
-    flat_names = []
-    flat_hetero = []
-
-    sequence_tokens = []
-    token_to_atoms = []
-    chain_ids_per_token = []
-    confidence_scores = []
-
-    chain_lookup = {}
-    entity_info = {}
-
-    global_atom_cursor = 0
-    global_res_cursor = 0
-    atom_array_idx = 0
-
-    for chain in structure.chains:
-        chain_idx_numeric = chain["asym_id"]
-        chain_name_str = str(chain["name"])
-        mol_type = chain["mol_type"]
-
-        chain_lookup[chain_idx_numeric] = chain_name_str
-        entity_info[chain["entity_id"]] = (
-            "polymer" if mol_type != MOL_TYPE_NONPOLYMER else "non-polymer"
-        )
-
-        res_start = chain["res_idx"]
-        res_end = chain["res_idx"] + chain["res_num"]
-        residues = structure.residues[res_start:res_end]
-
-        for residue in residues:
-            res_name = str(residue["name"])
-
-            sequence_tokens.append(res_name)
-            chain_ids_per_token.append(chain_idx_numeric)
-
-            score = plddt[global_res_cursor].item()
-            confidence_scores.append(score)
-            token_start_idx = atom_array_idx
-
-            atom_start = residue["atom_idx"]
-            atom_end = residue["atom_idx"] + residue["atom_num"]
-            atoms = structure.atoms[atom_start:atom_end]
-
-            for atom in atoms:
-                if not atom["is_present"]:
-                    continue
-
-                pos = coords[global_atom_cursor].tolist()
-                flat_positions.append(pos)
-
-                elem = get_element_symbol(atom["element"].item())
-                flat_elements.append(elem)
-
-                raw_name = atom["name"]
-                if hasattr(raw_name, "tolist"):
-                    raw_name = raw_name.tolist()
-                name_str = "".join([chr(c + 32) for c in raw_name if c != 0])
-                flat_names.append(name_str)
-
-                flat_hetero.append(mol_type == MOL_TYPE_NONPOLYMER)
-
-                global_atom_cursor += 1
-                atom_array_idx += 1
-
-            token_to_atoms.append([token_start_idx, atom_array_idx])
-            global_res_cursor += 1
-
-    return MolecularComplex(
-        id=complex_id,
-        sequence=sequence_tokens,
-        atom_positions=np.array(flat_positions, dtype=np.float32),
-        atom_elements=np.array(flat_elements, dtype=object),
-        token_to_atoms=np.array(token_to_atoms, dtype=np.int32),
-        chain_id=np.array(chain_ids_per_token, dtype=np.int64),
-        plddt=np.array(confidence_scores, dtype=np.float32),
-        atom_names=np.array(flat_names, dtype=object),
-        atom_hetero=np.array(flat_hetero, dtype=bool),
-        metadata=MolecularComplexMetadata(
-            entity_lookup=entity_info,
-            chain_lookup=chain_lookup,
-            assembly_composition=None,
-        ),
-    )
+from itertools import groupby
+from typing import Any
+
+import numpy as np
+import torch
+
+from .esmfold2_constants import ELEMENT_NUMBER_TO_SYMBOL, MOL_TYPE_NONPOLYMER
+from .esmfold2_molecular_complex import (
+    MolecularComplex,
+    MolecularComplexMetadata,
+)
+
+
+def get_element_symbol(atomic_num: int) -> str:
+    return ELEMENT_NUMBER_TO_SYMBOL.get(atomic_num, "X")
+
+
+def build_molecular_complex_from_features(
+    coords: torch.Tensor,
+    plddt: torch.Tensor,
+    atom_mask: torch.Tensor,
+    ref_element: torch.Tensor,
+    ref_atom_name_chars: torch.Tensor,
+    chain_infos: list,
+    complex_id: str,
+) -> MolecularComplex:
+    """Construct a MolecularComplex from feature-dict tensors and chain metadata.
+
+    Non-polymer chains (ligands) collapse all per-atom tokens into a single
+    residue token whose pLDDT is the per-token average and whose hetero flag
+    is True.
+    """
+    mask_np = atom_mask.bool().cpu().numpy()
+    coords_np = coords.float().cpu().numpy()
+    name_chars_np = ref_atom_name_chars.cpu().numpy()
+    elements_np = ref_element.cpu().numpy()
+    plddt_np = plddt.float().cpu().numpy()
+
+    sequence_tokens: list[str] = []
+    chain_ids_per_token: list[int] = []
+    token_to_atoms: list[list[int]] = []
+    confidence: list[float] = []
+    flat_positions: list[list[float]] = []
+    flat_elements: list[str] = []
+    flat_names: list[str] = []
+    flat_hetero: list[bool] = []
+
+    chain_lookup: dict[int, str] = {}
+    entity_info: dict[int, str] = {}
+    out_atom_cursor = 0
+
+    for ci in chain_infos:
+        chain_lookup[ci.asym_id] = ci.chain_id
+        is_nonpolymer = ci.mol_type == MOL_TYPE_NONPOLYMER
+        entity_info[ci.entity_id] = "non-polymer" if is_nonpolymer else "polymer"
+
+        if is_nonpolymer:
+            residue_name = ci.tokens[0].residue_name if ci.tokens else "LIG"
+            sequence_tokens.append(residue_name)
+            chain_ids_per_token.append(ci.asym_id)
+            avg_plddt = (
+                float(np.mean([plddt_np[ti.token_index] for ti in ci.tokens]))
+                if ci.tokens
+                else 0.0
+            )
+            confidence.append(avg_plddt)
+            token_atom_start = out_atom_cursor
+            for ti in ci.tokens:
+                for atom_idx in range(ti.atom_start, ti.atom_start + ti.atom_count):
+                    if not mask_np[atom_idx]:
+                        continue
+                    flat_positions.append(coords_np[atom_idx].tolist())
+                    flat_elements.append(get_element_symbol(int(elements_np[atom_idx])))
+                    chars = name_chars_np[atom_idx]
+                    name = "".join(
+                        chr(int(c) + 32) for c in chars if int(c) != 0
+                    ).strip()
+                    flat_names.append(name)
+                    flat_hetero.append(True)
+                    out_atom_cursor += 1
+            token_to_atoms.append([token_atom_start, out_atom_cursor])
+            continue
+
+        # Atom-tokenized modified residues (HYP, MSE, ...) span multiple
+        # tokens per residue; collapse them back to one mmCIF residue.
+        for _residue_index, ti_iter in groupby(
+            ci.tokens, key=lambda t: t.residue_index
+        ):
+            ti_group = list(ti_iter)
+            sequence_tokens.append(ti_group[0].residue_name)
+            chain_ids_per_token.append(ci.asym_id)
+            confidence.append(
+                float(np.mean([plddt_np[ti.token_index] for ti in ti_group]))
+            )
+            token_atom_start = out_atom_cursor
+            for ti in ti_group:
+                for atom_idx in range(ti.atom_start, ti.atom_start + ti.atom_count):
+                    if not mask_np[atom_idx]:
+                        continue
+                    flat_positions.append(coords_np[atom_idx].tolist())
+                    flat_elements.append(get_element_symbol(int(elements_np[atom_idx])))
+                    chars = name_chars_np[atom_idx]
+                    name = "".join(
+                        chr(int(c) + 32) for c in chars if int(c) != 0
+                    ).strip()
+                    flat_names.append(name)
+                    flat_hetero.append(False)
+                    out_atom_cursor += 1
+            token_to_atoms.append([token_atom_start, out_atom_cursor])
+
+    return MolecularComplex(
+        id=complex_id,
+        sequence=sequence_tokens,
+        atom_positions=np.array(flat_positions, dtype=np.float32).reshape(-1, 3),
+        atom_elements=np.array(flat_elements, dtype=object),
+        token_to_atoms=np.array(token_to_atoms, dtype=np.int32).reshape(-1, 2),
+        chain_id=np.array(chain_ids_per_token, dtype=np.int64),
+        plddt=np.array(confidence, dtype=np.float32),
+        atom_names=np.array(flat_names, dtype=object),
+        atom_hetero=np.array(flat_hetero, dtype=bool),
+        metadata=MolecularComplexMetadata(
+            entity_lookup=entity_info,
+            chain_lookup=chain_lookup,
+            assembly_composition=None,
+        ),
+    )
+
+
+def build_molecular_complex(
+    structure: Any, coords: torch.Tensor, plddt: torch.Tensor, complex_id: str
+) -> MolecularComplex:
+    """Directly constructs a MolecularComplex from model outputs without intermediate files.
+
+    Args:
+        structure: Object with .chains, .residues, .atoms numpy structured arrays.
+        coords: [N_atoms, 3] predicted atom coordinates.
+        plddt: [N_residues] per-residue confidence scores.
+        complex_id: Identifier string for the resulting complex.
+    """
+    flat_positions = []
+    flat_elements = []
+    flat_names = []
+    flat_hetero = []
+
+    sequence_tokens = []
+    token_to_atoms = []
+    chain_ids_per_token = []
+    confidence_scores = []
+
+    chain_lookup = {}
+    entity_info = {}
+
+    global_atom_cursor = 0
+    global_res_cursor = 0
+    atom_array_idx = 0
+
+    for chain in structure.chains:
+        chain_idx_numeric = chain["asym_id"]
+        chain_name_str = str(chain["name"])
+        mol_type = chain["mol_type"]
+
+        chain_lookup[chain_idx_numeric] = chain_name_str
+        entity_info[chain["entity_id"]] = (
+            "polymer" if mol_type != MOL_TYPE_NONPOLYMER else "non-polymer"
+        )
+
+        res_start = chain["res_idx"]
+        res_end = chain["res_idx"] + chain["res_num"]
+        residues = structure.residues[res_start:res_end]
+
+        for residue in residues:
+            res_name = str(residue["name"])
+
+            sequence_tokens.append(res_name)
+            chain_ids_per_token.append(chain_idx_numeric)
+
+            score = plddt[global_res_cursor].item()
+            confidence_scores.append(score)
+            token_start_idx = atom_array_idx
+
+            atom_start = residue["atom_idx"]
+            atom_end = residue["atom_idx"] + residue["atom_num"]
+            atoms = structure.atoms[atom_start:atom_end]
+
+            for atom in atoms:
+                if not atom["is_present"]:
+                    continue
+
+                pos = coords[global_atom_cursor].tolist()
+                flat_positions.append(pos)
+
+                elem = get_element_symbol(atom["element"].item())
+                flat_elements.append(elem)
+
+                raw_name = atom["name"]
+                if hasattr(raw_name, "tolist"):
+                    raw_name = raw_name.tolist()
+                name_str = "".join([chr(c + 32) for c in raw_name if c != 0])
+                flat_names.append(name_str)
+
+                flat_hetero.append(mol_type == MOL_TYPE_NONPOLYMER)
+
+                global_atom_cursor += 1
+                atom_array_idx += 1
+
+            token_to_atoms.append([token_start_idx, atom_array_idx])
+            global_res_cursor += 1
+
+    return MolecularComplex(
+        id=complex_id,
+        sequence=sequence_tokens,
+        atom_positions=np.array(flat_positions, dtype=np.float32),
+        atom_elements=np.array(flat_elements, dtype=object),
+        token_to_atoms=np.array(token_to_atoms, dtype=np.int32),
+        chain_id=np.array(chain_ids_per_token, dtype=np.int64),
+        plddt=np.array(confidence_scores, dtype=np.float32),
+        atom_names=np.array(flat_names, dtype=object),
+        atom_hetero=np.array(flat_hetero, dtype=bool),
+        metadata=MolecularComplexMetadata(
+            entity_lookup=entity_info,
+            chain_lookup=chain_lookup,
+            assembly_composition=None,
+        ),
+    )

esmfold2_paired_msa.py

@@ -1,245 +1,245 @@
-"""Taxonomy-paired MSA construction for ESMFold2 inference.
-
-Taxonomy IDs are read from FASTA headers as ``key=N`` tokens. Rows
-where any chain has ``key=-1`` (or no ``key=`` at all) are treated as
-unpaired and assigned to that chain's block-diagonal section after
-the paired rows.
-"""
-
-import re
-
-import numpy as np
-
-from .esmfold2_constants import (
-    MSA_GAP_TOKEN_ID,
-    PROTEIN_3TO1,
-    PROTEIN_RESIDUE_TO_RES_TYPE,
-    PROTEIN_UNK_RES_TYPE,
-)
-from .esmfold2_msa import MSA
-
-_KEY_RE = re.compile(r"key=(-?\d+)")
-
-
-def protein_letter_to_res_type() -> dict[str, int]:
-    """Return the protein 1-letter → res_type mapping used by the MSA encoder."""
-    mapping: dict[str, int] = {}
-    for three, one in PROTEIN_3TO1.items():
-        if three in PROTEIN_RESIDUE_TO_RES_TYPE:
-            mapping[one] = PROTEIN_RESIDUE_TO_RES_TYPE[three]
-    mapping["-"] = MSA_GAP_TOKEN_ID
-    mapping["X"] = PROTEIN_UNK_RES_TYPE
-    return mapping
-
-
-def _taxonomy_from_header(header: str) -> int:
-    if not header:
-        return -1
-    m = _KEY_RE.search(header)
-    return int(m.group(1)) if m else -1
-
-
-def msa_to_res_type_and_deletions(
-    msa: MSA, letter_to_res_type: dict[str, int]
-) -> tuple[np.ndarray, np.ndarray]:
-    """Convert an :class:`MSA` to ``(res_type[M, L], deletion_count[M, L])``.
-
-    Handles a3m insertion convention: lowercase letters and ``.`` are
-    insertions and are not emitted; their count is accumulated into the
-    next non-insertion position's deletion value. ``L`` is the query
-    length after stripping insertions from row 0.
-    """
-    query = msa.entries[0].sequence
-    L = sum(1 for ch in query if not (ch.islower() or ch == "."))
-    M = msa.depth
-
-    res_type = np.full((M, L), MSA_GAP_TOKEN_ID, dtype=np.int64)
-    deletions = np.zeros((M, L), dtype=np.float32)
-
-    for r, entry in enumerate(msa.entries):
-        col = 0
-        ins = 0
-        for ch in entry.sequence:
-            if ch == "." or (ch.islower() and ch != "-"):
-                ins += 1
-                continue
-            if col >= L:
-                break
-            if ch == "-":
-                res_type[r, col] = MSA_GAP_TOKEN_ID
-            else:
-                res_type[r, col] = letter_to_res_type.get(
-                    ch.upper(), PROTEIN_UNK_RES_TYPE
-                )
-            if ins > 0:
-                deletions[r, col] = float(ins)
-                ins = 0
-            col += 1
-    return res_type, deletions
-
-
-def _dummy_msa_residues(query_res_types: np.ndarray) -> np.ndarray:
-    """Single-row 'MSA' for chains without one — just the query."""
-    return query_res_types[None, :]  # [1, L]
-
-
-def construct_paired_msa(
-    chain_msas: dict[int, MSA | None],
-    chain_query_res_types: dict[int, np.ndarray],
-    token_asym_ids: np.ndarray,
-    token_res_ids: np.ndarray,
-    letter_to_res_type: dict[str, int] | None = None,
-    *,
-    max_pairs: int = 8192,
-    max_total: int = 16384,
-    max_seqs: int = 16384,
-) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """Build paired MSA features.
-
-    Parameters
-    ----------
-    chain_msas
-        ``asym_id -> MSA`` (or ``None`` for chains without an MSA).
-    chain_query_res_types
-        ``asym_id -> np.ndarray[L_c]`` of res-type ids for the chain's
-        query. Used to build dummy MSAs when a chain has no MSA.
-    token_asym_ids
-        Per-token asym_id, length ``T``. Must be non-decreasing.
-    token_res_ids
-        Per-token residue index within chain, length ``T``.
-    letter_to_res_type
-        1-letter → res-type mapping. Defaults to
-        :func:`protein_letter_to_res_type`.
-
-    Returns
-    -------
-    msa_residues : ``np.ndarray[M, T]`` int64
-    deletion_value : ``np.ndarray[M, T]`` float32 (raw deletion counts; the
-        ``arctan(/3) * pi/2`` transform is applied by the caller)
-    is_paired : ``np.ndarray[M, T]`` float32 broadcast of per-row,
-        per-chain paired flags.
-    """
-    if letter_to_res_type is None:
-        letter_to_res_type = protein_letter_to_res_type()
-
-    chain_ids: list[int] = sorted(chain_msas.keys())
-
-    # Build per-chain (res_type, deletions, taxonomy) tables.
-    chain_res_type: dict[int, np.ndarray] = {}
-    chain_deletions: dict[int, np.ndarray] = {}
-    chain_taxonomies: dict[int, list[int]] = {}
-    for c in chain_ids:
-        m = chain_msas.get(c)
-        if m is None or m.depth == 0:
-            qres = chain_query_res_types[c]
-            chain_res_type[c] = _dummy_msa_residues(qres)
-            chain_deletions[c] = np.zeros((1, qres.shape[0]), dtype=np.float32)
-            chain_taxonomies[c] = [-1]
-            continue
-        rt, dl = msa_to_res_type_and_deletions(m, letter_to_res_type)
-        chain_res_type[c] = rt
-        chain_deletions[c] = dl
-        chain_taxonomies[c] = [_taxonomy_from_header(e.header) for e in m.entries]
-
-    # Group by taxonomy, skip query row and unpaired (-1) entries.
-    taxonomy_map: dict[int, list[tuple[int, int]]] = {}
-    for c in chain_ids:
-        for seq_idx, taxon in enumerate(chain_taxonomies[c]):
-            if seq_idx == 0 or taxon == -1:
-                continue
-            taxonomy_map.setdefault(taxon, []).append((c, seq_idx))
-    taxonomy_map = {k: v for k, v in taxonomy_map.items() if len(v) > 1}
-    # Order taxonomies by number of distinct chains, descending.
-    sorted_taxa = sorted(
-        taxonomy_map.items(), key=lambda kv: len({c for c, _ in kv[1]}), reverse=True
-    )
-
-    visited = {s for _, items in taxonomy_map.items() for s in items}
-    available: dict[int, list[int]] = {
-        c: [i for i in range(1, len(chain_taxonomies[c])) if (c, i) not in visited]
-        for c in chain_ids
-    }
-
-    pairing: list[dict[int, int]] = [{c: 0 for c in chain_ids}]
-    is_paired: list[dict[int, int]] = [{c: 1 for c in chain_ids}]
-
-    for _, pairs in sorted_taxa:
-        per_chain: dict[int, list[int]] = {}
-        for c, seq_idx in pairs:
-            per_chain.setdefault(c, []).append(seq_idx)
-        max_occ = max(len(v) for v in per_chain.values())
-        for i in range(max_occ):
-            row_pairing: dict[int, int] = {}
-            row_is_paired: dict[int, int] = {}
-            for c, seq_idxs in per_chain.items():
-                row_pairing[c] = seq_idxs[i % len(seq_idxs)]
-                row_is_paired[c] = 1
-            for c in chain_ids:
-                if c in row_pairing:
-                    continue
-                row_is_paired[c] = 0
-                if available[c]:
-                    row_pairing[c] = available[c].pop(0)
-                else:
-                    row_pairing[c] = -1
-            pairing.append(row_pairing)
-            is_paired.append(row_is_paired)
-            if len(pairing) >= max_pairs:
-                break
-        if len(pairing) >= max_pairs:
-            break
-
-    max_left = max((len(v) for v in available.values()), default=0)
-    for _ in range(min(max_total - len(pairing), max_left)):
-        row_pairing = {}
-        row_is_paired = {}
-        for c in chain_ids:
-            row_is_paired[c] = 0
-            if available[c]:
-                row_pairing[c] = available[c].pop(0)
-            else:
-                row_pairing[c] = -1
-        pairing.append(row_pairing)
-        is_paired.append(row_is_paired)
-        if len(pairing) >= max_total:
-            break
-
-    pairing = pairing[:max_seqs]
-    is_paired = is_paired[:max_seqs]
-    M = len(pairing)
-    T = len(token_asym_ids)
-
-    msa_residues = np.full((M, T), MSA_GAP_TOKEN_ID, dtype=np.int64)
-    deletion_value = np.zeros((M, T), dtype=np.float32)
-    paired_mask = np.zeros((M, T), dtype=np.float32)
-
-    # Vectorize per chain: gather chain rows according to pairing[c], then
-    # index into them by the chain's token residue ids.
-    for c in chain_ids:
-        rt = chain_res_type[c]
-        dl = chain_deletions[c]
-        Lc = rt.shape[1]
-        chain_pairing = np.array([row[c] for row in pairing], dtype=np.int64)
-        chain_paired = np.array([row[c] for row in is_paired], dtype=np.float32)
-
-        token_mask = token_asym_ids == c
-        if not token_mask.any():
-            continue
-        token_res_in_chain = token_res_ids[token_mask]
-        # Clamp residue indices to the MSA's column range. Modified-residue
-        # tokens that exceed the query length fall back to the last column.
-        cols = np.minimum(token_res_in_chain, Lc - 1)
-
-        # Rows where pairing == -1 fall back to gap (already initialized).
-        valid_rows = chain_pairing >= 0
-        if valid_rows.any():
-            gathered_rt = rt[chain_pairing[valid_rows]][:, cols]
-            gathered_dl = dl[chain_pairing[valid_rows]][:, cols]
-            valid_idx = np.where(valid_rows)[0]
-            token_idx = np.where(token_mask)[0]
-            msa_residues[np.ix_(valid_idx, token_idx)] = gathered_rt
-            deletion_value[np.ix_(valid_idx, token_idx)] = gathered_dl
-
-        paired_mask[:, token_mask] = chain_paired[:, None]
-
-    return msa_residues, deletion_value, paired_mask
+"""Taxonomy-paired MSA construction for ESMFold2 inference.
+
+Taxonomy IDs are read from FASTA headers as ``key=N`` tokens. Rows
+where any chain has ``key=-1`` (or no ``key=`` at all) are treated as
+unpaired and assigned to that chain's block-diagonal section after
+the paired rows.
+"""
+
+import re
+
+import numpy as np
+
+from .esmfold2_constants import (
+    MSA_GAP_TOKEN_ID,
+    PROTEIN_3TO1,
+    PROTEIN_RESIDUE_TO_RES_TYPE,
+    PROTEIN_UNK_RES_TYPE,
+)
+from .esmfold2_msa import MSA
+
+_KEY_RE = re.compile(r"key=(-?\d+)")
+
+
+def protein_letter_to_res_type() -> dict[str, int]:
+    """Return the protein 1-letter → res_type mapping used by the MSA encoder."""
+    mapping: dict[str, int] = {}
+    for three, one in PROTEIN_3TO1.items():
+        if three in PROTEIN_RESIDUE_TO_RES_TYPE:
+            mapping[one] = PROTEIN_RESIDUE_TO_RES_TYPE[three]
+    mapping["-"] = MSA_GAP_TOKEN_ID
+    mapping["X"] = PROTEIN_UNK_RES_TYPE
+    return mapping
+
+
+def _taxonomy_from_header(header: str) -> int:
+    if not header:
+        return -1
+    m = _KEY_RE.search(header)
+    return int(m.group(1)) if m else -1
+
+
+def msa_to_res_type_and_deletions(
+    msa: MSA, letter_to_res_type: dict[str, int]
+) -> tuple[np.ndarray, np.ndarray]:
+    """Convert an :class:`MSA` to ``(res_type[M, L], deletion_count[M, L])``.
+
+    Handles a3m insertion convention: lowercase letters and ``.`` are
+    insertions and are not emitted; their count is accumulated into the
+    next non-insertion position's deletion value. ``L`` is the query
+    length after stripping insertions from row 0.
+    """
+    query = msa.entries[0].sequence
+    L = sum(1 for ch in query if not (ch.islower() or ch == "."))
+    M = msa.depth
+
+    res_type = np.full((M, L), MSA_GAP_TOKEN_ID, dtype=np.int64)
+    deletions = np.zeros((M, L), dtype=np.float32)
+
+    for r, entry in enumerate(msa.entries):
+        col = 0
+        ins = 0
+        for ch in entry.sequence:
+            if ch == "." or (ch.islower() and ch != "-"):
+                ins += 1
+                continue
+            if col >= L:
+                break
+            if ch == "-":
+                res_type[r, col] = MSA_GAP_TOKEN_ID
+            else:
+                res_type[r, col] = letter_to_res_type.get(
+                    ch.upper(), PROTEIN_UNK_RES_TYPE
+                )
+            if ins > 0:
+                deletions[r, col] = float(ins)
+                ins = 0
+            col += 1
+    return res_type, deletions
+
+
+def _dummy_msa_residues(query_res_types: np.ndarray) -> np.ndarray:
+    """Single-row 'MSA' for chains without one — just the query."""
+    return query_res_types[None, :]  # [1, L]
+
+
+def construct_paired_msa(
+    chain_msas: dict[int, MSA | None],
+    chain_query_res_types: dict[int, np.ndarray],
+    token_asym_ids: np.ndarray,
+    token_res_ids: np.ndarray,
+    letter_to_res_type: dict[str, int] | None = None,
+    *,
+    max_pairs: int = 8192,
+    max_total: int = 16384,
+    max_seqs: int = 16384,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Build paired MSA features.
+
+    Parameters
+    ----------
+    chain_msas
+        ``asym_id -> MSA`` (or ``None`` for chains without an MSA).
+    chain_query_res_types
+        ``asym_id -> np.ndarray[L_c]`` of res-type ids for the chain's
+        query. Used to build dummy MSAs when a chain has no MSA.
+    token_asym_ids
+        Per-token asym_id, length ``T``. Must be non-decreasing.
+    token_res_ids
+        Per-token residue index within chain, length ``T``.
+    letter_to_res_type
+        1-letter → res-type mapping. Defaults to
+        :func:`protein_letter_to_res_type`.
+
+    Returns
+    -------
+    msa_residues : ``np.ndarray[M, T]`` int64
+    deletion_value : ``np.ndarray[M, T]`` float32 (raw deletion counts; the
+        ``arctan(/3) * pi/2`` transform is applied by the caller)
+    is_paired : ``np.ndarray[M, T]`` float32 broadcast of per-row,
+        per-chain paired flags.
+    """
+    if letter_to_res_type is None:
+        letter_to_res_type = protein_letter_to_res_type()
+
+    chain_ids: list[int] = sorted(chain_msas.keys())
+
+    # Build per-chain (res_type, deletions, taxonomy) tables.
+    chain_res_type: dict[int, np.ndarray] = {}
+    chain_deletions: dict[int, np.ndarray] = {}
+    chain_taxonomies: dict[int, list[int]] = {}
+    for c in chain_ids:
+        m = chain_msas.get(c)
+        if m is None or m.depth == 0:
+            qres = chain_query_res_types[c]
+            chain_res_type[c] = _dummy_msa_residues(qres)
+            chain_deletions[c] = np.zeros((1, qres.shape[0]), dtype=np.float32)
+            chain_taxonomies[c] = [-1]
+            continue
+        rt, dl = msa_to_res_type_and_deletions(m, letter_to_res_type)
+        chain_res_type[c] = rt
+        chain_deletions[c] = dl
+        chain_taxonomies[c] = [_taxonomy_from_header(e.header) for e in m.entries]
+
+    # Group by taxonomy, skip query row and unpaired (-1) entries.
+    taxonomy_map: dict[int, list[tuple[int, int]]] = {}
+    for c in chain_ids:
+        for seq_idx, taxon in enumerate(chain_taxonomies[c]):
+            if seq_idx == 0 or taxon == -1:
+                continue
+            taxonomy_map.setdefault(taxon, []).append((c, seq_idx))
+    taxonomy_map = {k: v for k, v in taxonomy_map.items() if len(v) > 1}
+    # Order taxonomies by number of distinct chains, descending.
+    sorted_taxa = sorted(
+        taxonomy_map.items(), key=lambda kv: len({c for c, _ in kv[1]}), reverse=True
+    )
+
+    visited = {s for _, items in taxonomy_map.items() for s in items}
+    available: dict[int, list[int]] = {
+        c: [i for i in range(1, len(chain_taxonomies[c])) if (c, i) not in visited]
+        for c in chain_ids
+    }
+
+    pairing: list[dict[int, int]] = [{c: 0 for c in chain_ids}]
+    is_paired: list[dict[int, int]] = [{c: 1 for c in chain_ids}]
+
+    for _, pairs in sorted_taxa:
+        per_chain: dict[int, list[int]] = {}
+        for c, seq_idx in pairs:
+            per_chain.setdefault(c, []).append(seq_idx)
+        max_occ = max(len(v) for v in per_chain.values())
+        for i in range(max_occ):
+            row_pairing: dict[int, int] = {}
+            row_is_paired: dict[int, int] = {}
+            for c, seq_idxs in per_chain.items():
+                row_pairing[c] = seq_idxs[i % len(seq_idxs)]
+                row_is_paired[c] = 1
+            for c in chain_ids:
+                if c in row_pairing:
+                    continue
+                row_is_paired[c] = 0
+                if available[c]:
+                    row_pairing[c] = available[c].pop(0)
+                else:
+                    row_pairing[c] = -1
+            pairing.append(row_pairing)
+            is_paired.append(row_is_paired)
+            if len(pairing) >= max_pairs:
+                break
+        if len(pairing) >= max_pairs:
+            break
+
+    max_left = max((len(v) for v in available.values()), default=0)
+    for _ in range(min(max_total - len(pairing), max_left)):
+        row_pairing = {}
+        row_is_paired = {}
+        for c in chain_ids:
+            row_is_paired[c] = 0
+            if available[c]:
+                row_pairing[c] = available[c].pop(0)
+            else:
+                row_pairing[c] = -1
+        pairing.append(row_pairing)
+        is_paired.append(row_is_paired)
+        if len(pairing) >= max_total:
+            break
+
+    pairing = pairing[:max_seqs]
+    is_paired = is_paired[:max_seqs]
+    M = len(pairing)
+    T = len(token_asym_ids)
+
+    msa_residues = np.full((M, T), MSA_GAP_TOKEN_ID, dtype=np.int64)
+    deletion_value = np.zeros((M, T), dtype=np.float32)
+    paired_mask = np.zeros((M, T), dtype=np.float32)
+
+    # Vectorize per chain: gather chain rows according to pairing[c], then
+    # index into them by the chain's token residue ids.
+    for c in chain_ids:
+        rt = chain_res_type[c]
+        dl = chain_deletions[c]
+        Lc = rt.shape[1]
+        chain_pairing = np.array([row[c] for row in pairing], dtype=np.int64)
+        chain_paired = np.array([row[c] for row in is_paired], dtype=np.float32)
+
+        token_mask = token_asym_ids == c
+        if not token_mask.any():
+            continue
+        token_res_in_chain = token_res_ids[token_mask]
+        # Clamp residue indices to the MSA's column range. Modified-residue
+        # tokens that exceed the query length fall back to the last column.
+        cols = np.minimum(token_res_in_chain, Lc - 1)
+
+        # Rows where pairing == -1 fall back to gap (already initialized).
+        valid_rows = chain_pairing >= 0
+        if valid_rows.any():
+            gathered_rt = rt[chain_pairing[valid_rows]][:, cols]
+            gathered_dl = dl[chain_pairing[valid_rows]][:, cols]
+            valid_idx = np.where(valid_rows)[0]
+            token_idx = np.where(token_mask)[0]
+            msa_residues[np.ix_(valid_idx, token_idx)] = gathered_rt
+            deletion_value[np.ix_(valid_idx, token_idx)] = gathered_dl
+
+        paired_mask[:, token_mask] = chain_paired[:, None]
+
+    return msa_residues, deletion_value, paired_mask

esmfold2_parsing.py

@@ -1,112 +1,112 @@
-import io
-from pathlib import Path
-from typing import Generator, Iterable, NamedTuple
-
-PathOrBuffer = str | Path | io.TextIOBase
-FastaEntry = NamedTuple("FastaEntry", [("header", str), ("sequence", str)])
-
-
-def parse_fasta(fasta_string: str) -> Generator[FastaEntry, None, None]:
-    """
-    Parses a fasta file and yields FastaEntry objects
-
-    Args:
-        fasta_string: The fasta file as a string
-    Returns:
-        A generator of FastaEntry objects
-    """
-    header = None
-    seq = []
-    num_sequences = 0
-    for line in fasta_string.splitlines():
-        if not line or line[0] == "#":
-            continue
-        if line.startswith(">"):
-            if header is not None:
-                yield FastaEntry(header, "".join(seq))
-                seq = []
-            header = line[1:].strip()
-        else:
-            seq.append(line)
-    if header is not None:
-        num_sequences += 1
-        yield FastaEntry(header, "".join(seq))
-
-    if num_sequences == 0:
-        raise ValueError("Found no sequences in input")
-
-
-def read_sequences(path: PathOrBuffer) -> Generator[FastaEntry, None, None]:
-    # Uses duck typing to try and call the right method
-    # Doesn't use explicit isinstance check to support
-    # inputs that are not explicitly str/Path/TextIOBase but
-    # may support similar functionality
-    data = None  # type: ignore
-    try:
-        if str(path).endswith(".gz"):
-            import gzip
-
-            data = gzip.open(path, "rt")  # type: ignore
-        else:
-            try:
-                data = open(path)  # type: ignore
-            except TypeError:
-                data: io.TextIOBase = path  # type: ignore
-
-        yield from parse_fasta(data.read())
-    finally:
-        if data is not None:
-            data.close()
-
-
-def read_first_sequence(path: PathOrBuffer) -> FastaEntry:
-    return next(iter(read_sequences(path)))
-
-
-def count_fasta_sequences(path: str | Path) -> int:
-    """Count sequences in a FASTA file by counting header lines.
-
-    Faster than parsing the full file — only scans for '>' prefixes.
-    Returns 0 if the file does not exist.
-    """
-    path = Path(path)
-    if not path.exists():
-        return 0
-    with open(path) as f:
-        return sum(1 for line in f if line.startswith(">"))
-
-
-def append_fasta_sequence(header: str, sequence: str, path: str | Path) -> None:
-    """Append a single sequence to a FASTA file (creating it if needed)."""
-    path = Path(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    # The existing file may not end with a newline (e.g., write_sequences()
-    # explicitly avoids writing a newline at the end), so we insert one before
-    # appending to avoid merging with the last line.
-    needs_newline = (
-        path.exists() and path.stat().st_size > 0 and path.read_bytes()[-1:] != b"\n"
-    )
-    with open(path, "a") as f:
-        if needs_newline:
-            f.write("\n")
-        f.write(f">{header}\n{sequence}\n")
-
-
-def write_sequences(sequences: Iterable[tuple[str, str]], path: PathOrBuffer) -> None:
-    needs_closing = False
-    handle = None
-    try:
-        try:
-            handle = open(path, "w")  # type: ignore
-            needs_closing = True
-        except TypeError:
-            handle = path
-        has_prev = False
-        for header, seq in sequences:
-            if has_prev:
-                handle.write("\n")  # type: ignore
-            handle.write(f">{header}\n{seq}")  # type: ignore
-            has_prev = True
-    finally:
-        if needs_closing:
-            handle.close()  # type: ignore
+import io
+from pathlib import Path
+from typing import Generator, Iterable, NamedTuple
+
+PathOrBuffer = str | Path | io.TextIOBase
+FastaEntry = NamedTuple("FastaEntry", [("header", str), ("sequence", str)])
+
+
+def parse_fasta(fasta_string: str) -> Generator[FastaEntry, None, None]:
+    """
+    Parses a fasta file and yields FastaEntry objects
+
+    Args:
+        fasta_string: The fasta file as a string
+    Returns:
+        A generator of FastaEntry objects
+    """
+    header = None
+    seq = []
+    num_sequences = 0
+    for line in fasta_string.splitlines():
+        if not line or line[0] == "#":
+            continue
+        if line.startswith(">"):
+            if header is not None:
+                yield FastaEntry(header, "".join(seq))
+                seq = []
+            header = line[1:].strip()
+        else:
+            seq.append(line)
+    if header is not None:
+        num_sequences += 1
+        yield FastaEntry(header, "".join(seq))
+
+    if num_sequences == 0:
+        raise ValueError("Found no sequences in input")
+
+
+def read_sequences(path: PathOrBuffer) -> Generator[FastaEntry, None, None]:
+    # Uses duck typing to try and call the right method
+    # Doesn't use explicit isinstance check to support
+    # inputs that are not explicitly str/Path/TextIOBase but
+    # may support similar functionality
+    data = None  # type: ignore
+    try:
+        if str(path).endswith(".gz"):
+            import gzip
+
+            data = gzip.open(path, "rt")  # type: ignore
+        else:
+            try:
+                data = open(path)  # type: ignore
+            except TypeError:
+                data: io.TextIOBase = path  # type: ignore
+
+        yield from parse_fasta(data.read())
+    finally:
+        if data is not None:
+            data.close()
+
+
+def read_first_sequence(path: PathOrBuffer) -> FastaEntry:
+    return next(iter(read_sequences(path)))
+
+
+def count_fasta_sequences(path: str | Path) -> int:
+    """Count sequences in a FASTA file by counting header lines.
+
+    Faster than parsing the full file — only scans for '>' prefixes.
+    Returns 0 if the file does not exist.
+    """
+    path = Path(path)
+    if not path.exists():
+        return 0
+    with open(path) as f:
+        return sum(1 for line in f if line.startswith(">"))
+
+
+def append_fasta_sequence(header: str, sequence: str, path: str | Path) -> None:
+    """Append a single sequence to a FASTA file (creating it if needed)."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    # The existing file may not end with a newline (e.g., write_sequences()
+    # explicitly avoids writing a newline at the end), so we insert one before
+    # appending to avoid merging with the last line.
+    needs_newline = (
+        path.exists() and path.stat().st_size > 0 and path.read_bytes()[-1:] != b"\n"
+    )
+    with open(path, "a") as f:
+        if needs_newline:
+            f.write("\n")
+        f.write(f">{header}\n{sequence}\n")
+
+
+def write_sequences(sequences: Iterable[tuple[str, str]], path: PathOrBuffer) -> None:
+    needs_closing = False
+    handle = None
+    try:
+        try:
+            handle = open(path, "w")  # type: ignore
+            needs_closing = True
+        except TypeError:
+            handle = path
+        has_prev = False
+        for header, seq in sequences:
+            if has_prev:
+                handle.write("\n")  # type: ignore
+            handle.write(f">{header}\n{seq}")  # type: ignore
+            has_prev = True
+    finally:
+        if needs_closing:
+            handle.close()  # type: ignore

esmfold2_predicted_aligned_error.py

@@ -1,104 +1,104 @@
-import torch
-import torch.nn.functional as F
-
-from .esmfold2_affine3d import Affine3D
-
-
-def masked_mean(
-    mask: torch.Tensor,
-    value: torch.Tensor,
-    dim: int | None | tuple[int, ...] = None,
-    eps=1e-10,
-) -> torch.Tensor:
-    """Compute the mean of `value` where only positions where `mask == true` are
-    counted.
-    """
-    mask = mask.expand(*value.shape)
-    return torch.sum(mask * value, dim=dim) / (eps + torch.sum(mask, dim=dim))
-
-
-def _pae_bins(
-    max_bin: float = 31, num_bins: int = 64, device: torch.device = torch.device("cpu")
-):
-    bins = torch.linspace(0, max_bin, steps=(num_bins - 1), device=device)
-    step = max_bin / (num_bins - 2)
-    bin_centers = bins + step / 2
-    bin_centers = torch.cat(
-        [bin_centers, (bin_centers[-1] + step).unsqueeze(-1)], dim=0
-    )
-    return bin_centers
-
-
-def _compute_pae_masks(mask: torch.Tensor):
-    square_mask = (mask.unsqueeze(-1) * mask.unsqueeze(-2)).bool()
-    return square_mask
-
-
-def compute_predicted_aligned_error(
-    logits: torch.Tensor,
-    aa_mask: torch.Tensor,
-    sequence_id: torch.Tensor | None = None,
-    max_bin: float = 31,
-) -> torch.Tensor:
-    bins = _pae_bins(max_bin, logits.shape[-1], logits.device)
-    square_mask = _compute_pae_masks(aa_mask)
-    min_v = torch.finfo(logits.dtype).min
-    probs = logits.masked_fill(~square_mask.unsqueeze(-1), min_v).softmax(dim=-1)
-
-    return (probs * bins).sum(dim=-1)
-
-
-@torch.no_grad
-def compute_tm(logits: torch.Tensor, aa_mask: torch.Tensor, max_bin: float = 31.0):
-    square_mask = _compute_pae_masks(aa_mask)
-    seqlens = aa_mask.sum(-1, keepdim=True)
-    bins = _pae_bins(max_bin, logits.shape[-1], logits.device)
-    d0 = 1.24 * (seqlens.clamp_min(19) - 15) ** (1 / 3) - 1.8
-    f_d = 1.0 / (1 + (bins / d0.unsqueeze(-1)) ** 2)
-
-    min_v = torch.finfo(logits.dtype).min
-    probs = logits.masked_fill(~square_mask.unsqueeze(-1), min_v).softmax(dim=-1)
-    # This is the sum over bins
-    ptm = (probs * f_d.unsqueeze(-2)).sum(dim=-1)
-    # This is the mean over residues j
-    ptm = masked_mean(square_mask, ptm, dim=-1)
-    # The we do a max over residues i
-    return ptm.max(dim=-1).values
-
-
-def tm_loss(
-    logits: torch.Tensor,
-    pred_affine: torch.Tensor,
-    targ_affine: torch.Tensor,
-    targ_mask: torch.Tensor,
-    tm_mask: torch.Tensor | None = None,
-    sequence_id: torch.Tensor | None = None,
-    max_bin: float = 31,
-):
-    pred = Affine3D.from_tensor(pred_affine)
-    targ = Affine3D.from_tensor(targ_affine)
-
-    def transform(affine: Affine3D):
-        pts = affine.trans[..., None, :, :]
-        return affine.invert()[..., None].apply(pts)
-
-    with torch.no_grad():
-        sq_diff = (transform(pred) - transform(targ)).square().sum(dim=-1)
-
-        num_bins = logits.shape[-1]
-        sq_bins = torch.linspace(
-            0, max_bin, num_bins - 1, device=logits.device
-        ).square()
-        # Gets the bin id by using a sum.
-        true_bins = (sq_diff[..., None] > sq_bins).sum(dim=-1).long()
-
-    errors = F.cross_entropy(logits.movedim(3, 1), true_bins, reduction="none")
-    square_mask = _compute_pae_masks(targ_mask)
-    loss = masked_mean(square_mask, errors, dim=(-1, -2))
-
-    if tm_mask is not None:
-        loss = masked_mean(tm_mask, loss, dim=None)
-    else:
-        loss = loss.mean()
-
-    return loss
+import torch
+import torch.nn.functional as F
+
+from .esmfold2_affine3d import Affine3D
+
+
+def masked_mean(
+    mask: torch.Tensor,
+    value: torch.Tensor,
+    dim: int | None | tuple[int, ...] = None,
+    eps=1e-10,
+) -> torch.Tensor:
+    """Compute the mean of `value` where only positions where `mask == true` are
+    counted.
+    """
+    mask = mask.expand(*value.shape)
+    return torch.sum(mask * value, dim=dim) / (eps + torch.sum(mask, dim=dim))
+
+
+def _pae_bins(
+    max_bin: float = 31, num_bins: int = 64, device: torch.device = torch.device("cpu")
+):
+    bins = torch.linspace(0, max_bin, steps=(num_bins - 1), device=device)
+    step = max_bin / (num_bins - 2)
+    bin_centers = bins + step / 2
+    bin_centers = torch.cat(
+        [bin_centers, (bin_centers[-1] + step).unsqueeze(-1)], dim=0
+    )
+    return bin_centers
+
+
+def _compute_pae_masks(mask: torch.Tensor):
+    square_mask = (mask.unsqueeze(-1) * mask.unsqueeze(-2)).bool()
+    return square_mask
+
+
+def compute_predicted_aligned_error(
+    logits: torch.Tensor,
+    aa_mask: torch.Tensor,
+    sequence_id: torch.Tensor | None = None,
+    max_bin: float = 31,
+) -> torch.Tensor:
+    bins = _pae_bins(max_bin, logits.shape[-1], logits.device)
+    square_mask = _compute_pae_masks(aa_mask)
+    min_v = torch.finfo(logits.dtype).min
+    probs = logits.masked_fill(~square_mask.unsqueeze(-1), min_v).softmax(dim=-1)
+
+    return (probs * bins).sum(dim=-1)
+
+
+@torch.no_grad
+def compute_tm(logits: torch.Tensor, aa_mask: torch.Tensor, max_bin: float = 31.0):
+    square_mask = _compute_pae_masks(aa_mask)
+    seqlens = aa_mask.sum(-1, keepdim=True)
+    bins = _pae_bins(max_bin, logits.shape[-1], logits.device)
+    d0 = 1.24 * (seqlens.clamp_min(19) - 15) ** (1 / 3) - 1.8
+    f_d = 1.0 / (1 + (bins / d0.unsqueeze(-1)) ** 2)
+
+    min_v = torch.finfo(logits.dtype).min
+    probs = logits.masked_fill(~square_mask.unsqueeze(-1), min_v).softmax(dim=-1)
+    # This is the sum over bins
+    ptm = (probs * f_d.unsqueeze(-2)).sum(dim=-1)
+    # This is the mean over residues j
+    ptm = masked_mean(square_mask, ptm, dim=-1)
+    # The we do a max over residues i
+    return ptm.max(dim=-1).values
+
+
+def tm_loss(
+    logits: torch.Tensor,
+    pred_affine: torch.Tensor,
+    targ_affine: torch.Tensor,
+    targ_mask: torch.Tensor,
+    tm_mask: torch.Tensor | None = None,
+    sequence_id: torch.Tensor | None = None,
+    max_bin: float = 31,
+):
+    pred = Affine3D.from_tensor(pred_affine)
+    targ = Affine3D.from_tensor(targ_affine)
+
+    def transform(affine: Affine3D):
+        pts = affine.trans[..., None, :, :]
+        return affine.invert()[..., None].apply(pts)
+
+    with torch.no_grad():
+        sq_diff = (transform(pred) - transform(targ)).square().sum(dim=-1)
+
+        num_bins = logits.shape[-1]
+        sq_bins = torch.linspace(
+            0, max_bin, num_bins - 1, device=logits.device
+        ).square()
+        # Gets the bin id by using a sum.
+        true_bins = (sq_diff[..., None] > sq_bins).sum(dim=-1).long()
+
+    errors = F.cross_entropy(logits.movedim(3, 1), true_bins, reduction="none")
+    square_mask = _compute_pae_masks(targ_mask)
+    loss = masked_mean(square_mask, errors, dim=(-1, -2))
+
+    if tm_mask is not None:
+        loss = masked_mean(tm_mask, loss, dim=None)
+    else:
+        loss = loss.mean()
+
+    return loss

esmfold2_processor.py

@@ -1,355 +1,355 @@
-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"]
+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"]

esmfold2_protein_structure.py

@@ -1,306 +1,306 @@
-from __future__ import annotations
-
-from typing import Tuple, TypeVar
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import Tensor
-from torch.amp import autocast  # type: ignore
-
-from . import esmfold2_residue_constants as residue_constants
-from .esmfold2_misc import unbinpack
-from .esmfold2_affine3d import Affine3D
-
-ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor)
-
-
-def index_by_atom_name(
-    atom37: ArrayOrTensor, atom_names: str | list[str], dim: int = -2
-) -> ArrayOrTensor:
-    squeeze = False
-    if isinstance(atom_names, str):
-        atom_names = [atom_names]
-        squeeze = True
-    indices = [residue_constants.atom_order[atom_name] for atom_name in atom_names]
-    dim = dim % atom37.ndim
-    index = tuple(slice(None) if dim != i else indices for i in range(atom37.ndim))
-    result = atom37[index]  # type: ignore
-    if squeeze:
-        result = result.squeeze(dim)
-    return result
-
-
-def infer_cbeta_from_atom37(
-    atom37: ArrayOrTensor, L: float = 1.522, A: float = 1.927, D: float = -2.143
-):
-    """
-    Inspired by a util in trDesign:
-    https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L92
-
-    input:  atom37, (L)ength, (A)ngle, and (D)ihedral
-    output: 4th coord
-    """
-    N = index_by_atom_name(atom37, "N", dim=-2)
-    CA = index_by_atom_name(atom37, "CA", dim=-2)
-    C = index_by_atom_name(atom37, "C", dim=-2)
-
-    if isinstance(atom37, np.ndarray):
-
-        def normalize(x: ArrayOrTensor):
-            return x / np.linalg.norm(x, axis=-1, keepdims=True)
-
-        cross = np.cross
-    else:
-        normalize = F.normalize  # type: ignore
-        cross = torch.cross
-
-    with np.errstate(invalid="ignore"):  # inf - inf = nan is ok here
-        vec_nca = N - CA
-        vec_nc = N - C
-    nca = normalize(vec_nca)
-    n = normalize(cross(vec_nc, nca))  # type: ignore
-    m = [nca, cross(n, nca), n]
-    d = [L * np.cos(A), L * np.sin(A) * np.cos(D), -L * np.sin(A) * np.sin(D)]
-    return CA + sum([m * d for m, d in zip(m, d)])
-
-
-@torch.no_grad()
-@autocast("cuda", enabled=False)
-def compute_alignment_tensors(
-    mobile: torch.Tensor,
-    target: torch.Tensor,
-    atom_exists_mask: torch.Tensor | None = None,
-    sequence_id: torch.Tensor | None = None,
-):
-    """
-    Align two batches of structures with support for masking invalid atoms using PyTorch.
-
-    Args:
-    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
-    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
-    - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N)
-    - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking.
-
-    Returns:
-    - centered_mobile (torch.Tensor): Batch of coordinates of structure centered mobile (B, N, 3)
-    - centroid_mobile (torch.Tensor): Batch of coordinates of mobile centeroid (B, 3)
-    - centered_target (torch.Tensor): Batch of coordinates of structure centered target (B, N, 3)
-    - centroid_target (torch.Tensor): Batch of coordinates of target centeroid (B, 3)
-    - rotation_matrix (torch.Tensor): Batch of coordinates of rotation matrix (B, 3, 3)
-    - num_valid_atoms (torch.Tensor): Batch of number of valid atoms for alignment (B,)
-    """
-
-    # Ensure both batches have the same number of structures, atoms, and dimensions
-    if sequence_id is not None:
-        mobile = unbinpack(mobile, sequence_id, pad_value=torch.nan)
-        target = unbinpack(target, sequence_id, pad_value=torch.nan)
-        if atom_exists_mask is not None:
-            atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=0)
-        else:
-            atom_exists_mask = torch.isfinite(target).all(-1)
-
-    assert mobile.shape == target.shape, "Batch structure shapes do not match!"
-
-    # Number of structures in the batch
-    batch_size = mobile.shape[0]
-
-    # if [B, Nres, Natom, 3], resize
-    if mobile.dim() == 4:
-        mobile = mobile.view(batch_size, -1, 3)
-    if target.dim() == 4:
-        target = target.view(batch_size, -1, 3)
-    if atom_exists_mask is not None and atom_exists_mask.dim() == 3:
-        atom_exists_mask = atom_exists_mask.view(batch_size, -1)
-
-    # Number of atoms
-    num_atoms = mobile.shape[1]
-
-    # Apply masks if provided
-    if atom_exists_mask is not None:
-        mobile = mobile.masked_fill(~atom_exists_mask.unsqueeze(-1), 0)
-        target = target.masked_fill(~atom_exists_mask.unsqueeze(-1), 0)
-    else:
-        atom_exists_mask = torch.ones(
-            batch_size, num_atoms, dtype=torch.bool, device=mobile.device
-        )
-
-    num_valid_atoms = atom_exists_mask.sum(dim=-1, keepdim=True)
-    # Compute centroids for each batch
-    centroid_mobile = mobile.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1)
-    centroid_target = target.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1)
-
-    # Handle potential division by zero if all atoms are invalid in a structure
-    centroid_mobile[num_valid_atoms == 0] = 0
-    centroid_target[num_valid_atoms == 0] = 0
-
-    # Center structures by subtracting centroids
-    centered_mobile = mobile - centroid_mobile
-    centered_target = target - centroid_target
-
-    centered_mobile = centered_mobile.masked_fill(~atom_exists_mask.unsqueeze(-1), 0)
-    centered_target = centered_target.masked_fill(~atom_exists_mask.unsqueeze(-1), 0)
-
-    # Compute covariance matrix for each batch
-    covariance_matrix = torch.matmul(centered_mobile.transpose(1, 2), centered_target)
-
-    # Singular Value Decomposition for each batch
-    u, _, v = torch.svd(covariance_matrix)
-
-    # Calculate rotation matrices for each batch
-    rotation_matrix = torch.matmul(u, v.transpose(1, 2))
-
-    return (
-        centered_mobile,
-        centroid_mobile,
-        centered_target,
-        centroid_target,
-        rotation_matrix,
-        num_valid_atoms,
-    )
-
-
-@torch.no_grad()
-@autocast("cuda", enabled=False)
-def compute_rmsd_no_alignment(
-    aligned: torch.Tensor,
-    target: torch.Tensor,
-    num_valid_atoms: torch.Tensor,
-    reduction: str = "batch",
-) -> torch.Tensor:
-    """
-    Compute RMSD between two batches of structures without alignment.
-
-    Args:
-    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
-    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
-    - num_valid_atoms (torch.Tensor): Batch of number of valid atoms for alignment (B,)
-    - reduction (str): One of "batch", "per_sample", "per_residue".
-
-    Returns:
-
-    If reduction == "batch":
-        (torch.Tensor): 0-dim, Average Root Mean Square Deviation between the structures for each batch
-    If reduction == "per_sample":
-        (torch.Tensor): (B,)-dim, Root Mean Square Deviation between the structures for each batch
-    If reduction == "per_residue":
-        (torch.Tensor): (B, N)-dim, Root Mean Square Deviation between the structures for residue in the batch
-    """
-    if reduction not in ("per_residue", "per_sample", "batch"):
-        raise ValueError("Unrecognized reduction: '{reduction}'")
-    # Compute RMSD for each batch
-    diff = aligned - target
-    if reduction == "per_residue":
-        mean_squared_error = diff.square().view(diff.size(0), -1, 9).mean(dim=-1)
-    else:
-        mean_squared_error = diff.square().sum(dim=(1, 2)) / (
-            num_valid_atoms.squeeze(-1)
-        )
-
-    rmsd = torch.sqrt(mean_squared_error)
-    if reduction in ("per_sample", "per_residue"):
-        return rmsd
-    elif reduction == "batch":
-        avg_rmsd = rmsd.masked_fill(num_valid_atoms.squeeze(-1) == 0, 0).sum() / (
-            (num_valid_atoms > 0).sum() + 1e-8
-        )
-        return avg_rmsd
-    else:
-        raise ValueError(reduction)
-
-
-@torch.no_grad()
-@autocast("cuda", enabled=False)
-def compute_affine_and_rmsd(
-    mobile: torch.Tensor,
-    target: torch.Tensor,
-    atom_exists_mask: torch.Tensor | None = None,
-    sequence_id: torch.Tensor | None = None,
-) -> Tuple[Affine3D, torch.Tensor]:
-    """
-    Compute RMSD between two batches of structures with support for masking invalid atoms using PyTorch.
-
-    Args:
-    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
-    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
-    - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N)
-    - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking.
-
-    Returns:
-    - affine (Affine3D): Transformation between mobile and target structure
-    - avg_rmsd (torch.Tensor): Average Root Mean Square Deviation between the structures for each batch
-    """
-
-    (
-        centered_mobile,
-        centroid_mobile,
-        centered_target,
-        centroid_target,
-        rotation_matrix,
-        num_valid_atoms,
-    ) = compute_alignment_tensors(
-        mobile=mobile,
-        target=target,
-        atom_exists_mask=atom_exists_mask,
-        sequence_id=sequence_id,
-    )
-
-    # Apply rotation to mobile centroid
-    translation = torch.matmul(-centroid_mobile, rotation_matrix) + centroid_target
-    affine = Affine3D.from_tensor_pair(
-        translation, rotation_matrix.unsqueeze(dim=-3).transpose(-2, -1)
-    )
-
-    # Apply transformation to centered structure to compute rmsd
-    rotated_mobile = torch.matmul(centered_mobile, rotation_matrix)
-    avg_rmsd = compute_rmsd_no_alignment(
-        rotated_mobile, centered_target, num_valid_atoms, reduction="batch"
-    )
-
-    return affine, avg_rmsd
-
-
-def compute_gdt_ts_no_alignment(
-    aligned: torch.Tensor,
-    target: torch.Tensor,
-    atom_exists_mask: torch.Tensor,
-    reduction: str = "batch",
-) -> torch.Tensor:
-    """
-    Compute GDT_TS between two batches of structures without alignment.
-
-    Args:
-    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
-    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
-    - atom_exists_mask (torch.Tensor): Mask for Whether an atom exists of shape (B, N). noo
-    - reduction (str): One of "batch", "per_sample".
-
-    Returns:
-    If reduction == "batch":
-        (torch.Tensor): 0-dim, GDT_TS between the structures for each batch
-    If reduction == "per_sample":
-        (torch.Tensor): (B,)-dim, GDT_TS between the structures for each sample in the batch
-    """
-    if reduction not in ("per_sample", "batch"):
-        raise ValueError("Unrecognized reduction: '{reduction}'")
-
-    if atom_exists_mask is None:
-        atom_exists_mask = torch.isfinite(target).all(dim=-1)
-
-    deviation = torch.linalg.vector_norm(aligned - target, dim=-1)
-    num_valid_atoms = atom_exists_mask.sum(dim=-1)
-
-    # Compute GDT_TS
-    score = (
-        ((deviation < 1) * atom_exists_mask).sum(dim=-1) / num_valid_atoms
-        + ((deviation < 2) * atom_exists_mask).sum(dim=-1) / num_valid_atoms
-        + ((deviation < 4) * atom_exists_mask).sum(dim=-1) / num_valid_atoms
-        + ((deviation < 8) * atom_exists_mask).sum(dim=-1) / num_valid_atoms
-    ) * 0.25
-
-    if reduction == "batch":
-        return score.mean()
-    elif reduction == "per_sample":
-        return score
-    else:
-        raise ValueError("Unrecognized reduction: '{reduction}'")
+from __future__ import annotations
+
+from typing import Tuple, TypeVar
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.amp import autocast  # type: ignore
+
+from . import esmfold2_residue_constants as residue_constants
+from .esmfold2_misc import unbinpack
+from .esmfold2_affine3d import Affine3D
+
+ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor)
+
+
+def index_by_atom_name(
+    atom37: ArrayOrTensor, atom_names: str | list[str], dim: int = -2
+) -> ArrayOrTensor:
+    squeeze = False
+    if isinstance(atom_names, str):
+        atom_names = [atom_names]
+        squeeze = True
+    indices = [residue_constants.atom_order[atom_name] for atom_name in atom_names]
+    dim = dim % atom37.ndim
+    index = tuple(slice(None) if dim != i else indices for i in range(atom37.ndim))
+    result = atom37[index]  # type: ignore
+    if squeeze:
+        result = result.squeeze(dim)
+    return result
+
+
+def infer_cbeta_from_atom37(
+    atom37: ArrayOrTensor, L: float = 1.522, A: float = 1.927, D: float = -2.143
+):
+    """
+    Inspired by a util in trDesign:
+    https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L92
+
+    input:  atom37, (L)ength, (A)ngle, and (D)ihedral
+    output: 4th coord
+    """
+    N = index_by_atom_name(atom37, "N", dim=-2)
+    CA = index_by_atom_name(atom37, "CA", dim=-2)
+    C = index_by_atom_name(atom37, "C", dim=-2)
+
+    if isinstance(atom37, np.ndarray):
+
+        def normalize(x: ArrayOrTensor):
+            return x / np.linalg.norm(x, axis=-1, keepdims=True)
+
+        cross = np.cross
+    else:
+        normalize = F.normalize  # type: ignore
+        cross = torch.cross
+
+    with np.errstate(invalid="ignore"):  # inf - inf = nan is ok here
+        vec_nca = N - CA
+        vec_nc = N - C
+    nca = normalize(vec_nca)
+    n = normalize(cross(vec_nc, nca))  # type: ignore
+    m = [nca, cross(n, nca), n]
+    d = [L * np.cos(A), L * np.sin(A) * np.cos(D), -L * np.sin(A) * np.sin(D)]
+    return CA + sum([m * d for m, d in zip(m, d)])
+
+
+@torch.no_grad()
+@autocast("cuda", enabled=False)
+def compute_alignment_tensors(
+    mobile: torch.Tensor,
+    target: torch.Tensor,
+    atom_exists_mask: torch.Tensor | None = None,
+    sequence_id: torch.Tensor | None = None,
+):
+    """
+    Align two batches of structures with support for masking invalid atoms using PyTorch.
+
+    Args:
+    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
+    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
+    - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N)
+    - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking.
+
+    Returns:
+    - centered_mobile (torch.Tensor): Batch of coordinates of structure centered mobile (B, N, 3)
+    - centroid_mobile (torch.Tensor): Batch of coordinates of mobile centeroid (B, 3)
+    - centered_target (torch.Tensor): Batch of coordinates of structure centered target (B, N, 3)
+    - centroid_target (torch.Tensor): Batch of coordinates of target centeroid (B, 3)
+    - rotation_matrix (torch.Tensor): Batch of coordinates of rotation matrix (B, 3, 3)
+    - num_valid_atoms (torch.Tensor): Batch of number of valid atoms for alignment (B,)
+    """
+
+    # Ensure both batches have the same number of structures, atoms, and dimensions
+    if sequence_id is not None:
+        mobile = unbinpack(mobile, sequence_id, pad_value=torch.nan)
+        target = unbinpack(target, sequence_id, pad_value=torch.nan)
+        if atom_exists_mask is not None:
+            atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=0)
+        else:
+            atom_exists_mask = torch.isfinite(target).all(-1)
+
+    assert mobile.shape == target.shape, "Batch structure shapes do not match!"
+
+    # Number of structures in the batch
+    batch_size = mobile.shape[0]
+
+    # if [B, Nres, Natom, 3], resize
+    if mobile.dim() == 4:
+        mobile = mobile.view(batch_size, -1, 3)
+    if target.dim() == 4:
+        target = target.view(batch_size, -1, 3)
+    if atom_exists_mask is not None and atom_exists_mask.dim() == 3:
+        atom_exists_mask = atom_exists_mask.view(batch_size, -1)
+
+    # Number of atoms
+    num_atoms = mobile.shape[1]
+
+    # Apply masks if provided
+    if atom_exists_mask is not None:
+        mobile = mobile.masked_fill(~atom_exists_mask.unsqueeze(-1), 0)
+        target = target.masked_fill(~atom_exists_mask.unsqueeze(-1), 0)
+    else:
+        atom_exists_mask = torch.ones(
+            batch_size, num_atoms, dtype=torch.bool, device=mobile.device
+        )
+
+    num_valid_atoms = atom_exists_mask.sum(dim=-1, keepdim=True)
+    # Compute centroids for each batch
+    centroid_mobile = mobile.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1)
+    centroid_target = target.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1)
+
+    # Handle potential division by zero if all atoms are invalid in a structure
+    centroid_mobile[num_valid_atoms == 0] = 0
+    centroid_target[num_valid_atoms == 0] = 0
+
+    # Center structures by subtracting centroids
+    centered_mobile = mobile - centroid_mobile
+    centered_target = target - centroid_target
+
+    centered_mobile = centered_mobile.masked_fill(~atom_exists_mask.unsqueeze(-1), 0)
+    centered_target = centered_target.masked_fill(~atom_exists_mask.unsqueeze(-1), 0)
+
+    # Compute covariance matrix for each batch
+    covariance_matrix = torch.matmul(centered_mobile.transpose(1, 2), centered_target)
+
+    # Singular Value Decomposition for each batch
+    u, _, v = torch.svd(covariance_matrix)
+
+    # Calculate rotation matrices for each batch
+    rotation_matrix = torch.matmul(u, v.transpose(1, 2))
+
+    return (
+        centered_mobile,
+        centroid_mobile,
+        centered_target,
+        centroid_target,
+        rotation_matrix,
+        num_valid_atoms,
+    )
+
+
+@torch.no_grad()
+@autocast("cuda", enabled=False)
+def compute_rmsd_no_alignment(
+    aligned: torch.Tensor,
+    target: torch.Tensor,
+    num_valid_atoms: torch.Tensor,
+    reduction: str = "batch",
+) -> torch.Tensor:
+    """
+    Compute RMSD between two batches of structures without alignment.
+
+    Args:
+    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
+    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
+    - num_valid_atoms (torch.Tensor): Batch of number of valid atoms for alignment (B,)
+    - reduction (str): One of "batch", "per_sample", "per_residue".
+
+    Returns:
+
+    If reduction == "batch":
+        (torch.Tensor): 0-dim, Average Root Mean Square Deviation between the structures for each batch
+    If reduction == "per_sample":
+        (torch.Tensor): (B,)-dim, Root Mean Square Deviation between the structures for each batch
+    If reduction == "per_residue":
+        (torch.Tensor): (B, N)-dim, Root Mean Square Deviation between the structures for residue in the batch
+    """
+    if reduction not in ("per_residue", "per_sample", "batch"):
+        raise ValueError("Unrecognized reduction: '{reduction}'")
+    # Compute RMSD for each batch
+    diff = aligned - target
+    if reduction == "per_residue":
+        mean_squared_error = diff.square().view(diff.size(0), -1, 9).mean(dim=-1)
+    else:
+        mean_squared_error = diff.square().sum(dim=(1, 2)) / (
+            num_valid_atoms.squeeze(-1)
+        )
+
+    rmsd = torch.sqrt(mean_squared_error)
+    if reduction in ("per_sample", "per_residue"):
+        return rmsd
+    elif reduction == "batch":
+        avg_rmsd = rmsd.masked_fill(num_valid_atoms.squeeze(-1) == 0, 0).sum() / (
+            (num_valid_atoms > 0).sum() + 1e-8
+        )
+        return avg_rmsd
+    else:
+        raise ValueError(reduction)
+
+
+@torch.no_grad()
+@autocast("cuda", enabled=False)
+def compute_affine_and_rmsd(
+    mobile: torch.Tensor,
+    target: torch.Tensor,
+    atom_exists_mask: torch.Tensor | None = None,
+    sequence_id: torch.Tensor | None = None,
+) -> Tuple[Affine3D, torch.Tensor]:
+    """
+    Compute RMSD between two batches of structures with support for masking invalid atoms using PyTorch.
+
+    Args:
+    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
+    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
+    - atom_exists_mask (torch.Tensor, optional): Mask for Whether an atom exists of shape (B, N)
+    - sequence_id (torch.Tensor, optional): Sequence id tensor for binpacking.
+
+    Returns:
+    - affine (Affine3D): Transformation between mobile and target structure
+    - avg_rmsd (torch.Tensor): Average Root Mean Square Deviation between the structures for each batch
+    """
+
+    (
+        centered_mobile,
+        centroid_mobile,
+        centered_target,
+        centroid_target,
+        rotation_matrix,
+        num_valid_atoms,
+    ) = compute_alignment_tensors(
+        mobile=mobile,
+        target=target,
+        atom_exists_mask=atom_exists_mask,
+        sequence_id=sequence_id,
+    )
+
+    # Apply rotation to mobile centroid
+    translation = torch.matmul(-centroid_mobile, rotation_matrix) + centroid_target
+    affine = Affine3D.from_tensor_pair(
+        translation, rotation_matrix.unsqueeze(dim=-3).transpose(-2, -1)
+    )
+
+    # Apply transformation to centered structure to compute rmsd
+    rotated_mobile = torch.matmul(centered_mobile, rotation_matrix)
+    avg_rmsd = compute_rmsd_no_alignment(
+        rotated_mobile, centered_target, num_valid_atoms, reduction="batch"
+    )
+
+    return affine, avg_rmsd
+
+
+def compute_gdt_ts_no_alignment(
+    aligned: torch.Tensor,
+    target: torch.Tensor,
+    atom_exists_mask: torch.Tensor,
+    reduction: str = "batch",
+) -> torch.Tensor:
+    """
+    Compute GDT_TS between two batches of structures without alignment.
+
+    Args:
+    - mobile (torch.Tensor): Batch of coordinates of structure to be superimposed in shape (B, N, 3)
+    - target (torch.Tensor): Batch of coordinates of structure that is fixed in shape (B, N, 3)
+    - atom_exists_mask (torch.Tensor): Mask for Whether an atom exists of shape (B, N). noo
+    - reduction (str): One of "batch", "per_sample".
+
+    Returns:
+    If reduction == "batch":
+        (torch.Tensor): 0-dim, GDT_TS between the structures for each batch
+    If reduction == "per_sample":
+        (torch.Tensor): (B,)-dim, GDT_TS between the structures for each sample in the batch
+    """
+    if reduction not in ("per_sample", "batch"):
+        raise ValueError("Unrecognized reduction: '{reduction}'")
+
+    if atom_exists_mask is None:
+        atom_exists_mask = torch.isfinite(target).all(dim=-1)
+
+    deviation = torch.linalg.vector_norm(aligned - target, dim=-1)
+    num_valid_atoms = atom_exists_mask.sum(dim=-1)
+
+    # Compute GDT_TS
+    score = (
+        ((deviation < 1) * atom_exists_mask).sum(dim=-1) / num_valid_atoms
+        + ((deviation < 2) * atom_exists_mask).sum(dim=-1) / num_valid_atoms
+        + ((deviation < 4) * atom_exists_mask).sum(dim=-1) / num_valid_atoms
+        + ((deviation < 8) * atom_exists_mask).sum(dim=-1) / num_valid_atoms
+    ) * 0.25
+
+    if reduction == "batch":
+        return score.mean()
+    elif reduction == "per_sample":
+        return score
+    else:
+        raise ValueError("Unrecognized reduction: '{reduction}'")

esmfold2_residue_constants.py

@@ -1,1223 +1,1223 @@
-# Copyright 2025 EvolutionaryScale
-# Copyright 2021 AlQuraishi Laboratory
-# Copyright 2021 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Constants used in AlphaFold."""
-
-import collections
-import functools
-from pathlib import Path
-from typing import List, Mapping, Tuple
-
-import numpy as np
-
-# import tree
-
-# Internal import (35fd).
-
-
-# Distance from one CA to next CA [trans configuration: omega = 180].
-ca_ca = 3.80209737096
-
-# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
-# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
-# chi angles so their chi angle lists are empty.
-chi_angles_atoms = {
-    "ALA": [],
-    # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
-    "ARG": [
-        ["N", "CA", "CB", "CG"],
-        ["CA", "CB", "CG", "CD"],
-        ["CB", "CG", "CD", "NE"],
-        ["CG", "CD", "NE", "CZ"],
-    ],
-    "ASN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
-    "ASP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
-    "CYS": [["N", "CA", "CB", "SG"]],
-    "GLN": [
-        ["N", "CA", "CB", "CG"],
-        ["CA", "CB", "CG", "CD"],
-        ["CB", "CG", "CD", "OE1"],
-    ],
-    "GLU": [
-        ["N", "CA", "CB", "CG"],
-        ["CA", "CB", "CG", "CD"],
-        ["CB", "CG", "CD", "OE1"],
-    ],
-    "GLY": [],
-    "HIS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "ND1"]],
-    "ILE": [["N", "CA", "CB", "CG1"], ["CA", "CB", "CG1", "CD1"]],
-    "LEU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
-    "LYS": [
-        ["N", "CA", "CB", "CG"],
-        ["CA", "CB", "CG", "CD"],
-        ["CB", "CG", "CD", "CE"],
-        ["CG", "CD", "CE", "NZ"],
-    ],
-    "MET": [
-        ["N", "CA", "CB", "CG"],
-        ["CA", "CB", "CG", "SD"],
-        ["CB", "CG", "SD", "CE"],
-    ],
-    "PHE": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
-    "PRO": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"]],
-    "SER": [["N", "CA", "CB", "OG"]],
-    "THR": [["N", "CA", "CB", "OG1"]],
-    "TRP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
-    "TYR": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
-    "VAL": [["N", "CA", "CB", "CG1"]],
-    "UNK": [],
-}
-
-# If chi angles given in fixed-length array, this matrix determines how to mask
-# them for each AA type. The order is as per restype_order (see below).
-chi_angles_mask = [
-    [0.0, 0.0, 0.0, 0.0],  # ALA
-    [1.0, 1.0, 1.0, 1.0],  # ARG
-    [1.0, 1.0, 0.0, 0.0],  # ASN
-    [1.0, 1.0, 0.0, 0.0],  # ASP
-    [1.0, 0.0, 0.0, 0.0],  # CYS
-    [1.0, 1.0, 1.0, 0.0],  # GLN
-    [1.0, 1.0, 1.0, 0.0],  # GLU
-    [0.0, 0.0, 0.0, 0.0],  # GLY
-    [1.0, 1.0, 0.0, 0.0],  # HIS
-    [1.0, 1.0, 0.0, 0.0],  # ILE
-    [1.0, 1.0, 0.0, 0.0],  # LEU
-    [1.0, 1.0, 1.0, 1.0],  # LYS
-    [1.0, 1.0, 1.0, 0.0],  # MET
-    [1.0, 1.0, 0.0, 0.0],  # PHE
-    [1.0, 1.0, 0.0, 0.0],  # PRO
-    [1.0, 0.0, 0.0, 0.0],  # SER
-    [1.0, 0.0, 0.0, 0.0],  # THR
-    [1.0, 1.0, 0.0, 0.0],  # TRP
-    [1.0, 1.0, 0.0, 0.0],  # TYR
-    [1.0, 0.0, 0.0, 0.0],  # VAL
-    [0.0, 0.0, 0.0, 0.0],  # UNK
-]
-
-# The following chi angles are pi periodic: they can be rotated by a multiple
-# of pi without affecting the structure.
-chi_pi_periodic = [
-    [0.0, 0.0, 0.0, 0.0],  # ALA
-    [0.0, 0.0, 0.0, 0.0],  # ARG
-    [0.0, 0.0, 0.0, 0.0],  # ASN
-    [0.0, 1.0, 0.0, 0.0],  # ASP
-    [0.0, 0.0, 0.0, 0.0],  # CYS
-    [0.0, 0.0, 0.0, 0.0],  # GLN
-    [0.0, 0.0, 1.0, 0.0],  # GLU
-    [0.0, 0.0, 0.0, 0.0],  # GLY
-    [0.0, 0.0, 0.0, 0.0],  # HIS
-    [0.0, 0.0, 0.0, 0.0],  # ILE
-    [0.0, 0.0, 0.0, 0.0],  # LEU
-    [0.0, 0.0, 0.0, 0.0],  # LYS
-    [0.0, 0.0, 0.0, 0.0],  # MET
-    [0.0, 1.0, 0.0, 0.0],  # PHE
-    [0.0, 0.0, 0.0, 0.0],  # PRO
-    [0.0, 0.0, 0.0, 0.0],  # SER
-    [0.0, 0.0, 0.0, 0.0],  # THR
-    [0.0, 0.0, 0.0, 0.0],  # TRP
-    [0.0, 1.0, 0.0, 0.0],  # TYR
-    [0.0, 0.0, 0.0, 0.0],  # VAL
-    [0.0, 0.0, 0.0, 0.0],  # UNK
-]
-
-# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi,
-# psi and chi angles:
-# 0: 'backbone group',
-# 1: 'pre-omega-group', (empty)
-# 2: 'phi-group', (currently empty, because it defines only hydrogens)
-# 3: 'psi-group',
-# 4,5,6,7: 'chi1,2,3,4-group'
-# The atom positions are relative to the axis-end-atom of the corresponding
-# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis
-# is defined such that the dihedral-angle-definiting atom (the last entry in
-# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate).
-# format: [atomname, group_idx, rel_position]
-rigid_group_atom_positions = {
-    "ALA": [
-        ["N", 0, (-0.525, 1.363, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.526, -0.000, -0.000)],
-        ["CB", 0, (-0.529, -0.774, -1.205)],
-        ["O", 3, (0.627, 1.062, 0.000)],
-    ],
-    "ARG": [
-        ["N", 0, (-0.524, 1.362, -0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.525, -0.000, -0.000)],
-        ["CB", 0, (-0.524, -0.778, -1.209)],
-        ["O", 3, (0.626, 1.062, 0.000)],
-        ["CG", 4, (0.616, 1.390, -0.000)],
-        ["CD", 5, (0.564, 1.414, 0.000)],
-        ["NE", 6, (0.539, 1.357, -0.000)],
-        ["NH1", 7, (0.206, 2.301, 0.000)],
-        ["NH2", 7, (2.078, 0.978, -0.000)],
-        ["CZ", 7, (0.758, 1.093, -0.000)],
-    ],
-    "ASN": [
-        ["N", 0, (-0.536, 1.357, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.526, -0.000, -0.000)],
-        ["CB", 0, (-0.531, -0.787, -1.200)],
-        ["O", 3, (0.625, 1.062, 0.000)],
-        ["CG", 4, (0.584, 1.399, 0.000)],
-        ["ND2", 5, (0.593, -1.188, 0.001)],
-        ["OD1", 5, (0.633, 1.059, 0.000)],
-    ],
-    "ASP": [
-        ["N", 0, (-0.525, 1.362, -0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.527, 0.000, -0.000)],
-        ["CB", 0, (-0.526, -0.778, -1.208)],
-        ["O", 3, (0.626, 1.062, -0.000)],
-        ["CG", 4, (0.593, 1.398, -0.000)],
-        ["OD1", 5, (0.610, 1.091, 0.000)],
-        ["OD2", 5, (0.592, -1.101, -0.003)],
-    ],
-    "CYS": [
-        ["N", 0, (-0.522, 1.362, -0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.524, 0.000, 0.000)],
-        ["CB", 0, (-0.519, -0.773, -1.212)],
-        ["O", 3, (0.625, 1.062, -0.000)],
-        ["SG", 4, (0.728, 1.653, 0.000)],
-    ],
-    "GLN": [
-        ["N", 0, (-0.526, 1.361, -0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.526, 0.000, 0.000)],
-        ["CB", 0, (-0.525, -0.779, -1.207)],
-        ["O", 3, (0.626, 1.062, -0.000)],
-        ["CG", 4, (0.615, 1.393, 0.000)],
-        ["CD", 5, (0.587, 1.399, -0.000)],
-        ["NE2", 6, (0.593, -1.189, -0.001)],
-        ["OE1", 6, (0.634, 1.060, 0.000)],
-    ],
-    "GLU": [
-        ["N", 0, (-0.528, 1.361, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.526, -0.000, -0.000)],
-        ["CB", 0, (-0.526, -0.781, -1.207)],
-        ["O", 3, (0.626, 1.062, 0.000)],
-        ["CG", 4, (0.615, 1.392, 0.000)],
-        ["CD", 5, (0.600, 1.397, 0.000)],
-        ["OE1", 6, (0.607, 1.095, -0.000)],
-        ["OE2", 6, (0.589, -1.104, -0.001)],
-    ],
-    "GLY": [
-        ["N", 0, (-0.572, 1.337, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.517, -0.000, -0.000)],
-        ["O", 3, (0.626, 1.062, -0.000)],
-    ],
-    "HIS": [
-        ["N", 0, (-0.527, 1.360, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.525, 0.000, 0.000)],
-        ["CB", 0, (-0.525, -0.778, -1.208)],
-        ["O", 3, (0.625, 1.063, 0.000)],
-        ["CG", 4, (0.600, 1.370, -0.000)],
-        ["CD2", 5, (0.889, -1.021, 0.003)],
-        ["ND1", 5, (0.744, 1.160, -0.000)],
-        ["CE1", 5, (2.030, 0.851, 0.002)],
-        ["NE2", 5, (2.145, -0.466, 0.004)],
-    ],
-    "ILE": [
-        ["N", 0, (-0.493, 1.373, -0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.527, -0.000, -0.000)],
-        ["CB", 0, (-0.536, -0.793, -1.213)],
-        ["O", 3, (0.627, 1.062, -0.000)],
-        ["CG1", 4, (0.534, 1.437, -0.000)],
-        ["CG2", 4, (0.540, -0.785, -1.199)],
-        ["CD1", 5, (0.619, 1.391, 0.000)],
-    ],
-    "LEU": [
-        ["N", 0, (-0.520, 1.363, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.525, -0.000, -0.000)],
-        ["CB", 0, (-0.522, -0.773, -1.214)],
-        ["O", 3, (0.625, 1.063, -0.000)],
-        ["CG", 4, (0.678, 1.371, 0.000)],
-        ["CD1", 5, (0.530, 1.430, -0.000)],
-        ["CD2", 5, (0.535, -0.774, 1.200)],
-    ],
-    "LYS": [
-        ["N", 0, (-0.526, 1.362, -0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.526, 0.000, 0.000)],
-        ["CB", 0, (-0.524, -0.778, -1.208)],
-        ["O", 3, (0.626, 1.062, -0.000)],
-        ["CG", 4, (0.619, 1.390, 0.000)],
-        ["CD", 5, (0.559, 1.417, 0.000)],
-        ["CE", 6, (0.560, 1.416, 0.000)],
-        ["NZ", 7, (0.554, 1.387, 0.000)],
-    ],
-    "MET": [
-        ["N", 0, (-0.521, 1.364, -0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.525, 0.000, 0.000)],
-        ["CB", 0, (-0.523, -0.776, -1.210)],
-        ["O", 3, (0.625, 1.062, -0.000)],
-        ["CG", 4, (0.613, 1.391, -0.000)],
-        ["SD", 5, (0.703, 1.695, 0.000)],
-        ["CE", 6, (0.320, 1.786, -0.000)],
-    ],
-    "PHE": [
-        ["N", 0, (-0.518, 1.363, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.524, 0.000, -0.000)],
-        ["CB", 0, (-0.525, -0.776, -1.212)],
-        ["O", 3, (0.626, 1.062, -0.000)],
-        ["CG", 4, (0.607, 1.377, 0.000)],
-        ["CD1", 5, (0.709, 1.195, -0.000)],
-        ["CD2", 5, (0.706, -1.196, 0.000)],
-        ["CE1", 5, (2.102, 1.198, -0.000)],
-        ["CE2", 5, (2.098, -1.201, -0.000)],
-        ["CZ", 5, (2.794, -0.003, -0.001)],
-    ],
-    "PRO": [
-        ["N", 0, (-0.566, 1.351, -0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.527, -0.000, 0.000)],
-        ["CB", 0, (-0.546, -0.611, -1.293)],
-        ["O", 3, (0.621, 1.066, 0.000)],
-        ["CG", 4, (0.382, 1.445, 0.0)],
-        # ['CD', 5, (0.427, 1.440, 0.0)],
-        ["CD", 5, (0.477, 1.424, 0.0)],  # manually made angle 2 degrees larger
-    ],
-    "SER": [
-        ["N", 0, (-0.529, 1.360, -0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.525, -0.000, -0.000)],
-        ["CB", 0, (-0.518, -0.777, -1.211)],
-        ["O", 3, (0.626, 1.062, -0.000)],
-        ["OG", 4, (0.503, 1.325, 0.000)],
-    ],
-    "THR": [
-        ["N", 0, (-0.517, 1.364, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.526, 0.000, -0.000)],
-        ["CB", 0, (-0.516, -0.793, -1.215)],
-        ["O", 3, (0.626, 1.062, 0.000)],
-        ["CG2", 4, (0.550, -0.718, -1.228)],
-        ["OG1", 4, (0.472, 1.353, 0.000)],
-    ],
-    "TRP": [
-        ["N", 0, (-0.521, 1.363, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.525, -0.000, 0.000)],
-        ["CB", 0, (-0.523, -0.776, -1.212)],
-        ["O", 3, (0.627, 1.062, 0.000)],
-        ["CG", 4, (0.609, 1.370, -0.000)],
-        ["CD1", 5, (0.824, 1.091, 0.000)],
-        ["CD2", 5, (0.854, -1.148, -0.005)],
-        ["CE2", 5, (2.186, -0.678, -0.007)],
-        ["CE3", 5, (0.622, -2.530, -0.007)],
-        ["NE1", 5, (2.140, 0.690, -0.004)],
-        ["CH2", 5, (3.028, -2.890, -0.013)],
-        ["CZ2", 5, (3.283, -1.543, -0.011)],
-        ["CZ3", 5, (1.715, -3.389, -0.011)],
-    ],
-    "TYR": [
-        ["N", 0, (-0.522, 1.362, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.524, -0.000, -0.000)],
-        ["CB", 0, (-0.522, -0.776, -1.213)],
-        ["O", 3, (0.627, 1.062, -0.000)],
-        ["CG", 4, (0.607, 1.382, -0.000)],
-        ["CD1", 5, (0.716, 1.195, -0.000)],
-        ["CD2", 5, (0.713, -1.194, -0.001)],
-        ["CE1", 5, (2.107, 1.200, -0.002)],
-        ["CE2", 5, (2.104, -1.201, -0.003)],
-        ["OH", 5, (4.168, -0.002, -0.005)],
-        ["CZ", 5, (2.791, -0.001, -0.003)],
-    ],
-    "VAL": [
-        ["N", 0, (-0.494, 1.373, -0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.527, -0.000, -0.000)],
-        ["CB", 0, (-0.533, -0.795, -1.213)],
-        ["O", 3, (0.627, 1.062, -0.000)],
-        ["CG1", 4, (0.540, 1.429, -0.000)],
-        ["CG2", 4, (0.533, -0.776, 1.203)],
-    ],
-    # Assume alanine positions for unknown AA
-    "UNK": [
-        ["N", 0, (-0.525, 1.363, 0.000)],
-        ["CA", 0, (0.000, 0.000, 0.000)],
-        ["C", 0, (1.526, -0.000, -0.000)],
-    ],
-}
-
-# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
-residue_atoms = {
-    "ALA": ["C", "CA", "CB", "N", "O"],
-    "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"],
-    "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"],
-    "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"],
-    "CYS": ["C", "CA", "CB", "N", "O", "SG"],
-    "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"],
-    "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"],
-    "GLY": ["C", "CA", "N", "O"],
-    "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"],
-    "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"],
-    "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"],
-    "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"],
-    "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"],
-    "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"],
-    "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"],
-    "SER": ["C", "CA", "CB", "N", "O", "OG"],
-    "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"],
-    "TRP": [
-        "C",
-        "CA",
-        "CB",
-        "CG",
-        "CD1",
-        "CD2",
-        "CE2",
-        "CE3",
-        "CZ2",
-        "CZ3",
-        "CH2",
-        "N",
-        "NE1",
-        "O",
-    ],
-    "TYR": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O", "OH"],
-    "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"],
-    "UNK": ["C", "CA", "N"],
-}
-
-# Naming swaps for ambiguous atom names.
-# Due to symmetries in the amino acids the naming of atoms is ambiguous in
-# 4 of the 20 amino acids.
-# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities
-# in LEU, VAL and ARG can be resolved by using the 3d constellations of
-# the 'ambiguous' atoms and their neighbours)
-# TODO: ^ interpret this
-residue_atom_renaming_swaps = {
-    "ASP": {"OD1": "OD2"},
-    "GLU": {"OE1": "OE2"},
-    "PHE": {"CD1": "CD2", "CE1": "CE2"},
-    "TYR": {"CD1": "CD2", "CE1": "CE2"},
-}
-
-# Van der Waals radii [Angstroem] of the atoms (from Wikipedia)
-van_der_waals_radius = {"C": 1.7, "N": 1.55, "O": 1.52, "S": 1.8}
-
-Bond = collections.namedtuple("Bond", ["atom1_name", "atom2_name", "length", "stddev"])
-BondAngle = collections.namedtuple(
-    "BondAngle", ["atom1_name", "atom2_name", "atom3name", "angle_rad", "stddev"]
-)
-
-
-@functools.lru_cache(maxsize=None)
-def load_stereo_chemical_props() -> (
-    Tuple[
-        Mapping[str, List[Bond]],
-        Mapping[str, List[Bond]],
-        Mapping[str, List[BondAngle]],
-    ]
-):
-    """Load stereo_chemical_props.txt into a nice structure.
-
-    Load literature values for bond lengths and bond angles and translate
-    bond angles into the length of the opposite edge of the triangle
-    ("residue_virtual_bonds").
-
-    Returns:
-      residue_bonds:  dict that maps resname --> list of Bond tuples
-      residue_virtual_bonds: dict that maps resname --> list of Bond tuples
-      residue_bond_angles: dict that maps resname --> list of BondAngle tuples
-    """
-    stereo_chemical_props = Path(
-        "evolutionaryscale/structure/stereo_chemical_props.txt"
-    ).read_text()
-
-    lines_iter = iter(stereo_chemical_props.splitlines())
-    # Load bond lengths.
-    residue_bonds = {}
-    next(lines_iter)  # Skip header line.
-    for line in lines_iter:
-        if line.strip() == "-":
-            break
-        bond, resname, length, stddev = line.split()
-        atom1, atom2 = bond.split("-")
-        if resname not in residue_bonds:
-            residue_bonds[resname] = []
-        residue_bonds[resname].append(Bond(atom1, atom2, float(length), float(stddev)))
-    residue_bonds["UNK"] = []
-
-    # Load bond angles.
-    residue_bond_angles = {}
-    next(lines_iter)  # Skip empty line.
-    next(lines_iter)  # Skip header line.
-    for line in lines_iter:
-        if line.strip() == "-":
-            break
-        bond, resname, angle_degree, stddev_degree = line.split()
-        atom1, atom2, atom3 = bond.split("-")
-        if resname not in residue_bond_angles:
-            residue_bond_angles[resname] = []
-        residue_bond_angles[resname].append(
-            BondAngle(
-                atom1,
-                atom2,
-                atom3,
-                float(angle_degree) / 180.0 * np.pi,
-                float(stddev_degree) / 180.0 * np.pi,
-            )
-        )
-    residue_bond_angles["UNK"] = []
-
-    def make_bond_key(atom1_name, atom2_name):
-        """Unique key to lookup bonds."""
-        return "-".join(sorted([atom1_name, atom2_name]))
-
-    # Translate bond angles into distances ("virtual bonds").
-    residue_virtual_bonds = {}
-    for resname, bond_angles in residue_bond_angles.items():
-        # Create a fast lookup dict for bond lengths.
-        bond_cache = {}
-        for b in residue_bonds[resname]:
-            bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
-        residue_virtual_bonds[resname] = []
-        for ba in bond_angles:
-            bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
-            bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]
-
-            # Compute distance between atom1 and atom3 using the law of cosines
-            # c^2 = a^2 + b^2 - 2ab*cos(gamma).
-            gamma = ba.angle_rad
-            length = np.sqrt(
-                bond1.length**2
-                + bond2.length**2
-                - 2 * bond1.length * bond2.length * np.cos(gamma)
-            )
-
-            # Propagation of uncertainty assuming uncorrelated errors.
-            dl_outer = 0.5 / length
-            dl_dgamma = (2 * bond1.length * bond2.length * np.sin(gamma)) * dl_outer
-            dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer
-            dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer
-            stddev = np.sqrt(
-                (dl_dgamma * ba.stddev) ** 2
-                + (dl_db1 * bond1.stddev) ** 2
-                + (dl_db2 * bond2.stddev) ** 2
-            )
-            residue_virtual_bonds[resname].append(
-                Bond(ba.atom1_name, ba.atom3name, length, stddev)
-            )
-
-    return (residue_bonds, residue_virtual_bonds, residue_bond_angles)
-
-
-# Between-residue bond lengths for general bonds (first element) and for Proline
-# (second element).
-between_res_bond_length_c_n = [1.329, 1.341]
-between_res_bond_length_stddev_c_n = [0.014, 0.016]
-
-# Between-residue cos_angles.
-between_res_cos_angles_c_n_ca = [-0.5203, 0.0353]  # degrees: 121.352 +- 2.315
-between_res_cos_angles_ca_c_n = [-0.4473, 0.0311]  # degrees: 116.568 +- 1.995
-
-# This mapping is used when we need to store atom data in a format that requires
-# fixed atom data size for every residue (e.g. a numpy array).
-atom_types = [
-    "N",
-    "CA",
-    "C",
-    "CB",
-    "O",
-    "CG",
-    "CG1",
-    "CG2",
-    "OG",
-    "OG1",
-    "SG",
-    "CD",
-    "CD1",
-    "CD2",
-    "ND1",
-    "ND2",
-    "OD1",
-    "OD2",
-    "SD",
-    "CE",
-    "CE1",
-    "CE2",
-    "CE3",
-    "NE",
-    "NE1",
-    "NE2",
-    "OE1",
-    "OE2",
-    "CH2",
-    "NH1",
-    "NH2",
-    "OH",
-    "CZ",
-    "CZ2",
-    "CZ3",
-    "NZ",
-    "OXT",
-]
-atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
-atom_type_num = len(atom_types)  # := 37.
-
-# A compact atom encoding with 14 columns
-# pylint: disable=line-too-long
-# pylint: disable=bad-whitespace
-restype_name_to_atom14_names = {
-    "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""],
-    "ARG": [
-        "N",
-        "CA",
-        "C",
-        "O",
-        "CB",
-        "CG",
-        "CD",
-        "NE",
-        "CZ",
-        "NH1",
-        "NH2",
-        "",
-        "",
-        "",
-    ],
-    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2", "", "", "", "", "", ""],
-    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2", "", "", "", "", "", ""],
-    "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""],
-    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2", "", "", "", "", ""],
-    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2", "", "", "", "", ""],
-    "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""],
-    "HIS": [
-        "N",
-        "CA",
-        "C",
-        "O",
-        "CB",
-        "CG",
-        "ND1",
-        "CD2",
-        "CE1",
-        "NE2",
-        "",
-        "",
-        "",
-        "",
-    ],
-    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1", "", "", "", "", "", ""],
-    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "", "", "", "", "", ""],
-    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ", "", "", "", "", ""],
-    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE", "", "", "", "", "", ""],
-    "PHE": [
-        "N",
-        "CA",
-        "C",
-        "O",
-        "CB",
-        "CG",
-        "CD1",
-        "CD2",
-        "CE1",
-        "CE2",
-        "CZ",
-        "",
-        "",
-        "",
-    ],
-    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""],
-    "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""],
-    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2", "", "", "", "", "", "", ""],
-    "TRP": [
-        "N",
-        "CA",
-        "C",
-        "O",
-        "CB",
-        "CG",
-        "CD1",
-        "CD2",
-        "NE1",
-        "CE2",
-        "CE3",
-        "CZ2",
-        "CZ3",
-        "CH2",
-    ],
-    "TYR": [
-        "N",
-        "CA",
-        "C",
-        "O",
-        "CB",
-        "CG",
-        "CD1",
-        "CD2",
-        "CE1",
-        "CE2",
-        "CZ",
-        "OH",
-        "",
-        "",
-    ],
-    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "", "", "", "", "", "", ""],
-    "UNK": ["N", "CA", "C", "", "", "", "", "", "", "", "", "", "", ""],
-}
-# pylint: enable=line-too-long
-# pylint: enable=bad-whitespace
-
-
-# This is the standard residue order when coding AA type as a number.
-# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
-restypes = [
-    "A",
-    "R",
-    "N",
-    "D",
-    "C",
-    "Q",
-    "E",
-    "G",
-    "H",
-    "I",
-    "L",
-    "K",
-    "M",
-    "F",
-    "P",
-    "S",
-    "T",
-    "W",
-    "Y",
-    "V",
-]
-restype_order = {restype: i for i, restype in enumerate(restypes)}
-restype_num = len(restypes)  # := 20.
-unk_restype_index = restype_num  # Catch-all index for unknown restypes.
-
-restypes_with_x = restypes + ["X"]
-restype_order_with_x = {restype: i for i, restype in enumerate(restypes_with_x)}
-
-bb_atoms = ["N", "CA", "C", "O"]
-
-# Hydrophobicity by residue (positive values are hydrophobic). Derived from Black & Mould (1991), normalized by subtracting 0.5.
-hydrophobicity = {
-    "ALA": 0.116,
-    "ARG": -0.5,
-    "ASN": -0.264,
-    "ASP": -0.472,
-    "CYS": 0.18,
-    "GLN": -0.249,
-    "GLU": -0.457,
-    "GLY": 0.001,
-    "HIS": -0.335,
-    "ILE": 0.443,
-    "LEU": 0.443,
-    "LYS": -0.217,
-    "MET": 0.238,
-    "PHE": 0.5,
-    "PRO": 0.211,
-    "SER": -0.141,
-    "THR": -0.05,
-    "TRP": 0.378,
-    "TYR": 0.38,
-    "VAL": 0.325,
-}
-
-# Side chain max accessible surface area in Ala-X-Ala tripeptide (from Chennamsetty et al. 2010).
-side_chain_asa = {
-    "ALA": 64.7809,
-    "ARG": 210.02,
-    "ASN": 113.187,
-    "ASP": 110.209,
-    "CYS": 95.2439,
-    "GLN": 147.855,
-    "GLU": 143.924,
-    "GLY": 23.1338,
-    "HIS": 146.449,
-    "ILE": 151.242,
-    "LEU": 139.524,
-    "LYS": 177.366,
-    "MET": 164.674,
-    "PHE": 186.7,
-    "PRO": 111.533,
-    "SER": 81.2159,
-    "THR": 111.597,
-    "TRP": 229.619,
-    "TYR": 200.306,
-    "VAL": 124.237,
-}
-
-# Approximate Volumes of amino acids in cubic angstroms.
-# https://www.imgt.org/IMGTeducation/Aide-memoire/_UK/aminoacids/abbreviation.html
-amino_acid_volumes = {
-    "A": 88.6,  # Alanine
-    "R": 173.4,  # Arginine
-    "N": 114.1,  # Asparagine
-    "D": 111.1,  # Aspartic acid
-    "C": 108.5,  # Cysteine
-    "Q": 143.8,  # Glutamine
-    "E": 138.4,  # Glutamic acid
-    "G": 60.1,  # Glycine
-    "H": 153.2,  # Histidine
-    "I": 166.7,  # Isoleucine
-    "L": 166.7,  # Leucine
-    "K": 168.6,  # Lysine
-    "M": 162.9,  # Methionine
-    "F": 189.9,  # Phenylalanine
-    "P": 112.7,  # Proline
-    "S": 89.0,  # Serine
-    "T": 116.1,  # Threonine
-    "W": 227.8,  # Tryptophan
-    "Y": 193.6,  # Tyrosine
-    "V": 140.0,  # Valine
-    "X": 88.6,  # Unknown, use Alanine as approximation
-}
-
-
-def sequence_to_onehot(
-    sequence: str, mapping: Mapping[str, int], map_unknown_to_x: bool = False
-) -> np.ndarray:
-    """Maps the given sequence into a one-hot encoded matrix.
-
-    Args:
-      sequence: An amino acid sequence.
-      mapping: A dictionary mapping amino acids to integers.
-      map_unknown_to_x: If True, any amino acid that is not in the mapping will be
-        mapped to the unknown amino acid 'X'. If the mapping doesn't contain
-        amino acid 'X', an error will be thrown. If False, any amino acid not in
-        the mapping will throw an error.
-
-    Returns:
-      A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of
-      the sequence.
-
-    Raises:
-      ValueError: If the mapping doesn't contain values from 0 to
-        num_unique_aas - 1 without any gaps.
-    """
-    num_entries = max(mapping.values()) + 1
-
-    if sorted(set(mapping.values())) != list(range(num_entries)):
-        raise ValueError(
-            "The mapping must have values from 0 to num_unique_aas-1 "
-            "without any gaps. Got: %s" % sorted(mapping.values())
-        )
-
-    one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)
-
-    for aa_index, aa_type in enumerate(sequence):
-        if map_unknown_to_x:
-            if aa_type.isalpha() and aa_type.isupper():
-                aa_id = mapping.get(aa_type, mapping["X"])
-            else:
-                raise ValueError(f"Invalid character in the sequence: {aa_type}")
-        else:
-            aa_id = mapping[aa_type]
-        one_hot_arr[aa_index, aa_id] = 1
-
-    return one_hot_arr
-
-
-restype_1to3 = {
-    "A": "ALA",
-    "R": "ARG",
-    "N": "ASN",
-    "D": "ASP",
-    "C": "CYS",
-    "Q": "GLN",
-    "E": "GLU",
-    "G": "GLY",
-    "H": "HIS",
-    "I": "ILE",
-    "L": "LEU",
-    "K": "LYS",
-    "M": "MET",
-    "F": "PHE",
-    "P": "PRO",
-    "S": "SER",
-    "T": "THR",
-    "W": "TRP",
-    "Y": "TYR",
-    "V": "VAL",
-    "X": "UNK",
-}
-
-
-# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple
-# 1-to-1 mapping of 3 letter names to one letter names. The latter contains
-# many more, and less common, three letter names as keys and maps many of these
-# to the same one letter name (including 'X' and 'U' which we don't use here).
-restype_3to1 = {v: k for k, v in restype_1to3.items()}
-
-# Define a restype name for all unknown residues.
-unk_restype = "UNK"
-
-resnames = [restype_1to3[r] for r in restypes] + [unk_restype]
-resname_to_idx = {resname: i for i, resname in enumerate(resnames)}
-
-hydrophobic_resnames = {"VAL", "ILE", "LEU", "PHE", "MET", "TRP"}
-
-# The mapping here uses hhblits convention, so that B is mapped to D, J and O
-# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the
-# remaining 20 amino acids are kept in alphabetical order.
-# There are 2 non-amino acid codes, X (representing any amino acid) and
-# "-" representing a missing amino acid in an alignment.  The id for these
-# codes is put at the end (20 and 21) so that they can easily be ignored if
-# desired.
-HHBLITS_AA_TO_ID = {
-    "A": 0,
-    "B": 2,
-    "C": 1,
-    "D": 2,
-    "E": 3,
-    "F": 4,
-    "G": 5,
-    "H": 6,
-    "I": 7,
-    "J": 20,
-    "K": 8,
-    "L": 9,
-    "M": 10,
-    "N": 11,
-    "O": 20,
-    "P": 12,
-    "Q": 13,
-    "R": 14,
-    "S": 15,
-    "T": 16,
-    "U": 1,
-    "V": 17,
-    "W": 18,
-    "X": 20,
-    "Y": 19,
-    "Z": 3,
-    "-": 21,
-}
-
-# Partial inversion of HHBLITS_AA_TO_ID.
-ID_TO_HHBLITS_AA = {
-    0: "A",
-    1: "C",  # Also U.
-    2: "D",  # Also B.
-    3: "E",  # Also Z.
-    4: "F",
-    5: "G",
-    6: "H",
-    7: "I",
-    8: "K",
-    9: "L",
-    10: "M",
-    11: "N",
-    12: "P",
-    13: "Q",
-    14: "R",
-    15: "S",
-    16: "T",
-    17: "V",
-    18: "W",
-    19: "Y",
-    20: "X",  # Includes J and O.
-    21: "-",
-}
-
-restypes_with_x_and_gap = restypes + ["X", "-"]
-MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
-    restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i])
-    for i in range(len(restypes_with_x_and_gap))
-)
-
-
-def _make_standard_atom_mask() -> np.ndarray:
-    """Returns [num_res_types, num_atom_types] mask array."""
-    # +1 to account for unknown (all 0s).
-    mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32)
-    for restype, restype_letter in enumerate(restypes):
-        restype_name = restype_1to3[restype_letter]
-        atom_names = residue_atoms[restype_name]
-        for atom_name in atom_names:
-            atom_type = atom_order[atom_name]
-            mask[restype, atom_type] = 1
-    return mask
-
-
-STANDARD_ATOM_MASK = _make_standard_atom_mask()
-
-
-# A one hot representation for the first and second atoms defining the axis
-# of rotation for each chi-angle in each residue.
-def chi_angle_atom(atom_index: int) -> np.ndarray:
-    """Define chi-angle rigid groups via one-hot representations."""
-    chi_angles_index = {}
-    one_hots = []
-
-    for k, v in chi_angles_atoms.items():
-        indices = [atom_types.index(s[atom_index]) for s in v]
-        indices.extend([-1] * (4 - len(indices)))
-        chi_angles_index[k] = indices
-
-    for r in restypes:
-        res3 = restype_1to3[r]
-        one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
-        one_hots.append(one_hot)
-
-    one_hots.append(np.zeros([4, atom_type_num]))  # Add zeros for residue `X`.
-    one_hot = np.stack(one_hots, axis=0)
-    one_hot = np.transpose(one_hot, [0, 2, 1])
-
-    return one_hot
-
-
-chi_atom_1_one_hot = chi_angle_atom(1)
-chi_atom_2_one_hot = chi_angle_atom(2)
-
-# An array like chi_angles_atoms but using indices rather than names.
-chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes]
-# chi_angles_atom_indices = tree.map_structure(
-#     lambda atom_name: atom_order[atom_name], chi_angles_atom_indices
-# )
-chi_angles_atom_indices = np.array(
-    [
-        chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms)))
-        for chi_atoms in chi_angles_atom_indices
-    ]
-)
-
-# Mapping from (res_name, atom_name) pairs to the atom's chi group index
-# and atom index within that group.
-chi_groups_for_atom = collections.defaultdict(list)
-for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items():
-    for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
-        for atom_i, atom in enumerate(chi_group):
-            chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i))
-chi_groups_for_atom = dict(chi_groups_for_atom)
-
-
-def _make_rigid_transformation_4x4(ex, ey, translation):
-    """Create a rigid 4x4 transformation matrix from two axes and transl."""
-    # Normalize ex.
-    ex_normalized = ex / np.linalg.norm(ex)
-
-    # make ey perpendicular to ex
-    ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized
-    ey_normalized /= np.linalg.norm(ey_normalized)
-
-    # compute ez as cross product
-    eznorm = np.cross(ex_normalized, ey_normalized)
-    m = np.stack([ex_normalized, ey_normalized, eznorm, translation]).transpose()
-    m = np.concatenate([m, [[0.0, 0.0, 0.0, 1.0]]], axis=0)
-    return m
-
-
-# create an array with (restype, atomtype) --> rigid_group_idx
-# and an array with (restype, atomtype, coord) for the atom positions
-# and compute affine transformation matrices (4,4) from one rigid group to the
-# previous group
-restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=int)
-restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
-restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32)
-restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=int)
-restype_atom14_mask = np.zeros([21, 14], dtype=np.float32)
-restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32)
-restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32)
-
-
-def _make_rigid_group_constants():
-    """Fill the arrays above."""
-    for restype, restype_letter in enumerate(restypes_with_x):
-        resname = restype_1to3[restype_letter]
-        for atomname, group_idx, atom_position in rigid_group_atom_positions[resname]:
-            atomtype = atom_order[atomname]
-            restype_atom37_to_rigid_group[restype, atomtype] = group_idx
-            restype_atom37_mask[restype, atomtype] = 1
-            restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position
-
-            atom14idx = restype_name_to_atom14_names[resname].index(atomname)
-            restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
-            restype_atom14_mask[restype, atom14idx] = 1
-            restype_atom14_rigid_group_positions[restype, atom14idx, :] = atom_position
-
-    for restype, restype_letter in enumerate(restypes_with_x):
-        resname = restype_1to3[restype_letter]
-        atom_positions = {
-            name: np.array(pos) for name, _, pos in rigid_group_atom_positions[resname]
-        }
-
-        # backbone to backbone is the identity transform
-        restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)
-
-        # pre-omega-frame to backbone (currently dummy identity matrix)
-        restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)
-
-        # phi-frame to backbone
-        mat = _make_rigid_transformation_4x4(
-            ex=atom_positions["N"] - atom_positions["CA"],
-            ey=np.array([1.0, 0.0, 0.0]),
-            translation=atom_positions["N"],
-        )
-        restype_rigid_group_default_frame[restype, 2, :, :] = mat
-
-        # psi-frame to backbone
-        mat = _make_rigid_transformation_4x4(
-            ex=atom_positions["C"] - atom_positions["CA"],
-            ey=atom_positions["CA"] - atom_positions["N"],
-            translation=atom_positions["C"],
-        )
-        restype_rigid_group_default_frame[restype, 3, :, :] = mat
-
-        # chi1-frame to backbone
-        if chi_angles_mask[restype][0]:
-            base_atom_names = chi_angles_atoms[resname][0]
-            base_atom_positions = [atom_positions[name] for name in base_atom_names]
-            mat = _make_rigid_transformation_4x4(
-                ex=base_atom_positions[2] - base_atom_positions[1],
-                ey=base_atom_positions[0] - base_atom_positions[1],
-                translation=base_atom_positions[2],
-            )
-            restype_rigid_group_default_frame[restype, 4, :, :] = mat
-
-        # chi2-frame to chi1-frame
-        # chi3-frame to chi2-frame
-        # chi4-frame to chi3-frame
-        # luckily all rotation axes for the next frame start at (0,0,0) of the
-        # previous frame
-        for chi_idx in range(1, 4):
-            if chi_angles_mask[restype][chi_idx]:
-                axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
-                axis_end_atom_position = atom_positions[axis_end_atom_name]
-                mat = _make_rigid_transformation_4x4(
-                    ex=axis_end_atom_position,
-                    ey=np.array([-1.0, 0.0, 0.0]),
-                    translation=axis_end_atom_position,
-                )
-                restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat
-
-
-_make_rigid_group_constants()
-
-
-def make_atom14_dists_bounds(overlap_tolerance=1.5, bond_length_tolerance_factor=15.0):
-    """compute upper and lower bounds for bonds to assess violations."""
-    restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32)
-    restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32)
-    restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32)
-    residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props()
-    for restype, restype_letter in enumerate(restypes):
-        resname = restype_1to3[restype_letter]
-        atom_list = restype_name_to_atom14_names[resname]
-
-        # create lower and upper bounds for clashes
-        for atom1_idx, atom1_name in enumerate(atom_list):
-            if not atom1_name:
-                continue
-            atom1_radius = van_der_waals_radius[atom1_name[0]]
-            for atom2_idx, atom2_name in enumerate(atom_list):
-                if (not atom2_name) or atom1_idx == atom2_idx:
-                    continue
-                atom2_radius = van_der_waals_radius[atom2_name[0]]
-                lower = atom1_radius + atom2_radius - overlap_tolerance
-                upper = 1e10
-                restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
-                restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
-                restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
-                restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
-
-        # overwrite lower and upper bounds for bonds and angles
-        for b in residue_bonds[resname] + residue_virtual_bonds[resname]:
-            atom1_idx = atom_list.index(b.atom1_name)
-            atom2_idx = atom_list.index(b.atom2_name)
-            lower = b.length - bond_length_tolerance_factor * b.stddev
-            upper = b.length + bond_length_tolerance_factor * b.stddev
-            restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
-            restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
-            restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
-            restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
-            restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev
-            restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev
-    return {
-        "lower_bound": restype_atom14_bond_lower_bound,  # shape (21,14,14)
-        "upper_bound": restype_atom14_bond_upper_bound,  # shape (21,14,14)
-        "stddev": restype_atom14_bond_stddev,  # shape (21,14,14)
-    }
-
-
-restype_atom14_ambiguous_atoms = np.zeros((21, 14), dtype=np.float32)
-restype_atom14_ambiguous_atoms_swap_idx = np.tile(np.arange(14, dtype=int), (21, 1))
-
-
-def _make_atom14_ambiguity_feats():
-    for res, pairs in residue_atom_renaming_swaps.items():
-        res_idx = restype_order[restype_3to1[res]]
-        for atom1, atom2 in pairs.items():
-            atom1_idx = restype_name_to_atom14_names[res].index(atom1)
-            atom2_idx = restype_name_to_atom14_names[res].index(atom2)
-            restype_atom14_ambiguous_atoms[res_idx, atom1_idx] = 1
-            restype_atom14_ambiguous_atoms[res_idx, atom2_idx] = 1
-            restype_atom14_ambiguous_atoms_swap_idx[res_idx, atom1_idx] = atom2_idx
-            restype_atom14_ambiguous_atoms_swap_idx[res_idx, atom2_idx] = atom1_idx
-
-
-_make_atom14_ambiguity_feats()
-
-
-def aatype_to_str_sequence(aatype):
-    return "".join([restypes_with_x[aatype[i]] for i in range(len(aatype))])
-
-
-# NOTE(thayes): These are computed based on the average CA->C and CA->N norm from rigid_group_atom_positions
-CA_TO_N_NORM = 1.4591
-CA_TO_C_NORM = 1.5252
-
-
-def _make_restype_atom37_to_atom14():
-    """Map from atom37 to atom14 per residue type."""
-    restype_atom37_to_atom14 = []  # mapping (restype, atom37) --> atom14
-    for rt in restypes:
-        atom_names = restype_name_to_atom14_names[restype_1to3[rt]]
-        atom_name_to_idx14 = {name: i for i, name in enumerate(atom_names)}
-        restype_atom37_to_atom14.append(
-            [
-                (atom_name_to_idx14[name] if name in atom_name_to_idx14 else 0)
-                for name in atom_types
-            ]
-        )
-
-    restype_atom37_to_atom14.append([0] * 37)
-    restype_atom37_to_atom14 = np.array(restype_atom37_to_atom14, dtype=np.int32)
-    return restype_atom37_to_atom14
-
-
-def _make_restype_atom14_to_atom37():
-    """Map from atom14 to atom37 per residue type."""
-    restype_atom14_to_atom37 = []  # mapping (restype, atom14) --> atom37
-    for rt in restypes:
-        atom_names = restype_name_to_atom14_names[restype_1to3[rt]]
-        restype_atom14_to_atom37.append(
-            [(atom_order[name] if name else 0) for name in atom_names]
-        )
-    # Add dummy mapping for restype 'UNK'
-    restype_atom14_to_atom37.append([0] * 14)
-    restype_atom14_to_atom37 = np.array(restype_atom14_to_atom37, dtype=np.int32)
-    return restype_atom14_to_atom37
-
-
-RESTYPE_ATOM14_TO_ATOM37 = _make_restype_atom14_to_atom37()
-RESTYPE_ATOM37_TO_ATOM14 = _make_restype_atom37_to_atom14()
-CHAIN_BREAK_TOKEN = "|"
+# Copyright 2025 EvolutionaryScale
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Constants used in AlphaFold."""
+
+import collections
+import functools
+from pathlib import Path
+from typing import List, Mapping, Tuple
+
+import numpy as np
+
+# import tree
+
+# Internal import (35fd).
+
+
+# Distance from one CA to next CA [trans configuration: omega = 180].
+ca_ca = 3.80209737096
+
+# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
+# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
+# chi angles so their chi angle lists are empty.
+chi_angles_atoms = {
+    "ALA": [],
+    # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
+    "ARG": [
+        ["N", "CA", "CB", "CG"],
+        ["CA", "CB", "CG", "CD"],
+        ["CB", "CG", "CD", "NE"],
+        ["CG", "CD", "NE", "CZ"],
+    ],
+    "ASN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
+    "ASP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
+    "CYS": [["N", "CA", "CB", "SG"]],
+    "GLN": [
+        ["N", "CA", "CB", "CG"],
+        ["CA", "CB", "CG", "CD"],
+        ["CB", "CG", "CD", "OE1"],
+    ],
+    "GLU": [
+        ["N", "CA", "CB", "CG"],
+        ["CA", "CB", "CG", "CD"],
+        ["CB", "CG", "CD", "OE1"],
+    ],
+    "GLY": [],
+    "HIS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "ND1"]],
+    "ILE": [["N", "CA", "CB", "CG1"], ["CA", "CB", "CG1", "CD1"]],
+    "LEU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+    "LYS": [
+        ["N", "CA", "CB", "CG"],
+        ["CA", "CB", "CG", "CD"],
+        ["CB", "CG", "CD", "CE"],
+        ["CG", "CD", "CE", "NZ"],
+    ],
+    "MET": [
+        ["N", "CA", "CB", "CG"],
+        ["CA", "CB", "CG", "SD"],
+        ["CB", "CG", "SD", "CE"],
+    ],
+    "PHE": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+    "PRO": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"]],
+    "SER": [["N", "CA", "CB", "OG"]],
+    "THR": [["N", "CA", "CB", "OG1"]],
+    "TRP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+    "TYR": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+    "VAL": [["N", "CA", "CB", "CG1"]],
+    "UNK": [],
+}
+
+# If chi angles given in fixed-length array, this matrix determines how to mask
+# them for each AA type. The order is as per restype_order (see below).
+chi_angles_mask = [
+    [0.0, 0.0, 0.0, 0.0],  # ALA
+    [1.0, 1.0, 1.0, 1.0],  # ARG
+    [1.0, 1.0, 0.0, 0.0],  # ASN
+    [1.0, 1.0, 0.0, 0.0],  # ASP
+    [1.0, 0.0, 0.0, 0.0],  # CYS
+    [1.0, 1.0, 1.0, 0.0],  # GLN
+    [1.0, 1.0, 1.0, 0.0],  # GLU
+    [0.0, 0.0, 0.0, 0.0],  # GLY
+    [1.0, 1.0, 0.0, 0.0],  # HIS
+    [1.0, 1.0, 0.0, 0.0],  # ILE
+    [1.0, 1.0, 0.0, 0.0],  # LEU
+    [1.0, 1.0, 1.0, 1.0],  # LYS
+    [1.0, 1.0, 1.0, 0.0],  # MET
+    [1.0, 1.0, 0.0, 0.0],  # PHE
+    [1.0, 1.0, 0.0, 0.0],  # PRO
+    [1.0, 0.0, 0.0, 0.0],  # SER
+    [1.0, 0.0, 0.0, 0.0],  # THR
+    [1.0, 1.0, 0.0, 0.0],  # TRP
+    [1.0, 1.0, 0.0, 0.0],  # TYR
+    [1.0, 0.0, 0.0, 0.0],  # VAL
+    [0.0, 0.0, 0.0, 0.0],  # UNK
+]
+
+# The following chi angles are pi periodic: they can be rotated by a multiple
+# of pi without affecting the structure.
+chi_pi_periodic = [
+    [0.0, 0.0, 0.0, 0.0],  # ALA
+    [0.0, 0.0, 0.0, 0.0],  # ARG
+    [0.0, 0.0, 0.0, 0.0],  # ASN
+    [0.0, 1.0, 0.0, 0.0],  # ASP
+    [0.0, 0.0, 0.0, 0.0],  # CYS
+    [0.0, 0.0, 0.0, 0.0],  # GLN
+    [0.0, 0.0, 1.0, 0.0],  # GLU
+    [0.0, 0.0, 0.0, 0.0],  # GLY
+    [0.0, 0.0, 0.0, 0.0],  # HIS
+    [0.0, 0.0, 0.0, 0.0],  # ILE
+    [0.0, 0.0, 0.0, 0.0],  # LEU
+    [0.0, 0.0, 0.0, 0.0],  # LYS
+    [0.0, 0.0, 0.0, 0.0],  # MET
+    [0.0, 1.0, 0.0, 0.0],  # PHE
+    [0.0, 0.0, 0.0, 0.0],  # PRO
+    [0.0, 0.0, 0.0, 0.0],  # SER
+    [0.0, 0.0, 0.0, 0.0],  # THR
+    [0.0, 0.0, 0.0, 0.0],  # TRP
+    [0.0, 1.0, 0.0, 0.0],  # TYR
+    [0.0, 0.0, 0.0, 0.0],  # VAL
+    [0.0, 0.0, 0.0, 0.0],  # UNK
+]
+
+# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi,
+# psi and chi angles:
+# 0: 'backbone group',
+# 1: 'pre-omega-group', (empty)
+# 2: 'phi-group', (currently empty, because it defines only hydrogens)
+# 3: 'psi-group',
+# 4,5,6,7: 'chi1,2,3,4-group'
+# The atom positions are relative to the axis-end-atom of the corresponding
+# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis
+# is defined such that the dihedral-angle-definiting atom (the last entry in
+# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate).
+# format: [atomname, group_idx, rel_position]
+rigid_group_atom_positions = {
+    "ALA": [
+        ["N", 0, (-0.525, 1.363, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.526, -0.000, -0.000)],
+        ["CB", 0, (-0.529, -0.774, -1.205)],
+        ["O", 3, (0.627, 1.062, 0.000)],
+    ],
+    "ARG": [
+        ["N", 0, (-0.524, 1.362, -0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.525, -0.000, -0.000)],
+        ["CB", 0, (-0.524, -0.778, -1.209)],
+        ["O", 3, (0.626, 1.062, 0.000)],
+        ["CG", 4, (0.616, 1.390, -0.000)],
+        ["CD", 5, (0.564, 1.414, 0.000)],
+        ["NE", 6, (0.539, 1.357, -0.000)],
+        ["NH1", 7, (0.206, 2.301, 0.000)],
+        ["NH2", 7, (2.078, 0.978, -0.000)],
+        ["CZ", 7, (0.758, 1.093, -0.000)],
+    ],
+    "ASN": [
+        ["N", 0, (-0.536, 1.357, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.526, -0.000, -0.000)],
+        ["CB", 0, (-0.531, -0.787, -1.200)],
+        ["O", 3, (0.625, 1.062, 0.000)],
+        ["CG", 4, (0.584, 1.399, 0.000)],
+        ["ND2", 5, (0.593, -1.188, 0.001)],
+        ["OD1", 5, (0.633, 1.059, 0.000)],
+    ],
+    "ASP": [
+        ["N", 0, (-0.525, 1.362, -0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.527, 0.000, -0.000)],
+        ["CB", 0, (-0.526, -0.778, -1.208)],
+        ["O", 3, (0.626, 1.062, -0.000)],
+        ["CG", 4, (0.593, 1.398, -0.000)],
+        ["OD1", 5, (0.610, 1.091, 0.000)],
+        ["OD2", 5, (0.592, -1.101, -0.003)],
+    ],
+    "CYS": [
+        ["N", 0, (-0.522, 1.362, -0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.524, 0.000, 0.000)],
+        ["CB", 0, (-0.519, -0.773, -1.212)],
+        ["O", 3, (0.625, 1.062, -0.000)],
+        ["SG", 4, (0.728, 1.653, 0.000)],
+    ],
+    "GLN": [
+        ["N", 0, (-0.526, 1.361, -0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.526, 0.000, 0.000)],
+        ["CB", 0, (-0.525, -0.779, -1.207)],
+        ["O", 3, (0.626, 1.062, -0.000)],
+        ["CG", 4, (0.615, 1.393, 0.000)],
+        ["CD", 5, (0.587, 1.399, -0.000)],
+        ["NE2", 6, (0.593, -1.189, -0.001)],
+        ["OE1", 6, (0.634, 1.060, 0.000)],
+    ],
+    "GLU": [
+        ["N", 0, (-0.528, 1.361, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.526, -0.000, -0.000)],
+        ["CB", 0, (-0.526, -0.781, -1.207)],
+        ["O", 3, (0.626, 1.062, 0.000)],
+        ["CG", 4, (0.615, 1.392, 0.000)],
+        ["CD", 5, (0.600, 1.397, 0.000)],
+        ["OE1", 6, (0.607, 1.095, -0.000)],
+        ["OE2", 6, (0.589, -1.104, -0.001)],
+    ],
+    "GLY": [
+        ["N", 0, (-0.572, 1.337, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.517, -0.000, -0.000)],
+        ["O", 3, (0.626, 1.062, -0.000)],
+    ],
+    "HIS": [
+        ["N", 0, (-0.527, 1.360, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.525, 0.000, 0.000)],
+        ["CB", 0, (-0.525, -0.778, -1.208)],
+        ["O", 3, (0.625, 1.063, 0.000)],
+        ["CG", 4, (0.600, 1.370, -0.000)],
+        ["CD2", 5, (0.889, -1.021, 0.003)],
+        ["ND1", 5, (0.744, 1.160, -0.000)],
+        ["CE1", 5, (2.030, 0.851, 0.002)],
+        ["NE2", 5, (2.145, -0.466, 0.004)],
+    ],
+    "ILE": [
+        ["N", 0, (-0.493, 1.373, -0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.527, -0.000, -0.000)],
+        ["CB", 0, (-0.536, -0.793, -1.213)],
+        ["O", 3, (0.627, 1.062, -0.000)],
+        ["CG1", 4, (0.534, 1.437, -0.000)],
+        ["CG2", 4, (0.540, -0.785, -1.199)],
+        ["CD1", 5, (0.619, 1.391, 0.000)],
+    ],
+    "LEU": [
+        ["N", 0, (-0.520, 1.363, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.525, -0.000, -0.000)],
+        ["CB", 0, (-0.522, -0.773, -1.214)],
+        ["O", 3, (0.625, 1.063, -0.000)],
+        ["CG", 4, (0.678, 1.371, 0.000)],
+        ["CD1", 5, (0.530, 1.430, -0.000)],
+        ["CD2", 5, (0.535, -0.774, 1.200)],
+    ],
+    "LYS": [
+        ["N", 0, (-0.526, 1.362, -0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.526, 0.000, 0.000)],
+        ["CB", 0, (-0.524, -0.778, -1.208)],
+        ["O", 3, (0.626, 1.062, -0.000)],
+        ["CG", 4, (0.619, 1.390, 0.000)],
+        ["CD", 5, (0.559, 1.417, 0.000)],
+        ["CE", 6, (0.560, 1.416, 0.000)],
+        ["NZ", 7, (0.554, 1.387, 0.000)],
+    ],
+    "MET": [
+        ["N", 0, (-0.521, 1.364, -0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.525, 0.000, 0.000)],
+        ["CB", 0, (-0.523, -0.776, -1.210)],
+        ["O", 3, (0.625, 1.062, -0.000)],
+        ["CG", 4, (0.613, 1.391, -0.000)],
+        ["SD", 5, (0.703, 1.695, 0.000)],
+        ["CE", 6, (0.320, 1.786, -0.000)],
+    ],
+    "PHE": [
+        ["N", 0, (-0.518, 1.363, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.524, 0.000, -0.000)],
+        ["CB", 0, (-0.525, -0.776, -1.212)],
+        ["O", 3, (0.626, 1.062, -0.000)],
+        ["CG", 4, (0.607, 1.377, 0.000)],
+        ["CD1", 5, (0.709, 1.195, -0.000)],
+        ["CD2", 5, (0.706, -1.196, 0.000)],
+        ["CE1", 5, (2.102, 1.198, -0.000)],
+        ["CE2", 5, (2.098, -1.201, -0.000)],
+        ["CZ", 5, (2.794, -0.003, -0.001)],
+    ],
+    "PRO": [
+        ["N", 0, (-0.566, 1.351, -0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.527, -0.000, 0.000)],
+        ["CB", 0, (-0.546, -0.611, -1.293)],
+        ["O", 3, (0.621, 1.066, 0.000)],
+        ["CG", 4, (0.382, 1.445, 0.0)],
+        # ['CD', 5, (0.427, 1.440, 0.0)],
+        ["CD", 5, (0.477, 1.424, 0.0)],  # manually made angle 2 degrees larger
+    ],
+    "SER": [
+        ["N", 0, (-0.529, 1.360, -0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.525, -0.000, -0.000)],
+        ["CB", 0, (-0.518, -0.777, -1.211)],
+        ["O", 3, (0.626, 1.062, -0.000)],
+        ["OG", 4, (0.503, 1.325, 0.000)],
+    ],
+    "THR": [
+        ["N", 0, (-0.517, 1.364, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.526, 0.000, -0.000)],
+        ["CB", 0, (-0.516, -0.793, -1.215)],
+        ["O", 3, (0.626, 1.062, 0.000)],
+        ["CG2", 4, (0.550, -0.718, -1.228)],
+        ["OG1", 4, (0.472, 1.353, 0.000)],
+    ],
+    "TRP": [
+        ["N", 0, (-0.521, 1.363, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.525, -0.000, 0.000)],
+        ["CB", 0, (-0.523, -0.776, -1.212)],
+        ["O", 3, (0.627, 1.062, 0.000)],
+        ["CG", 4, (0.609, 1.370, -0.000)],
+        ["CD1", 5, (0.824, 1.091, 0.000)],
+        ["CD2", 5, (0.854, -1.148, -0.005)],
+        ["CE2", 5, (2.186, -0.678, -0.007)],
+        ["CE3", 5, (0.622, -2.530, -0.007)],
+        ["NE1", 5, (2.140, 0.690, -0.004)],
+        ["CH2", 5, (3.028, -2.890, -0.013)],
+        ["CZ2", 5, (3.283, -1.543, -0.011)],
+        ["CZ3", 5, (1.715, -3.389, -0.011)],
+    ],
+    "TYR": [
+        ["N", 0, (-0.522, 1.362, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.524, -0.000, -0.000)],
+        ["CB", 0, (-0.522, -0.776, -1.213)],
+        ["O", 3, (0.627, 1.062, -0.000)],
+        ["CG", 4, (0.607, 1.382, -0.000)],
+        ["CD1", 5, (0.716, 1.195, -0.000)],
+        ["CD2", 5, (0.713, -1.194, -0.001)],
+        ["CE1", 5, (2.107, 1.200, -0.002)],
+        ["CE2", 5, (2.104, -1.201, -0.003)],
+        ["OH", 5, (4.168, -0.002, -0.005)],
+        ["CZ", 5, (2.791, -0.001, -0.003)],
+    ],
+    "VAL": [
+        ["N", 0, (-0.494, 1.373, -0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.527, -0.000, -0.000)],
+        ["CB", 0, (-0.533, -0.795, -1.213)],
+        ["O", 3, (0.627, 1.062, -0.000)],
+        ["CG1", 4, (0.540, 1.429, -0.000)],
+        ["CG2", 4, (0.533, -0.776, 1.203)],
+    ],
+    # Assume alanine positions for unknown AA
+    "UNK": [
+        ["N", 0, (-0.525, 1.363, 0.000)],
+        ["CA", 0, (0.000, 0.000, 0.000)],
+        ["C", 0, (1.526, -0.000, -0.000)],
+    ],
+}
+
+# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
+residue_atoms = {
+    "ALA": ["C", "CA", "CB", "N", "O"],
+    "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"],
+    "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"],
+    "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"],
+    "CYS": ["C", "CA", "CB", "N", "O", "SG"],
+    "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"],
+    "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"],
+    "GLY": ["C", "CA", "N", "O"],
+    "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"],
+    "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"],
+    "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"],
+    "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"],
+    "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"],
+    "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"],
+    "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"],
+    "SER": ["C", "CA", "CB", "N", "O", "OG"],
+    "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"],
+    "TRP": [
+        "C",
+        "CA",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "CE2",
+        "CE3",
+        "CZ2",
+        "CZ3",
+        "CH2",
+        "N",
+        "NE1",
+        "O",
+    ],
+    "TYR": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O", "OH"],
+    "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"],
+    "UNK": ["C", "CA", "N"],
+}
+
+# Naming swaps for ambiguous atom names.
+# Due to symmetries in the amino acids the naming of atoms is ambiguous in
+# 4 of the 20 amino acids.
+# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities
+# in LEU, VAL and ARG can be resolved by using the 3d constellations of
+# the 'ambiguous' atoms and their neighbours)
+# TODO: ^ interpret this
+residue_atom_renaming_swaps = {
+    "ASP": {"OD1": "OD2"},
+    "GLU": {"OE1": "OE2"},
+    "PHE": {"CD1": "CD2", "CE1": "CE2"},
+    "TYR": {"CD1": "CD2", "CE1": "CE2"},
+}
+
+# Van der Waals radii [Angstroem] of the atoms (from Wikipedia)
+van_der_waals_radius = {"C": 1.7, "N": 1.55, "O": 1.52, "S": 1.8}
+
+Bond = collections.namedtuple("Bond", ["atom1_name", "atom2_name", "length", "stddev"])
+BondAngle = collections.namedtuple(
+    "BondAngle", ["atom1_name", "atom2_name", "atom3name", "angle_rad", "stddev"]
+)
+
+
+@functools.lru_cache(maxsize=None)
+def load_stereo_chemical_props() -> (
+    Tuple[
+        Mapping[str, List[Bond]],
+        Mapping[str, List[Bond]],
+        Mapping[str, List[BondAngle]],
+    ]
+):
+    """Load stereo_chemical_props.txt into a nice structure.
+
+    Load literature values for bond lengths and bond angles and translate
+    bond angles into the length of the opposite edge of the triangle
+    ("residue_virtual_bonds").
+
+    Returns:
+      residue_bonds:  dict that maps resname --> list of Bond tuples
+      residue_virtual_bonds: dict that maps resname --> list of Bond tuples
+      residue_bond_angles: dict that maps resname --> list of BondAngle tuples
+    """
+    stereo_chemical_props = Path(
+        "evolutionaryscale/structure/stereo_chemical_props.txt"
+    ).read_text()
+
+    lines_iter = iter(stereo_chemical_props.splitlines())
+    # Load bond lengths.
+    residue_bonds = {}
+    next(lines_iter)  # Skip header line.
+    for line in lines_iter:
+        if line.strip() == "-":
+            break
+        bond, resname, length, stddev = line.split()
+        atom1, atom2 = bond.split("-")
+        if resname not in residue_bonds:
+            residue_bonds[resname] = []
+        residue_bonds[resname].append(Bond(atom1, atom2, float(length), float(stddev)))
+    residue_bonds["UNK"] = []
+
+    # Load bond angles.
+    residue_bond_angles = {}
+    next(lines_iter)  # Skip empty line.
+    next(lines_iter)  # Skip header line.
+    for line in lines_iter:
+        if line.strip() == "-":
+            break
+        bond, resname, angle_degree, stddev_degree = line.split()
+        atom1, atom2, atom3 = bond.split("-")
+        if resname not in residue_bond_angles:
+            residue_bond_angles[resname] = []
+        residue_bond_angles[resname].append(
+            BondAngle(
+                atom1,
+                atom2,
+                atom3,
+                float(angle_degree) / 180.0 * np.pi,
+                float(stddev_degree) / 180.0 * np.pi,
+            )
+        )
+    residue_bond_angles["UNK"] = []
+
+    def make_bond_key(atom1_name, atom2_name):
+        """Unique key to lookup bonds."""
+        return "-".join(sorted([atom1_name, atom2_name]))
+
+    # Translate bond angles into distances ("virtual bonds").
+    residue_virtual_bonds = {}
+    for resname, bond_angles in residue_bond_angles.items():
+        # Create a fast lookup dict for bond lengths.
+        bond_cache = {}
+        for b in residue_bonds[resname]:
+            bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
+        residue_virtual_bonds[resname] = []
+        for ba in bond_angles:
+            bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
+            bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]
+
+            # Compute distance between atom1 and atom3 using the law of cosines
+            # c^2 = a^2 + b^2 - 2ab*cos(gamma).
+            gamma = ba.angle_rad
+            length = np.sqrt(
+                bond1.length**2
+                + bond2.length**2
+                - 2 * bond1.length * bond2.length * np.cos(gamma)
+            )
+
+            # Propagation of uncertainty assuming uncorrelated errors.
+            dl_outer = 0.5 / length
+            dl_dgamma = (2 * bond1.length * bond2.length * np.sin(gamma)) * dl_outer
+            dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer
+            dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer
+            stddev = np.sqrt(
+                (dl_dgamma * ba.stddev) ** 2
+                + (dl_db1 * bond1.stddev) ** 2
+                + (dl_db2 * bond2.stddev) ** 2
+            )
+            residue_virtual_bonds[resname].append(
+                Bond(ba.atom1_name, ba.atom3name, length, stddev)
+            )
+
+    return (residue_bonds, residue_virtual_bonds, residue_bond_angles)
+
+
+# Between-residue bond lengths for general bonds (first element) and for Proline
+# (second element).
+between_res_bond_length_c_n = [1.329, 1.341]
+between_res_bond_length_stddev_c_n = [0.014, 0.016]
+
+# Between-residue cos_angles.
+between_res_cos_angles_c_n_ca = [-0.5203, 0.0353]  # degrees: 121.352 +- 2.315
+between_res_cos_angles_ca_c_n = [-0.4473, 0.0311]  # degrees: 116.568 +- 1.995
+
+# This mapping is used when we need to store atom data in a format that requires
+# fixed atom data size for every residue (e.g. a numpy array).
+atom_types = [
+    "N",
+    "CA",
+    "C",
+    "CB",
+    "O",
+    "CG",
+    "CG1",
+    "CG2",
+    "OG",
+    "OG1",
+    "SG",
+    "CD",
+    "CD1",
+    "CD2",
+    "ND1",
+    "ND2",
+    "OD1",
+    "OD2",
+    "SD",
+    "CE",
+    "CE1",
+    "CE2",
+    "CE3",
+    "NE",
+    "NE1",
+    "NE2",
+    "OE1",
+    "OE2",
+    "CH2",
+    "NH1",
+    "NH2",
+    "OH",
+    "CZ",
+    "CZ2",
+    "CZ3",
+    "NZ",
+    "OXT",
+]
+atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
+atom_type_num = len(atom_types)  # := 37.
+
+# A compact atom encoding with 14 columns
+# pylint: disable=line-too-long
+# pylint: disable=bad-whitespace
+restype_name_to_atom14_names = {
+    "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""],
+    "ARG": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD",
+        "NE",
+        "CZ",
+        "NH1",
+        "NH2",
+        "",
+        "",
+        "",
+    ],
+    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2", "", "", "", "", "", ""],
+    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2", "", "", "", "", "", ""],
+    "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""],
+    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2", "", "", "", "", ""],
+    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2", "", "", "", "", ""],
+    "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""],
+    "HIS": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "ND1",
+        "CD2",
+        "CE1",
+        "NE2",
+        "",
+        "",
+        "",
+        "",
+    ],
+    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1", "", "", "", "", "", ""],
+    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "", "", "", "", "", ""],
+    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ", "", "", "", "", ""],
+    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE", "", "", "", "", "", ""],
+    "PHE": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "CE1",
+        "CE2",
+        "CZ",
+        "",
+        "",
+        "",
+    ],
+    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""],
+    "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""],
+    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2", "", "", "", "", "", "", ""],
+    "TRP": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "NE1",
+        "CE2",
+        "CE3",
+        "CZ2",
+        "CZ3",
+        "CH2",
+    ],
+    "TYR": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "CE1",
+        "CE2",
+        "CZ",
+        "OH",
+        "",
+        "",
+    ],
+    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "", "", "", "", "", "", ""],
+    "UNK": ["N", "CA", "C", "", "", "", "", "", "", "", "", "", "", ""],
+}
+# pylint: enable=line-too-long
+# pylint: enable=bad-whitespace
+
+
+# This is the standard residue order when coding AA type as a number.
+# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
+restypes = [
+    "A",
+    "R",
+    "N",
+    "D",
+    "C",
+    "Q",
+    "E",
+    "G",
+    "H",
+    "I",
+    "L",
+    "K",
+    "M",
+    "F",
+    "P",
+    "S",
+    "T",
+    "W",
+    "Y",
+    "V",
+]
+restype_order = {restype: i for i, restype in enumerate(restypes)}
+restype_num = len(restypes)  # := 20.
+unk_restype_index = restype_num  # Catch-all index for unknown restypes.
+
+restypes_with_x = restypes + ["X"]
+restype_order_with_x = {restype: i for i, restype in enumerate(restypes_with_x)}
+
+bb_atoms = ["N", "CA", "C", "O"]
+
+# Hydrophobicity by residue (positive values are hydrophobic). Derived from Black & Mould (1991), normalized by subtracting 0.5.
+hydrophobicity = {
+    "ALA": 0.116,
+    "ARG": -0.5,
+    "ASN": -0.264,
+    "ASP": -0.472,
+    "CYS": 0.18,
+    "GLN": -0.249,
+    "GLU": -0.457,
+    "GLY": 0.001,
+    "HIS": -0.335,
+    "ILE": 0.443,
+    "LEU": 0.443,
+    "LYS": -0.217,
+    "MET": 0.238,
+    "PHE": 0.5,
+    "PRO": 0.211,
+    "SER": -0.141,
+    "THR": -0.05,
+    "TRP": 0.378,
+    "TYR": 0.38,
+    "VAL": 0.325,
+}
+
+# Side chain max accessible surface area in Ala-X-Ala tripeptide (from Chennamsetty et al. 2010).
+side_chain_asa = {
+    "ALA": 64.7809,
+    "ARG": 210.02,
+    "ASN": 113.187,
+    "ASP": 110.209,
+    "CYS": 95.2439,
+    "GLN": 147.855,
+    "GLU": 143.924,
+    "GLY": 23.1338,
+    "HIS": 146.449,
+    "ILE": 151.242,
+    "LEU": 139.524,
+    "LYS": 177.366,
+    "MET": 164.674,
+    "PHE": 186.7,
+    "PRO": 111.533,
+    "SER": 81.2159,
+    "THR": 111.597,
+    "TRP": 229.619,
+    "TYR": 200.306,
+    "VAL": 124.237,
+}
+
+# Approximate Volumes of amino acids in cubic angstroms.
+# https://www.imgt.org/IMGTeducation/Aide-memoire/_UK/aminoacids/abbreviation.html
+amino_acid_volumes = {
+    "A": 88.6,  # Alanine
+    "R": 173.4,  # Arginine
+    "N": 114.1,  # Asparagine
+    "D": 111.1,  # Aspartic acid
+    "C": 108.5,  # Cysteine
+    "Q": 143.8,  # Glutamine
+    "E": 138.4,  # Glutamic acid
+    "G": 60.1,  # Glycine
+    "H": 153.2,  # Histidine
+    "I": 166.7,  # Isoleucine
+    "L": 166.7,  # Leucine
+    "K": 168.6,  # Lysine
+    "M": 162.9,  # Methionine
+    "F": 189.9,  # Phenylalanine
+    "P": 112.7,  # Proline
+    "S": 89.0,  # Serine
+    "T": 116.1,  # Threonine
+    "W": 227.8,  # Tryptophan
+    "Y": 193.6,  # Tyrosine
+    "V": 140.0,  # Valine
+    "X": 88.6,  # Unknown, use Alanine as approximation
+}
+
+
+def sequence_to_onehot(
+    sequence: str, mapping: Mapping[str, int], map_unknown_to_x: bool = False
+) -> np.ndarray:
+    """Maps the given sequence into a one-hot encoded matrix.
+
+    Args:
+      sequence: An amino acid sequence.
+      mapping: A dictionary mapping amino acids to integers.
+      map_unknown_to_x: If True, any amino acid that is not in the mapping will be
+        mapped to the unknown amino acid 'X'. If the mapping doesn't contain
+        amino acid 'X', an error will be thrown. If False, any amino acid not in
+        the mapping will throw an error.
+
+    Returns:
+      A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of
+      the sequence.
+
+    Raises:
+      ValueError: If the mapping doesn't contain values from 0 to
+        num_unique_aas - 1 without any gaps.
+    """
+    num_entries = max(mapping.values()) + 1
+
+    if sorted(set(mapping.values())) != list(range(num_entries)):
+        raise ValueError(
+            "The mapping must have values from 0 to num_unique_aas-1 "
+            "without any gaps. Got: %s" % sorted(mapping.values())
+        )
+
+    one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)
+
+    for aa_index, aa_type in enumerate(sequence):
+        if map_unknown_to_x:
+            if aa_type.isalpha() and aa_type.isupper():
+                aa_id = mapping.get(aa_type, mapping["X"])
+            else:
+                raise ValueError(f"Invalid character in the sequence: {aa_type}")
+        else:
+            aa_id = mapping[aa_type]
+        one_hot_arr[aa_index, aa_id] = 1
+
+    return one_hot_arr
+
+
+restype_1to3 = {
+    "A": "ALA",
+    "R": "ARG",
+    "N": "ASN",
+    "D": "ASP",
+    "C": "CYS",
+    "Q": "GLN",
+    "E": "GLU",
+    "G": "GLY",
+    "H": "HIS",
+    "I": "ILE",
+    "L": "LEU",
+    "K": "LYS",
+    "M": "MET",
+    "F": "PHE",
+    "P": "PRO",
+    "S": "SER",
+    "T": "THR",
+    "W": "TRP",
+    "Y": "TYR",
+    "V": "VAL",
+    "X": "UNK",
+}
+
+
+# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple
+# 1-to-1 mapping of 3 letter names to one letter names. The latter contains
+# many more, and less common, three letter names as keys and maps many of these
+# to the same one letter name (including 'X' and 'U' which we don't use here).
+restype_3to1 = {v: k for k, v in restype_1to3.items()}
+
+# Define a restype name for all unknown residues.
+unk_restype = "UNK"
+
+resnames = [restype_1to3[r] for r in restypes] + [unk_restype]
+resname_to_idx = {resname: i for i, resname in enumerate(resnames)}
+
+hydrophobic_resnames = {"VAL", "ILE", "LEU", "PHE", "MET", "TRP"}
+
+# The mapping here uses hhblits convention, so that B is mapped to D, J and O
+# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the
+# remaining 20 amino acids are kept in alphabetical order.
+# There are 2 non-amino acid codes, X (representing any amino acid) and
+# "-" representing a missing amino acid in an alignment.  The id for these
+# codes is put at the end (20 and 21) so that they can easily be ignored if
+# desired.
+HHBLITS_AA_TO_ID = {
+    "A": 0,
+    "B": 2,
+    "C": 1,
+    "D": 2,
+    "E": 3,
+    "F": 4,
+    "G": 5,
+    "H": 6,
+    "I": 7,
+    "J": 20,
+    "K": 8,
+    "L": 9,
+    "M": 10,
+    "N": 11,
+    "O": 20,
+    "P": 12,
+    "Q": 13,
+    "R": 14,
+    "S": 15,
+    "T": 16,
+    "U": 1,
+    "V": 17,
+    "W": 18,
+    "X": 20,
+    "Y": 19,
+    "Z": 3,
+    "-": 21,
+}
+
+# Partial inversion of HHBLITS_AA_TO_ID.
+ID_TO_HHBLITS_AA = {
+    0: "A",
+    1: "C",  # Also U.
+    2: "D",  # Also B.
+    3: "E",  # Also Z.
+    4: "F",
+    5: "G",
+    6: "H",
+    7: "I",
+    8: "K",
+    9: "L",
+    10: "M",
+    11: "N",
+    12: "P",
+    13: "Q",
+    14: "R",
+    15: "S",
+    16: "T",
+    17: "V",
+    18: "W",
+    19: "Y",
+    20: "X",  # Includes J and O.
+    21: "-",
+}
+
+restypes_with_x_and_gap = restypes + ["X", "-"]
+MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
+    restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i])
+    for i in range(len(restypes_with_x_and_gap))
+)
+
+
+def _make_standard_atom_mask() -> np.ndarray:
+    """Returns [num_res_types, num_atom_types] mask array."""
+    # +1 to account for unknown (all 0s).
+    mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32)
+    for restype, restype_letter in enumerate(restypes):
+        restype_name = restype_1to3[restype_letter]
+        atom_names = residue_atoms[restype_name]
+        for atom_name in atom_names:
+            atom_type = atom_order[atom_name]
+            mask[restype, atom_type] = 1
+    return mask
+
+
+STANDARD_ATOM_MASK = _make_standard_atom_mask()
+
+
+# A one hot representation for the first and second atoms defining the axis
+# of rotation for each chi-angle in each residue.
+def chi_angle_atom(atom_index: int) -> np.ndarray:
+    """Define chi-angle rigid groups via one-hot representations."""
+    chi_angles_index = {}
+    one_hots = []
+
+    for k, v in chi_angles_atoms.items():
+        indices = [atom_types.index(s[atom_index]) for s in v]
+        indices.extend([-1] * (4 - len(indices)))
+        chi_angles_index[k] = indices
+
+    for r in restypes:
+        res3 = restype_1to3[r]
+        one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
+        one_hots.append(one_hot)
+
+    one_hots.append(np.zeros([4, atom_type_num]))  # Add zeros for residue `X`.
+    one_hot = np.stack(one_hots, axis=0)
+    one_hot = np.transpose(one_hot, [0, 2, 1])
+
+    return one_hot
+
+
+chi_atom_1_one_hot = chi_angle_atom(1)
+chi_atom_2_one_hot = chi_angle_atom(2)
+
+# An array like chi_angles_atoms but using indices rather than names.
+chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes]
+# chi_angles_atom_indices = tree.map_structure(
+#     lambda atom_name: atom_order[atom_name], chi_angles_atom_indices
+# )
+chi_angles_atom_indices = np.array(
+    [
+        chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms)))
+        for chi_atoms in chi_angles_atom_indices
+    ]
+)
+
+# Mapping from (res_name, atom_name) pairs to the atom's chi group index
+# and atom index within that group.
+chi_groups_for_atom = collections.defaultdict(list)
+for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items():
+    for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
+        for atom_i, atom in enumerate(chi_group):
+            chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i))
+chi_groups_for_atom = dict(chi_groups_for_atom)
+
+
+def _make_rigid_transformation_4x4(ex, ey, translation):
+    """Create a rigid 4x4 transformation matrix from two axes and transl."""
+    # Normalize ex.
+    ex_normalized = ex / np.linalg.norm(ex)
+
+    # make ey perpendicular to ex
+    ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized
+    ey_normalized /= np.linalg.norm(ey_normalized)
+
+    # compute ez as cross product
+    eznorm = np.cross(ex_normalized, ey_normalized)
+    m = np.stack([ex_normalized, ey_normalized, eznorm, translation]).transpose()
+    m = np.concatenate([m, [[0.0, 0.0, 0.0, 1.0]]], axis=0)
+    return m
+
+
+# create an array with (restype, atomtype) --> rigid_group_idx
+# and an array with (restype, atomtype, coord) for the atom positions
+# and compute affine transformation matrices (4,4) from one rigid group to the
+# previous group
+restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=int)
+restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
+restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32)
+restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=int)
+restype_atom14_mask = np.zeros([21, 14], dtype=np.float32)
+restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32)
+restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32)
+
+
+def _make_rigid_group_constants():
+    """Fill the arrays above."""
+    for restype, restype_letter in enumerate(restypes_with_x):
+        resname = restype_1to3[restype_letter]
+        for atomname, group_idx, atom_position in rigid_group_atom_positions[resname]:
+            atomtype = atom_order[atomname]
+            restype_atom37_to_rigid_group[restype, atomtype] = group_idx
+            restype_atom37_mask[restype, atomtype] = 1
+            restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position
+
+            atom14idx = restype_name_to_atom14_names[resname].index(atomname)
+            restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
+            restype_atom14_mask[restype, atom14idx] = 1
+            restype_atom14_rigid_group_positions[restype, atom14idx, :] = atom_position
+
+    for restype, restype_letter in enumerate(restypes_with_x):
+        resname = restype_1to3[restype_letter]
+        atom_positions = {
+            name: np.array(pos) for name, _, pos in rigid_group_atom_positions[resname]
+        }
+
+        # backbone to backbone is the identity transform
+        restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)
+
+        # pre-omega-frame to backbone (currently dummy identity matrix)
+        restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)
+
+        # phi-frame to backbone
+        mat = _make_rigid_transformation_4x4(
+            ex=atom_positions["N"] - atom_positions["CA"],
+            ey=np.array([1.0, 0.0, 0.0]),
+            translation=atom_positions["N"],
+        )
+        restype_rigid_group_default_frame[restype, 2, :, :] = mat
+
+        # psi-frame to backbone
+        mat = _make_rigid_transformation_4x4(
+            ex=atom_positions["C"] - atom_positions["CA"],
+            ey=atom_positions["CA"] - atom_positions["N"],
+            translation=atom_positions["C"],
+        )
+        restype_rigid_group_default_frame[restype, 3, :, :] = mat
+
+        # chi1-frame to backbone
+        if chi_angles_mask[restype][0]:
+            base_atom_names = chi_angles_atoms[resname][0]
+            base_atom_positions = [atom_positions[name] for name in base_atom_names]
+            mat = _make_rigid_transformation_4x4(
+                ex=base_atom_positions[2] - base_atom_positions[1],
+                ey=base_atom_positions[0] - base_atom_positions[1],
+                translation=base_atom_positions[2],
+            )
+            restype_rigid_group_default_frame[restype, 4, :, :] = mat
+
+        # chi2-frame to chi1-frame
+        # chi3-frame to chi2-frame
+        # chi4-frame to chi3-frame
+        # luckily all rotation axes for the next frame start at (0,0,0) of the
+        # previous frame
+        for chi_idx in range(1, 4):
+            if chi_angles_mask[restype][chi_idx]:
+                axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
+                axis_end_atom_position = atom_positions[axis_end_atom_name]
+                mat = _make_rigid_transformation_4x4(
+                    ex=axis_end_atom_position,
+                    ey=np.array([-1.0, 0.0, 0.0]),
+                    translation=axis_end_atom_position,
+                )
+                restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat
+
+
+_make_rigid_group_constants()
+
+
+def make_atom14_dists_bounds(overlap_tolerance=1.5, bond_length_tolerance_factor=15.0):
+    """compute upper and lower bounds for bonds to assess violations."""
+    restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32)
+    restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32)
+    restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32)
+    residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props()
+    for restype, restype_letter in enumerate(restypes):
+        resname = restype_1to3[restype_letter]
+        atom_list = restype_name_to_atom14_names[resname]
+
+        # create lower and upper bounds for clashes
+        for atom1_idx, atom1_name in enumerate(atom_list):
+            if not atom1_name:
+                continue
+            atom1_radius = van_der_waals_radius[atom1_name[0]]
+            for atom2_idx, atom2_name in enumerate(atom_list):
+                if (not atom2_name) or atom1_idx == atom2_idx:
+                    continue
+                atom2_radius = van_der_waals_radius[atom2_name[0]]
+                lower = atom1_radius + atom2_radius - overlap_tolerance
+                upper = 1e10
+                restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
+                restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
+                restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
+                restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
+
+        # overwrite lower and upper bounds for bonds and angles
+        for b in residue_bonds[resname] + residue_virtual_bonds[resname]:
+            atom1_idx = atom_list.index(b.atom1_name)
+            atom2_idx = atom_list.index(b.atom2_name)
+            lower = b.length - bond_length_tolerance_factor * b.stddev
+            upper = b.length + bond_length_tolerance_factor * b.stddev
+            restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
+            restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
+            restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
+            restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
+            restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev
+            restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev
+    return {
+        "lower_bound": restype_atom14_bond_lower_bound,  # shape (21,14,14)
+        "upper_bound": restype_atom14_bond_upper_bound,  # shape (21,14,14)
+        "stddev": restype_atom14_bond_stddev,  # shape (21,14,14)
+    }
+
+
+restype_atom14_ambiguous_atoms = np.zeros((21, 14), dtype=np.float32)
+restype_atom14_ambiguous_atoms_swap_idx = np.tile(np.arange(14, dtype=int), (21, 1))
+
+
+def _make_atom14_ambiguity_feats():
+    for res, pairs in residue_atom_renaming_swaps.items():
+        res_idx = restype_order[restype_3to1[res]]
+        for atom1, atom2 in pairs.items():
+            atom1_idx = restype_name_to_atom14_names[res].index(atom1)
+            atom2_idx = restype_name_to_atom14_names[res].index(atom2)
+            restype_atom14_ambiguous_atoms[res_idx, atom1_idx] = 1
+            restype_atom14_ambiguous_atoms[res_idx, atom2_idx] = 1
+            restype_atom14_ambiguous_atoms_swap_idx[res_idx, atom1_idx] = atom2_idx
+            restype_atom14_ambiguous_atoms_swap_idx[res_idx, atom2_idx] = atom1_idx
+
+
+_make_atom14_ambiguity_feats()
+
+
+def aatype_to_str_sequence(aatype):
+    return "".join([restypes_with_x[aatype[i]] for i in range(len(aatype))])
+
+
+# NOTE(thayes): These are computed based on the average CA->C and CA->N norm from rigid_group_atom_positions
+CA_TO_N_NORM = 1.4591
+CA_TO_C_NORM = 1.5252
+
+
+def _make_restype_atom37_to_atom14():
+    """Map from atom37 to atom14 per residue type."""
+    restype_atom37_to_atom14 = []  # mapping (restype, atom37) --> atom14
+    for rt in restypes:
+        atom_names = restype_name_to_atom14_names[restype_1to3[rt]]
+        atom_name_to_idx14 = {name: i for i, name in enumerate(atom_names)}
+        restype_atom37_to_atom14.append(
+            [
+                (atom_name_to_idx14[name] if name in atom_name_to_idx14 else 0)
+                for name in atom_types
+            ]
+        )
+
+    restype_atom37_to_atom14.append([0] * 37)
+    restype_atom37_to_atom14 = np.array(restype_atom37_to_atom14, dtype=np.int32)
+    return restype_atom37_to_atom14
+
+
+def _make_restype_atom14_to_atom37():
+    """Map from atom14 to atom37 per residue type."""
+    restype_atom14_to_atom37 = []  # mapping (restype, atom14) --> atom37
+    for rt in restypes:
+        atom_names = restype_name_to_atom14_names[restype_1to3[rt]]
+        restype_atom14_to_atom37.append(
+            [(atom_order[name] if name else 0) for name in atom_names]
+        )
+    # Add dummy mapping for restype 'UNK'
+    restype_atom14_to_atom37.append([0] * 14)
+    restype_atom14_to_atom37 = np.array(restype_atom14_to_atom37, dtype=np.int32)
+    return restype_atom14_to_atom37
+
+
+RESTYPE_ATOM14_TO_ATOM37 = _make_restype_atom14_to_atom37()
+RESTYPE_ATOM37_TO_ATOM14 = _make_restype_atom37_to_atom14()
+CHAIN_BREAK_TOKEN = "|"

esmfold2_sequential_dataclass.py

@@ -1,157 +1,157 @@
-from abc import ABC, abstractmethod
-from dataclasses import dataclass, fields, replace
-from typing import TypeVar
-
-import numpy as np
-
-from .esmfold2_misc import concat_objects, slice_any_object
-
-T = TypeVar("T")
-
-
-@dataclass(frozen=True)
-class SequentialDataclass(ABC):
-    """
-    This is a builder on a dataclass that allows for automatic slicing and concatenation.
-
-    When representing multimodal data, we often have multiple datatypes which have sequence dimensions that are the same (e.g. the length of the protein).
-
-    When applying a transformation like a crop, we want to apply this to all tensors at the same time (e.g. crop the sequence, structure, and function).
-
-    We also have some fields that are not sequential (like an id, or data source), which we don't want to crop.
-
-    The SequentialDataclass abstracts this cropping away, allowing you to define dataclasses that implement `__len__`, `__getitem__` and `concat` automatically.
-
-    This is done through the `metadata` field, which can take 3 values:
-        `sequence` (bool): True or False, tells the dataclass whether this field is a sequential type. Default: False.
-        `sequence_dim` (int): Which dimension is the sequential dimension (e.g. for a list of inverse folded sequences, we want to index each sequence in the list, not the list itself). Default: 0.
-        `join_token` (Any): What token to use to join when concatenating elements. Default: None.
-
-
-    Example:
-
-        @dataclass(frozen=True)
-        class Foo(SequentialDataclass):
-            id: str
-            sequence: str = field(metadata={"sequence": True, "join_token": "|"})
-            tensor: torch.Tensor = field(metadata={"sequence": True, "join_token": torch.nan})
-
-            def __len__(self):
-                # Must implement the __len__ method
-                return len(self.sequence)
-
-        >>> foo = Foo(id="foo", sequence="ABCDE", tensor=torch.randn(5))
-        Foo(id='foo', sequence='ABCDE', tensor=tensor([ 0.0252, -0.3335, -0.5143,  0.0251, -1.0717]))
-
-        >>> foo[1:4]
-        Foo(id='foo', sequence='BCD', tensor=tensor([-0.3335, -0.5143,  0.0251]))
-
-        >>> foo[np.arange(5) < 3]
-        Foo(id='foo', sequence='ABC', tensor=tensor([ 0.0252, -0.3335, -0.5143]))
-
-        >>> Foo.concat([foo[:2], foo[3:]])
-        Foo(id='foo', sequence='AB|DE', tensor=tensor([ 0.0252, -0.3335,     nan,  0.0251, -1.0717]))
-
-        # Trying to create a type where the sequence lengths do not match raises an error
-        >>> foo = Foo(id="foo", sequence="ABCDE", tensor=torch.randn(6))
-        ValueError: Mismatch in sequence length for field: tensor. Expected 5, received 6
-
-    """
-
-    def __post_init__(self):
-        self._check_sequence_lengths_match()
-
-    @abstractmethod
-    def __len__(self):
-        raise NotImplementedError
-
-    def __getitem__(self, idx: int | list[int] | slice | np.ndarray):
-        updated_fields = {}
-        if isinstance(idx, int):
-            # make it so that things remain sequential
-            idx = [idx]
-
-        for fld in fields(self):
-            if fld.metadata.get("sequence", False):
-                # this is a sequence, should be the same length as all other sequences
-                sequence_dim = fld.metadata.get("sequence_dim", 0)
-                value = getattr(self, fld.name)
-                if value is None:
-                    continue
-                match sequence_dim:
-                    case 0:
-                        # sequence is first dimension
-                        value = getattr(self, fld.name)
-                        value = slice_any_object(value, idx)
-                        updated_fields[fld.name] = value
-                    case 1:
-                        new_value = [slice_any_object(item, idx) for item in value]
-                        updated_fields[fld.name] = value.__class__(new_value)
-                    case _:
-                        raise NotImplementedError(
-                            "Arbitrary slicing for different sequence length fields is not implemented"
-                        )
-
-        return replace(self, **updated_fields)
-
-    def _check_sequence_lengths_match(self):
-        """Checks if sequence lengths of all "sequence" fields match."""
-        for fld in fields(self):
-            if fld.metadata.get("sequence", False) and fld.name != "complex":
-                # this is a sequence, should be the same length as all other sequences
-                sequence_dim = fld.metadata.get("sequence_dim", 0)
-                value = getattr(self, fld.name)
-                if value is None:
-                    continue
-                match sequence_dim:
-                    case 0:
-                        # sequence is first dimension
-                        value = getattr(self, fld.name)
-                        if len(value) != len(self):
-                            raise ValueError(
-                                f"Mismatch in sequence length for field: {fld.name}. Expected {len(self)}, received {len(value)}"
-                            )
-                    case 1:
-                        for item in value:
-                            if len(item) != len(self):
-                                raise ValueError(
-                                    f"Mismatch in sequence length for field: {fld.name}. Expected {len(self)}, received {len(item)}"
-                                )
-                    case _:
-                        raise NotImplementedError(
-                            "Arbitrary matching for different sequence length fields is not implemented"
-                        )
-
-    @classmethod
-    def concat(cls, items: list[T], **kwargs) -> T:
-        updated_fields = {}
-        for fld in fields(cls):
-            if fld.metadata.get("sequence", False):
-                # this is a sequence, should be the same length as all other sequences
-                sequence_dim = fld.metadata.get("sequence_dim", 0)
-                join_value = fld.metadata.get("join_token", None)
-                if getattr(items[0], fld.name) is None:
-                    continue
-                values = [getattr(item, fld.name) for item in items]
-                match sequence_dim:
-                    case 0:
-                        # sequence is first dimension
-                        value = concat_objects(values, join_value)
-                        updated_fields[fld.name] = value
-                    case 1:
-                        new_value = [
-                            concat_objects(item, join_value) for item in zip(*values)
-                        ]
-                        updated_fields[fld.name] = getattr(
-                            items[0], fld.name
-                        ).__class__(new_value)
-                    case _:
-                        raise NotImplementedError(
-                            "Arbitrary joining for different sequence length fields is not implemented"
-                        )
-        updated_fields.update(kwargs)
-
-        return replace(
-            items[0],  # type: ignore
-            **updated_fields,
-        )
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, fields, replace
+from typing import TypeVar
+
+import numpy as np
+
+from .esmfold2_misc import concat_objects, slice_any_object
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True)
+class SequentialDataclass(ABC):
+    """
+    This is a builder on a dataclass that allows for automatic slicing and concatenation.
+
+    When representing multimodal data, we often have multiple datatypes which have sequence dimensions that are the same (e.g. the length of the protein).
+
+    When applying a transformation like a crop, we want to apply this to all tensors at the same time (e.g. crop the sequence, structure, and function).
+
+    We also have some fields that are not sequential (like an id, or data source), which we don't want to crop.
+
+    The SequentialDataclass abstracts this cropping away, allowing you to define dataclasses that implement `__len__`, `__getitem__` and `concat` automatically.
+
+    This is done through the `metadata` field, which can take 3 values:
+        `sequence` (bool): True or False, tells the dataclass whether this field is a sequential type. Default: False.
+        `sequence_dim` (int): Which dimension is the sequential dimension (e.g. for a list of inverse folded sequences, we want to index each sequence in the list, not the list itself). Default: 0.
+        `join_token` (Any): What token to use to join when concatenating elements. Default: None.
+
+
+    Example:
+
+        @dataclass(frozen=True)
+        class Foo(SequentialDataclass):
+            id: str
+            sequence: str = field(metadata={"sequence": True, "join_token": "|"})
+            tensor: torch.Tensor = field(metadata={"sequence": True, "join_token": torch.nan})
+
+            def __len__(self):
+                # Must implement the __len__ method
+                return len(self.sequence)
+
+        >>> foo = Foo(id="foo", sequence="ABCDE", tensor=torch.randn(5))
+        Foo(id='foo', sequence='ABCDE', tensor=tensor([ 0.0252, -0.3335, -0.5143,  0.0251, -1.0717]))
+
+        >>> foo[1:4]
+        Foo(id='foo', sequence='BCD', tensor=tensor([-0.3335, -0.5143,  0.0251]))
+
+        >>> foo[np.arange(5) < 3]
+        Foo(id='foo', sequence='ABC', tensor=tensor([ 0.0252, -0.3335, -0.5143]))
+
+        >>> Foo.concat([foo[:2], foo[3:]])
+        Foo(id='foo', sequence='AB|DE', tensor=tensor([ 0.0252, -0.3335,     nan,  0.0251, -1.0717]))
+
+        # Trying to create a type where the sequence lengths do not match raises an error
+        >>> foo = Foo(id="foo", sequence="ABCDE", tensor=torch.randn(6))
+        ValueError: Mismatch in sequence length for field: tensor. Expected 5, received 6
+
+    """
+
+    def __post_init__(self):
+        self._check_sequence_lengths_match()
+
+    @abstractmethod
+    def __len__(self):
+        raise NotImplementedError
+
+    def __getitem__(self, idx: int | list[int] | slice | np.ndarray):
+        updated_fields = {}
+        if isinstance(idx, int):
+            # make it so that things remain sequential
+            idx = [idx]
+
+        for fld in fields(self):
+            if fld.metadata.get("sequence", False):
+                # this is a sequence, should be the same length as all other sequences
+                sequence_dim = fld.metadata.get("sequence_dim", 0)
+                value = getattr(self, fld.name)
+                if value is None:
+                    continue
+                match sequence_dim:
+                    case 0:
+                        # sequence is first dimension
+                        value = getattr(self, fld.name)
+                        value = slice_any_object(value, idx)
+                        updated_fields[fld.name] = value
+                    case 1:
+                        new_value = [slice_any_object(item, idx) for item in value]
+                        updated_fields[fld.name] = value.__class__(new_value)
+                    case _:
+                        raise NotImplementedError(
+                            "Arbitrary slicing for different sequence length fields is not implemented"
+                        )
+
+        return replace(self, **updated_fields)
+
+    def _check_sequence_lengths_match(self):
+        """Checks if sequence lengths of all "sequence" fields match."""
+        for fld in fields(self):
+            if fld.metadata.get("sequence", False) and fld.name != "complex":
+                # this is a sequence, should be the same length as all other sequences
+                sequence_dim = fld.metadata.get("sequence_dim", 0)
+                value = getattr(self, fld.name)
+                if value is None:
+                    continue
+                match sequence_dim:
+                    case 0:
+                        # sequence is first dimension
+                        value = getattr(self, fld.name)
+                        if len(value) != len(self):
+                            raise ValueError(
+                                f"Mismatch in sequence length for field: {fld.name}. Expected {len(self)}, received {len(value)}"
+                            )
+                    case 1:
+                        for item in value:
+                            if len(item) != len(self):
+                                raise ValueError(
+                                    f"Mismatch in sequence length for field: {fld.name}. Expected {len(self)}, received {len(item)}"
+                                )
+                    case _:
+                        raise NotImplementedError(
+                            "Arbitrary matching for different sequence length fields is not implemented"
+                        )
+
+    @classmethod
+    def concat(cls, items: list[T], **kwargs) -> T:
+        updated_fields = {}
+        for fld in fields(cls):
+            if fld.metadata.get("sequence", False):
+                # this is a sequence, should be the same length as all other sequences
+                sequence_dim = fld.metadata.get("sequence_dim", 0)
+                join_value = fld.metadata.get("join_token", None)
+                if getattr(items[0], fld.name) is None:
+                    continue
+                values = [getattr(item, fld.name) for item in items]
+                match sequence_dim:
+                    case 0:
+                        # sequence is first dimension
+                        value = concat_objects(values, join_value)
+                        updated_fields[fld.name] = value
+                    case 1:
+                        new_value = [
+                            concat_objects(item, join_value) for item in zip(*values)
+                        ]
+                        updated_fields[fld.name] = getattr(
+                            items[0], fld.name
+                        ).__class__(new_value)
+                    case _:
+                        raise NotImplementedError(
+                            "Arbitrary joining for different sequence length fields is not implemented"
+                        )
+        updated_fields.update(kwargs)
+
+        return replace(
+            items[0],  # type: ignore
+            **updated_fields,
+        )

esmfold2_system.py

@@ -1,45 +1,45 @@
-import io
-import subprocess
-import typing as T
-from pathlib import Path
-
-PathLike = T.Union[str, Path]
-PathOrBuffer = T.Union[PathLike, io.StringIO]
-
-
-def run_subprocess_with_errorcheck(
-    *popenargs,
-    capture_output: bool = False,
-    quiet: bool = False,
-    env: dict[str, str] | None = None,
-    shell: bool = False,
-    executable: str | None = None,
-    **kws,
-) -> subprocess.CompletedProcess:
-    """A command similar to subprocess.run, however the errormessage will
-    contain the stderr when using this function. This makes it significantly
-    easier to diagnose issues.
-    """
-    try:
-        if capture_output:
-            stdout = subprocess.PIPE
-        elif quiet:
-            stdout = subprocess.DEVNULL
-        else:
-            stdout = None
-
-        p = subprocess.run(
-            *popenargs,
-            stderr=subprocess.PIPE,
-            stdout=stdout,
-            check=True,
-            env=env,
-            shell=shell,
-            executable=executable,
-            **kws,
-        )
-    except subprocess.CalledProcessError as e:
-        raise RuntimeError(
-            f"Command failed with errorcode {e.returncode}." f"\n\n{e.stderr.decode()}"
-        )
-    return p
+import io
+import subprocess
+import typing as T
+from pathlib import Path
+
+PathLike = T.Union[str, Path]
+PathOrBuffer = T.Union[PathLike, io.StringIO]
+
+
+def run_subprocess_with_errorcheck(
+    *popenargs,
+    capture_output: bool = False,
+    quiet: bool = False,
+    env: dict[str, str] | None = None,
+    shell: bool = False,
+    executable: str | None = None,
+    **kws,
+) -> subprocess.CompletedProcess:
+    """A command similar to subprocess.run, however the errormessage will
+    contain the stderr when using this function. This makes it significantly
+    easier to diagnose issues.
+    """
+    try:
+        if capture_output:
+            stdout = subprocess.PIPE
+        elif quiet:
+            stdout = subprocess.DEVNULL
+        else:
+            stdout = None
+
+        p = subprocess.run(
+            *popenargs,
+            stderr=subprocess.PIPE,
+            stdout=stdout,
+            check=True,
+            env=env,
+            shell=shell,
+            executable=executable,
+            **kws,
+        )
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(
+            f"Command failed with errorcode {e.returncode}." f"\n\n{e.stderr.decode()}"
+        )
+    return p

esmfold2_types.py

@@ -1,33 +1,33 @@
-"""Re-exports of the canonical SPI dataclasses from input_builder.
-
-This module exists so the HF processor and downstream code can import the
-ESMFold2 input types from a single namespace without picking up internal-only
-sibling utilities. The actual definitions live in
-``esm.utils.structure.input_builder``.
-"""
-
-from .esmfold2_msa import MSA
-from .esmfold2_parsing import FastaEntry
-from .esmfold2_input_builder import (
-    CovalentBond,
-    DistogramConditioning,
-    DNAInput,
-    LigandInput,
-    Modification,
-    ProteinInput,
-    RNAInput,
-    StructurePredictionInput,
-)
-
-__all__ = [
-    "FastaEntry",
-    "MSA",
-    "Modification",
-    "ProteinInput",
-    "RNAInput",
-    "DNAInput",
-    "LigandInput",
-    "DistogramConditioning",
-    "CovalentBond",
-    "StructurePredictionInput",
-]
+"""Re-exports of the canonical SPI dataclasses from input_builder.
+
+This module exists so the HF processor and downstream code can import the
+ESMFold2 input types from a single namespace without picking up internal-only
+sibling utilities. The actual definitions live in
+``esm.utils.structure.input_builder``.
+"""
+
+from .esmfold2_msa import MSA
+from .esmfold2_parsing import FastaEntry
+from .esmfold2_input_builder import (
+    CovalentBond,
+    DistogramConditioning,
+    DNAInput,
+    LigandInput,
+    Modification,
+    ProteinInput,
+    RNAInput,
+    StructurePredictionInput,
+)
+
+__all__ = [
+    "FastaEntry",
+    "MSA",
+    "Modification",
+    "ProteinInput",
+    "RNAInput",
+    "DNAInput",
+    "LigandInput",
+    "DistogramConditioning",
+    "CovalentBond",
+    "StructurePredictionInput",
+]

esmfold2_utils_types.py

@@ -1,33 +1,33 @@
-from __future__ import annotations
-
-import io
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Union
-
-from cloudpathlib import CloudPath
-
-PathLike = Union[str, Path, CloudPath]
-PathOrBuffer = Union[PathLike, io.StringIO]
-
-
-@dataclass
-class FunctionAnnotation:
-    """Represents an annotation of a protein's function over a range of residues.
-
-    Fields:
-        label (str): An entry in either the function_tokens or residue_annotations tokenizer vocabs
-        start (int): Start index of this annotation.  1-indexed, inclusive.
-        end (int): End index of this annotation.  1-indexed, inclusive.
-    """
-
-    label: str
-    start: int
-    end: int
-
-    def to_tuple(self) -> tuple[str, int, int]:
-        return self.label, self.start, self.end
-
-    def __len__(self) -> int:
-        """Length of the annotation."""
-        return self.end - self.start + 1
+from __future__ import annotations
+
+import io
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Union
+
+from cloudpathlib import CloudPath
+
+PathLike = Union[str, Path, CloudPath]
+PathOrBuffer = Union[PathLike, io.StringIO]
+
+
+@dataclass
+class FunctionAnnotation:
+    """Represents an annotation of a protein's function over a range of residues.
+
+    Fields:
+        label (str): An entry in either the function_tokens or residue_annotations tokenizer vocabs
+        start (int): Start index of this annotation.  1-indexed, inclusive.
+        end (int): End index of this annotation.  1-indexed, inclusive.
+    """
+
+    label: str
+    start: int
+    end: int
+
+    def to_tuple(self) -> tuple[str, int, int]:
+        return self.label, self.start, self.end
+
+    def __len__(self) -> int:
+        """Length of the annotation."""
+        return self.end - self.start + 1

modeling_esmfold2_experimental.py

@@ -1,997 +1,1001 @@
-"""FastPLMs ESMFold2 experimental architecture.
-
-This module supports Biohub's experimental binder-design checkpoints. The
-released ESMFold2 architecture in ``modeling_esmfold2.py`` intentionally
-rejects those configs because the experimental trunk uses explicit pair-loop
-re-injection and a different confidence/MSA stack.
-"""
-
-from __future__ import annotations
-
-import math
-from pathlib import Path
-from typing import Any, cast
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch import Tensor
-from transformers.modeling_utils import PreTrainedModel
-
-from .configuration_esmfold2 import ESMFold2Config
-from .modeling_esmfold2 import (
-    _convert_te_modules_to_fp8_inplace,
-    _lm_precision_context,
-)
-from .modeling_esmfold2_common import (
-    CHAR_VOCAB_SIZE,
-    MAX_ATOMIC_NUMBER,
-    NUM_RES_TYPES,
-    DiffusionModule,
-    DiffusionStructureHead,
-    DiffusionTransformer,
-    FoldingTrunk,
-    InputsEmbedder,
-    LanguageModelShim,
-    MSAPairWeightedAveraging,
-    OuterProductMean,
-    PairUpdateBlock,
-    ResIdxAsymIdSymIdEntityIdEncoding,
-    RowAttentionPooling,
-    SwiGLUMLP,
-    TriangleMultiplicativeUpdate,
-    _categorical_mean,
-    _compute_intra_token_idx,
-    _seed_context,
-    compute_lm_hidden_states,
-    gather_rep_atom_coords,
-    gather_token_to_atom,
-)
-
-_EPS = 1e-5
-_NONPOLYMER_ID = 3
-
-
-class ConfidenceHead(nn.Module):
-    """Experimental confidence head predicting pLDDT, PAE, pTM, and ipTM."""
-
-    boundaries: Tensor
-
-    def __init__(self, config: ESMFold2Config) -> None:
-        super().__init__()
-        ch = config.confidence_head
-        d_single = config.d_single
-        d_pair = config.d_pair
-        d_inputs = config.inputs.d_inputs
-
-        boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1)
-        self.register_buffer("boundaries", boundaries)
-        self.dist_bin_pairwise_embed = nn.Embedding(ch.distogram_bins, d_pair)
-
-        self.s_norm = nn.LayerNorm(d_single)
-        self.s_inputs_to_single = nn.Linear(d_inputs, d_single, bias=False)
-        self.s_to_z = nn.Linear(d_inputs, d_pair, bias=False)
-        self.s_to_z_transpose = nn.Linear(d_inputs, d_pair, bias=False)
-        self.s_to_z_prod_in1 = nn.Linear(d_inputs, d_pair, bias=False)
-        self.s_to_z_prod_in2 = nn.Linear(d_inputs, d_pair, bias=False)
-        self.s_to_z_prod_out = nn.Linear(d_pair, d_pair, bias=False)
-        self.s_input_to_s = nn.Linear(d_inputs, d_single, bias=False)
-        self.s_inputs_norm = nn.LayerNorm(d_inputs)
-        self.z_norm = nn.LayerNorm(d_pair)
-        self.row_attention_pooling = RowAttentionPooling(
-            d_pair=d_pair, d_single=d_single
-        )
-
-        pf = ch.folding_trunk
-        self.folding_trunk = FoldingTrunk(
-            n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4
-        )
-
-        self.plddt_ln = nn.LayerNorm(d_single)
-        max_atoms_per_token = 23
-        self.plddt_weight = nn.Parameter(
-            torch.zeros(max_atoms_per_token, d_single, ch.num_plddt_bins)
-        )
-        self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False)
-
-    def set_kernel_backend(self, backend: str | None) -> None:
-        self.folding_trunk.set_kernel_backend(backend)
-
-    def set_chunk_size(self, chunk_size: int | None) -> None:
-        self.folding_trunk.set_chunk_size(chunk_size)
-
-    @staticmethod
-    def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor:
-        if num_diffusion_samples == 1:
-            return x
-        return x.repeat_interleave(num_diffusion_samples, 0)
-
-    @staticmethod
-    def _flatten_sample_axis(x: Tensor) -> Tensor:
-        if x.ndim == 4:
-            b, mult, n, c = x.shape
-            return x.reshape(b * mult, n, c)
-        return x
-
-    def forward(
-        self,
-        s_inputs: Tensor,
-        z: Tensor,
-        x_pred: Tensor,
-        distogram_atom_idx: Tensor,
-        token_attention_mask: Tensor,
-        atom_to_token: Tensor,
-        atom_attention_mask: Tensor,
-        asym_id: Tensor,
-        mol_type: Tensor,
-        num_diffusion_samples: int = 1,
-        relative_position_encoding: Tensor | None = None,
-        token_bonds_encoding: Tensor | None = None,
-    ) -> dict[str, Tensor]:
-        s_inputs_normed = self.s_inputs_norm(s_inputs)
-        z_base = self.z_norm(z)
-        if relative_position_encoding is not None:
-            z_base = z_base + relative_position_encoding
-        if token_bonds_encoding is not None:
-            z_base = z_base + token_bonds_encoding
-        z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2)
-        z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1)
-        z_base = z_base + self.s_to_z_prod_out(
-            self.s_to_z_prod_in1(s_inputs_normed)[:, :, None, :]
-            * self.s_to_z_prod_in2(s_inputs_normed)[:, None, :, :]
-        )
-
-        pair = self._repeat_batch(z_base, num_diffusion_samples)
-        x_pred_flat = self._flatten_sample_axis(x_pred)
-        atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples)
-        atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples)
-        rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long()
-        mask = self._repeat_batch(token_attention_mask, num_diffusion_samples)
-        batch_mult = pair.shape[0]
-
-        rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m)
-        rep_distances = torch.cdist(
-            rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist"
-        )
-        distogram_bins = (
-            (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()
-        )
-        pair = pair + self.dist_bin_pairwise_embed(distogram_bins)
-
-        pair_mask = mask[:, :, None].float() * mask[:, None, :].float()
-        pair = pair + self.folding_trunk(pair, pair_attention_mask=pair_mask)
-        single = self.row_attention_pooling(pair, mask)
-
-        atom_mask_f = atom_mask_m.float()
-        s_at_atoms = gather_token_to_atom(single, atom_to_token_m)
-        s_at_atoms = self.plddt_ln(s_at_atoms)
-        intra_idx = _compute_intra_token_idx(atom_to_token_m)
-        intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1)
-        plddt_weight = self.plddt_weight[intra_idx]
-        plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms, plddt_weight)
-        plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0)
-
-        length = single.shape[1]
-        plddt_sum = torch.zeros(
-            batch_mult, length, device=single.device, dtype=plddt_per_atom.dtype
-        )
-        atom_count = torch.zeros(
-            batch_mult, length, device=single.device, dtype=plddt_per_atom.dtype
-        )
-        atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype)
-        plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t)
-        atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t)
-        plddt = plddt_sum / atom_count.clamp(min=1e-6)
-
-        complex_plddt = (plddt_per_atom * atom_mask_f).sum(dim=-1) / (
-            atom_mask_f.sum(dim=-1) + _EPS
-        )
-
-        expanded_type = self._repeat_batch(mol_type, num_diffusion_samples)
-        expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples)
-        is_ligand = (expanded_type == _NONPOLYMER_ID).float()
-        inter_chain = (
-            expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
-        ).float()
-        near_contact = (rep_distances < 8).float()
-        interface_per_token = (
-            near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1)
-        ).amax(dim=-1)
-        iplddt_weight = torch.where(
-            is_ligand.bool(),
-            torch.full_like(interface_per_token, 2.0),
-            interface_per_token,
-        )
-        iplddt_weight_atoms = gather_token_to_atom(
-            iplddt_weight.unsqueeze(-1), atom_to_token_m
-        ).squeeze(-1)
-        atom_iplddt_w = atom_mask_f * iplddt_weight_atoms
-        complex_iplddt = (plddt_per_atom * atom_iplddt_w).sum(dim=-1) / (
-            atom_iplddt_w.sum(dim=-1) + _EPS
-        )
-        plddt_ca = plddt_per_atom.gather(1, rep_idx_m)
-
-        pae_logits = self.pae_head(pair)
-        pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach()
-
-        n_bins = pae_logits.shape[-1]
-        bin_width = 32.0 / n_bins
-        bin_centers = torch.arange(
-            0.5 * bin_width, 32.0, bin_width, device=pae_logits.device
-        )
-        mask_f = mask.float()
-        n_res = mask_f.sum(dim=-1, keepdim=True)
-        d0 = 1.24 * (n_res.clamp(min=19) - 15) ** (1 / 3) - 1.8
-        tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2)
-        pae_probs = F.softmax(pae_logits, dim=-1)
-        tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1)
-
-        pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2)
-        ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (
-            pair_mask_2d.sum(dim=-1) + _EPS
-        )
-        ptm = ptm_per_row.max(dim=-1).values
-
-        inter_chain_mask = (
-            expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
-        ).float() * pair_mask_2d
-        iptm_per_row = (tm_expected * inter_chain_mask).sum(dim=-1) / (
-            inter_chain_mask.sum(dim=-1) + _EPS
-        )
-        iptm = iptm_per_row.max(dim=-1).values
-
-        max_chain_id = int(expanded_asym.max().item()) if batch_mult > 0 else 0
-        n_chains = max_chain_id + 1
-        pair_chains_iptm = torch.zeros(
-            batch_mult,
-            n_chains,
-            n_chains,
-            device=tm_expected.device,
-            dtype=tm_expected.dtype,
-        )
-        for c1 in range(n_chains):
-            chain_c1 = (expanded_asym == c1).float() * mask_f
-            if chain_c1.sum() == 0:
-                continue
-            for c2 in range(n_chains):
-                chain_c2 = (expanded_asym == c2).float() * mask_f
-                pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2)
-                denom = pair_m.sum(dim=(-1, -2)) + _EPS
-                pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(
-                    dim=(-1, -2)
-                ) / denom
-
-        return {
-            "plddt_logits": plddt_logits,
-            "plddt": plddt.detach(),
-            "plddt_per_atom": plddt_per_atom.detach(),
-            "plddt_ca": plddt_ca.detach(),
-            "complex_plddt": complex_plddt.detach(),
-            "complex_iplddt": complex_iplddt.detach(),
-            "pae_logits": pae_logits,
-            "pae": pae,
-            "ptm": ptm.detach(),
-            "iptm": iptm.detach(),
-            "pair_chains_iptm": pair_chains_iptm.detach(),
-        }
-
-
-class _TransitionFFN(nn.Module):
-    def __init__(self, d_model: int, expansion_ratio: int = 4) -> None:
-        super().__init__()
-        self.norm = nn.LayerNorm(d_model)
-        self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.ffn(self.norm(x))
-
-
-class MSAEncoderBlock(nn.Module):
-    """One experimental MSA update block."""
-
-    def __init__(
-        self,
-        d_msa: int,
-        d_pair: int,
-        d_hidden: int = 32,
-        n_heads_msa: int = 8,
-        msa_head_width: int = 32,
-    ) -> None:
-        super().__init__()
-        self.outer_product_mean = OuterProductMean(
-            d_msa, d_hidden, d_pair, divide_outer_before_proj=True
-        )
-        self.msa_pair_weighted_averaging = MSAPairWeightedAveraging(
-            d_msa, d_pair, n_heads_msa, msa_head_width
-        )
-        self.msa_transition = _TransitionFFN(d_msa, expansion_ratio=4)
-        self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True)
-        self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False)
-        self.pair_transition = _TransitionFFN(d_pair, expansion_ratio=4)
-
-    def set_chunk_size(self, chunk_size: int | None) -> None:
-        self.outer_product_mean.set_chunk_size(chunk_size)
-        self.tri_mul_out.set_chunk_size(chunk_size)
-        self.tri_mul_in.set_chunk_size(chunk_size)
-
-    def forward(
-        self,
-        msa_repr: Tensor,
-        pair_repr: Tensor,
-        msa_attention_mask: Tensor,
-        pair_attention_mask: Tensor,
-        msa_track_mask: Tensor | None = None,
-    ) -> tuple[Tensor, Tensor]:
-        mask4d = (
-            msa_track_mask[:, None, None, None].to(dtype=msa_repr.dtype)
-            if msa_track_mask is not None
-            else None
-        )
-
-        pair_mask4d = mask4d[:, :, :1] if mask4d is not None else None
-
-        msa_update = self.msa_pair_weighted_averaging(
-            msa_repr, pair_repr, pair_attention_mask
-        )
-        if mask4d is not None:
-            msa_update = msa_update * mask4d
-        msa_repr = msa_repr + msa_update
-
-        msa_transition = self.msa_transition(msa_repr)
-        if mask4d is not None:
-            msa_transition = msa_transition * mask4d
-        msa_repr = msa_repr + msa_transition
-
-        pair_opm = self.outer_product_mean(msa_repr, msa_attention_mask)
-        if pair_mask4d is not None:
-            pair_opm = pair_opm * pair_mask4d
-        pair_repr = pair_repr + pair_opm
-
-        pair_out = self.tri_mul_out(pair_repr, mask=pair_attention_mask)
-        if pair_mask4d is not None:
-            pair_out = pair_out * pair_mask4d
-        pair_repr = pair_repr + pair_out
-
-        pair_in = self.tri_mul_in(pair_repr, mask=pair_attention_mask)
-        if pair_mask4d is not None:
-            pair_in = pair_in * pair_mask4d
-        pair_repr = pair_repr + pair_in
-
-        pair_transition = self.pair_transition(pair_repr)
-        if pair_mask4d is not None:
-            pair_transition = pair_transition * pair_mask4d
-        pair_repr = pair_repr + pair_transition
-        return msa_repr, pair_repr
-
-
-class MSAEncoder(nn.Module):
-    def __init__(
-        self,
-        d_msa: int,
-        d_pair: int,
-        d_inputs: int,
-        d_hidden: int = 32,
-        n_layers: int = 4,
-        n_heads_msa: int = 8,
-        msa_head_width: int = 32,
-    ) -> None:
-        super().__init__()
-        self.embed = nn.Linear(35, d_msa, bias=False)
-        self.project_inputs = nn.Linear(d_inputs, d_msa, bias=False)
-        self.blocks = nn.ModuleList(
-            [
-                MSAEncoderBlock(
-                    d_msa=d_msa,
-                    d_pair=d_pair,
-                    d_hidden=d_hidden,
-                    n_heads_msa=n_heads_msa,
-                    msa_head_width=msa_head_width,
-                )
-                for _ in range(n_layers)
-            ]
-        )
-
-    def set_chunk_size(self, chunk_size: int | None) -> None:
-        for block in self.blocks:
-            cast(MSAEncoderBlock, block).set_chunk_size(chunk_size)
-
-    def forward(
-        self,
-        x_pair: Tensor,
-        x_inputs: Tensor,
-        msa_oh: Tensor,
-        has_deletion: Tensor,
-        deletion_value: Tensor,
-        msa_attention_mask: Tensor,
-    ) -> Tensor:
-        batch_size, _, depth = msa_attention_mask.shape
-        m_feat = torch.cat(
-            [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)],
-            dim=-1,
-        )
-        m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2)
-        if depth > 1:
-            msa_track_mask = msa_attention_mask[:, :, 1:].any(dim=(1, 2))
-        else:
-            msa_track_mask = torch.zeros(
-                batch_size, dtype=torch.bool, device=x_pair.device
-            )
-        tok_mask = msa_attention_mask[:, :, 0]
-        pair_attention_mask = tok_mask.unsqueeze(2) * tok_mask.unsqueeze(1)
-        for block in self.blocks:
-            m, x_pair = cast(MSAEncoderBlock, block)(
-                m,
-                x_pair,
-                msa_attention_mask,
-                pair_attention_mask,
-                msa_track_mask,
-            )
-        return x_pair * msa_track_mask[:, None, None, None].to(dtype=x_pair.dtype)
-
-
-class ESMFold2ExperimentalModel(PreTrainedModel):
-    """Experimental ESMFold2 architecture used by binder-design checkpoints."""
-
-    config_class = ESMFold2Config
-    _keys_to_ignore_on_load_unexpected = [r"\._extra_state$"]
-
-    def __init__(self, config: ESMFold2Config) -> None:
-        super().__init__(config)
-        d_inputs = config.inputs.d_inputs
-        d_pair = config.d_pair
-
-        self.inputs_embedder = InputsEmbedder(config)
-        self.z_init_1 = nn.Linear(d_inputs, d_pair, bias=False)
-        self.z_init_2 = nn.Linear(d_inputs, d_pair, bias=False)
-        self.rel_pos = ResIdxAsymIdSymIdEntityIdEncoding(
-            n_relative_residx_bins=config.n_relative_residx_bins,
-            n_relative_chain_bins=config.n_relative_chain_bins,
-            d_pair=d_pair,
-        )
-        self.token_bonds = nn.Linear(1, d_pair, bias=False)
-        self.language_model = LanguageModelShim(
-            d_z=d_pair, d_model=config.lm_d_model, num_layers=config.lm_num_layers
-        )
-        self._esmc: nn.Module | None = None
-        self._esmc_fp8 = False
-        self._esmfold2_input_builder: Any | None = None
-
-        pf = config.folding_trunk
-        self.folding_trunk = FoldingTrunk(
-            n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4
-        )
-        self.pair_loop_proj = nn.Sequential(
-            nn.LayerNorm(d_pair), nn.Linear(d_pair, d_pair, bias=False)
-        )
-        nn.init.zeros_(cast(nn.Linear, self.pair_loop_proj[1]).weight)
-
-        self.structure_head = DiffusionStructureHead(config)
-        self.distogram_head = nn.Linear(
-            d_pair, config.structure_head.distogram_bins, bias=True
-        )
-        self.confidence_head: ConfidenceHead | None = (
-            ConfidenceHead(config) if config.confidence_head.enabled else None
-        )
-
-        msa_cfg = config.msa_encoder
-        self.msa_encoder: MSAEncoder | None = None
-        if msa_cfg.enabled:
-            self.msa_encoder = MSAEncoder(
-                d_msa=msa_cfg.d_msa,
-                d_pair=d_pair,
-                d_inputs=d_inputs,
-                d_hidden=msa_cfg.d_hidden,
-                n_layers=msa_cfg.n_layers,
-                n_heads_msa=msa_cfg.n_heads_msa,
-                msa_head_width=msa_cfg.msa_head_width,
-            )
-
-        self.post_init()
-
-    @property
-    def device(self) -> torch.device:
-        return next(self.parameters()).device
-
-    def set_kernel_backend(self, backend: str | None) -> None:
-        self.folding_trunk.set_kernel_backend(backend)
-        if self.confidence_head is not None:
-            self.confidence_head.set_kernel_backend(backend)
-        self.structure_head.set_kernel_backend(backend)
-
-    def set_chunk_size(self, chunk_size: int | None) -> None:
-        self.folding_trunk.set_chunk_size(chunk_size)
-        if self.confidence_head is not None:
-            self.confidence_head.set_chunk_size(chunk_size)
-        if self.msa_encoder is not None:
-            self.msa_encoder.set_chunk_size(chunk_size)
-
-    def configure_lm_dropout(
-        self,
-        lm_dropout: float,
-        *,
-        force_lm_dropout_during_inference: bool = True,
-    ) -> None:
-        self.config.lm_dropout = lm_dropout
-        self.config.force_lm_dropout_during_inference = (
-            force_lm_dropout_during_inference
-        )
-
-    def load_esmc(self, esmc_model_path: str, precision: str = "bf16") -> None:
-        from .modeling_esmc import ESMCModel
-
-        dtype_map = {
-            "bf16": torch.bfloat16,
-            "fp32": torch.float32,
-            "fp8": torch.bfloat16,
-        }
+"""FastPLMs ESMFold2 experimental architecture.
+
+This module supports Biohub's experimental binder-design checkpoints. The
+released ESMFold2 architecture in ``modeling_esmfold2.py`` intentionally
+rejects those configs because the experimental trunk uses explicit pair-loop
+re-injection and a different confidence/MSA stack.
+"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import Any, cast
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_esmfold2 import ESMFold2Config
+from .modeling_esmfold2 import (
+    _load_fastplms_esmplusplus_for_esmfold2,
+    _lm_precision_context,
+)
+from .modeling_esmfold2_common import (
+    CHAR_VOCAB_SIZE,
+    MAX_ATOMIC_NUMBER,
+    NUM_RES_TYPES,
+    DiffusionModule,
+    DiffusionStructureHead,
+    DiffusionTransformer,
+    FoldingTrunk,
+    InputsEmbedder,
+    LanguageModelShim,
+    MSAPairWeightedAveraging,
+    OuterProductMean,
+    PairUpdateBlock,
+    ResIdxAsymIdSymIdEntityIdEncoding,
+    RowAttentionPooling,
+    SwiGLUMLP,
+    TriangleMultiplicativeUpdate,
+    _categorical_mean,
+    _compute_intra_token_idx,
+    _seed_context,
+    compute_lm_hidden_states,
+    gather_rep_atom_coords,
+    gather_token_to_atom,
+)
+
+_EPS = 1e-5
+_NONPOLYMER_ID = 3
+
+
+class ConfidenceHead(nn.Module):
+    """Experimental confidence head predicting pLDDT, PAE, pTM, and ipTM."""
+
+    boundaries: Tensor
+
+    def __init__(self, config: ESMFold2Config) -> None:
+        super().__init__()
+        ch = config.confidence_head
+        d_single = config.d_single
+        d_pair = config.d_pair
+        d_inputs = config.inputs.d_inputs
+
+        boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1)
+        self.register_buffer("boundaries", boundaries)
+        self.dist_bin_pairwise_embed = nn.Embedding(ch.distogram_bins, d_pair)
+
+        self.s_norm = nn.LayerNorm(d_single)
+        self.s_inputs_to_single = nn.Linear(d_inputs, d_single, bias=False)
+        self.s_to_z = nn.Linear(d_inputs, d_pair, bias=False)
+        self.s_to_z_transpose = nn.Linear(d_inputs, d_pair, bias=False)
+        self.s_to_z_prod_in1 = nn.Linear(d_inputs, d_pair, bias=False)
+        self.s_to_z_prod_in2 = nn.Linear(d_inputs, d_pair, bias=False)
+        self.s_to_z_prod_out = nn.Linear(d_pair, d_pair, bias=False)
+        self.s_input_to_s = nn.Linear(d_inputs, d_single, bias=False)
+        self.s_inputs_norm = nn.LayerNorm(d_inputs)
+        self.z_norm = nn.LayerNorm(d_pair)
+        self.row_attention_pooling = RowAttentionPooling(
+            d_pair=d_pair, d_single=d_single
+        )
+
+        pf = ch.folding_trunk
+        self.folding_trunk = FoldingTrunk(
+            n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4
+        )
+
+        self.plddt_ln = nn.LayerNorm(d_single)
+        max_atoms_per_token = 23
+        self.plddt_weight = nn.Parameter(
+            torch.zeros(max_atoms_per_token, d_single, ch.num_plddt_bins)
+        )
+        self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False)
+
+    def set_kernel_backend(self, backend: str | None) -> None:
+        self.folding_trunk.set_kernel_backend(backend)
+
+    def set_chunk_size(self, chunk_size: int | None) -> None:
+        self.folding_trunk.set_chunk_size(chunk_size)
+
+    @staticmethod
+    def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor:
+        if num_diffusion_samples == 1:
+            return x
+        return x.repeat_interleave(num_diffusion_samples, 0)
+
+    @staticmethod
+    def _flatten_sample_axis(x: Tensor) -> Tensor:
+        if x.ndim == 4:
+            b, mult, n, c = x.shape
+            return x.reshape(b * mult, n, c)
+        return x
+
+    def forward(
+        self,
+        s_inputs: Tensor,
+        z: Tensor,
+        x_pred: Tensor,
+        distogram_atom_idx: Tensor,
+        token_attention_mask: Tensor,
+        atom_to_token: Tensor,
+        atom_attention_mask: Tensor,
+        asym_id: Tensor,
+        mol_type: Tensor,
+        num_diffusion_samples: int = 1,
+        relative_position_encoding: Tensor | None = None,
+        token_bonds_encoding: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        s_inputs_normed = self.s_inputs_norm(s_inputs)
+        z_base = self.z_norm(z)
+        if relative_position_encoding is not None:
+            z_base = z_base + relative_position_encoding
+        if token_bonds_encoding is not None:
+            z_base = z_base + token_bonds_encoding
+        z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2)
+        z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1)
+        z_base = z_base + self.s_to_z_prod_out(
+            self.s_to_z_prod_in1(s_inputs_normed)[:, :, None, :]
+            * self.s_to_z_prod_in2(s_inputs_normed)[:, None, :, :]
+        )
+
+        pair = self._repeat_batch(z_base, num_diffusion_samples)
+        x_pred_flat = self._flatten_sample_axis(x_pred)
+        atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples)
+        atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples)
+        rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long()
+        mask = self._repeat_batch(token_attention_mask, num_diffusion_samples)
+        batch_mult = pair.shape[0]
+
+        rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m)
+        rep_distances = torch.cdist(
+            rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist"
+        )
+        distogram_bins = (
+            (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()
+        )
+        pair = pair + self.dist_bin_pairwise_embed(distogram_bins)
+
+        pair_mask = mask[:, :, None].float() * mask[:, None, :].float()
+        pair = pair + self.folding_trunk(pair, pair_attention_mask=pair_mask)
+        single = self.row_attention_pooling(pair, mask)
+
+        atom_mask_f = atom_mask_m.float()
+        s_at_atoms = gather_token_to_atom(single, atom_to_token_m)
+        s_at_atoms = self.plddt_ln(s_at_atoms)
+        intra_idx = _compute_intra_token_idx(atom_to_token_m)
+        intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1)
+        plddt_weight = self.plddt_weight[intra_idx]
+        plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms, plddt_weight)
+        plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0)
+
+        length = single.shape[1]
+        plddt_sum = torch.zeros(
+            batch_mult, length, device=single.device, dtype=plddt_per_atom.dtype
+        )
+        atom_count = torch.zeros(
+            batch_mult, length, device=single.device, dtype=plddt_per_atom.dtype
+        )
+        atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype)
+        plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t)
+        atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t)
+        plddt = plddt_sum / atom_count.clamp(min=1e-6)
+
+        complex_plddt = (plddt_per_atom * atom_mask_f).sum(dim=-1) / (
+            atom_mask_f.sum(dim=-1) + _EPS
+        )
+
+        expanded_type = self._repeat_batch(mol_type, num_diffusion_samples)
+        expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples)
+        is_ligand = (expanded_type == _NONPOLYMER_ID).float()
+        inter_chain = (
+            expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
+        ).float()
+        near_contact = (rep_distances < 8).float()
+        interface_per_token = (
+            near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1)
+        ).amax(dim=-1)
+        iplddt_weight = torch.where(
+            is_ligand.bool(),
+            torch.full_like(interface_per_token, 2.0),
+            interface_per_token,
+        )
+        iplddt_weight_atoms = gather_token_to_atom(
+            iplddt_weight.unsqueeze(-1), atom_to_token_m
+        ).squeeze(-1)
+        atom_iplddt_w = atom_mask_f * iplddt_weight_atoms
+        complex_iplddt = (plddt_per_atom * atom_iplddt_w).sum(dim=-1) / (
+            atom_iplddt_w.sum(dim=-1) + _EPS
+        )
+        plddt_ca = plddt_per_atom.gather(1, rep_idx_m)
+
+        pae_logits = self.pae_head(pair)
+        pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach()
+
+        n_bins = pae_logits.shape[-1]
+        bin_width = 32.0 / n_bins
+        bin_centers = torch.arange(
+            0.5 * bin_width, 32.0, bin_width, device=pae_logits.device
+        )
+        mask_f = mask.float()
+        n_res = mask_f.sum(dim=-1, keepdim=True)
+        d0 = 1.24 * (n_res.clamp(min=19) - 15) ** (1 / 3) - 1.8
+        tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2)
+        pae_probs = F.softmax(pae_logits, dim=-1)
+        tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1)
+
+        pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2)
+        ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (
+            pair_mask_2d.sum(dim=-1) + _EPS
+        )
+        ptm = ptm_per_row.max(dim=-1).values
+
+        inter_chain_mask = (
+            expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
+        ).float() * pair_mask_2d
+        iptm_per_row = (tm_expected * inter_chain_mask).sum(dim=-1) / (
+            inter_chain_mask.sum(dim=-1) + _EPS
+        )
+        iptm = iptm_per_row.max(dim=-1).values
+
+        max_chain_id = int(expanded_asym.max().item()) if batch_mult > 0 else 0
+        n_chains = max_chain_id + 1
+        pair_chains_iptm = torch.zeros(
+            batch_mult,
+            n_chains,
+            n_chains,
+            device=tm_expected.device,
+            dtype=tm_expected.dtype,
+        )
+        for c1 in range(n_chains):
+            chain_c1 = (expanded_asym == c1).float() * mask_f
+            if chain_c1.sum() == 0:
+                continue
+            for c2 in range(n_chains):
+                chain_c2 = (expanded_asym == c2).float() * mask_f
+                pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2)
+                denom = pair_m.sum(dim=(-1, -2)) + _EPS
+                pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(
+                    dim=(-1, -2)
+                ) / denom
+
+        return {
+            "plddt_logits": plddt_logits,
+            "plddt": plddt.detach(),
+            "plddt_per_atom": plddt_per_atom.detach(),
+            "plddt_ca": plddt_ca.detach(),
+            "complex_plddt": complex_plddt.detach(),
+            "complex_iplddt": complex_iplddt.detach(),
+            "pae_logits": pae_logits,
+            "pae": pae,
+            "ptm": ptm.detach(),
+            "iptm": iptm.detach(),
+            "pair_chains_iptm": pair_chains_iptm.detach(),
+        }
+
+
+class _TransitionFFN(nn.Module):
+    def __init__(self, d_model: int, expansion_ratio: int = 4) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(d_model)
+        self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.ffn(self.norm(x))
+
+
+class MSAEncoderBlock(nn.Module):
+    """One experimental MSA update block."""
+
+    def __init__(
+        self,
+        d_msa: int,
+        d_pair: int,
+        d_hidden: int = 32,
+        n_heads_msa: int = 8,
+        msa_head_width: int = 32,
+    ) -> None:
+        super().__init__()
+        self.outer_product_mean = OuterProductMean(
+            d_msa, d_hidden, d_pair, divide_outer_before_proj=True
+        )
+        self.msa_pair_weighted_averaging = MSAPairWeightedAveraging(
+            d_msa, d_pair, n_heads_msa, msa_head_width
+        )
+        self.msa_transition = _TransitionFFN(d_msa, expansion_ratio=4)
+        self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True)
+        self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False)
+        self.pair_transition = _TransitionFFN(d_pair, expansion_ratio=4)
+
+    def set_chunk_size(self, chunk_size: int | None) -> None:
+        self.outer_product_mean.set_chunk_size(chunk_size)
+        self.tri_mul_out.set_chunk_size(chunk_size)
+        self.tri_mul_in.set_chunk_size(chunk_size)
+
+    def forward(
+        self,
+        msa_repr: Tensor,
+        pair_repr: Tensor,
+        msa_attention_mask: Tensor,
+        pair_attention_mask: Tensor,
+        msa_track_mask: Tensor | None = None,
+    ) -> tuple[Tensor, Tensor]:
+        mask4d = (
+            msa_track_mask[:, None, None, None].to(dtype=msa_repr.dtype)
+            if msa_track_mask is not None
+            else None
+        )
+
+        pair_mask4d = mask4d[:, :, :1] if mask4d is not None else None
+
+        msa_update = self.msa_pair_weighted_averaging(
+            msa_repr, pair_repr, pair_attention_mask
+        )
+        if mask4d is not None:
+            msa_update = msa_update * mask4d
+        msa_repr = msa_repr + msa_update
+
+        msa_transition = self.msa_transition(msa_repr)
+        if mask4d is not None:
+            msa_transition = msa_transition * mask4d
+        msa_repr = msa_repr + msa_transition
+
+        pair_opm = self.outer_product_mean(msa_repr, msa_attention_mask)
+        if pair_mask4d is not None:
+            pair_opm = pair_opm * pair_mask4d
+        pair_repr = pair_repr + pair_opm
+
+        pair_out = self.tri_mul_out(pair_repr, mask=pair_attention_mask)
+        if pair_mask4d is not None:
+            pair_out = pair_out * pair_mask4d
+        pair_repr = pair_repr + pair_out
+
+        pair_in = self.tri_mul_in(pair_repr, mask=pair_attention_mask)
+        if pair_mask4d is not None:
+            pair_in = pair_in * pair_mask4d
+        pair_repr = pair_repr + pair_in
+
+        pair_transition = self.pair_transition(pair_repr)
+        if pair_mask4d is not None:
+            pair_transition = pair_transition * pair_mask4d
+        pair_repr = pair_repr + pair_transition
+        return msa_repr, pair_repr
+
+
+class MSAEncoder(nn.Module):
+    def __init__(
+        self,
+        d_msa: int,
+        d_pair: int,
+        d_inputs: int,
+        d_hidden: int = 32,
+        n_layers: int = 4,
+        n_heads_msa: int = 8,
+        msa_head_width: int = 32,
+    ) -> None:
+        super().__init__()
+        self.embed = nn.Linear(35, d_msa, bias=False)
+        self.project_inputs = nn.Linear(d_inputs, d_msa, bias=False)
+        self.blocks = nn.ModuleList(
+            [
+                MSAEncoderBlock(
+                    d_msa=d_msa,
+                    d_pair=d_pair,
+                    d_hidden=d_hidden,
+                    n_heads_msa=n_heads_msa,
+                    msa_head_width=msa_head_width,
+                )
+                for _ in range(n_layers)
+            ]
+        )
+
+    def set_chunk_size(self, chunk_size: int | None) -> None:
+        for block in self.blocks:
+            cast(MSAEncoderBlock, block).set_chunk_size(chunk_size)
+
+    def forward(
+        self,
+        x_pair: Tensor,
+        x_inputs: Tensor,
+        msa_oh: Tensor,
+        has_deletion: Tensor,
+        deletion_value: Tensor,
+        msa_attention_mask: Tensor,
+    ) -> Tensor:
+        batch_size, _, depth = msa_attention_mask.shape
+        m_feat = torch.cat(
+            [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)],
+            dim=-1,
+        )
+        m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2)
+        if depth > 1:
+            msa_track_mask = msa_attention_mask[:, :, 1:].any(dim=(1, 2))
+        else:
+            msa_track_mask = torch.zeros(
+                batch_size, dtype=torch.bool, device=x_pair.device
+            )
+        tok_mask = msa_attention_mask[:, :, 0]
+        pair_attention_mask = tok_mask.unsqueeze(2) * tok_mask.unsqueeze(1)
+        for block in self.blocks:
+            m, x_pair = cast(MSAEncoderBlock, block)(
+                m,
+                x_pair,
+                msa_attention_mask,
+                pair_attention_mask,
+                msa_track_mask,
+            )
+        return x_pair * msa_track_mask[:, None, None, None].to(dtype=x_pair.dtype)
+
+
+class ESMFold2ExperimentalModel(PreTrainedModel):
+    """Experimental ESMFold2 architecture used by binder-design checkpoints."""
+
+    config_class = ESMFold2Config
+    _keys_to_ignore_on_load_unexpected = [r"\._extra_state$"]
+
+    def __init__(self, config: ESMFold2Config) -> None:
+        super().__init__(config)
+        d_inputs = config.inputs.d_inputs
+        d_pair = config.d_pair
+
+        self.inputs_embedder = InputsEmbedder(config)
+        self.z_init_1 = nn.Linear(d_inputs, d_pair, bias=False)
+        self.z_init_2 = nn.Linear(d_inputs, d_pair, bias=False)
+        self.rel_pos = ResIdxAsymIdSymIdEntityIdEncoding(
+            n_relative_residx_bins=config.n_relative_residx_bins,
+            n_relative_chain_bins=config.n_relative_chain_bins,
+            d_pair=d_pair,
+        )
+        self.token_bonds = nn.Linear(1, d_pair, bias=False)
+        self.language_model = LanguageModelShim(
+            d_z=d_pair, d_model=config.lm_d_model, num_layers=config.lm_num_layers
+        )
+        self._esmc: nn.Module | None = None
+        self._esmc_fp8 = False
+        self._esmfold2_input_builder: Any | None = None
+
+        pf = config.folding_trunk
+        self.folding_trunk = FoldingTrunk(
+            n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4
+        )
+        self.pair_loop_proj = nn.Sequential(
+            nn.LayerNorm(d_pair), nn.Linear(d_pair, d_pair, bias=False)
+        )
+        nn.init.zeros_(cast(nn.Linear, self.pair_loop_proj[1]).weight)
+
+        self.structure_head = DiffusionStructureHead(config)
+        self.distogram_head = nn.Linear(
+            d_pair, config.structure_head.distogram_bins, bias=True
+        )
+        self.confidence_head: ConfidenceHead | None = (
+            ConfidenceHead(config) if config.confidence_head.enabled else None
+        )
+
+        msa_cfg = config.msa_encoder
+        self.msa_encoder: MSAEncoder | None = None
+        if msa_cfg.enabled:
+            self.msa_encoder = MSAEncoder(
+                d_msa=msa_cfg.d_msa,
+                d_pair=d_pair,
+                d_inputs=d_inputs,
+                d_hidden=msa_cfg.d_hidden,
+                n_layers=msa_cfg.n_layers,
+                n_heads_msa=msa_cfg.n_heads_msa,
+                msa_head_width=msa_cfg.msa_head_width,
+            )
+
+        self.post_init()
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def set_kernel_backend(self, backend: str | None) -> None:
+        self.folding_trunk.set_kernel_backend(backend)
+        if self.confidence_head is not None:
+            self.confidence_head.set_kernel_backend(backend)
+        self.structure_head.set_kernel_backend(backend)
+
+    def set_chunk_size(self, chunk_size: int | None) -> None:
+        self.folding_trunk.set_chunk_size(chunk_size)
+        if self.confidence_head is not None:
+            self.confidence_head.set_chunk_size(chunk_size)
+        if self.msa_encoder is not None:
+            self.msa_encoder.set_chunk_size(chunk_size)
+
+    def configure_lm_dropout(
+        self,
+        lm_dropout: float,
+        *,
+        force_lm_dropout_during_inference: bool = True,
+    ) -> None:
+        self.config.lm_dropout = lm_dropout
+        self.config.force_lm_dropout_during_inference = (
+            force_lm_dropout_during_inference
+        )
+
+    def load_esmc(self, esmc_model_path: str, precision: str = "bf16") -> None:
+        dtype_map = {
+            "bf16": torch.bfloat16,
+            "fp32": torch.float32,
+        }
         if precision not in dtype_map:
-            raise ValueError(
-                f"precision must be one of {list(dtype_map)}, got {precision!r}"
-            )
-        esmc = (
-            ESMCModel.from_pretrained(esmc_model_path)
-            .to(device=self.device, dtype=dtype_map[precision])
-            .eval()
-        )
-        for parameter in esmc.parameters():
-            parameter.requires_grad_(False)
-        if precision == "fp8":
-            with torch.no_grad():
-                _convert_te_modules_to_fp8_inplace(esmc)
-            self._esmc_fp8 = True
-        else:
-            self._esmc_fp8 = False
-        self._esmc = esmc
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path,
-        *model_args,
-        load_esmc: bool = True,
-        **kwargs,
-    ):
-        if "config" not in kwargs:
-            kwargs["config"] = ESMFold2Config.from_pretrained(
-                pretrained_model_name_or_path, **kwargs
-            )
-        esmc_precision = kwargs.pop("esmc_precision", "bf16")
-        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
-        if load_esmc:
-            model.load_esmc(model.config.esmc_id, precision=esmc_precision)
-        return model
-
-    def apply_torch_compile(
-        self, mode: str = "fixed_seqlen", dynamic: bool | None = None
-    ) -> None:
-        import torch._dynamo
-
-        torch._dynamo.config.cache_size_limit = 512
-        torch._dynamo.config.accumulated_cache_size_limit = 512
-        torch._dynamo.config.capture_scalar_outputs = True
-
-        if dynamic is None:
-            dynamic = mode == "dynamic_seqlen"
-        compile_kwargs: dict[str, bool] = {"dynamic": dynamic}
-        compile_targets = (
-            PairUpdateBlock,
-            DiffusionTransformer,
-            DiffusionModule,
-            MSAEncoderBlock,
-        )
-
-        def _maybe_compile(module: nn.Module) -> None:
-            if isinstance(module, compile_targets):
-                module.forward = torch.compile(module.forward, **compile_kwargs)
-
-        self.apply(_maybe_compile)
-
-    def _compute_lm_hidden_states(
-        self,
-        input_ids: Tensor,
-        asym_id: Tensor,
-        residue_index: Tensor,
-        mol_type: Tensor,
-        tok_mask: Tensor,
-    ) -> Tensor:
-        assert self._esmc is not None
-        pad_to = 8 if self._esmc_fp8 else None
-        with _lm_precision_context(self._esmc_fp8):
-            return compute_lm_hidden_states(
-                self._esmc,
-                input_ids,
-                asym_id,
-                residue_index,
-                mol_type,
-                tok_mask,
-                pad_to_multiple=pad_to,
-            )
-
-    def forward(
-        self,
-        token_index: Tensor,
-        residue_index: Tensor,
-        asym_id: Tensor,
-        sym_id: Tensor,
-        entity_id: Tensor,
-        mol_type: Tensor,
-        res_type: Tensor,
-        token_bonds: Tensor,
-        token_attention_mask: Tensor,
-        ref_pos: Tensor,
-        ref_element: Tensor,
-        ref_charge: Tensor,
-        ref_atom_name_chars: Tensor,
-        ref_space_uid: Tensor,
-        atom_attention_mask: Tensor,
-        atom_to_token: Tensor,
-        distogram_atom_idx: Tensor,
-        deletion_mean: Tensor | None = None,
-        msa: Tensor | None = None,
-        has_deletion: Tensor | None = None,
-        deletion_value: Tensor | None = None,
-        msa_attention_mask: Tensor | None = None,
-        input_ids: Tensor | None = None,
-        lm_hidden_states: Tensor | None = None,
-        res_type_soft: Tensor | None = None,
-        num_loops: int | None = None,
-        num_diffusion_samples: int | None = None,
-        num_sampling_steps: int | None = None,
-        early_exit: bool = False,
-        seed: int | None = None,
-        calculate_confidence: bool = True,
-        provide_soft_sequence_to_msa_and_profile: bool = True,
-        noise_scale: float | None = None,
-        step_scale: float | None = None,
-        max_inference_sigma: int | None = None,
-    ) -> dict[str, Tensor]:
-        del noise_scale, step_scale, max_inference_sigma
-        tok_mask = token_attention_mask
-        atm_mask = atom_attention_mask
-        n_loops = num_loops if num_loops is not None else self.config.num_loops
-        n_samples = (
-            num_diffusion_samples
-            if num_diffusion_samples is not None
-            else self.config.num_diffusion_samples
-        )
-
-        if res_type.dim() == 2:
-            res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float()
-            res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float()
-        else:
-            res_type_oh = res_type.float()
-
-        if msa is not None:
-            msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float()
-            if msa_attention_mask is not None:
-                mask_f = msa_attention_mask.float().unsqueeze(-1)
-                msa_oh_profile = msa_oh_profile * mask_f
-                valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1)
-                profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1)
-            else:
-                profile = msa_oh_profile.mean(dim=1)
-        else:
-            profile = res_type_oh
-
-        if res_type_soft is not None:
-            res_type_oh = res_type_soft.float()
-            if (
-                not self.config.disable_msa_features
-                and provide_soft_sequence_to_msa_and_profile
-            ):
-                profile = res_type_oh
-                msa = res_type_oh.unsqueeze(1)
-                msa_attention_mask = tok_mask.unsqueeze(1)
-
-        if deletion_mean is None:
-            deletion_mean = torch.zeros(
-                res_type.shape[0], res_type.shape[1], device=res_type.device
-            )
-        if self.config.disable_msa_features:
-            profile = torch.zeros_like(profile)
-            deletion_mean = torch.zeros_like(deletion_mean)
-
-        ref_element_oh = F.one_hot(
-            ref_element.long(), num_classes=MAX_ATOMIC_NUMBER
-        ).float()
-        ref_atom_name_chars_oh = F.one_hot(
-            ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE
-        ).float()
-        atm_mask_f = atm_mask.float()
-        ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1)
-        ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(
-            -1
-        ).unsqueeze(-1)
-        atom_to_token = atom_to_token * atm_mask.long()
-
-        use_amp = ref_pos.device.type == "cuda"
-        with (
-            torch.set_grad_enabled(res_type_soft is not None),
-            torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16),
-        ):
-            x_inputs = self.inputs_embedder(
-                aatype=res_type_oh,
-                profile=profile.float(),
-                deletion_mean=deletion_mean.float(),
-                ref_pos=ref_pos,
-                atom_attention_mask=atm_mask,
-                ref_space_uid=ref_space_uid,
-                ref_charge=ref_charge,
-                ref_element=ref_element_oh,
-                ref_atom_name_chars=ref_atom_name_chars_oh,
-                atom_to_token=atom_to_token,
-            )
-
-            z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(
-                x_inputs
-            ).unsqueeze(1)
-            relative_position_encoding = self.rel_pos(
-                residue_index=residue_index,
-                asym_id=asym_id,
-                sym_id=sym_id,
-                entity_id=entity_id,
-                token_index=token_index,
-            )
-            token_bonds_encoding = self.token_bonds(token_bonds.float())
-            z_init = z_init + relative_position_encoding + token_bonds_encoding
-
-            if (
-                lm_hidden_states is None
-                and input_ids is not None
-                and self._esmc is not None
-            ):
-                lm_hidden_states = self._compute_lm_hidden_states(
-                    input_ids, asym_id, residue_index, mol_type, tok_mask
-                )
-            if lm_hidden_states is not None:
-                lm_dropout = (
-                    self.config.lm_dropout
-                    if self.config.force_lm_dropout_during_inference or self.training
-                    else 0.0
-                )
-                lm_z = self.language_model(
-                    lm_hidden_states.detach(), lm_dropout=lm_dropout
-                )
-                z_init = z_init + lm_z.to(z_init.dtype)
-
-            msa_kwargs: dict[str, Tensor] | None = None
-            if self.msa_encoder is not None and msa is not None:
-                if msa.dim() == 4:
-                    batch_msa, depth, length_msa, _ = msa.shape
-                    msa_oh = msa.permute(0, 2, 1, 3).float()
-                else:
-                    batch_msa, depth, length_msa = msa.shape
-                    msa_oh = F.one_hot(
-                        msa.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES
-                    ).float()
-                msa_attn = (
-                    msa_attention_mask.permute(0, 2, 1).float()
-                    if msa_attention_mask is not None
-                    else tok_mask[:, :, None].expand(-1, -1, depth).float()
+            if precision == "fp8":
+                raise RuntimeError(
+                    "esmc_precision='fp8' is supported only by the standard "
+                    "released ESMFold2 model. The experimental binder-design "
+                    "model keeps the FastPLMs ESM++ backbone in bf16 or fp32."
                 )
-                msa_oh = msa_oh * msa_attn.unsqueeze(-1)
-                hd = (
-                    has_deletion.permute(0, 2, 1).float()
-                    if has_deletion is not None
-                    else torch.zeros(batch_msa, length_msa, depth, device=msa.device)
-                )
-                dv = (
-                    deletion_value.permute(0, 2, 1).float()
-                    if deletion_value is not None
-                    else torch.zeros(batch_msa, length_msa, depth, device=msa.device)
-                )
-                msa_kwargs = {
-                    "x_inputs": x_inputs,
-                    "msa_oh": msa_oh,
-                    "has_deletion": hd,
-                    "deletion_value": dv,
-                    "msa_attention_mask": msa_attn,
-                }
-
-            pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float()
-            z = torch.zeros_like(z_init)
-            prev_pair: Tensor | None = None
-            prev_disto_probs: Tensor | None = None
-            for loop_num in range(n_loops + 1):
-                z = z_init + self.pair_loop_proj(z)
-                if msa_kwargs is not None and self.msa_encoder is not None:
-                    z = z + self.msa_encoder(x_pair=z, **msa_kwargs).to(z.dtype)
-                z = self.folding_trunk(z, pair_attention_mask=pair_mask)
-
-                if early_exit and loop_num < n_loops:
-                    l2_converged = False
-                    if prev_pair is not None and loop_num > 0:
-                        rel_l2 = (z.float() - prev_pair.float()).norm() / prev_pair.float().norm().clamp(
-                            min=1e-8
-                        )
-                        l2_converged = rel_l2.item() < 0.25
-                    prev_pair = z.detach().clone()
-                    sym_z = z.float() + z.float().transpose(-2, -3)
-                    cur_probs = F.softmax(self.distogram_head(sym_z).float(), dim=-1)
-                    if prev_disto_probs is not None and loop_num > 0:
-                        kl_per_pair = (
-                            cur_probs
-                            * (
-                                cur_probs.clamp(min=1e-8)
-                                / prev_disto_probs.clamp(min=1e-8)
-                            ).log()
-                        ).sum(-1)
-                        kl = (kl_per_pair + kl_per_pair.transpose(-1, -2)).mean() / 2
-                        if l2_converged or kl.item() < 0.05:
-                            break
-                    prev_disto_probs = cur_probs.detach()
-
-            distogram_logits = self.distogram_head(z + z.transpose(-2, -3))
-
-        with torch.no_grad(), _seed_context(seed):
-            structure_output = self.structure_head.sample(
-                z_trunk=z.float(),
-                s_inputs=x_inputs,
-                s_trunk=None,
-                relative_position_encoding=relative_position_encoding,
-                ref_pos=ref_pos,
-                ref_charge=ref_charge,
-                ref_mask=atm_mask,
-                ref_element=ref_element_oh,
-                ref_atom_name_chars=ref_atom_name_chars_oh,
-                ref_space_uid=ref_space_uid,
-                tok_idx=atom_to_token,
-                asym_id=asym_id,
-                residue_index=residue_index,
-                entity_id=entity_id,
-                token_index=token_index,
-                sym_id=sym_id,
-                token_attention_mask=tok_mask,
-                num_diffusion_samples=n_samples,
-                num_sampling_steps=num_sampling_steps,
-                return_atom_repr=False,
-                denoising_early_exit_rmsd=(0.10 if early_exit else None),
-            )
-        sample_coords = structure_output["sample_atom_coords"]
-        assert sample_coords is not None
-
-        output: dict[str, Tensor] = {
-            "distogram_logits": distogram_logits,
-            "sample_atom_coords": sample_coords,
-        }
-        if calculate_confidence and self.confidence_head is not None:
-            confidence_output = self.confidence_head(
-                s_inputs=x_inputs.detach(),
-                z=z.detach().float(),
-                x_pred=sample_coords.detach(),
-                distogram_atom_idx=distogram_atom_idx,
-                token_attention_mask=tok_mask,
-                atom_to_token=atom_to_token,
-                atom_attention_mask=atm_mask,
-                asym_id=asym_id,
-                mol_type=mol_type,
-                num_diffusion_samples=n_samples,
-                relative_position_encoding=relative_position_encoding.detach(),
-                token_bonds_encoding=token_bonds_encoding.detach(),
-            )
-            output.update(confidence_output)
-        output["atom_pad_mask"] = (
-            atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask
-        )
-        output["residue_index"] = residue_index
-        output["entity_id"] = entity_id
-        return output
-
-    @property
-    def input_builder(self):
-        if self._esmfold2_input_builder is None:
-            from .esmfold2_processor import ESMFold2InputBuilder
-
-            self._esmfold2_input_builder = ESMFold2InputBuilder()
-        return self._esmfold2_input_builder
-
-    @property
-    def input_types(self):
-        from . import esmfold2_types
-
-        return esmfold2_types
-
-    def prepare_structure_input(self, input, seed: int | None = None):
-        return self.input_builder.prepare_input(input, seed=seed, device=self.device)
-
-    @torch.no_grad()
-    def infer_protein(self, seq: str, **forward_kwargs) -> dict[str, Tensor]:
-        from .protein_utils import prepare_protein_features
-
-        features = prepare_protein_features(seq)
-        features = {name: tensor.to(self.device) for name, tensor in features.items()}
-        output = self(**features, **forward_kwargs)
-        for name in (
-            "res_type",
-            "atom_to_token",
-            "ref_atom_name_chars",
-            "atom_attention_mask",
-            "token_attention_mask",
-            "residue_index",
-        ):
-            output[name] = features[name]
-        return output
-
-    def fold(
-        self,
-        input,
-        *,
-        num_loops: int = 3,
-        num_sampling_steps: int = 50,
-        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",
-    ):
-        return self.input_builder.fold(
-            self,
-            input,
-            num_loops=num_loops,
-            num_sampling_steps=num_sampling_steps,
-            num_diffusion_samples=num_diffusion_samples,
-            seed=seed,
-            noise_scale=noise_scale,
-            step_scale=step_scale,
-            max_inference_sigma=max_inference_sigma,
-            early_exit=early_exit,
-            complex_id=complex_id,
-        )
-
-    def fold_protein(
-        self,
-        sequence: str,
-        *,
-        chain_id: str = "A",
-        num_loops: int = 3,
-        num_sampling_steps: int = 50,
-        num_diffusion_samples: int = 1,
-        seed: int | None = None,
-        complex_id: str = "pred",
-    ):
-        from .esmfold2_types import ProteinInput, StructurePredictionInput
-
-        input = StructurePredictionInput(
-            sequences=[ProteinInput(id=chain_id, sequence=sequence)]
-        )
-        return self.fold(
-            input,
-            num_loops=num_loops,
-            num_sampling_steps=num_sampling_steps,
-            num_diffusion_samples=num_diffusion_samples,
-            seed=seed,
-            complex_id=complex_id,
-        )
-
-    @staticmethod
-    def result_to_cif(result) -> str:
-        assert not isinstance(result, list), "Pass one MolecularComplexResult at a time."
-        return result.complex.to_mmcif()
-
-    @staticmethod
-    def result_to_pdb(result) -> str:
-        assert not isinstance(result, list), "Pass one MolecularComplexResult at a time."
-        return result.complex.to_protein_complex().to_pdb_string()
-
-    def save_as_cif(self, result, output_path: str | Path) -> None:
-        Path(output_path).write_text(self.result_to_cif(result))
-
-    def save_as_pdb(self, result, output_path: str | Path) -> None:
-        Path(output_path).write_text(self.result_to_pdb(result))
-
-    def infer_protein_as_cif(self, seq: str, **forward_kwargs) -> str:
-        return self.result_to_cif(self.fold_protein(seq, **forward_kwargs))
-
-    def infer_protein_as_pdb(self, seq: str, **forward_kwargs) -> str:
-        return self.result_to_pdb(self.fold_protein(seq, **forward_kwargs))
-
-
-__all__ = [
-    "ConfidenceHead",
-    "MSAEncoder",
-    "MSAEncoderBlock",
-    "ESMFold2ExperimentalModel",
-]
+            raise ValueError(f"precision must be one of {list(dtype_map)}, got {precision!r}")
+        esmc = _load_fastplms_esmplusplus_for_esmfold2(
+            esmc_model_path=esmc_model_path,
+            attn_backend=self.config.esmc_attn_backend,
+            device=self.device,
+            dtype=dtype_map[precision],
+        )
+        assert esmc.config.hidden_size == self.config.lm_d_model, (
+            f"ESMFold2 expected lm_d_model={self.config.lm_d_model}, "
+            f"but loaded ESM++ hidden_size={esmc.config.hidden_size}."
+        )
+        assert esmc.config.num_hidden_layers == self.config.lm_num_layers, (
+            f"ESMFold2 expected lm_num_layers={self.config.lm_num_layers}, "
+            f"but loaded ESM++ num_hidden_layers={esmc.config.num_hidden_layers}."
+        )
+        for parameter in esmc.parameters():
+            parameter.requires_grad_(False)
+        self._esmc_fp8 = False
+        self._esmc = esmc
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *model_args,
+        load_esmc: bool = True,
+        **kwargs,
+    ):
+        if "config" not in kwargs:
+            kwargs["config"] = ESMFold2Config.from_pretrained(
+                pretrained_model_name_or_path, **kwargs
+            )
+        esmc_precision = kwargs.pop("esmc_precision", "bf16")
+        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        if load_esmc:
+            model.load_esmc(model.config.esmc_id, precision=esmc_precision)
+        return model
+
+    def apply_torch_compile(
+        self, mode: str = "fixed_seqlen", dynamic: bool | None = None
+    ) -> None:
+        import torch._dynamo
+
+        torch._dynamo.config.cache_size_limit = 512
+        torch._dynamo.config.accumulated_cache_size_limit = 512
+        torch._dynamo.config.capture_scalar_outputs = True
+
+        if dynamic is None:
+            dynamic = mode == "dynamic_seqlen"
+        compile_kwargs: dict[str, bool] = {"dynamic": dynamic}
+        compile_targets = (
+            PairUpdateBlock,
+            DiffusionTransformer,
+            DiffusionModule,
+            MSAEncoderBlock,
+        )
+
+        def _maybe_compile(module: nn.Module) -> None:
+            if isinstance(module, compile_targets):
+                module.forward = torch.compile(module.forward, **compile_kwargs)
+
+        self.apply(_maybe_compile)
+
+    def _compute_lm_hidden_states(
+        self,
+        input_ids: Tensor,
+        asym_id: Tensor,
+        residue_index: Tensor,
+        mol_type: Tensor,
+        tok_mask: Tensor,
+    ) -> Tensor:
+        assert self._esmc is not None
+        pad_to = 8 if self._esmc_fp8 else None
+        with _lm_precision_context(self._esmc_fp8):
+            return compute_lm_hidden_states(
+                self._esmc,
+                input_ids,
+                asym_id,
+                residue_index,
+                mol_type,
+                tok_mask,
+                pad_to_multiple=pad_to,
+            )
+
+    def forward(
+        self,
+        token_index: Tensor,
+        residue_index: Tensor,
+        asym_id: Tensor,
+        sym_id: Tensor,
+        entity_id: Tensor,
+        mol_type: Tensor,
+        res_type: Tensor,
+        token_bonds: Tensor,
+        token_attention_mask: Tensor,
+        ref_pos: Tensor,
+        ref_element: Tensor,
+        ref_charge: Tensor,
+        ref_atom_name_chars: Tensor,
+        ref_space_uid: Tensor,
+        atom_attention_mask: Tensor,
+        atom_to_token: Tensor,
+        distogram_atom_idx: Tensor,
+        deletion_mean: Tensor | None = None,
+        msa: Tensor | None = None,
+        has_deletion: Tensor | None = None,
+        deletion_value: Tensor | None = None,
+        msa_attention_mask: Tensor | None = None,
+        input_ids: Tensor | None = None,
+        lm_hidden_states: Tensor | None = None,
+        res_type_soft: Tensor | None = None,
+        num_loops: int | None = None,
+        num_diffusion_samples: int | None = None,
+        num_sampling_steps: int | None = None,
+        early_exit: bool = False,
+        seed: int | None = None,
+        calculate_confidence: bool = True,
+        provide_soft_sequence_to_msa_and_profile: bool = True,
+        noise_scale: float | None = None,
+        step_scale: float | None = None,
+        max_inference_sigma: int | None = None,
+    ) -> dict[str, Tensor]:
+        del noise_scale, step_scale, max_inference_sigma
+        tok_mask = token_attention_mask
+        atm_mask = atom_attention_mask
+        n_loops = num_loops if num_loops is not None else self.config.num_loops
+        n_samples = (
+            num_diffusion_samples
+            if num_diffusion_samples is not None
+            else self.config.num_diffusion_samples
+        )
+
+        if res_type.dim() == 2:
+            res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float()
+            res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float()
+        else:
+            res_type_oh = res_type.float()
+
+        if msa is not None:
+            msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float()
+            if msa_attention_mask is not None:
+                mask_f = msa_attention_mask.float().unsqueeze(-1)
+                msa_oh_profile = msa_oh_profile * mask_f
+                valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1)
+                profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1)
+            else:
+                profile = msa_oh_profile.mean(dim=1)
+        else:
+            profile = res_type_oh
+
+        if res_type_soft is not None:
+            res_type_oh = res_type_soft.float()
+            if (
+                not self.config.disable_msa_features
+                and provide_soft_sequence_to_msa_and_profile
+            ):
+                profile = res_type_oh
+                msa = res_type_oh.unsqueeze(1)
+                msa_attention_mask = tok_mask.unsqueeze(1)
+
+        if deletion_mean is None:
+            deletion_mean = torch.zeros(
+                res_type.shape[0], res_type.shape[1], device=res_type.device
+            )
+        if self.config.disable_msa_features:
+            profile = torch.zeros_like(profile)
+            deletion_mean = torch.zeros_like(deletion_mean)
+
+        ref_element_oh = F.one_hot(
+            ref_element.long(), num_classes=MAX_ATOMIC_NUMBER
+        ).float()
+        ref_atom_name_chars_oh = F.one_hot(
+            ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE
+        ).float()
+        atm_mask_f = atm_mask.float()
+        ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1)
+        ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(
+            -1
+        ).unsqueeze(-1)
+        atom_to_token = atom_to_token * atm_mask.long()
+
+        use_amp = ref_pos.device.type == "cuda"
+        with (
+            torch.set_grad_enabled(res_type_soft is not None),
+            torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16),
+        ):
+            x_inputs = self.inputs_embedder(
+                aatype=res_type_oh,
+                profile=profile.float(),
+                deletion_mean=deletion_mean.float(),
+                ref_pos=ref_pos,
+                atom_attention_mask=atm_mask,
+                ref_space_uid=ref_space_uid,
+                ref_charge=ref_charge,
+                ref_element=ref_element_oh,
+                ref_atom_name_chars=ref_atom_name_chars_oh,
+                atom_to_token=atom_to_token,
+            )
+
+            z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(
+                x_inputs
+            ).unsqueeze(1)
+            relative_position_encoding = self.rel_pos(
+                residue_index=residue_index,
+                asym_id=asym_id,
+                sym_id=sym_id,
+                entity_id=entity_id,
+                token_index=token_index,
+            )
+            token_bonds_encoding = self.token_bonds(token_bonds.float())
+            z_init = z_init + relative_position_encoding + token_bonds_encoding
+
+            if (
+                lm_hidden_states is None
+                and input_ids is not None
+                and self._esmc is not None
+            ):
+                lm_hidden_states = self._compute_lm_hidden_states(
+                    input_ids, asym_id, residue_index, mol_type, tok_mask
+                )
+            if lm_hidden_states is not None:
+                lm_dropout = (
+                    self.config.lm_dropout
+                    if self.config.force_lm_dropout_during_inference or self.training
+                    else 0.0
+                )
+                lm_z = self.language_model(
+                    lm_hidden_states.detach(), lm_dropout=lm_dropout
+                )
+                z_init = z_init + lm_z.to(z_init.dtype)
+
+            msa_kwargs: dict[str, Tensor] | None = None
+            if self.msa_encoder is not None and msa is not None:
+                if msa.dim() == 4:
+                    batch_msa, depth, length_msa, _ = msa.shape
+                    msa_oh = msa.permute(0, 2, 1, 3).float()
+                else:
+                    batch_msa, depth, length_msa = msa.shape
+                    msa_oh = F.one_hot(
+                        msa.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES
+                    ).float()
+                msa_attn = (
+                    msa_attention_mask.permute(0, 2, 1).float()
+                    if msa_attention_mask is not None
+                    else tok_mask[:, :, None].expand(-1, -1, depth).float()
+                )
+                msa_oh = msa_oh * msa_attn.unsqueeze(-1)
+                hd = (
+                    has_deletion.permute(0, 2, 1).float()
+                    if has_deletion is not None
+                    else torch.zeros(batch_msa, length_msa, depth, device=msa.device)
+                )
+                dv = (
+                    deletion_value.permute(0, 2, 1).float()
+                    if deletion_value is not None
+                    else torch.zeros(batch_msa, length_msa, depth, device=msa.device)
+                )
+                msa_kwargs = {
+                    "x_inputs": x_inputs,
+                    "msa_oh": msa_oh,
+                    "has_deletion": hd,
+                    "deletion_value": dv,
+                    "msa_attention_mask": msa_attn,
+                }
+
+            pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float()
+            z = torch.zeros_like(z_init)
+            prev_pair: Tensor | None = None
+            prev_disto_probs: Tensor | None = None
+            for loop_num in range(n_loops + 1):
+                z = z_init + self.pair_loop_proj(z)
+                if msa_kwargs is not None and self.msa_encoder is not None:
+                    z = z + self.msa_encoder(x_pair=z, **msa_kwargs).to(z.dtype)
+                z = self.folding_trunk(z, pair_attention_mask=pair_mask)
+
+                if early_exit and loop_num < n_loops:
+                    l2_converged = False
+                    if prev_pair is not None and loop_num > 0:
+                        rel_l2 = (z.float() - prev_pair.float()).norm() / prev_pair.float().norm().clamp(
+                            min=1e-8
+                        )
+                        l2_converged = rel_l2.item() < 0.25
+                    prev_pair = z.detach().clone()
+                    sym_z = z.float() + z.float().transpose(-2, -3)
+                    cur_probs = F.softmax(self.distogram_head(sym_z).float(), dim=-1)
+                    if prev_disto_probs is not None and loop_num > 0:
+                        kl_per_pair = (
+                            cur_probs
+                            * (
+                                cur_probs.clamp(min=1e-8)
+                                / prev_disto_probs.clamp(min=1e-8)
+                            ).log()
+                        ).sum(-1)
+                        kl = (kl_per_pair + kl_per_pair.transpose(-1, -2)).mean() / 2
+                        if l2_converged or kl.item() < 0.05:
+                            break
+                    prev_disto_probs = cur_probs.detach()
+
+            distogram_logits = self.distogram_head(z + z.transpose(-2, -3))
+
+        with torch.no_grad(), _seed_context(seed):
+            structure_output = self.structure_head.sample(
+                z_trunk=z.float(),
+                s_inputs=x_inputs,
+                s_trunk=None,
+                relative_position_encoding=relative_position_encoding,
+                ref_pos=ref_pos,
+                ref_charge=ref_charge,
+                ref_mask=atm_mask,
+                ref_element=ref_element_oh,
+                ref_atom_name_chars=ref_atom_name_chars_oh,
+                ref_space_uid=ref_space_uid,
+                tok_idx=atom_to_token,
+                asym_id=asym_id,
+                residue_index=residue_index,
+                entity_id=entity_id,
+                token_index=token_index,
+                sym_id=sym_id,
+                token_attention_mask=tok_mask,
+                num_diffusion_samples=n_samples,
+                num_sampling_steps=num_sampling_steps,
+                return_atom_repr=False,
+                denoising_early_exit_rmsd=(0.10 if early_exit else None),
+            )
+        sample_coords = structure_output["sample_atom_coords"]
+        assert sample_coords is not None
+
+        output: dict[str, Tensor] = {
+            "distogram_logits": distogram_logits,
+            "sample_atom_coords": sample_coords,
+        }
+        if calculate_confidence and self.confidence_head is not None:
+            confidence_output = self.confidence_head(
+                s_inputs=x_inputs.detach(),
+                z=z.detach().float(),
+                x_pred=sample_coords.detach(),
+                distogram_atom_idx=distogram_atom_idx,
+                token_attention_mask=tok_mask,
+                atom_to_token=atom_to_token,
+                atom_attention_mask=atm_mask,
+                asym_id=asym_id,
+                mol_type=mol_type,
+                num_diffusion_samples=n_samples,
+                relative_position_encoding=relative_position_encoding.detach(),
+                token_bonds_encoding=token_bonds_encoding.detach(),
+            )
+            output.update(confidence_output)
+        output["atom_pad_mask"] = (
+            atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask
+        )
+        output["residue_index"] = residue_index
+        output["entity_id"] = entity_id
+        return output
+
+    @property
+    def input_builder(self):
+        if self._esmfold2_input_builder is None:
+            from .esmfold2_processor import ESMFold2InputBuilder
+
+            self._esmfold2_input_builder = ESMFold2InputBuilder()
+        return self._esmfold2_input_builder
+
+    @property
+    def input_types(self):
+        from . import esmfold2_types
+
+        return esmfold2_types
+
+    def prepare_structure_input(self, input, seed: int | None = None):
+        return self.input_builder.prepare_input(input, seed=seed, device=self.device)
+
+    @torch.no_grad()
+    def infer_protein(self, seq: str, **forward_kwargs) -> dict[str, Tensor]:
+        from .protein_utils import prepare_protein_features
+
+        features = prepare_protein_features(seq)
+        features = {name: tensor.to(self.device) for name, tensor in features.items()}
+        output = self(**features, **forward_kwargs)
+        for name in (
+            "res_type",
+            "atom_to_token",
+            "ref_atom_name_chars",
+            "atom_attention_mask",
+            "token_attention_mask",
+            "residue_index",
+        ):
+            output[name] = features[name]
+        return output
+
+    def fold(
+        self,
+        input,
+        *,
+        num_loops: int = 3,
+        num_sampling_steps: int = 50,
+        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",
+    ):
+        return self.input_builder.fold(
+            self,
+            input,
+            num_loops=num_loops,
+            num_sampling_steps=num_sampling_steps,
+            num_diffusion_samples=num_diffusion_samples,
+            seed=seed,
+            noise_scale=noise_scale,
+            step_scale=step_scale,
+            max_inference_sigma=max_inference_sigma,
+            early_exit=early_exit,
+            complex_id=complex_id,
+        )
+
+    def fold_protein(
+        self,
+        sequence: str,
+        *,
+        chain_id: str = "A",
+        num_loops: int = 3,
+        num_sampling_steps: int = 50,
+        num_diffusion_samples: int = 1,
+        seed: int | None = None,
+        complex_id: str = "pred",
+    ):
+        from .esmfold2_types import ProteinInput, StructurePredictionInput
+
+        input = StructurePredictionInput(
+            sequences=[ProteinInput(id=chain_id, sequence=sequence)]
+        )
+        return self.fold(
+            input,
+            num_loops=num_loops,
+            num_sampling_steps=num_sampling_steps,
+            num_diffusion_samples=num_diffusion_samples,
+            seed=seed,
+            complex_id=complex_id,
+        )
+
+    @staticmethod
+    def result_to_cif(result) -> str:
+        assert not isinstance(result, list), "Pass one MolecularComplexResult at a time."
+        return result.complex.to_mmcif()
+
+    @staticmethod
+    def result_to_pdb(result) -> str:
+        assert not isinstance(result, list), "Pass one MolecularComplexResult at a time."
+        return result.complex.to_protein_complex().to_pdb_string()
+
+    def save_as_cif(self, result, output_path: str | Path) -> None:
+        Path(output_path).write_text(self.result_to_cif(result))
+
+    def save_as_pdb(self, result, output_path: str | Path) -> None:
+        Path(output_path).write_text(self.result_to_pdb(result))
+
+    def infer_protein_as_cif(self, seq: str, **forward_kwargs) -> str:
+        return self.result_to_cif(self.fold_protein(seq, **forward_kwargs))
+
+    def infer_protein_as_pdb(self, seq: str, **forward_kwargs) -> str:
+        return self.result_to_pdb(self.fold_protein(seq, **forward_kwargs))
+
 
+__all__ = [
+    "ConfidenceHead",
+    "MSAEncoder",
+    "MSAEncoderBlock",
+    "ESMFold2ExperimentalModel",
+]

protein_utils.py

@@ -1,488 +1,488 @@
-# coding=utf-8
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Self-contained protein featurization for ESMFold2 inference.
-
-Lets ``ESMFold2ExperimentalModel.infer_protein_as_pdb`` fold a protein sequence
-ESMFold-style without the ``esm`` companion package. The featurization
-mirrors ``ESMFold2InputBuilder.prepare_input`` for the protein-only path —
-``test_prepare_protein_features.py`` enforces tensor-exact parity.
-"""
-
-from __future__ import annotations
-
-import math
-
-import torch
-from torch import Tensor
-
-MOL_TYPE_PROTEIN = 0
-PROTEIN_UNK_RES_TYPE = 22
-MSA_GAP_TOKEN_ID = 1
-
-PROTEIN_RESIDUE_TO_RES_TYPE: dict[str, int] = {
-    "ALA": 2,
-    "ARG": 3,
-    "ASN": 4,
-    "ASP": 5,
-    "CYS": 6,
-    "GLN": 7,
-    "GLU": 8,
-    "GLY": 9,
-    "HIS": 10,
-    "ILE": 11,
-    "LEU": 12,
-    "LYS": 13,
-    "MET": 14,
-    "PHE": 15,
-    "PRO": 16,
-    "SER": 17,
-    "THR": 18,
-    "TRP": 19,
-    "TYR": 20,
-    "VAL": 21,
-}
-
-PROTEIN_1TO3: dict[str, str] = {
-    "A": "ALA",
-    "R": "ARG",
-    "N": "ASN",
-    "D": "ASP",
-    "C": "CYS",
-    "Q": "GLN",
-    "E": "GLU",
-    "G": "GLY",
-    "H": "HIS",
-    "I": "ILE",
-    "L": "LEU",
-    "K": "LYS",
-    "M": "MET",
-    "F": "PHE",
-    "P": "PRO",
-    "S": "SER",
-    "T": "THR",
-    "W": "TRP",
-    "Y": "TYR",
-    "V": "VAL",
-    "X": "UNK",
-}
-
-ESM_PROTEIN_VOCAB: dict[str, int] = {
-    "L": 4,
-    "A": 5,
-    "G": 6,
-    "V": 7,
-    "S": 8,
-    "E": 9,
-    "R": 10,
-    "T": 11,
-    "I": 12,
-    "D": 13,
-    "P": 14,
-    "K": 15,
-    "Q": 16,
-    "N": 17,
-    "F": 18,
-    "Y": 19,
-    "M": 20,
-    "H": 21,
-    "W": 22,
-    "C": 23,
-    "X": 3,
-}
-
-# Heavy atoms per canonical residue, in training-time order.
-PROTEIN_HEAVY_ATOMS: dict[str, list[str]] = {
-    "ALA": ["N", "CA", "C", "O", "CB"],
-    "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"],
-    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"],
-    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"],
-    "CYS": ["N", "CA", "C", "O", "CB", "SG"],
-    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"],
-    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"],
-    "GLY": ["N", "CA", "C", "O"],
-    "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"],
-    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"],
-    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"],
-    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"],
-    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
-    "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"],
-    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"],
-    "SER": ["N", "CA", "C", "O", "CB", "OG"],
-    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"],
-    "TRP": [
-        "N",
-        "CA",
-        "C",
-        "O",
-        "CB",
-        "CG",
-        "CD1",
-        "CD2",
-        "NE1",
-        "CE2",
-        "CE3",
-        "CZ2",
-        "CZ3",
-        "CH2",
-    ],
-    "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"],
-    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"],
-    "UNK": ["N", "CA", "C", "O"],
-}
-
-PROTEIN_REF_POS: dict[str, dict[str, tuple[float, float, float]]] = {
-    "ALA": {
-        "N": (-0.01003183238208294, -1.2073018550872803, -1.0555061101913452),
-        "CA": (-0.04190138354897499, 0.17447763681411743, -0.5729365348815918),
-        "C": (1.2127548456192017, 0.4737588167190552, 0.19521640241146088),
-        "O": (1.9390329122543335, 1.4484562873840332, -0.13759790360927582),
-        "CB": (-1.276943325996399, 0.4288230538368225, 0.29937705397605896),
-    },
-    "ARG": {
-        "N": (-2.0170421600341797, 0.6717798113822937, -1.1794233322143555),
-        "CA": (-2.0503084659576416, -0.5735036730766296, -0.4097220301628113),
-        "C": (-3.469440460205078, -1.0612813234329224, -0.2755832374095917),
-        "O": (-3.8218462467193604, -2.1369943618774414, -0.8294969797134399),
-        "CB": (-1.4193516969680786, -0.3735991418361664, 0.9852858781814575),
-        "CG": (0.11878877878189087, -0.3112654983997345, 0.963895857334137),
-        "CD": (0.6643245816230774, 1.0068185329437256, 0.3963329493999481),
-        "NE": (2.1090238094329834, 1.0977025032043457, 0.6120952367782593),
-        "CZ": (3.098905324935913, 0.3215920031070709, -0.09047172218561172),
-        "NH1": (4.461230278015137, 0.3844667971134186, 0.34141138195991516),
-        "NH2": (2.7856509685516357, -0.4166366159915924, -1.1148239374160767),
-    },
-    "ASN": {
-        "N": (-0.7595629096031189, 0.7503494620323181, 1.1369825601577759),
-        "CA": (-0.76087886095047, 0.23876343667507172, -0.23573364317417145),
-        "C": (-1.9211044311523438, -0.6982439160346985, -0.42196929454803467),
-        "O": (-2.677666187286377, -0.5753439664840698, -1.4223182201385498),
-        "CB": (0.5504899024963379, -0.5078350305557251, -0.5390339493751526),
-        "CG": (1.7250099182128906, 0.4264017939567566, -0.5778228640556335),
-        "OD1": (1.9470350742340088, 1.1086392402648926, -1.613560438156128),
-        "ND2": (2.57365345954895, 0.5730618834495544, 0.5608599781990051),
-    },
-    "ASP": {
-        "N": (-1.8452696800231934, -1.2169504165649414, 0.19437327980995178),
-        "CA": (-0.6379959583282471, -0.41974392533302307, 0.41681644320487976),
-        "C": (-0.9431572556495667, 1.0356197357177734, 0.18555717170238495),
-        "O": (-1.5183608531951904, 1.4045922756195068, -0.8739855885505676),
-        "CB": (0.48594576120376587, -0.8970447778701782, -0.5209363698959351),
-        "CG": (1.780342936515808, -0.19918935000896454, -0.2310730367898941),
-        "OD1": (2.5202910900115967, -0.6044584512710571, 0.7049641013145447),
-        "OD2": (2.1454880237579346, 0.9208861589431763, -0.9712985157966614),
-    },
-    "CYS": {
-        "N": (0.0469963513314724, 1.190075159072876, -1.1607273817062378),
-        "CA": (0.11344368755817413, -0.09400428831577301, -0.45952197909355164),
-        "C": (-1.2652032375335693, -0.6832379698753357, -0.3594406247138977),
-        "O": (-1.4631439447402954, -1.8851220607757568, -0.6826791763305664),
-        "CB": (0.6919880509376526, 0.09034398198127747, 0.952482283115387),
-        "SG": (2.4619927406311035, 0.5235707759857178, 0.9020372629165649),
-    },
-    "GLN": {
-        "N": (-2.370004653930664, -0.9637529850006104, -0.7942749261856079),
-        "CA": (-1.370002269744873, -0.6000258922576904, 0.2103111445903778),
-        "C": (-1.7545503377914429, 0.7091967463493347, 0.8433493971824646),
-        "O": (-1.8520662784576416, 0.7999289631843567, 2.0964975357055664),
-        "CB": (0.02040259726345539, -0.5004461407661438, -0.44764479994773865),
-        "CG": (1.1377512216567993, -0.28680720925331116, 0.582992434501648),
-        "CD": (2.4745187759399414, -0.24800164997577667, -0.09364881366491318),
-        "OE1": (3.1685523986816406, -1.2966246604919434, -0.1717153936624527),
-        "NE2": (2.947425603866577, 0.9601329565048218, -0.6888364553451538),
-    },
-    "GLU": {
-        "N": (-1.5850872993469238, -1.337684154510498, 0.9490851163864136),
-        "CA": (-1.0560977458953857, 0.027459044009447098, 1.0306966304779053),
-        "C": (-1.7741456031799316, 0.9664392471313477, 0.09259600937366486),
-        "O": (-1.9012441635131836, 2.181349992752075, 0.402479350566864),
-        "CB": (0.4706551432609558, 0.048803869634866714, 0.8114414811134338),
-        "CG": (0.9133604764938354, -0.4219329059123993, -0.5830985307693481),
-        "CD": (2.398822069168091, -0.3097084164619446, -0.7210537791252136),
-        "OE1": (3.1389315128326416, -1.274524450302124, -0.39029765129089355),
-        "OE2": (2.9647817611694336, 0.8781346082687378, -1.1732689142227173),
-    },
-    "GLY": {
-        "N": (-1.3942985534667969, -0.39875128865242004, -0.3370324671268463),
-        "CA": (-0.39974430203437805, 0.5488945245742798, 0.15242962539196014),
-        "C": (0.9440054893493652, -0.10314033925533295, 0.19859643280506134),
-        "O": (1.3352899551391602, -0.669218122959137, 1.2541258335113525),
-    },
-    "HIS": {
-        "N": (-1.4532867670059204, -1.0689626932144165, 0.881072461605072),
-        "CA": (-1.3396095037460327, 0.24797579646110535, 0.24960045516490936),
-        "C": (-2.675257921218872, 0.6571555733680725, -0.30441102385520935),
-        "O": (-3.1311378479003906, 1.8079776763916016, -0.06785715371370316),
-        "CB": (-0.3041955828666687, 0.21721023321151733, -0.8885309100151062),
-        "CG": (1.0887513160705566, 0.028941065073013306, -0.36419469118118286),
-        "ND1": (1.840459942817688, 1.0411773920059204, 0.29804590344429016),
-        "CD2": (1.780855417251587, -1.1011489629745483, -0.3814258575439453),
-        "CE1": (2.9566943645477295, 0.4924798905849457, 0.6477115750312805),
-        "NE2": (3.0280203819274902, -0.8751969337463379, 0.26084381341934204),
-    },
-    "ILE": {
-        "N": (-0.7167549729347229, -1.5426139831542969, -0.9983330368995667),
-        "CA": (-1.0636085271835327, -0.35169270634651184, -0.21393552422523499),
-        "C": (-1.3896740674972534, 0.8142145276069641, -1.1164065599441528),
-        "O": (-1.2377792596817017, 0.7302915453910828, -2.3656840324401855),
-        "CB": (0.061667006462812424, 0.01599610224366188, 0.8057394623756409),
-        "CG1": (1.502519965171814, -0.08899776637554169, 0.24154816567897797),
-        "CG2": (-0.053174979984760284, -0.8521055579185486, 2.0702083110809326),
-        "CD1": (1.7929610013961792, 0.899773120880127, -0.8863027691841125),
-    },
-    "LEU": {
-        "N": (1.9657520055770874, -1.9763224124908447, -0.18391533195972443),
-        "CA": (1.3077669143676758, -0.6677430868148804, -0.19492436945438385),
-        "C": (1.9905058145523071, 0.24182087182998657, 0.7879968285560608),
-        "O": (2.06896710395813, -0.07880014181137085, 2.0048046112060547),
-        "CB": (-0.20306941866874695, -0.8093230128288269, 0.11243502795696259),
-        "CG": (-0.9916267395019531, 0.5234957337379456, 0.06723011285066605),
-        "CD1": (-2.4228057861328125, 0.29949337244033813, 0.573042094707489),
-        "CD2": (-1.0282856225967407, 1.1250264644622803, -1.346014380455017),
-    },
-    "LYS": {
-        "N": (2.4221372604370117, -0.6473312377929688, 0.6370573043823242),
-        "CA": (2.0314927101135254, 0.2786507308483124, -0.4298512041568756),
-        "C": (2.7168593406677246, 1.595757246017456, -0.20924785733222961),
-        "O": (3.397681713104248, 2.116427421569824, -1.1332510709762573),
-        "CB": (0.5018402934074402, 0.4873858690261841, -0.49062973260879517),
-        "CG": (-0.25062066316604614, -0.7894009947776794, -0.9055535793304443),
-        "CD": (-1.769762635231018, -0.5552700161933899, -1.040329933166504),
-        "CE": (-2.576533555984497, -1.0221366882324219, 0.18493641912937164),
-        "NZ": (-2.269151210784912, -0.24293844401836395, 1.3849012851715088),
-    },
-    "MET": {
-        "N": (1.8903918266296387, -1.5252995491027832, -0.42638593912124634),
-        "CA": (1.2630571126937866, -0.24417810142040253, -0.7626462578773499),
-        "C": (2.30391001701355, 0.8367712497711182, -0.7254616618156433),
-        "O": (2.465414524078369, 1.5928632020950317, -1.7207728624343872),
-        "CB": (0.10567972809076309, 0.10861825942993164, 0.19741646945476532),
-        "CG": (-1.0658042430877686, -0.8736631274223328, 0.08811883628368378),
-        "SD": (-2.4557132720947266, -0.3332225978374481, 1.1461700201034546),
-        "CE": (-3.265165090560913, 0.7033554911613464, -0.11588376015424728),
-    },
-    "PHE": {
-        "N": (-2.8484435081481934, -1.525790810585022, 0.01789816841483116),
-        "CA": (-1.591969609260559, -0.8545162677764893, 0.35214468836784363),
-        "C": (-1.8900631666183472, 0.45833414793014526, 1.0232222080230713),
-        "O": (-1.3424992561340332, 0.74432373046875, 2.121629476547241),
-        "CB": (-0.760358452796936, -0.6342853307723999, -0.9257160425186157),
-        "CG": (0.604112982749939, -0.07200468331575394, -0.6148118376731873),
-        "CD1": (0.8468314409255981, 1.2480632066726685, -0.7146694660186768),
-        "CD2": (1.6827683448791504, -0.9758077263832092, -0.1423054188489914),
-        "CE1": (2.1801748275756836, 1.7875733375549316, -0.3744623064994812),
-        "CE2": (2.888307809829712, -0.48277512192726135, 0.16804970800876617),
-        "CZ": (3.149812936782837, 0.9656873941421509, 0.04440271109342575),
-    },
-    "PRO": {
-        "N": (-0.836250364780426, -0.9899801015853882, 0.5561304688453674),
-        "CA": (0.32722190022468567, -0.6164458394050598, -0.25072571635246277),
-        "C": (1.6121541261672974, -1.1711241006851196, 0.31082412600517273),
-        "O": (1.6127740144729614, -2.2771971225738525, 0.9156193733215332),
-        "CB": (0.3248198926448822, 0.9028244018554688, -0.33368146419525146),
-        "CG": (-1.1425083875656128, 1.2730128765106201, -0.2590600252151489),
-        "CD": (-1.8495968580245972, 0.026575811207294464, 0.2681289613246918),
-    },
-    "SER": {
-        "N": (0.674650251865387, 1.5018702745437622, -0.5367295145988464),
-        "CA": (0.00013792862591799349, 0.4966467022895813, 0.28510504961013794),
-        "C": (0.9941009879112244, -0.5374617576599121, 0.73505038022995),
-        "O": (1.0545241832733154, -0.8683545589447021, 1.9495396614074707),
-        "CB": (-1.1279288530349731, -0.1659376323223114, -0.5160963535308838),
-        "OG": (-1.8135979175567627, -1.085249662399292, 0.28947514295578003),
-    },
-    "THR": {
-        "N": (-1.325830340385437, -1.3728225231170654, 0.6882233023643494),
-        "CA": (-0.5433306097984314, -0.16364754736423492, 0.41697052121162415),
-        "C": (-1.294381856918335, 0.7077372074127197, -0.5549946427345276),
-        "O": (-1.6939635276794434, 0.23654410243034363, -1.6540418863296509),
-        "CB": (0.853203296661377, -0.5363803505897522, -0.14109353721141815),
-        "OG1": (1.5220820903778076, -1.379003643989563, 0.7635167837142944),
-        "CG2": (1.7225933074951172, 0.7054727077484131, -0.3651331067085266),
-    },
-    "TRP": {
-        "N": (3.686030864715576, 0.7599999904632568, 0.496155709028244),
-        "CA": (2.384092092514038, 0.09079249948263168, 0.5325262546539307),
-        "C": (2.1113572120666504, -0.6121063232421875, -0.7733646035194397),
-        "O": (1.796526312828064, -1.8323148488998413, -0.7775964140892029),
-        "CB": (1.281521201133728, 1.1139036417007446, 0.8559791445732117),
-        "CG": (-0.04292375594377518, 0.44645074009895325, 1.0942792892456055),
-        "CD1": (-0.42329534888267517, -0.15470874309539795, 2.2227554321289062),
-        "CD2": (-1.1023900508880615, 0.2158389836549759, 0.11529432237148285),
-        "NE1": (-1.7030320167541504, -0.7665823101997375, 2.0595016479492188),
-        "CE2": (-2.045644998550415, -0.4881173074245453, 0.710669219493866),
-        "CE3": (-1.2173502445220947, 0.6102271676063538, -1.300106406211853),
-        "CZ2": (-3.256009340286255, -0.9164394736289978, -0.00984987337142229),
-        "CZ3": (-2.315925121307373, 0.2306906282901764, -1.9776310920715332),
-        "CH2": (-3.3817875385284424, -0.5677337646484375, -1.3032053709030151),
-    },
-    "TYR": {
-        "N": (-1.7900604009628296, -0.8409399390220642, 1.3180142641067505),
-        "CA": (-1.913882851600647, 0.23552845418453217, 0.330669641494751),
-        "C": (-3.347280740737915, 0.3588399887084961, -0.09830684959888458),
-        "O": (-3.967811346054077, -0.6449354290962219, -0.5423302054405212),
-        "CB": (-1.0093992948532104, 0.0004731413209810853, -0.8981552124023438),
-        "CG": (0.4520410895347595, 0.021162061020731926, -0.5305932760238647),
-        "CD1": (1.0992432832717896, 1.1877919435501099, -0.3579142987728119),
-        "CD2": (1.1803174018859863, -1.253401279449463, -0.31122180819511414),
-        "CE1": (2.5253450870513916, 1.1990256309509277, 0.029804613441228867),
-        "CE2": (2.471151113510132, -1.240687608718872, 0.043534230440855026),
-        "CZ": (3.180687665939331, 0.04672492295503616, 0.2214856892824173),
-        "OH": (4.523719787597656, 0.0671030730009079, 0.5877485871315002),
-    },
-    "VAL": {
-        "N": (0.5987519025802612, -1.569443702697754, -0.7379124760627747),
-        "CA": (0.6014357209205627, -0.10503966361284256, -0.6336286664009094),
-        "C": (1.8391697406768799, 0.4067850410938263, 0.06351757049560547),
-        "O": (2.3952062129974365, -0.2666190266609192, 0.9731166958808899),
-        "CB": (-0.694736897945404, 0.4259096384048462, 0.03581475466489792),
-        "CG1": (-1.9276031255722046, 0.09515828639268875, -0.8172357082366943),
-        "CG2": (-0.8938426971435547, -0.08640842139720917, 1.472349762916565),
-    },
-    "UNK": {
-        "N": (0.0, 0.0, 0.0),
-        "CA": (0.0, 0.0, 0.0),
-        "C": (0.0, 0.0, 0.0),
-        "O": (0.0, 0.0, 0.0),
-    },
-}
-
-# Protonated nitrogens at physiological pH (matches CHARGED_ATOMS in the
-# opensource constants for the protein subset).
-PROTEIN_CHARGED_ATOMS: dict[tuple[str, str], int] = {
-    ("LYS", "NZ"): 1,
-    ("ARG", "NH2"): 1,
-    ("HIS", "ND1"): 1,
-}
-
-# Only the elements that appear in canonical protein heavy atoms.
-_PROTEIN_ELEMENT_TO_ATOMIC_NUM: dict[str, int] = {"C": 6, "N": 7, "O": 8, "S": 16}
-
-
-def _encode_atom_name(name: str) -> list[int]:
-    padded = name.ljust(4)[:4]
-    return [ord(c) - 32 if c != " " else 0 for c in padded]
-
-
-def prepare_protein_features(sequence: str) -> dict[str, Tensor]:
-    """Featurize a single protein sequence for ESMFold2ExperimentalModel.forward.
-
-    Returns the same keys with the same dtypes/shapes as
-    ``ESMFold2InputBuilder.prepare_input(StructurePredictionInput(...))``
-    restricted to a single-chain protein with no MSA, modifications,
-    distogram conditioning, or covalent bonds. All tensors have a
-    leading batch dim of 1; the caller is responsible for moving them
-    to the model device.
-    """
-    if not sequence:
-        raise ValueError("sequence must be non-empty")
-
-    res_3letter = [PROTEIN_1TO3.get(c, "UNK") for c in sequence]
-    L = len(sequence)
-
-    token_atom_starts: list[int] = []
-    atom_records: list[tuple[int, str, str, int, tuple[float, float, float]]] = []
-    res_type_vals: list[int] = []
-    input_id_vals: list[int] = []
-    distogram_rep_atom_idx: list[int] = []
-
-    atom_cursor = 0
-    for t_idx, (letter, res_3) in enumerate(zip(sequence, res_3letter)):
-        atom_names = PROTEIN_HEAVY_ATOMS[res_3]
-        res_type = PROTEIN_RESIDUE_TO_RES_TYPE.get(res_3, PROTEIN_UNK_RES_TYPE)
-        input_id = ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"])
-
-        token_atom_starts.append(atom_cursor)
-        for name in atom_names:
-            charge = PROTEIN_CHARGED_ATOMS.get((res_3, name), 0)
-            element = name[0]  # protein heavy atoms are all single-letter C/N/O/S
-            ref_pos = PROTEIN_REF_POS[res_3][name]
-            atom_records.append((t_idx, name, element, charge, ref_pos))
-            atom_cursor += 1
-
-        rep_name = "CB" if "CB" in atom_names else "CA"
-        distogram_rep_atom_idx.append(
-            token_atom_starts[t_idx] + atom_names.index(rep_name)
-        )
-
-        res_type_vals.append(res_type)
-        input_id_vals.append(input_id)
-
-    n_real_atoms = len(atom_records)
-    n_atoms = math.ceil(n_real_atoms / 32) * 32 if n_real_atoms > 0 else 32
-
-    ref_pos = torch.zeros(n_atoms, 3, dtype=torch.float32)
-    ref_element = torch.zeros(n_atoms, dtype=torch.int64)
-    ref_charge = torch.zeros(n_atoms, dtype=torch.int8)
-    ref_atom_name_chars = torch.zeros(n_atoms, 4, dtype=torch.int64)
-    ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64)
-    atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool)
-    atom_to_token = torch.zeros(n_atoms, dtype=torch.int64)
-
-    for i, (t_idx, name, element, charge, pos) in enumerate(atom_records):
-        ref_pos[i] = torch.tensor(pos, dtype=torch.float32)
-        ref_element[i] = _PROTEIN_ELEMENT_TO_ATOMIC_NUM[element]
-        ref_charge[i] = charge
-        ref_atom_name_chars[i] = torch.tensor(
-            _encode_atom_name(name), dtype=torch.int64
-        )
-        ref_space_uid[i] = t_idx
-        atom_attention_mask[i] = True
-        atom_to_token[i] = t_idx
-
-    token_index = torch.arange(L, dtype=torch.int64)
-    residue_index = torch.arange(L, dtype=torch.int64)
-    asym_id = torch.zeros(L, dtype=torch.int64)
-    sym_id = torch.zeros(L, dtype=torch.int64)
-    entity_id = torch.ones(L, dtype=torch.int64)
-    mol_type = torch.full((L,), MOL_TYPE_PROTEIN, dtype=torch.int64)
-    res_type = torch.tensor(res_type_vals, dtype=torch.int64)
-    input_ids = torch.tensor(input_id_vals, dtype=torch.int64)
-    token_bonds = torch.zeros(L, L, 1, dtype=torch.float32)
-    token_attention_mask = torch.ones(L, dtype=torch.bool)
-    distogram_atom_idx = torch.tensor(distogram_rep_atom_idx, dtype=torch.int64)
-
-    # Single-sequence MSA: depth 1, row 0 is the sequence itself.
-    msa = res_type.unsqueeze(0)
-    msa_attention_mask = torch.ones(1, L, dtype=torch.bool)
-    has_deletion = torch.zeros(1, L, dtype=torch.bool)
-    deletion_value = torch.zeros(1, L, dtype=torch.float32)
-    deletion_mean = torch.zeros(L, dtype=torch.float32)
-
-    features = {
-        "token_index": token_index,
-        "residue_index": residue_index,
-        "asym_id": asym_id,
-        "sym_id": sym_id,
-        "entity_id": entity_id,
-        "mol_type": mol_type,
-        "res_type": res_type,
-        "input_ids": input_ids,
-        "token_bonds": token_bonds,
-        "token_attention_mask": token_attention_mask,
-        "ref_pos": ref_pos,
-        "ref_element": ref_element,
-        "ref_charge": ref_charge,
-        "ref_atom_name_chars": ref_atom_name_chars,
-        "ref_space_uid": ref_space_uid,
-        "atom_attention_mask": atom_attention_mask,
-        "atom_to_token": atom_to_token,
-        "distogram_atom_idx": distogram_atom_idx,
-        "msa": msa,
-        "msa_attention_mask": msa_attention_mask,
-        "has_deletion": has_deletion,
-        "deletion_value": deletion_value,
-        "deletion_mean": deletion_mean,
-    }
-    return {k: v.unsqueeze(0) for k, v in features.items()}
+# coding=utf-8
+# Copyright 2026 Biohub. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Self-contained protein featurization for ESMFold2 inference.
+
+Lets ``ESMFold2ExperimentalModel.infer_protein_as_pdb`` fold a protein sequence
+ESMFold-style without the ``esm`` companion package. The featurization
+mirrors ``ESMFold2InputBuilder.prepare_input`` for the protein-only path —
+``test_prepare_protein_features.py`` enforces tensor-exact parity.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor
+
+MOL_TYPE_PROTEIN = 0
+PROTEIN_UNK_RES_TYPE = 22
+MSA_GAP_TOKEN_ID = 1
+
+PROTEIN_RESIDUE_TO_RES_TYPE: dict[str, int] = {
+    "ALA": 2,
+    "ARG": 3,
+    "ASN": 4,
+    "ASP": 5,
+    "CYS": 6,
+    "GLN": 7,
+    "GLU": 8,
+    "GLY": 9,
+    "HIS": 10,
+    "ILE": 11,
+    "LEU": 12,
+    "LYS": 13,
+    "MET": 14,
+    "PHE": 15,
+    "PRO": 16,
+    "SER": 17,
+    "THR": 18,
+    "TRP": 19,
+    "TYR": 20,
+    "VAL": 21,
+}
+
+PROTEIN_1TO3: dict[str, str] = {
+    "A": "ALA",
+    "R": "ARG",
+    "N": "ASN",
+    "D": "ASP",
+    "C": "CYS",
+    "Q": "GLN",
+    "E": "GLU",
+    "G": "GLY",
+    "H": "HIS",
+    "I": "ILE",
+    "L": "LEU",
+    "K": "LYS",
+    "M": "MET",
+    "F": "PHE",
+    "P": "PRO",
+    "S": "SER",
+    "T": "THR",
+    "W": "TRP",
+    "Y": "TYR",
+    "V": "VAL",
+    "X": "UNK",
+}
+
+ESM_PROTEIN_VOCAB: dict[str, int] = {
+    "L": 4,
+    "A": 5,
+    "G": 6,
+    "V": 7,
+    "S": 8,
+    "E": 9,
+    "R": 10,
+    "T": 11,
+    "I": 12,
+    "D": 13,
+    "P": 14,
+    "K": 15,
+    "Q": 16,
+    "N": 17,
+    "F": 18,
+    "Y": 19,
+    "M": 20,
+    "H": 21,
+    "W": 22,
+    "C": 23,
+    "X": 3,
+}
+
+# Heavy atoms per canonical residue, in training-time order.
+PROTEIN_HEAVY_ATOMS: dict[str, list[str]] = {
+    "ALA": ["N", "CA", "C", "O", "CB"],
+    "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"],
+    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"],
+    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"],
+    "CYS": ["N", "CA", "C", "O", "CB", "SG"],
+    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"],
+    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"],
+    "GLY": ["N", "CA", "C", "O"],
+    "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"],
+    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"],
+    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"],
+    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"],
+    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
+    "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"],
+    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"],
+    "SER": ["N", "CA", "C", "O", "CB", "OG"],
+    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"],
+    "TRP": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "NE1",
+        "CE2",
+        "CE3",
+        "CZ2",
+        "CZ3",
+        "CH2",
+    ],
+    "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"],
+    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"],
+    "UNK": ["N", "CA", "C", "O"],
+}
+
+PROTEIN_REF_POS: dict[str, dict[str, tuple[float, float, float]]] = {
+    "ALA": {
+        "N": (-0.01003183238208294, -1.2073018550872803, -1.0555061101913452),
+        "CA": (-0.04190138354897499, 0.17447763681411743, -0.5729365348815918),
+        "C": (1.2127548456192017, 0.4737588167190552, 0.19521640241146088),
+        "O": (1.9390329122543335, 1.4484562873840332, -0.13759790360927582),
+        "CB": (-1.276943325996399, 0.4288230538368225, 0.29937705397605896),
+    },
+    "ARG": {
+        "N": (-2.0170421600341797, 0.6717798113822937, -1.1794233322143555),
+        "CA": (-2.0503084659576416, -0.5735036730766296, -0.4097220301628113),
+        "C": (-3.469440460205078, -1.0612813234329224, -0.2755832374095917),
+        "O": (-3.8218462467193604, -2.1369943618774414, -0.8294969797134399),
+        "CB": (-1.4193516969680786, -0.3735991418361664, 0.9852858781814575),
+        "CG": (0.11878877878189087, -0.3112654983997345, 0.963895857334137),
+        "CD": (0.6643245816230774, 1.0068185329437256, 0.3963329493999481),
+        "NE": (2.1090238094329834, 1.0977025032043457, 0.6120952367782593),
+        "CZ": (3.098905324935913, 0.3215920031070709, -0.09047172218561172),
+        "NH1": (4.461230278015137, 0.3844667971134186, 0.34141138195991516),
+        "NH2": (2.7856509685516357, -0.4166366159915924, -1.1148239374160767),
+    },
+    "ASN": {
+        "N": (-0.7595629096031189, 0.7503494620323181, 1.1369825601577759),
+        "CA": (-0.76087886095047, 0.23876343667507172, -0.23573364317417145),
+        "C": (-1.9211044311523438, -0.6982439160346985, -0.42196929454803467),
+        "O": (-2.677666187286377, -0.5753439664840698, -1.4223182201385498),
+        "CB": (0.5504899024963379, -0.5078350305557251, -0.5390339493751526),
+        "CG": (1.7250099182128906, 0.4264017939567566, -0.5778228640556335),
+        "OD1": (1.9470350742340088, 1.1086392402648926, -1.613560438156128),
+        "ND2": (2.57365345954895, 0.5730618834495544, 0.5608599781990051),
+    },
+    "ASP": {
+        "N": (-1.8452696800231934, -1.2169504165649414, 0.19437327980995178),
+        "CA": (-0.6379959583282471, -0.41974392533302307, 0.41681644320487976),
+        "C": (-0.9431572556495667, 1.0356197357177734, 0.18555717170238495),
+        "O": (-1.5183608531951904, 1.4045922756195068, -0.8739855885505676),
+        "CB": (0.48594576120376587, -0.8970447778701782, -0.5209363698959351),
+        "CG": (1.780342936515808, -0.19918935000896454, -0.2310730367898941),
+        "OD1": (2.5202910900115967, -0.6044584512710571, 0.7049641013145447),
+        "OD2": (2.1454880237579346, 0.9208861589431763, -0.9712985157966614),
+    },
+    "CYS": {
+        "N": (0.0469963513314724, 1.190075159072876, -1.1607273817062378),
+        "CA": (0.11344368755817413, -0.09400428831577301, -0.45952197909355164),
+        "C": (-1.2652032375335693, -0.6832379698753357, -0.3594406247138977),
+        "O": (-1.4631439447402954, -1.8851220607757568, -0.6826791763305664),
+        "CB": (0.6919880509376526, 0.09034398198127747, 0.952482283115387),
+        "SG": (2.4619927406311035, 0.5235707759857178, 0.9020372629165649),
+    },
+    "GLN": {
+        "N": (-2.370004653930664, -0.9637529850006104, -0.7942749261856079),
+        "CA": (-1.370002269744873, -0.6000258922576904, 0.2103111445903778),
+        "C": (-1.7545503377914429, 0.7091967463493347, 0.8433493971824646),
+        "O": (-1.8520662784576416, 0.7999289631843567, 2.0964975357055664),
+        "CB": (0.02040259726345539, -0.5004461407661438, -0.44764479994773865),
+        "CG": (1.1377512216567993, -0.28680720925331116, 0.582992434501648),
+        "CD": (2.4745187759399414, -0.24800164997577667, -0.09364881366491318),
+        "OE1": (3.1685523986816406, -1.2966246604919434, -0.1717153936624527),
+        "NE2": (2.947425603866577, 0.9601329565048218, -0.6888364553451538),
+    },
+    "GLU": {
+        "N": (-1.5850872993469238, -1.337684154510498, 0.9490851163864136),
+        "CA": (-1.0560977458953857, 0.027459044009447098, 1.0306966304779053),
+        "C": (-1.7741456031799316, 0.9664392471313477, 0.09259600937366486),
+        "O": (-1.9012441635131836, 2.181349992752075, 0.402479350566864),
+        "CB": (0.4706551432609558, 0.048803869634866714, 0.8114414811134338),
+        "CG": (0.9133604764938354, -0.4219329059123993, -0.5830985307693481),
+        "CD": (2.398822069168091, -0.3097084164619446, -0.7210537791252136),
+        "OE1": (3.1389315128326416, -1.274524450302124, -0.39029765129089355),
+        "OE2": (2.9647817611694336, 0.8781346082687378, -1.1732689142227173),
+    },
+    "GLY": {
+        "N": (-1.3942985534667969, -0.39875128865242004, -0.3370324671268463),
+        "CA": (-0.39974430203437805, 0.5488945245742798, 0.15242962539196014),
+        "C": (0.9440054893493652, -0.10314033925533295, 0.19859643280506134),
+        "O": (1.3352899551391602, -0.669218122959137, 1.2541258335113525),
+    },
+    "HIS": {
+        "N": (-1.4532867670059204, -1.0689626932144165, 0.881072461605072),
+        "CA": (-1.3396095037460327, 0.24797579646110535, 0.24960045516490936),
+        "C": (-2.675257921218872, 0.6571555733680725, -0.30441102385520935),
+        "O": (-3.1311378479003906, 1.8079776763916016, -0.06785715371370316),
+        "CB": (-0.3041955828666687, 0.21721023321151733, -0.8885309100151062),
+        "CG": (1.0887513160705566, 0.028941065073013306, -0.36419469118118286),
+        "ND1": (1.840459942817688, 1.0411773920059204, 0.29804590344429016),
+        "CD2": (1.780855417251587, -1.1011489629745483, -0.3814258575439453),
+        "CE1": (2.9566943645477295, 0.4924798905849457, 0.6477115750312805),
+        "NE2": (3.0280203819274902, -0.8751969337463379, 0.26084381341934204),
+    },
+    "ILE": {
+        "N": (-0.7167549729347229, -1.5426139831542969, -0.9983330368995667),
+        "CA": (-1.0636085271835327, -0.35169270634651184, -0.21393552422523499),
+        "C": (-1.3896740674972534, 0.8142145276069641, -1.1164065599441528),
+        "O": (-1.2377792596817017, 0.7302915453910828, -2.3656840324401855),
+        "CB": (0.061667006462812424, 0.01599610224366188, 0.8057394623756409),
+        "CG1": (1.502519965171814, -0.08899776637554169, 0.24154816567897797),
+        "CG2": (-0.053174979984760284, -0.8521055579185486, 2.0702083110809326),
+        "CD1": (1.7929610013961792, 0.899773120880127, -0.8863027691841125),
+    },
+    "LEU": {
+        "N": (1.9657520055770874, -1.9763224124908447, -0.18391533195972443),
+        "CA": (1.3077669143676758, -0.6677430868148804, -0.19492436945438385),
+        "C": (1.9905058145523071, 0.24182087182998657, 0.7879968285560608),
+        "O": (2.06896710395813, -0.07880014181137085, 2.0048046112060547),
+        "CB": (-0.20306941866874695, -0.8093230128288269, 0.11243502795696259),
+        "CG": (-0.9916267395019531, 0.5234957337379456, 0.06723011285066605),
+        "CD1": (-2.4228057861328125, 0.29949337244033813, 0.573042094707489),
+        "CD2": (-1.0282856225967407, 1.1250264644622803, -1.346014380455017),
+    },
+    "LYS": {
+        "N": (2.4221372604370117, -0.6473312377929688, 0.6370573043823242),
+        "CA": (2.0314927101135254, 0.2786507308483124, -0.4298512041568756),
+        "C": (2.7168593406677246, 1.595757246017456, -0.20924785733222961),
+        "O": (3.397681713104248, 2.116427421569824, -1.1332510709762573),
+        "CB": (0.5018402934074402, 0.4873858690261841, -0.49062973260879517),
+        "CG": (-0.25062066316604614, -0.7894009947776794, -0.9055535793304443),
+        "CD": (-1.769762635231018, -0.5552700161933899, -1.040329933166504),
+        "CE": (-2.576533555984497, -1.0221366882324219, 0.18493641912937164),
+        "NZ": (-2.269151210784912, -0.24293844401836395, 1.3849012851715088),
+    },
+    "MET": {
+        "N": (1.8903918266296387, -1.5252995491027832, -0.42638593912124634),
+        "CA": (1.2630571126937866, -0.24417810142040253, -0.7626462578773499),
+        "C": (2.30391001701355, 0.8367712497711182, -0.7254616618156433),
+        "O": (2.465414524078369, 1.5928632020950317, -1.7207728624343872),
+        "CB": (0.10567972809076309, 0.10861825942993164, 0.19741646945476532),
+        "CG": (-1.0658042430877686, -0.8736631274223328, 0.08811883628368378),
+        "SD": (-2.4557132720947266, -0.3332225978374481, 1.1461700201034546),
+        "CE": (-3.265165090560913, 0.7033554911613464, -0.11588376015424728),
+    },
+    "PHE": {
+        "N": (-2.8484435081481934, -1.525790810585022, 0.01789816841483116),
+        "CA": (-1.591969609260559, -0.8545162677764893, 0.35214468836784363),
+        "C": (-1.8900631666183472, 0.45833414793014526, 1.0232222080230713),
+        "O": (-1.3424992561340332, 0.74432373046875, 2.121629476547241),
+        "CB": (-0.760358452796936, -0.6342853307723999, -0.9257160425186157),
+        "CG": (0.604112982749939, -0.07200468331575394, -0.6148118376731873),
+        "CD1": (0.8468314409255981, 1.2480632066726685, -0.7146694660186768),
+        "CD2": (1.6827683448791504, -0.9758077263832092, -0.1423054188489914),
+        "CE1": (2.1801748275756836, 1.7875733375549316, -0.3744623064994812),
+        "CE2": (2.888307809829712, -0.48277512192726135, 0.16804970800876617),
+        "CZ": (3.149812936782837, 0.9656873941421509, 0.04440271109342575),
+    },
+    "PRO": {
+        "N": (-0.836250364780426, -0.9899801015853882, 0.5561304688453674),
+        "CA": (0.32722190022468567, -0.6164458394050598, -0.25072571635246277),
+        "C": (1.6121541261672974, -1.1711241006851196, 0.31082412600517273),
+        "O": (1.6127740144729614, -2.2771971225738525, 0.9156193733215332),
+        "CB": (0.3248198926448822, 0.9028244018554688, -0.33368146419525146),
+        "CG": (-1.1425083875656128, 1.2730128765106201, -0.2590600252151489),
+        "CD": (-1.8495968580245972, 0.026575811207294464, 0.2681289613246918),
+    },
+    "SER": {
+        "N": (0.674650251865387, 1.5018702745437622, -0.5367295145988464),
+        "CA": (0.00013792862591799349, 0.4966467022895813, 0.28510504961013794),
+        "C": (0.9941009879112244, -0.5374617576599121, 0.73505038022995),
+        "O": (1.0545241832733154, -0.8683545589447021, 1.9495396614074707),
+        "CB": (-1.1279288530349731, -0.1659376323223114, -0.5160963535308838),
+        "OG": (-1.8135979175567627, -1.085249662399292, 0.28947514295578003),
+    },
+    "THR": {
+        "N": (-1.325830340385437, -1.3728225231170654, 0.6882233023643494),
+        "CA": (-0.5433306097984314, -0.16364754736423492, 0.41697052121162415),
+        "C": (-1.294381856918335, 0.7077372074127197, -0.5549946427345276),
+        "O": (-1.6939635276794434, 0.23654410243034363, -1.6540418863296509),
+        "CB": (0.853203296661377, -0.5363803505897522, -0.14109353721141815),
+        "OG1": (1.5220820903778076, -1.379003643989563, 0.7635167837142944),
+        "CG2": (1.7225933074951172, 0.7054727077484131, -0.3651331067085266),
+    },
+    "TRP": {
+        "N": (3.686030864715576, 0.7599999904632568, 0.496155709028244),
+        "CA": (2.384092092514038, 0.09079249948263168, 0.5325262546539307),
+        "C": (2.1113572120666504, -0.6121063232421875, -0.7733646035194397),
+        "O": (1.796526312828064, -1.8323148488998413, -0.7775964140892029),
+        "CB": (1.281521201133728, 1.1139036417007446, 0.8559791445732117),
+        "CG": (-0.04292375594377518, 0.44645074009895325, 1.0942792892456055),
+        "CD1": (-0.42329534888267517, -0.15470874309539795, 2.2227554321289062),
+        "CD2": (-1.1023900508880615, 0.2158389836549759, 0.11529432237148285),
+        "NE1": (-1.7030320167541504, -0.7665823101997375, 2.0595016479492188),
+        "CE2": (-2.045644998550415, -0.4881173074245453, 0.710669219493866),
+        "CE3": (-1.2173502445220947, 0.6102271676063538, -1.300106406211853),
+        "CZ2": (-3.256009340286255, -0.9164394736289978, -0.00984987337142229),
+        "CZ3": (-2.315925121307373, 0.2306906282901764, -1.9776310920715332),
+        "CH2": (-3.3817875385284424, -0.5677337646484375, -1.3032053709030151),
+    },
+    "TYR": {
+        "N": (-1.7900604009628296, -0.8409399390220642, 1.3180142641067505),
+        "CA": (-1.913882851600647, 0.23552845418453217, 0.330669641494751),
+        "C": (-3.347280740737915, 0.3588399887084961, -0.09830684959888458),
+        "O": (-3.967811346054077, -0.6449354290962219, -0.5423302054405212),
+        "CB": (-1.0093992948532104, 0.0004731413209810853, -0.8981552124023438),
+        "CG": (0.4520410895347595, 0.021162061020731926, -0.5305932760238647),
+        "CD1": (1.0992432832717896, 1.1877919435501099, -0.3579142987728119),
+        "CD2": (1.1803174018859863, -1.253401279449463, -0.31122180819511414),
+        "CE1": (2.5253450870513916, 1.1990256309509277, 0.029804613441228867),
+        "CE2": (2.471151113510132, -1.240687608718872, 0.043534230440855026),
+        "CZ": (3.180687665939331, 0.04672492295503616, 0.2214856892824173),
+        "OH": (4.523719787597656, 0.0671030730009079, 0.5877485871315002),
+    },
+    "VAL": {
+        "N": (0.5987519025802612, -1.569443702697754, -0.7379124760627747),
+        "CA": (0.6014357209205627, -0.10503966361284256, -0.6336286664009094),
+        "C": (1.8391697406768799, 0.4067850410938263, 0.06351757049560547),
+        "O": (2.3952062129974365, -0.2666190266609192, 0.9731166958808899),
+        "CB": (-0.694736897945404, 0.4259096384048462, 0.03581475466489792),
+        "CG1": (-1.9276031255722046, 0.09515828639268875, -0.8172357082366943),
+        "CG2": (-0.8938426971435547, -0.08640842139720917, 1.472349762916565),
+    },
+    "UNK": {
+        "N": (0.0, 0.0, 0.0),
+        "CA": (0.0, 0.0, 0.0),
+        "C": (0.0, 0.0, 0.0),
+        "O": (0.0, 0.0, 0.0),
+    },
+}
+
+# Protonated nitrogens at physiological pH (matches CHARGED_ATOMS in the
+# opensource constants for the protein subset).
+PROTEIN_CHARGED_ATOMS: dict[tuple[str, str], int] = {
+    ("LYS", "NZ"): 1,
+    ("ARG", "NH2"): 1,
+    ("HIS", "ND1"): 1,
+}
+
+# Only the elements that appear in canonical protein heavy atoms.
+_PROTEIN_ELEMENT_TO_ATOMIC_NUM: dict[str, int] = {"C": 6, "N": 7, "O": 8, "S": 16}
+
+
+def _encode_atom_name(name: str) -> list[int]:
+    padded = name.ljust(4)[:4]
+    return [ord(c) - 32 if c != " " else 0 for c in padded]
+
+
+def prepare_protein_features(sequence: str) -> dict[str, Tensor]:
+    """Featurize a single protein sequence for ESMFold2ExperimentalModel.forward.
+
+    Returns the same keys with the same dtypes/shapes as
+    ``ESMFold2InputBuilder.prepare_input(StructurePredictionInput(...))``
+    restricted to a single-chain protein with no MSA, modifications,
+    distogram conditioning, or covalent bonds. All tensors have a
+    leading batch dim of 1; the caller is responsible for moving them
+    to the model device.
+    """
+    if not sequence:
+        raise ValueError("sequence must be non-empty")
+
+    res_3letter = [PROTEIN_1TO3.get(c, "UNK") for c in sequence]
+    L = len(sequence)
+
+    token_atom_starts: list[int] = []
+    atom_records: list[tuple[int, str, str, int, tuple[float, float, float]]] = []
+    res_type_vals: list[int] = []
+    input_id_vals: list[int] = []
+    distogram_rep_atom_idx: list[int] = []
+
+    atom_cursor = 0
+    for t_idx, (letter, res_3) in enumerate(zip(sequence, res_3letter)):
+        atom_names = PROTEIN_HEAVY_ATOMS[res_3]
+        res_type = PROTEIN_RESIDUE_TO_RES_TYPE.get(res_3, PROTEIN_UNK_RES_TYPE)
+        input_id = ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"])
+
+        token_atom_starts.append(atom_cursor)
+        for name in atom_names:
+            charge = PROTEIN_CHARGED_ATOMS.get((res_3, name), 0)
+            element = name[0]  # protein heavy atoms are all single-letter C/N/O/S
+            ref_pos = PROTEIN_REF_POS[res_3][name]
+            atom_records.append((t_idx, name, element, charge, ref_pos))
+            atom_cursor += 1
+
+        rep_name = "CB" if "CB" in atom_names else "CA"
+        distogram_rep_atom_idx.append(
+            token_atom_starts[t_idx] + atom_names.index(rep_name)
+        )
+
+        res_type_vals.append(res_type)
+        input_id_vals.append(input_id)
+
+    n_real_atoms = len(atom_records)
+    n_atoms = math.ceil(n_real_atoms / 32) * 32 if n_real_atoms > 0 else 32
+
+    ref_pos = torch.zeros(n_atoms, 3, dtype=torch.float32)
+    ref_element = torch.zeros(n_atoms, dtype=torch.int64)
+    ref_charge = torch.zeros(n_atoms, dtype=torch.int8)
+    ref_atom_name_chars = torch.zeros(n_atoms, 4, dtype=torch.int64)
+    ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64)
+    atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool)
+    atom_to_token = torch.zeros(n_atoms, dtype=torch.int64)
+
+    for i, (t_idx, name, element, charge, pos) in enumerate(atom_records):
+        ref_pos[i] = torch.tensor(pos, dtype=torch.float32)
+        ref_element[i] = _PROTEIN_ELEMENT_TO_ATOMIC_NUM[element]
+        ref_charge[i] = charge
+        ref_atom_name_chars[i] = torch.tensor(
+            _encode_atom_name(name), dtype=torch.int64
+        )
+        ref_space_uid[i] = t_idx
+        atom_attention_mask[i] = True
+        atom_to_token[i] = t_idx
+
+    token_index = torch.arange(L, dtype=torch.int64)
+    residue_index = torch.arange(L, dtype=torch.int64)
+    asym_id = torch.zeros(L, dtype=torch.int64)
+    sym_id = torch.zeros(L, dtype=torch.int64)
+    entity_id = torch.ones(L, dtype=torch.int64)
+    mol_type = torch.full((L,), MOL_TYPE_PROTEIN, dtype=torch.int64)
+    res_type = torch.tensor(res_type_vals, dtype=torch.int64)
+    input_ids = torch.tensor(input_id_vals, dtype=torch.int64)
+    token_bonds = torch.zeros(L, L, 1, dtype=torch.float32)
+    token_attention_mask = torch.ones(L, dtype=torch.bool)
+    distogram_atom_idx = torch.tensor(distogram_rep_atom_idx, dtype=torch.int64)
+
+    # Single-sequence MSA: depth 1, row 0 is the sequence itself.
+    msa = res_type.unsqueeze(0)
+    msa_attention_mask = torch.ones(1, L, dtype=torch.bool)
+    has_deletion = torch.zeros(1, L, dtype=torch.bool)
+    deletion_value = torch.zeros(1, L, dtype=torch.float32)
+    deletion_mean = torch.zeros(L, dtype=torch.float32)
+
+    features = {
+        "token_index": token_index,
+        "residue_index": residue_index,
+        "asym_id": asym_id,
+        "sym_id": sym_id,
+        "entity_id": entity_id,
+        "mol_type": mol_type,
+        "res_type": res_type,
+        "input_ids": input_ids,
+        "token_bonds": token_bonds,
+        "token_attention_mask": token_attention_mask,
+        "ref_pos": ref_pos,
+        "ref_element": ref_element,
+        "ref_charge": ref_charge,
+        "ref_atom_name_chars": ref_atom_name_chars,
+        "ref_space_uid": ref_space_uid,
+        "atom_attention_mask": atom_attention_mask,
+        "atom_to_token": atom_to_token,
+        "distogram_atom_idx": distogram_atom_idx,
+        "msa": msa,
+        "msa_attention_mask": msa_attention_mask,
+        "has_deletion": has_deletion,
+        "deletion_value": deletion_value,
+        "deletion_mean": deletion_mean,
+    }
+    return {k: v.unsqueeze(0) for k, v in features.items()}