bytedance / bytedance/Protenix

[Feature Request] Add optional Amber relaxation to protenix pred (CUDA enabled)

Open
#236 1 comment 2 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2.1k
Forks
312
PR merge metrics
No merged PRs in 30d

Description

**I’d like to propose adding post-inference structure relaxation to Protenix-v1, with a simple one-flag UX.**

**Why this is useful**
Protenix currently outputs predicted structures, but users who want physically cleaner structures still need separate post-processing scripts/tools. Having built-in Amber relaxation would make out-cif files much useful and more reproducible, in multi-seed runs.

**Unrelaxed:**
- Steric clashes (atoms too close together)
- Unphysical bond lengths/angles

**Relaxed**: Better for MD simulations, docking or some quick comparison etc analysis
- Resolve clashes between atoms
- Fix bond lengths/angles to ideal values

Single flag in ```protenix pred --relax_amber true```
```bash
protenix pred \
-i ./design.json \
-o ./output_design \
-n protenix_base_default_v1.0.0 \
--seeds 101 \
--enable_cache true \
--enable_tf32 true \
--enable_fusion true \
--relax_amber true
```

Expected behavior when ```--relax_amber true```:
- Relax all predicted CIF structures ([*_sample_*.cif] for all seeds/results generated in the run.
- Prefer GPU/CUDA relaxation automatically if OpenMM CUDA platform is available.
- Fallback to CPU if no CUDA.
- Out [*_relaxed.cif] files.

When ```--relax_amber false``` (default), functions remains unchanged.

**Suggested implementation details**
_- I have already implemented by adding a dedicated relaxation module (for example [relax.py] and connect it into post-inference flow in [batch_inference.py]_

Core functions:
- validate_relaxation_runtime(use_amber_relax, use_gpu_relax)
- get_openmm_platform_names()
- is_openmm_cuda_available()
- run_amber_relaxation and save cif files

Notes:
- Use OpenMM Amber forcefield minimization.
- Accept CIF input directly via OpenMM PDBxFile.
- Optional PDBFixer path; fallback to OpenMM Modeller if not installed.
- Convert minimized structure back to CIF.
- Clean up intermediate files (e.g., remove temporary relaxed PDB after CIF write).

Dependencies:
- openmm (CUDA-enabled build recommended)
- gemmi
- optional: pdbfixer

Validation:
- I have already tested this feature in the current latest Protenix-v1
```bash
> ls
design_sample_0.cif design_sample_2_relaxed.cif design_summary_confidence_sample_0.json
design_sample_0_relaxed.cif design_sample_3.cif design_summary_confidence_sample_1.json
design_sample_1.cif design_sample_3_relaxed.cif design_summary_confidence_sample_2.json
design_sample_1_relaxed.cif design_sample_4.cif design_summary_confidence_sample_3.json
design_sample_2.cif design_sample_4_relaxed.cif design_summary_confidence_sample_4.json
```

Some quick energy analysis of relaxed out-cif files:
Sample 0: 39,720 → -10,424 kcal/mol ✓ good
Sample 2: 2,503 → -10,614 kcal/mol ✓ excellent (low initial energy)
Sample 3: 7.5M → -10,310 kcal/mol ✓ huge clash fixed!
Image
img caption: relaxed(in cyan) and unrelaxed(in green)

Let me know If you I can Open Pull request

Contributor guide

Open the contributing guide

Research direction

Start with the post-inference flow in batch_inference.py and the proposed relaxation module relax.py, then review how the current *_sample_*.cif outputs are produced. Run the documented protenix pred command with relaxation enabled and verify that every seed produces a corresponding *_relaxed.cif, with CUDA preferred and CPU fallback, while the default false flag preserves current behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.