❄️ diff-em: Differentiable Cryo-EM Fitting in 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:

• 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:

• diff-biophys — Core differentiable biophysics engine.
• diff-hdx — Differentiable HDX-MS prediction.
• synth-cryo-em — Cryo-EM simulation.

📖 Citation

@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