Skip to main content

Differentiable Cryo-EM map fitting in JAX

Project description

❄️ 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.

🏗️ 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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

diff_em-0.1.0.tar.gz (6.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

diff_em-0.1.0-py3-none-any.whl (4.7 kB view details)

Uploaded Python 3

File details

Details for the file diff_em-0.1.0.tar.gz.

File metadata

  • Download URL: diff_em-0.1.0.tar.gz
  • Upload date:
  • Size: 6.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.10

File hashes

Hashes for diff_em-0.1.0.tar.gz
Algorithm Hash digest
SHA256 dd538c5eaf59c1f2ce0a35b99d8932412ebebd39db4c09e6ce1bf5b9741b9d16
MD5 c3a9aff7db6d5ba310a90d2f33bb78a9
BLAKE2b-256 d4bf20266dc53c2c5735276794c4f28c3f5546dbf37fa33ebcbac263c787edfb

See more details on using hashes here.

File details

Details for the file diff_em-0.1.0-py3-none-any.whl.

File metadata

  • Download URL: diff_em-0.1.0-py3-none-any.whl
  • Upload date:
  • Size: 4.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.10

File hashes

Hashes for diff_em-0.1.0-py3-none-any.whl
Algorithm Hash digest
SHA256 360319d22eed04cbfd8f7f5b3a07f8a55de7f27cb9db3b53d2c9c991bd90984c
MD5 3d2c4a522d846f2b070a6c6447275798
BLAKE2b-256 bf839f6527040e73dcd80880ce0aab8c62d4b6fee407b4c396d9842611288908

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page