CryoJAX Ensemble Optimization

Cryo-electron microscopy ensemble optimization using individual particles and physical constraints

cryojax_eo is a module of the cryoJAX library, a JAX and Equinox-based library for differentiable cryo-EM forward models. The purpose of this library is to provide a framework for optimizing structural ensembles, defined as a weighted discrete set of atomic structures, given a set of cryo-EM images.

To do this, we implement an algorithm inspired by projected gradient descent, where the optimization step is performed by comparing the ensemble to the cryo-EM dataset, and the projection step is done through Steered Molecular Dynamics using the popular OpenMM library.

Details and results are available in our preprint.

Capabilities

• Ensemble Optimization — optimize a weighted ensemble of atomic structures against cryo-EM particle images using Steered Molecular Dynamics
• Ensemble Reweighting — compute optimal weights for an existing set of structures or volumes against cryo-EM images (no OpenMM required)
• Dataset Simulation — generate synthetic heterogeneous cryo-EM datasets from multiple atomic models
• Flexible Fitting — fit a single atomic model to a consensus density map using steered MD

Getting started

See the Installation page, then pick the workflow that matches your use case from the Usage section.

Reproducing paper results

All data, atomic models, config files, and instructions are available on Zenodo.

Contact

Please submit bug reports, feature requests, or general feedback as a GitHub issue.