https://github.com/elkins-lab/diff-em
Differentiable Cryo-EM map fitting in JAX
https://github.com/elkins-lab/diff-em
biophysics computational-biophysics computational-structural-biology cryo-em differentiable-programming jax structural-biology
Last synced: 16 days ago
JSON representation
Differentiable Cryo-EM map fitting in JAX
- Host: GitHub
- URL: https://github.com/elkins-lab/diff-em
- Owner: elkins-lab
- License: mit
- Created: 2026-06-05T00:59:46.000Z (about 1 month ago)
- Default Branch: main
- Last Pushed: 2026-06-30T21:11:29.000Z (16 days ago)
- Last Synced: 2026-06-30T22:24:04.602Z (16 days ago)
- Topics: biophysics, computational-biophysics, computational-structural-biology, cryo-em, differentiable-programming, jax, structural-biology
- Language: Python
- Homepage: https://elkins-lab.github.io/diff-em/
- Size: 718 KB
- Stars: 2
- Watchers: 0
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- Changelog: CHANGELOG.md
- License: LICENSE
Awesome Lists containing this project
README
# ❄️ diff-em: Differentiable Cryo-EM Fitting in JAX
[](https://github.com/elkins-lab/diff-em/actions/workflows/test.yml)
[](https://opensource.org/licenses/MIT)
[](https://github.com/google/jax)
**diff-em** provides high-performance, auto-differentiable kernels for fitting atomic structures into Cryo-EM density maps. Built on **JAX**, it enables gradient-based optimization of coordinates directly against 3D experimental data.
---
## 🎯 Features
- **Gaussian Mixture Volumes:** Represent atomic models as differentiable 3D density maps using sum-of-Gaussians (electrostatic potential approximation).
- **Cross-Correlation Kernels:** Differentiable computation of map-to-model correlation coefficients (CC) for structural refinement (Rossmann, 2000).
- **Optimization Strategy:** Compatible with multi-resolution fitting and neural density fields (Zhong et al., 2021).
- **Hardware Acceleration:** Optimized for GPU/TPU execution via XLA, enabling the fitting of large complexes in seconds.
---
## 📚 Tutorials
Experience **diff-em** directly in your browser:
- [](https://colab.research.google.com/github/elkins-lab/diff-em/blob/main/examples/interactive_tutorials/cryo_em_fitting_tutorial.ipynb) **Cryo-EM Density Fitting** — Learn how to optimize atomic coordinates directly against 3D density maps using cross-correlation.
---
## 🏗️ Technical Architecture
- **Backend:** JAX (XLA-compiled).
- **Physics:** 3D Gaussian placement with B-factor smoothing.
- **Optimization:** Pure JAX implementation compatible with `optax` for high-dimensional gradient descent.
---
## 🧪 Scientific Validation
- **Density Parity:** Simulated densities are verified against standard EM map generation tools (e.g., `gemmi` or `ChimeraX`).
- **CC Gradient Stability:** Verified numerically stable gradients for structural refinement in the presence of noise.
- **Resolution Limits:** Benchmarked against known high-resolution and low-resolution experimental maps.
---
## 🚀 Roadmap
- [x] Differentiable 3D Gaussian density kernels.
- [x] Cross-correlation (CC) loss functions.
- [ ] Integration with MRC map loaders.
- [ ] Automated multi-resolution refinement schedules.
---
## 🔗 Related Projects
diff-em is part of the **differentiable biophysics** ecosystem:
- [diff-biophys](https://github.com/elkins-lab/diff-biophys) — Core differentiable biophysics engine.
- [diff-hdx](https://github.com/elkins-lab/diff-hdx) — Differentiable HDX-MS prediction.
- [synth-cryo-em](https://github.com/elkins-lab/synth-cryo-em) — Cryo-EM simulation.
---
## 📖 Citation
```bibtex
@software{diff_em,
author = {Elkins, George},
title = {diff-em: Differentiable Cryo-EM map fitting in JAX},
year = {2026},
url = {https://github.com/elkins-lab/diff-em},
version = {0.1.2}
}
```
## ⚖️ License
MIT