foldsteer

Implementes basic chemical steering for OpenFold3-p2. Reimplements methods previously developed in Boltz-1x and Protenix-v2:

• physical guidance — gradient descent on a flat-bottom chemical energy, applied to the denoiser's x0 prediction at each diffusion step;
• Feynman-Kac steering — sample a particle population and resample it toward low-energy trajectories.

The package is set up in a model agnostic way. Adaptors can be written for new models.

Benchmark on OF3p2

OpenFold3-p2 (155k) on the 100 smallest Runs N' Poses post-2025 targets, 1 seed × 5 samples, guidance only (20 GD steps), scored with PXMeter. Rates are over the 103 ligand chains, each represented by its top-1 sample ranked by chain_pair_iptm. Ligand success = RMSD < 2 Å and lDDT-PLI > 0.8.

	baseline	steered
PoseBusters valid	66.0%	89.3%
ligand success	49.5%	50.5%
PB valid & ligand success	38.8%	47.6%
wall clock, 4× GH200	9m56s	17m11s (1.73×)

Steering fixes chemistry without moving placement: paired ligand by ligand it fixed 25 validity failures and broke 1 (McNemar p < 1e-4), while ligand success moved by a net +1 (p = 1.0). Nearly all of the gain is sterics — minimum_distance_to_protein failures fall 25 → 0; chirality is second (9 → 4).

Caveat: 10 of 161 multi-atom ligand chains were skipped because RDKit rejected the molecule rebuilt from atom_array (valence errors on quaternary nitrogen and boron, which need a formal charge the rebuild does not assign), leaving 7 of the 100 targets unsteered and still counted in the steered arm.

TODOs

• Test out the FK steering to see how much better / slower that is
• See how valency errors can be overcome (possibly some change to featureization is needed?)

Usage

Install

pip install -e .                 # engine only, no folding model needed
pip install -e '.[openfold3]'    # with the OpenFold3 adapter

Already have OpenFold3 installed? Install foldsteer into that same environment with pip install -e . --no-deps — it must share the interpreter with OF3 (see below), and --no-deps keeps pip from touching your pinned torch. Nothing about the OF3 install changes; foldsteer imports it lazily and patches at runtime. You need OF3 >= 0.4.5 (or any main carrying SampleDiffusion._sample_rollout; at tag 0.4.4 that loop is still inlined in forward and the adapter will raise AttributeError).

Use with OpenFold3

From the command line, examples/run_of3_steered.py does the patches OF3 inference to use the steering. It passes everything after -- straight to run_openfold predict, so the steered and unsteered arms take identical OF3 arguments:

python examples/run_of3_steered.py --steer --num-gd-steps 20 \
    --stats-json steering_stats.json -- \
    --query_json examples/query_protein_ligand.json \
    --inference_ckpt_path /path/to/of3-p2-155k.pt \
    --num_diffusion_samples 5 \
    --use_msa_server false --use_templates false \
    --output_dir out_steered

Swap --steer for --no-steer to get the baseline. Always check the [foldsteer] targets=... ligand_chains=... guided_steps=... line it prints: a run that found no steerable ligand produces output indistinguishable from an unsteered one.

Guidance only is the default: it is the cheaper half and carries the benefit measured below. Feynman-Kac steering has not been tested end-to-end against a real model — it is exercised only by unit tests against a mock sampler. It also reuses OF3's rollout-sample axis as the particle axis, so num_diffusion_samples must be a multiple of num_particles and fewer structures come back than were sampled.

A complete runnable example — query JSON, both arms, and how to confirm steering actually fired — is in examples/.

Use standalone

from foldsteer import SteeringEngine, default_config
from foldsteer.adapters.rdkit_source import context_from_mols

ctx = context_from_mols([mol], [{i: i for i in range(mol.GetNumAtoms())}],
                        n_atoms=coords.shape[-2])
engine = SteeringEngine(default_config(), ctx)
coords = coords + engine.guide(coords, t=0.5)   # t: 1 = noisy, 0 = clean

Potentials

BoundsMatrixPotential, VDWOverlapPotential, ChiralAtomPotential, StereoBondPotential, PlanarBondPotential, Sp2CenterPotential, Sp3CenterPotential, ConjugatedTorsionPotential, ConnectionsPotential.

All flat-bottom: zero energy and zero gradient on valid geometry, so steering is inert on structures the model already got right.

Tests

pytest tests/ -q      # 31 tests, no GPU or model weights required

Every analytic gradient is checked against autograd (max error ~1e-15). See DESIGN.md for the full analysis, the OpenFold3 integration points, and validation results.

Status

Prototype. The engine, potentials, and RDKit extraction are tested. The OpenFold3 AtomArray reconstruction path has now been exercised against real OF3 inference input — which turned up three silent failures in it (integer MoleculeType ids, Kekulé order on aromatic bonds, stereo perceived from ref_pos rather than the not-yet-predicted coord), all fixed and pinned by tests/test_of3_extraction.py. Guidance is now benchmarked against real OF3 inference (see Benchmark above); Feynman-Kac steering is not, and the formal charge gap in the AtomArray rebuild is the clearest outstanding fix.

Process

Using Claude Science, I wrote a delibrated about what the API should look like. Then Claude code implemented the plan.