esmfold2-atomworks

ESMFold2 on AtomWorks structures, with native ESMFold2 feature semantics preserved exactly.

Not a rewrite. The published model is held intact, fed from and returned to AtomWorks structures — then checked to be the same model: the same features, and the same outputs under deterministic kernels. An optional Foundry integration adds training, configuration and checkpointing on Foundry's trainer, without changing the native model.

AtomWorks AtomArray ──adapter──> ESMFold2 (unmodified) ──adapter──> AtomWorks AtomArray
                         │
                         └── feature parity, exact, on CPU, no weights

Status

part	state
AtomWorks ↔ ESMFold2 — adapter and its reverse, AtomWorks pipeline, supervision labels	done; feature parity exact on 5 fixtures
Inference — model wrapper, engine, CLI, a provenance record per fold	done; output parity exact on a GPU under deterministic kernels
Foundry integration (optional) — trainer, Hydra configs	contracts verified against the pinned Foundry checkout, the runtime ones also against rc-foundry 0.2.0; training steps run on a GPU (docs/06); the objective is the caller's (docs/05)

The core milestone is met: for monomer, multimer, metal, cofactor and modified-residue systems, the input the adapter derives from a structure featurizes to the same 29 tensors as an independently frozen reference input. Featurization is a pure function at a fixed seed, so that equality is exact. (Only a SMILES ligand makes the seed matter: its conformer is embedded at call time.)

The model's structure head is a diffusion sampler, and on a GPU with torch's deterministic algorithms off a seeded fold is not repeatable, so identical features alone say that the two paths condition the model identically. Under deterministic kernels they say more, and it is confirmed on a GPU: the two paths return identical structures, every coordinate and confidence value equal (docs/02).

Quickstart

pip install -e .                 # esm and atomworks resolve like any dependency
pip install -e ".[foundry]"      # optional: training through Foundry (Python 3.12)

esmfold2-atomworks doctor        # imports, weights; parity given a clone or AtomWorks' test data

esmfold2-atomworks parity structures/*.cif        # survey a corpus
esmfold2-atomworks fold input.cif --out-dir runs/ # intended for a GPU

from esmfold2_atomworks.data.loading import parse_structure
from esmfold2_atomworks.model.esmfold2 import AtomWorksESMFold2

atoms, chain_info = parse_structure("2hhb.cif.gz")
model = AtomWorksESMFold2()
structure, result = model.fold_atom_array(atoms, chain_info=chain_info)

What it refuses to do

The adapter fails loudly on the cases that otherwise produce a confident prediction of the wrong molecule:

• a ligand labelled LIG/UNL/UNK/UNX is refused, because those are real CCD codes and what a model writes on anything it was handed as SMILES;
• a declared expected_formula that disagrees with the heavy atoms present is an error;
• a sequence is taken from the full entity record, not from the residues that happen to be modelled, so unresolved loops are not silently deleted;
• a non-standard residue is declared by CCD code, not folded as its parent;
• a chain kind, sequence override, MSA or ligand declaration that does not bind to exactly the chain it names — a typo, the wrong kind of chain, an alignment built for another sequence — is an error, not a no-op;
• a chain that holds more than one molecule — a protein, its ligands and its waters under one author chain id, or the copies of a built assembly under one chain_id — is an error, whether its annotations say so or a declared kind is contradicted by the CCD;
• a fold with no ESMC backbone attached and no LM hidden states supplied is refused: the native model would fold anyway, without the LM prior the checkpoint was trained with;
• a chain that cannot be expressed, a covalent bond that cannot be placed, a chain kind that would have to be guessed, or a modification whose position is unknown each raises by default — accepted only by name, on every path, and recorded per output when accepted.

Each of these has a decision ID in docs/00_SCOPE.md and a test; the strictness policy is laid out in docs/01.

Layout

src/esmfold2_atomworks/
  data/                 the adapter, the reverse adapter, the AtomWorks pipeline, labels
  model/                AtomWorksESMFold2 — holds the native module
  inference/            engine: one resident model, many structures
  parity/               feature and output comparison
  metrics.py  paths.py  doctor.py  cli.py
  training/             optional Foundry integration: FabricTrainer subclass
configs/                optional Foundry integration: Hydra configs
tests/                  parity and adapter tests; CPU-only by default
docs/                   the design record — read docs/README.md first

Nothing outside training/ and configs/ imports Foundry, and tests/test_core_boundary.py holds it to that. The optional integration is described in docs/06.

Environments

pip install -e . is all the package needs: esm and atomworks are ordinary dependencies, resolved from PyPI (Python ≥ 3.12; the foundry extra needs 3.12, as rc-foundry does). The structure-based tests, feature parity among them, read test structures, four of which are in tests/data/structures; the others come from an atomworks checkout rather than with the package. The suite passes against the PyPI releases of the upstreams, and the tests that need the checkout skip without it. See docs/04.

The suite and the GPU parity checks also run in a pinned reference environment: Python 3.14, torch 2.14, atomworks as a package, and esm and foundry as source trees at the revisions in reproducibility/UPSTREAM.lock. reproducibility/ defines it (uv sync --project reproducibility, then source reproducibility/env.sh). It is not needed for ordinary use.

License

MIT — see LICENSE. This project wraps, but does not vendor, ESM (MIT) and AtomWorks / Foundry (BSD-3-Clause), each of which carries its own license and must be obtained separately.