Chemical steering: gradient guidance toward valid ligand chemistry - #395
Open
jnwei wants to merge 11 commits into
Open
Chemical steering: gradient guidance toward valid ligand chemistry#395jnwei wants to merge 11 commits into
jnwei wants to merge 11 commits into
Conversation
Chemical steering nudges the diffusion sampler toward valid ligand chemistry. This is the package skeleton: the torch-only `__init__` that the model stack imports, a README covering the architecture and the derived/original split, and THIRD_PARTY_NOTICES.md. The package merges two prior contributions, and the notices file records which files derive from what, file by file. The engine, registry and schedule design come from foldsteer (Etowah Adams, Apache 2.0), itself deriving from Boltz and Protenix; the RDKit constraint extraction and the acceptance test suite come from OpenFold3 PR #385 (Peter Obi). Boltz (MIT) supplies the flat-bottom formulation and the default parameters, Protenix the registry structure. `__init__` deliberately exports only the sampling-loop side. Featurization pulls in RDKit, biotite and the data pipeline, so it is imported directly by the data path rather than re-exported here -- otherwise every importer of this package, the model stack included, would pay for them. pyproject gains openfold3/steering/tests to testpaths, so the package's own tests run with the default suite. Co-authored-by: Etowah Adams <etowahadams@gmail.com> Co-authored-by: Peter Obi <peter.obi@psivant.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The vocabulary the rest of the package is written in. types.py holds four frozen dataclasses: RestraintSet (indices plus lower and upper bounds, either of which may be infinite), SteeringContext, StepState and SteeringUpdate. StepState.steering_t implements the Boltz/Protenix convention, 1 - step_index / num_steps, counting steps taken rather than sigmas visited. SteeringUpdate is explicitly an additive correction to x0, never coordinates. defaults.py is the single place derived numeric parameters live -- weights, buffers, the vdW cutoff offset, the gradient-step count -- so that a value adapted from Boltz cannot end up uncredited by being restated in a config model elsewhere. schedules.py carries the three foldsteer schedule types (Constant, ExponentialInterpolation, PiecewiseStepFunction). Only Constant is reachable from a runner yaml today; the others are implemented and unit-tested against the day a real schedule is wanted. Co-authored-by: Etowah Adams <etowahadams@gmail.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A potential is a restraint family: zero energy inside a permitted interval,
linear penalty outside it, so dU/dv is exactly -w below the window, +w above,
and zero within. Gradients are analytic rather than autograd -- the sampler
runs under torch.no_grad / inference mode, where taping the denoised tensor
would raise, and the guidance loop evaluates each term num_gd_steps times per
denoising step for what are closed-form derivatives anyway.
DistanceBoundsPotential is the only family implemented. Its energy_and_gradient
accepts any number of leading batch/sample axes; the reshape that makes that
work is load-bearing and regression-tested against shapes that collide with
the term's arity.
Registration takes an explicit snake_case name (`@register("distance_bounds_
potential")`) rather than deriving one from the class. That name is what a
user writes in a runner yaml, so it should not be CamelCase, should not
change when someone renames a class, and should not be produced by a
converter that turns VDWOverlapPotential into v_d_w_overlap_potential.
Collisions are rejected outright: a silent overwrite would make one potential
unreachable and misroute the other's restraints.
Co-authored-by: Etowah Adams <etowahadams@gmail.com>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ChemicalSteering.on_denoised runs num_gd_steps of gradient descent on the denoised coordinate estimate and returns the accumulated correction. Because the flat-bottom subgradient is exactly +-1 and a distance has |dv/dx| = 1, a term's weight is literally its step size: a violating atom moves `weight` Angstrom per step. One correction to the design as written: usable terms are resolved once per call and the loop then runs unconditionally. Deciding per step whether any term applies, and breaking out when none does, ends the whole descent at the first interval gap rather than skipping that one step. A test covers it. Adapted from foldsteer. See THIRD_PARTY_NOTICES.md. Co-authored-by: Etowah Adams <etowahadams@gmail.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Steering is a property of the run, not of a query: whether ligands are steered is set once in the runner yaml and applies to every query, which also makes a steered/unsteered ablation a yaml toggle rather than a second set of query JSONs. Query and inference_query_format are untouched, which in turn keeps this package free of any openfold3.projects import. The two config lines that make the setting reachable land here too -- a field on each of the two inference config models, and the line handing it to the data module -- so the routing test has the path it asserts on. batch_features.py is the wire format, both directions in one module, so a rename cannot desynchronize the producer from the consumer -- a failure that would surface as steering quietly doing nothing rather than as an error. A SteeringContext cannot ride the batch as an object: the collator does pad_sequence(...).squeeze(-1) and the model unsqueezes a sample axis into every leaf, and only two batch entries have hardcoded non-tensor escape hatches. Everything therefore crosses as plain tensors, with an explicit per-term _count as the shape authority -- orientation is never inferred from a tensor's stored shape. Disabled steering emits no keys at all rather than zeroed ones. That is what makes it a structural no-op: the sampler probes with .get, finds nothing, and runs the unmodified path. The tests put features through openfold_batch_collator itself rather than a stand-in, since the shapes that reach the sampler are the whole point, and round-trip them back through prepare_steering against an independent name-based conformer oracle. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The chemistry, and the only module in the package that imports RDKit -- which is what keeps RDKit off the model's import path. Bounds come from GetMoleculeBoundsMatrix on the intact, hydrogen-carrying reference molecule, exactly as Boltz computes them, with buffers applied per pair classification rather than uniformly: bonded pairs, 1-3 angle pairs, pairs that are both (which take the tighter of the two), and everything else, which keeps only a lower bound floored at the pair's van der Waals cutoff. The reference molecule is never rebuilt from an AtomArray. A round trip through RDKit valence perception drops molecules carrying formal charges it cannot reassign -- quaternary nitrogen, boron -- which is what left 10 of 161 ligand chains unsteered in foldsteer's benchmark. build_context maps RDKit-local indices onto the global atom axis by zipping each ligand residue against its in_crop_mask. In a complex that axis is renumbered relative to a ligand-only query, so the tests check the mapping against an oracle that is independent of it: reference conformer coordinates placed by matching atom names. If the two disagree, the conformer stops satisfying its own restraints. Five chain layouts are covered, and rotating every index by one atom is asserted to register, so the check cannot pass vacuously. Mismatched reference molecules raise MissingReferenceMoleculeError naming both counts, rather than surfacing an incidental "zip() argument is shorter" from further down. Constraint extraction and the acceptance tests derive from OpenFold3 PR #385. See THIRD_PARTY_NOTICES.md. Co-authored-by: Peter Obi <peter.obi@psivant.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The hook: 41 lines in diffusion_module.py, sitting between the denoiser call and the Euler step. Guidance corrects the denoised estimate rather than the noisy state, so it flows through that step instead of fighting it, and the dtype round trip lives here so the engine only ever sees float32. forward decides once, up front, and passes the decision down explicitly. _sample_rollout takes a use_steering flag alongside the prepared object, and the two must agree -- a True flag with nothing prepared raises rather than quietly running an unsteered rollout under a name that says otherwise. The primary rollout passes the run's setting; the pocket-refinement rollout passes False in the open, so guidance effects stay attributable in a benchmark rather than entangled with the second pass. prepare_steering returns None whenever a run is not steered -- steering off, no ligand to restrain, no term with restraints -- and validates before anything runs: one query per model batch, the emitted atom count against the model's atom axis, index dtype and range, element counts against _count. It is called before the first allocation, so a malformed batch fails before compute is spent. The property the design hangs on is tested directly: a batch carrying no steering features produces byte-identical output at rtol=0, atol=0 and leaves the RNG state untouched. Steering itself draws no randomness, so an enabled run cannot desynchronize any subsequent draw. bfloat16 rollouts are covered both ways -- the engine sees float32 whatever the rollout's precision, and no dtype changes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
One call in create_all_features, sitting beside its pocket_sampling counterpart, which is what turns the run-level setting into batch features. The entry point is named maybe_create_steering_features because three ordinary situations produce no features at all -- steering off, every term off, or a query with no ligand -- and the empty return is what makes a disabled run a structural no-op rather than one carrying zeroed restraints. The tests go through InferenceDataset.create_all_features itself rather than the featurizer, because that call is the thing that can be deleted while every other steering test stays green. Ethanol behind a four-residue protein chain: three heavy-atom restraints, small enough to enumerate by hand, identified by atom name so the ligand's global offset is checked rather than assumed, with bond windows asserted to bracket the real C-C and C-O lengths. mypy-baseline.txt loses one line: annotating the features dict in create_all_features resolves a pre-existing finding as well as the new one. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
An entry in the configuration reference, a steering block in the reference full config, and a runnable example runner yaml. Terms are keyed by their snake_case registry name, so what a user writes is `distance_bounds_potential` rather than a class name. The docs state what steering does and does not claim: it improves the internal chemistry of a predicted ligand, it applies to every query in the run, and only the primary rollout is steered. Derived numeric defaults are referenced from defaults.py rather than restated here -- restating them is how a value adapted from Boltz ends up uncredited outside the package boundary. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Answers "is this molecule built correctly" rather than "is it in the right place", from the prediction alone with no experimental reference -- which is what makes it complementary to metrics/alignment.py. worst_bounds_violation reports how far any atom pair sits outside its distance-geometry window, covering bond lengths, angles and internal clashes at once, the way PoseBusters' internal validity checks do. saturated_ring_atom_names and mean_ring_torsion measure whether a saturated ring kept its pucker: ~55 degrees at every ring bond for a chair, 0 for a flattened one. Ring detection is restricted to saturated all-carbon six-rings because the torsion metric is calibrated against a chair -- an aromatic ring is flat by right, a five-ring puckers to a different amplitude, and a ring fused to an arene is a half-chair (tetralin is a test case for that last one). Tested against coordinates whose answers are known by construction: dihedrals built at a chosen angle, one hexagon builder parametrized by pucker giving both an ideal chair and an explicitly flat ring. CPU only, no weights, a third of a second. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Runs one query through real inference twice, steering on and off, and
measures the predicted ligand's internal geometry. Nothing else differs
between the arms.
FKBP12 with rapamycin (PDB 1FKB), the ligand given by CCD code rather than
SMILES. Rapamycin is a 31-membered macrocycle, 65 heavy atoms, flexible
enough that a single-sequence prediction strains it; FKBP12 is 107 residues,
so both arms run in about 70 seconds.
Measured over two runs of 5 samples per arm, of3-ob-2025-06-30-174k on a CUDA
GB10:
worst distance-bounds violation off 0.74-1.21 A (mean 0.93)
on 0.000-0.02 A (mean ~0.00)
violated pairs per sample off 10-13, on 0-1
The violations are entirely nonbonded 1-4 and 1-5 pairs -- torsional strain,
atoms folded closer than any reachable torsion allows (C6-O4 at 3.11 A
against a 4.28 A floor). Bond lengths and angles were already correct in both
arms, so this term is not fixing those here.
Two candidate metrics were tried first and rejected on evidence, both
recorded in the module docstring. Aromatic ring planarity on foldsteer's own
example, SH2 with phosphotyrosine, comes out at 0.01 A in both arms -- this
checkpoint gets benzene right unaided. Violations against the raw, unbuffered
bounds matrix are not usable as an independent yardstick either: RDKit's own
generated conformer for rapamycin violates those by up to 2.12 A over 15
pairs, worse than the predictions do.
A second assertion covers the failure mode Tom Goddard's badchem survey
documents, where predictors return cyclohexane planar: rapamycin's saturated
six-ring must stay in a chair. Steered it is 53.7 +- 0.8 degrees across
samples, unsteered 50.3 +- 8.6, one sample having partially flattened to 35.
The primary metric is openly the quantity steering optimizes, so this
demonstrates that guidance reaches the sampler and changes the output
chemistry -- not independently that the chemistry is right. The ring check
and foldsteer's PoseBusters benchmark (66.0% -> 89.3% valid) are the
independent evidence.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
jnwei
marked this pull request as ready for review
September 4, 2026 09:55
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds stereochemistry steering to OpenFold3. In this PR the foundation is laid for adding more potentials, and also more test cases, in particular by introducing a new top level
steeringdirectory.Provenance
The work here is based on @peter-s-obi #385 and @etowahadams Foldsteer repo. Both of these works in turn borrowed formulations and code from Boltz and Protenix
Changes
InferenceDataset) and in the diffusion moduleSteeringSettingsopenfold3/steeringdirectoryTesting
Other Notes
The individual commits are organized by themes. It might be easiest to review this PR commit by commit.