https://github.com/elkins-lab/diff-ensemble
Differentiable VAE framework for predicting protein structural ensembles (IDPs) consistent with SAXS and NMR data. Built on JAX/Flax.
https://github.com/elkins-lab/diff-ensemble
biophysics computational-biophysics computational-structural-biology differentiable-programming flax intrinsically-disordered-proteins jax nmr-spectroscopy saxs science-education structural-biology vae variational-autoencoder
Last synced: 16 days ago
JSON representation
Differentiable VAE framework for predicting protein structural ensembles (IDPs) consistent with SAXS and NMR data. Built on JAX/Flax.
- Host: GitHub
- URL: https://github.com/elkins-lab/diff-ensemble
- Owner: elkins-lab
- License: mit
- Created: 2026-05-26T15:57:13.000Z (about 2 months ago)
- Default Branch: main
- Last Pushed: 2026-06-30T20:39:25.000Z (16 days ago)
- Last Synced: 2026-06-30T21:25:15.015Z (16 days ago)
- Topics: biophysics, computational-biophysics, computational-structural-biology, differentiable-programming, flax, intrinsically-disordered-proteins, jax, nmr-spectroscopy, saxs, science-education, structural-biology, vae, variational-autoencoder
- Language: Python
- Homepage: https://elkins.github.io/diff-ensemble
- Size: 997 KB
- Stars: 1
- Watchers: 0
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- Changelog: CHANGELOG.md
- License: LICENSE
Awesome Lists containing this project
README
# 🧬 DiffEnsemble
[](https://github.com/elkins-lab/diff-ensemble/actions/workflows/test.yml)
[](https://elkins-lab.github.io/diff-ensemble/)
[](https://codecov.io/gh/elkins-lab/diff-ensemble)
[](https://opensource.org/licenses/MIT)
[](https://pypi.org/project/diff-ensemble/)
[](https://github.com/google/jax)
[](https://github.com/astral-sh/ruff)
[](https://mypy-lang.org/)
**DiffEnsemble** is a high-performance, differentiable framework for predicting structural ensembles of Intrinsically Disordered Proteins (IDPs). By combining generative deep learning with hardware-accelerated biophysics, it bridges the gap between protein sequence and solution-state experimental data.
---
## 🚀 For Machine Learning Engineers
### Generative Architecture
DiffEnsemble utilizes a **Variational Autoencoder (VAE)** implemented in **Flax** to learn the conformational manifold of flexible proteins.
- **Encoder**: Maps sequence-derived features to a latent Gaussian distribution $(\mu, \sigma)$.
- **Latent Space**: Enables efficient sampling of diverse structural states.
- **Decoder**: Predicts a set of $(\phi, \psi)$ backbone torsions that define the protein's fold.
### End-to-End Differentiability
Built entirely on **JAX**, the entire pipeline—from the VAE weights to the final biophysical observable—is auto-differentiable.
- **Biophysical Loss**: We compute the gradient of the error between predicted ensemble-averaged spectra and experimental data (SAXS/NMR) to update model weights directly.
- **Vectorized Sampling**: Leveraging JAX's `vmap`, we generate and process ensembles of 100+ structures in parallel on GPUs/TPUs.
---
## 🧪 For Structural Biologists
### The "Ensemble" Concept
Unlike AlphaFold, which predicts a single static structure, IDPs exist as a "cloud" of interconverting conformations. DiffEnsemble predicts this **structural ensemble**, which is essential for understanding proteins that do not have a stable fold.
### Differentiable Physics Engine
We use the **NeRF (Natural Extension Reference Frame)** algorithm to convert predicted torsions into 3D Cartesian coordinates. These coordinates are then passed to **DiffBiophys** kernels to calculate:
- **SAXS**: Small-Angle X-ray Scattering profiles via the Debye formula.
- **NMR**: Residual Dipolar Couplings (RDCs) and Chemical Shifts.
### Scientific Validation
DiffEnsemble is rigorously validated against peer-reviewed standards:
- **Sic1 Benchmark**: Recapitulates the dimensions of the Sic1 IDP as determined by the **Forman-Kay Group** (*JACS 2020*).
- **CASP16 T1200**: Benchmarked against the **Montelione Group's** SpA domain-linker-domain challenge using RDC Q-factors.
- **Polymer Physics**: Obeys Flory's scaling laws ($R_g \propto N^{0.588}$) for random coils in a good solvent.
---
## 🛠️ Quick Start
```python
import jax
import jax.numpy as jnp
from diff_ensemble.model import EnsembleVAE
# Initialize model (90 residues, 32 latent dims, 100 models)
model = EnsembleVAE(seq_len=90, latent_dim=32, ensemble_size=100)
rng = jax.random.PRNGKey(0)
# Generate a structural ensemble from sequence features
batch_x = jnp.ones((1, 90, 4)) # Example features
torsions, mean, logvar = model.apply({"params": params}, batch_x, rng)
coords = model.generate_coordinates(torsions) # Shape: (100, 270, 3)
# Save the cloud to a multi-model PDB for visualization
from diff_ensemble.io import save_ensemble_to_pdb
save_ensemble_to_pdb(coords, "ensemble_cloud.pdb")
```
## 📚 References
1. **Kingma & Welling (2013)**: *Auto-Encoding Variational Bayes*.
2. **Gomes et al. (2020)**: *Conformational Ensembles of an IDP Consistent with NMR, SAXS, and smFRET* (Forman-Kay Lab).
3. **McBride et al. (2025)**: *Predicting Pose Distribution of Protein Domains* (Montelione Lab).
4. **Parsons et al. (2005)**: *Practical conversion from torsion space to Cartesian space for in silico protein synthesis*. J. Comput. Chem. 26(10), 1063–1068.
## 🛠 Software Architecture
The project is structured for modularity and high-performance execution:
* **`diff_ensemble/model.py`**: The Flax-based VAE architecture (Encoder/Decoder).
* **`diff_ensemble/observables.py`**: Forward biophysical kernels and multi-objective loss functions.
* **`diff_ensemble/train.py`**: The training orchestration and optimization loop using Optax.
* **`diff_ensemble/io.py`**: PDB trajectory export and multi-model stack management.
* **`diff_ensemble/ensemble.py`**: High-level API for population-weighted averaging.
---
## 🤝 Contributing & Support
We welcome contributions from both the Machine Learning and Structural Biology communities!
* **Bugs/Features:** Please open an issue on the GitHub repository.
* **Questions:** Visit our [Documentation](https://elkins-lab.github.io/diff-ensemble/) or reach out via GitHub Discussions.
---
## ⚖️ License
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
---
## 🔗 Related Projects
DiffEnsemble depends on and integrates with:
- [diff-biophys](https://github.com/elkins-lab/diff-biophys) — Differentiable JAX kernels for SAXS/NMR (core dependency)
- [synth-pdb](https://github.com/elkins-lab/synth-pdb) — Synthetic structure generation for training data
- [synth-nmr](https://github.com/elkins-lab/synth-nmr) — NMR observables for experimental targets
- [synth-saxs](https://github.com/elkins-lab/synth-saxs) — SAXS profile simulation
- [torsion-tuner](https://github.com/elkins-lab/torsion-tuner) — Single-structure refinement counterpart
---
## 📖 Citation
```bibtex
@software{diff_ensemble,
author = {Elkins, George},
title = {DiffEnsemble: Differentiable structural ensemble prediction for IDPs},
year = {2026},
url = {https://github.com/elkins-lab/diff-ensemble},
version = {0.1.3}
}
```
## 📖 Tutorials
Get started immediately with our interactive Jupyter notebooks:
* **[Quick Start: Differentiable IDP Ensemble Prediction](examples/quickstart_ensemble.ipynb)**: Train a VAE to predict structural ensembles constrained by SAXS data.
[](https://colab.research.google.com/github/elkins/diff-ensemble/blob/main/examples/quickstart_ensemble.ipynb)