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.

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.

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.

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.

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.

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.

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")

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.

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.

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 or reach out via GitHub Discussions.

This project is licensed under the MIT License - see the LICENSE file for details.

DiffEnsemble depends on and integrates with:

• diff-biophys — Differentiable JAX kernels for SAXS/NMR (core dependency)
• synth-pdb — Synthetic structure generation for training data
• synth-nmr — NMR observables for experimental targets
• synth-saxs — SAXS profile simulation
• torsion-tuner — Single-structure refinement counterpart

@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}
}

Get started immediately with our interactive Jupyter notebooks:

• Quick Start: Differentiable IDP Ensemble Prediction: Train a VAE to predict structural ensembles constrained by SAXS data.