Skip to content

Add TFG contact restraints for JSON constraints#310

Open
taivu1998 wants to merge 1 commit into
bytedance:mainfrom
taivu1998:tdv/issue-307-tfg-constraints
Open

Add TFG contact restraints for JSON constraints#310
taivu1998 wants to merge 1 commit into
bytedance:mainfrom
taivu1998:tdv/issue-307-tfg-constraints

Conversation

@taivu1998
Copy link
Copy Markdown

Summary

Addresses #307 by adding a supported Training-Free Guidance (TFG) path for JSON constraint.contact inputs on Protenix v1/v2 models without enabling untrained constraint embedders.

This PR:

  • adds UserDistanceRestraintPotential, a flat-bottom pairwise distance restraint term for user-provided contact bounds
  • converts inference JSON constraint.contact entries into atom-pair TFG restraint features when TFG feature extraction is enabled
  • supports residue-residue, atom-atom, and mixed atom-residue contact formats
  • adds a --tfg_constraint_weight CLI knob for tuning restraint strength
  • warns users when JSON constraints are present but the selected model/path will not consume them
  • documents the v0.5 learned-constraint path versus the v1/v2 TFG contact-restraint path
  • adds a runnable v1 contact-constraint example and focused tests

Root Cause

The v1/v2 checkpoints do not have trained ConstraintEmbedder modules enabled, so JSON constraints can be parsed without actually conditioning diffusion sampling. The existing TFG framework already supports differentiable geometry terms, but JSON contact constraints were not converted into TFG-compatible atom-pair features.

Implementation Notes

The new restraint path is intentionally narrow:

  • it does not turn on untrained constraint embedders for v1/v2 models
  • it does not reinterpret user contacts as RDKit pairwise geometry bounds
  • it preserves user-provided min/max distances as soft flat-bottom TFG penalties
  • it keeps constraint.pocket unsupported for v1/v2 TFG and documents that limitation

Validation

  • python3 -m py_compile protenix/tfg/potentials.py protenix/tfg/config.py protenix/data/constraint/constraint_featurizer.py protenix/data/inference/json_to_feature.py runner/batch_inference.py tests/test_tfg_user_restraints.py
  • git diff --check
  • git diff --cached --check
  • python3 -m json.tool examples/example_tfg_contact_v1.json
  • manual smoke test for UserDistanceRestraintPotential energy and gradient direction
  • python -m pytest tests/test_tfg_user_restraints.py -q

The focused pytest file collected successfully in the local environment, but all tests were skipped because the local Python environment is missing optional Protenix chemistry dependencies such as RDKit/Biotite. uv run pytest ... is also blocked locally by an existing invalid Anaconda llvmlite egg-info version, before test execution begins.

@taivu1998 taivu1998 marked this pull request as ready for review May 11, 2026 03:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant