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.1.tar.gz (6.6 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.1-py3-none-any.whl (4.7 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: diff_em-0.1.1.tar.gz
  • Upload date:
  • Size: 6.6 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.1.tar.gz
Algorithm Hash digest
SHA256 b1b0e819e882482e4446da6a1af4ee3b7333c9bfba5e8cf693a6d417374df04d
MD5 2faef22c05609299cc37d2ae51d8a430
BLAKE2b-256 28d74da59cd7989f585f468041723152c048965ea1a5fc09251cb9c376ac1d6a

See more details on using hashes here.

File details

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

File metadata

  • Download URL: diff_em-0.1.1-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.1-py3-none-any.whl
Algorithm Hash digest
SHA256 5abb35292654f62ad9c0feea68def3fdd3599afeb9a6d18c55540018f80c1b78
MD5 63c8b81423ae0f75cf2a816a242f69c9
BLAKE2b-256 d78b668831364c0588a2cfe92a2630d3120c62e2dace2fe5e574f1a6af1e10a2

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