https://github.com/elkins/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/diff-ensemble

biophysics differentiable-programming flax intrinsically-disordered-proteins jax nmr-spectroscopy saxs structural-biology variational-autoencoder

Last synced: about 2 months 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/diff-ensemble
• Owner: elkins
• License: mit
• Created: 2026-05-26T15:57:13.000Z (about 2 months ago)
• Default Branch: main
• Last Pushed: 2026-05-26T18:05:04.000Z (about 2 months ago)
• Last Synced: 2026-05-26T18:08:14.755Z (about 2 months ago)
• Topics: biophysics, differentiable-programming, flax, intrinsically-disordered-proteins, jax, nmr-spectroscopy, saxs, structural-biology, variational-autoencoder
• Language: Python
• Homepage: https://elkins.github.io/diff-ensemble
• Size: 605 KB
• Stars: 0
• Watchers: 0
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

README

          
# 🧬 DiffEnsemble: Differentiable IDP Ensemble Prediction

[![PyPI version](https://img.shields.io/pypi/v/diff-ensemble.svg)](https://pypi.org/project/diff-ensemble/)

[![Supported Python versions](https://img.shields.io/pypi/pyversions/diff-ensemble.svg)](https://pypi.org/project/diff-ensemble/)

[![Tests](https://github.com/elkins/diff-ensemble/actions/workflows/test.yml/badge.svg)](https://github.com/elkins/diff-ensemble/actions/workflows/test.yml)

[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

[![JAX](https://img.shields.io/badge/Accelerated_by-JAX-blue.svg)](https://github.com/google/jax)

[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)

DiffEnsemble is a JAX-powered framework for predicting structural ensembles of Intrinsically Disordered Proteins (IDPs) using a Variational Autoencoder (VAE) coupled with differentiable biophysical observables.

---

### 🧪 For Structural Biologists

*   **Ensemble Averaging:** Automatically calculates ensemble-averaged SAXS profiles and NMR observables.

*   **Disorder Recovery:** Specifically designed for proteins that don't have a single "fixed" structure, providing a statistical view of the conformational landscape.

### 🤖 For Machine Learning Geeks

*   **VAE-Physics Integration:** A latent-space generative model where the reconstruction loss is a combination of latent KLD and physical observables (SAXS/NMR).

*   **Differentiable Torsions:** Maps latent vectors to 3D coordinates via a differentiable NeRF (Natural Extension Reference Frame) implementation.

---

## 🚀 Key Features

*   **JAX-Accelerated VAE:** High-performance training of generative models for IDPs.

*   **Debye-Based SAXS Prediction:** Differentiable back-calculation of SAXS profiles from structural ensembles.

*   **Latent Space Exploration:** Sample new conformations from the learned disordered landscape.

## 📦 Installation

```bash

pip install diff-ensemble

```

## 📖 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.

    [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/elkins/diff-ensemble/blob/main/examples/quickstart_ensemble.ipynb)

## 📜 License

Distributed under the MIT License. See `LICENSE` for more information.