{"id":51310359,"url":"https://github.com/elkins-lab/diff-ensemble","last_synced_at":"2026-07-01T03:03:01.281Z","repository":{"id":360495811,"uuid":"1250417582","full_name":"elkins-lab/diff-ensemble","owner":"elkins-lab","description":"Differentiable VAE framework for predicting protein structural ensembles (IDPs) consistent with SAXS and NMR data. Built on JAX/Flax.","archived":false,"fork":false,"pushed_at":"2026-06-30T20:39:25.000Z","size":1021,"stargazers_count":1,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-06-30T21:25:15.015Z","etag":null,"topics":["biophysics","computational-biophysics","computational-structural-biology","differentiable-programming","flax","intrinsically-disordered-proteins","jax","nmr-spectroscopy","saxs","science-education","structural-biology","vae","variational-autoencoder"],"latest_commit_sha":null,"homepage":"https://elkins.github.io/diff-ensemble","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/elkins-lab.png","metadata":{"files":{"readme":"README.md","changelog":"CHANGELOG.md","contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-05-26T15:57:13.000Z","updated_at":"2026-06-30T20:24:19.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/elkins-lab/diff-ensemble","commit_stats":null,"previous_names":["elkins/diff-ensemble","elkins-lab/diff-ensemble"],"tags_count":2,"template":false,"template_full_name":null,"purl":"pkg:github/elkins-lab/diff-ensemble","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/elkins-lab%2Fdiff-ensemble","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/elkins-lab%2Fdiff-ensemble/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/elkins-lab%2Fdiff-ensemble/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/elkins-lab%2Fdiff-ensemble/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/elkins-lab","download_url":"https://codeload.github.com/elkins-lab/diff-ensemble/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/elkins-lab%2Fdiff-ensemble/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":34990846,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-07-01T02:00:05.325Z","response_time":130,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["biophysics","computational-biophysics","computational-structural-biology","differentiable-programming","flax","intrinsically-disordered-proteins","jax","nmr-spectroscopy","saxs","science-education","structural-biology","vae","variational-autoencoder"],"created_at":"2026-07-01T03:03:00.823Z","updated_at":"2026-07-01T03:03:01.276Z","avatar_url":"https://github.com/elkins-lab.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# 🧬 DiffEnsemble\n\n[![Tests](https://github.com/elkins-lab/diff-ensemble/actions/workflows/test.yml/badge.svg)](https://github.com/elkins-lab/diff-ensemble/actions/workflows/test.yml)\n[![Docs](https://github.com/elkins-lab/diff-ensemble/actions/workflows/docs.yml/badge.svg)](https://elkins-lab.github.io/diff-ensemble/)\n[![codecov](https://codecov.io/gh/elkins-lab/diff-ensemble/branch/main/graph/badge.svg)](https://codecov.io/gh/elkins-lab/diff-ensemble)\n[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)\n[![Python versions](https://img.shields.io/pypi/pyversions/diff-ensemble)](https://pypi.org/project/diff-ensemble/)\n[![JAX](https://img.shields.io/badge/backend-JAX-9cf.svg)](https://github.com/google/jax)\n[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)\n[![Checked with mypy](https://img.shields.io/badge/type%20checked-mypy-blue)](https://mypy-lang.org/)\n\n**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.\n\n---\n\n## 🚀 For Machine Learning Engineers\n\n### Generative Architecture\nDiffEnsemble utilizes a **Variational Autoencoder (VAE)** implemented in **Flax** to learn the conformational manifold of flexible proteins.\n- **Encoder**: Maps sequence-derived features to a latent Gaussian distribution $(\\mu, \\sigma)$.\n- **Latent Space**: Enables efficient sampling of diverse structural states.\n- **Decoder**: Predicts a set of $(\\phi, \\psi)$ backbone torsions that define the protein's fold.\n\n### End-to-End Differentiability\nBuilt entirely on **JAX**, the entire pipeline—from the VAE weights to the final biophysical observable—is auto-differentiable.\n- **Biophysical Loss**: We compute the gradient of the error between predicted ensemble-averaged spectra and experimental data (SAXS/NMR) to update model weights directly.\n- **Vectorized Sampling**: Leveraging JAX's `vmap`, we generate and process ensembles of 100+ structures in parallel on GPUs/TPUs.\n\n---\n\n## 🧪 For Structural Biologists\n\n### The \"Ensemble\" Concept\nUnlike 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.\n\n### Differentiable Physics Engine\nWe 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:\n- **SAXS**: Small-Angle X-ray Scattering profiles via the Debye formula.\n- **NMR**: Residual Dipolar Couplings (RDCs) and Chemical Shifts.\n\n### Scientific Validation\nDiffEnsemble is rigorously validated against peer-reviewed standards:\n- **Sic1 Benchmark**: Recapitulates the dimensions of the Sic1 IDP as determined by the **Forman-Kay Group** (*JACS 2020*).\n- **CASP16 T1200**: Benchmarked against the **Montelione Group's** SpA domain-linker-domain challenge using RDC Q-factors.\n- **Polymer Physics**: Obeys Flory's scaling laws ($R_g \\propto N^{0.588}$) for random coils in a good solvent.\n\n---\n\n## 🛠️ Quick Start\n\n```python\nimport jax\nimport jax.numpy as jnp\nfrom diff_ensemble.model import EnsembleVAE\n\n# Initialize model (90 residues, 32 latent dims, 100 models)\nmodel = EnsembleVAE(seq_len=90, latent_dim=32, ensemble_size=100)\nrng = jax.random.PRNGKey(0)\n\n# Generate a structural ensemble from sequence features\nbatch_x = jnp.ones((1, 90, 4)) # Example features\ntorsions, mean, logvar = model.apply({\"params\": params}, batch_x, rng)\ncoords = model.generate_coordinates(torsions) # Shape: (100, 270, 3)\n\n# Save the cloud to a multi-model PDB for visualization\nfrom diff_ensemble.io import save_ensemble_to_pdb\nsave_ensemble_to_pdb(coords, \"ensemble_cloud.pdb\")\n```\n\n## 📚 References\n\n1. **Kingma \u0026 Welling (2013)**: *Auto-Encoding Variational Bayes*.\n2. **Gomes et al. (2020)**: *Conformational Ensembles of an IDP Consistent with NMR, SAXS, and smFRET* (Forman-Kay Lab).\n3. **McBride et al. (2025)**: *Predicting Pose Distribution of Protein Domains* (Montelione Lab).\n4. **Parsons et al. (2005)**: *Practical conversion from torsion space to Cartesian space for in silico protein synthesis*. J. Comput. Chem. 26(10), 1063–1068.\n\n## 🛠 Software Architecture\n\nThe project is structured for modularity and high-performance execution:\n\n*   **`diff_ensemble/model.py`**: The Flax-based VAE architecture (Encoder/Decoder).\n*   **`diff_ensemble/observables.py`**: Forward biophysical kernels and multi-objective loss functions.\n*   **`diff_ensemble/train.py`**: The training orchestration and optimization loop using Optax.\n*   **`diff_ensemble/io.py`**: PDB trajectory export and multi-model stack management.\n*   **`diff_ensemble/ensemble.py`**: High-level API for population-weighted averaging.\n\n---\n\n## 🤝 Contributing \u0026 Support\n\nWe welcome contributions from both the Machine Learning and Structural Biology communities!\n*   **Bugs/Features:** Please open an issue on the GitHub repository.\n*   **Questions:** Visit our [Documentation](https://elkins-lab.github.io/diff-ensemble/) or reach out via GitHub Discussions.\n\n---\n\n## ⚖️ License\n\nThis project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.\n\n---\n\n## 🔗 Related Projects\n\nDiffEnsemble depends on and integrates with:\n\n- [diff-biophys](https://github.com/elkins-lab/diff-biophys) — Differentiable JAX kernels for SAXS/NMR (core dependency)\n- [synth-pdb](https://github.com/elkins-lab/synth-pdb) — Synthetic structure generation for training data\n- [synth-nmr](https://github.com/elkins-lab/synth-nmr) — NMR observables for experimental targets\n- [synth-saxs](https://github.com/elkins-lab/synth-saxs) — SAXS profile simulation\n- [torsion-tuner](https://github.com/elkins-lab/torsion-tuner) — Single-structure refinement counterpart\n\n---\n\n## 📖 Citation\n\n```bibtex\n@software{diff_ensemble,\n  author  = {Elkins, George},\n  title   = {DiffEnsemble: Differentiable structural ensemble prediction for IDPs},\n  year    = {2026},\n  url     = {https://github.com/elkins-lab/diff-ensemble},\n  version = {0.1.3}\n}\n```\n\n## 📖 Tutorials\n\nGet started immediately with our interactive Jupyter notebooks:\n\n*   **[Quick Start: Differentiable IDP Ensemble Prediction](examples/quickstart_ensemble.ipynb)**: Train a VAE to predict structural ensembles constrained by SAXS data.\n    [![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)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Felkins-lab%2Fdiff-ensemble","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Felkins-lab%2Fdiff-ensemble","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Felkins-lab%2Fdiff-ensemble/lists"}