Skip to content

elkins-lab/diff-em

Repository files navigation

❄️ diff-em: Differentiable Cryo-EM Fitting in JAX

Tests License: MIT 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:

  • Open In Colab 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

  • Differentiable 3D Gaussian density kernels.
  • 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:


📖 Citation

@software{diff_em,
  author  = {Elkins, George},
  title   = {diff-em: Differentiable Cryo-EM map fitting in JAX},
  year    = {2026},
  url     = {https://github.com/elkins/diff-em},
  version = {0.1.0}
}

⚖️ License

MIT

Packages

 
 
 

Contributors