Metadata-Version: 2.4
Name: jaxfvld
Version: 0.1.0
Summary: A small, fast, fully differentiable antibody Fv structure predictor in JAX/Flax.
Author: Fabricagen
License: BSD-3-Clause
Project-URL: Homepage, https://github.com/fabricagen/jaxfvld
Project-URL: Repository, https://github.com/fabricagen/jaxfvld
Keywords: jax,flax,antibody,protein-structure,fv,cdr,vhh,differentiable
Classifier: License :: OSI Approved :: BSD License
Classifier: Programming Language :: Python :: 3
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: jax>=0.4.20
Requires-Dist: jaxlib>=0.4.20
Requires-Dist: flax>=0.7.5
Requires-Dist: numpy>=1.24
Provides-Extra: ablanx
Requires-Dist: ablanx; extra == "ablanx"
Provides-Extra: igfold
Requires-Dist: igfold; extra == "igfold"
Requires-Dist: torch; extra == "igfold"
Provides-Extra: data
Requires-Dist: biopython; extra == "data"
Requires-Dist: requests; extra == "data"
Provides-Extra: train
Requires-Dist: optax; extra == "train"
Requires-Dist: ablanx; extra == "train"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: matplotlib; extra == "dev"
Requires-Dist: optax; extra == "dev"
Dynamic: license-file

# jaxfvld

jaxfvld is a small, fast, fully differentiable antibody Fv structure predictor in JAX/Flax. It folds an
antibody variable region (paired VH and VL, or a single VHH domain) from sequence to an all-atom backbone,
returning per-residue frames, side-chain torsions, atom14 coordinates, and a per-residue pLDDT. It takes no
antigen and no input structure.

It is ~0.76 M parameters, BSD-3-Clause, and matches IgFold on CDR-H3 (see [Accuracy](#accuracy)). Provenance
and citation are in [ATTRIBUTION.md](ATTRIBUTION.md) and [CITATION.cff](CITATION.cff).

![A gallery of jaxfvld antibody Fv and VHH folds](figures/gallery.png)

## Why use it

- Fully differentiable: the sequence enters as a soft PSSM, so gradients flow from any function of the
  predicted structure back to the sequence, for gradient-based design.
- Small and fast: ~0.76 M parameters; a single seed folds a ~230-residue Fv in about 0.8 s on CPU (the
  4-seed ensemble in about 2 s), no GPU required.
- Accurate on the hard loop: ~2.4 Å median CDR-H3 backbone RMSD on a leak-free, post-2021 test set.
- Permissive: BSD-3-Clause.
- JAX-native: `jit` / `vmap` / `grad`-compatible and composable inside a larger model.

## Quickstart

Install:

    pip install jaxfvld
    pip install "jaxfvld[ablanx]"    # + the AbLang2 feature path the released checkpoint uses
    # from source: pip install git+https://github.com/fabricagen/jaxfvld

Fold:

    from jaxfvld import fold_fv

    out = fold_fv(
        vh="QVQLVESGGGLVQPGGSLRLSCAAS...",
        vl="DIQMTQSPSSLSASVGDRVTITCRAS...",   # omit for a single-domain VHH
        weights="jaxfvld_weights.npz",
        ablanx_weights="ablang2_weights.npz",
    )
    out["atom14"]   # [N, 14, 3] all-atom coordinates
    out["plddt"]    # [N] per-residue confidence

Command line:

    jaxfvld fold --vh QVQL... --vl DIQM... --out fv.pdb \
        --weights jaxfvld_weights.npz --ablanx-weights ablang2_weights.npz
    # --single folds the best single seed instead of the 4-seed ensemble (faster, ~0.25 Å worse on CDR-H3)

The released checkpoint was trained with the AbLang2 feature path and depends on it: pass the ablang2 weights
(from the [ablanx](https://github.com/fabricagen/ablanx) release) to fold at the reported accuracy. Without
them CDR-H3 collapses (median ~18 Å vs ~2.5 Å), so the path is effectively required. Weights from `weights=`
/ `--weights` / `JAXFVLD_WEIGHTS`; ablanx weights from `ablanx_weights=` / `--ablanx-weights` /
`JAXFVLD_ABLANX_WEIGHTS` (or `ABLANG_WEIGHTS`).

## What it is

The forward pass is a featurizer, a structure module, and an all-atom decoder:

1. The soft PSSM is projected to a single representation with a per-chain embedding (VH = chain 0, VL =
   chain 1); the pair representation encodes relative position and a cross-chain indicator, mixed by a
   triangle-multiplicative update. An optional AbLang2 sequence-embedding + attention path feeds both; the
   released checkpoint uses it.
2. Invariant Point Attention iterates from a black-hole initialization, updating the single representation
   and a per-residue rigid frame. The block is weight-tied (`n_layer` = refinement iterations), and the
   trunk recycles `n_recycle` times with a stop-gradient.
3. The final frame and seven predicted torsions become atom14 coordinates via the vendored AlphaFold
   geometry tables; a head predicts pLDDT.

Training minimizes Frame-Aligned Point Error, so the loss is invariant to a global rigid transform and
independent of centering. The released architecture is `Fold(c_s=128, c_z=128, n_layer=4, n_recycle=3,
tri_pair=True)`, `max_len=280`. The weights archive carries this config, so the loader reconstructs the
exact model and the weight-tied `n_layer`/`n_recycle` trap cannot bite. See
[TECHNICAL_BRIEF.md](TECHNICAL_BRIEF.md) for the architecture diagram, the AbLang2 path, and the geometry
provenance.

## Accuracy

The shipped model is a 4-seed ensemble at the converged step. CDR-H3 backbone RMSD is the headline metric
(the hard loop that decides antibody structure quality), measured framework-aligned the standard
IgFold / ABodyBuilder2 way. The test set is leak-free: no test antibody shares a CDR-H3 cluster with
training, restricted to post-2021 depositions. Both predictors are raw (no OpenMM relaxation).

Per-region median backbone RMSD, set300 (n = 300):

| region     | jaxfvld (4-seed) | IgFold (raw) |
|------------|------------------|--------------|
| framework  | 0.62             | **0.56**     |
| CDR-H1     | 1.00             | **0.98**     |
| CDR-H2     | 0.77             | **0.73**     |
| **CDR-H3** | **2.39**         | 2.48         |

![jaxfvld CDR-H3 backbone accuracy on set300, IgFold as the baseline](figures/parity.png)

Paired CDR-H3 on the 153 PDBs both predictors share, bootstrapping the per-structure paired delta directly
(do not read small differences off overlapping CIs): median delta = **−0.067 Å**, 95% CI **[−0.16, +0.08]**
(crosses zero), win-rate **54%**. Both predictors are scored on the identical 153 PDBs, nothing dropped on
either side.

On CDR-H3 the paired delta straddles zero: a statistical tie with the IgFold baseline, not a beat, from a
0.76 M-parameter model; the baseline keeps a small edge on the framework. See [Validation](#validation) for
the reproduction, and reproduce the tables and figure from the committed fixtures:

    python -m jaxfvld.eval.parity_report          # prints the tables, writes figures/parity.png

## Speed and size

The trained folder is small and folds a variable region on CPU, no GPU. On one VH+VL Fv, one CPU, no
refinement:

| | jaxfvld | IgFold |
|------------------------------|--------------|----------|
| trained folder, per seed / model | **0.76 M**   | 1.56 M   |
| ensemble (4 members)             | 3.04 M       | 6.23 M   |
| fold time, 4-member ensemble     | **~2.1 s**   | ~7.1 s   |
| CDR-H3 median, set300            | **2.39 Å**   | 2.48 Å   |

The 4-seed ensemble folds about 3x faster than IgFold's 4-model ensemble on the same CPU, from a trained
folder about half the per-model size, at a statistical tie on CDR-H3; a single seed folds in about 0.8 s.
Both models pair the folder with a frozen antibody language model as a sequence prior (jaxfvld: AbLang2,
44 M; IgFold: AntiBERTy, ~26 M), which is shared rather than trained here and is not included in the
trained-folder rows above.

Reproduce the folder size and fold time (add IgFold with the `[igfold]` extra for the head-to-head):

    JAXFVLD_WEIGHTS=jaxfvld_weights.npz ABLANG_WEIGHTS=ablang2_weights.npz \
        python -m jaxfvld.eval.speed_size --igfold

## Validation

Folding the leak-free set300 through the released package reproduces the committed benchmark: the aggregate
per-region medians match to within 0.01 Å, and every per-PDB CDR-H3 matches its recorded value to a mean of
0.013 Å. Predicted confidence tracks error (pLDDT calibration). See [VALIDATION.md](VALIDATION.md) for the
receipt, the figures, and the exact reproduction commands.

## Differentiable sequence design

Because the sequence enters as a soft PSSM, gradients flow from any function of the predicted structure back
to the sequence logits. `examples/` has two worked scripts: design a soft sequence so the folded Fv matches
a target CDR-H3 geometry, and maximize predicted foldability. This is the main reason to use jaxfvld over a
non-differentiable folder.

## Weights

Trained weights are attached to the GitHub release as a single archive, `jaxfvld_weights.npz` (4 seeds + the
best single seed), which carries its own reconstruction config. The AbLang2 weights (`ablang2_weights.npz`)
come from the [ablanx](https://github.com/fabricagen/ablanx) release. `export_ensemble.py` builds the archive
from per-seed checkpoints.

## Tests

    pip install "jaxfvld[dev]"
    python -m pytest -q -k "not weights_reproduce"     # IPA geometry, inference, parity fixtures (no weights)
    # the frame/FAPE/fold/padding gates need parsed data and self-skip without it; run them where the data lives:
    JAXFVLD_TEST_DATA=<allatom shards> SABDAB_DIR=<sabdab2 dir> \
        python -m pytest -q tests/test_frames.py tests/test_fape.py tests/test_fold.py tests/test_padding_invariance.py
    # the weights-gated reproduction folds set300 through the package; run it on a GPU box, not a laptop:
    JAXFVLD_WEIGHTS=... ABLANG_WEIGHTS=... JAXFVLD_EVAL_DATA=... python -m pytest -q tests/test_parity.py

## Evaluate, train, and ingest data

`jaxfvld.eval` scores any predictor on the same metric (`cdr_rmsd`, `igfold_headtohead`, `plddt_calib`);
`jaxfvld.train.train_folder` trains the folder (`pip install "jaxfvld[train]"`); `jaxfvld.data` is the
self-contained SAbDab / SAbDab2 ingestion path. Data comes from SAbDab (Oxford OPIG), used under its own
terms and not redistributed. Locations come from environment variables with local defaults.

## Ecosystem

jaxfvld is the structure folder in the **seam** bundle, which couples it with the [ablanx](
https://github.com/fabricagen/ablanx) sequence prior. For antibody developability screening and repair, see
**sift**: https://sift.fabricagen.ai (coming soon, ca 08/26).

## License

BSD-3-Clause. See [LICENSE](LICENSE) and [ATTRIBUTION.md](ATTRIBUTION.md).
`jaxfvld/model/residue_constants.py` is vendored from AlphaFold (Copyright 2021 DeepMind, Apache-2.0) with
its original header retained; those terms are compatible with BSD-3-Clause redistribution.
