Metadata-Version: 2.4
Name: cryojax_eo
Version: 0.1.0
Summary: Cryo-EM Ensemble Optimization using JAX and OpenMM for Molecular Dynamics simulations.
Project-URL: repository, https://github.com/flatironinstitute/cryojax-ensemble-optimization
Author: Erik Henning Thiede, Pilar Cossio, Roy Lederman
Author-email: David Silva Sanchez <david.silva@yale.edu>
License: MIT License
        
        Copyright (c) 2023 David Silva-Sánchez
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Keywords: cryo-EM,electron-microscopy,equinox,jax,molecular-dynamics
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: GNU Lesser General Public License v2 (LGPLv2)
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.11
Requires-Dist: cryojax
Requires-Dist: dm-pix
Requires-Dist: equinox
Requires-Dist: jax
Requires-Dist: jax-dataloader
Requires-Dist: jaxlib
Requires-Dist: jaxopt
Requires-Dist: jaxtyping
Requires-Dist: mdtraj
Requires-Dist: natsort
Requires-Dist: numpy
Requires-Dist: openmm
Requires-Dist: optax
Requires-Dist: optimistix
Requires-Dist: pydantic
Requires-Dist: pyyaml
Requires-Dist: tqdm
Provides-Extra: dev
Requires-Dist: pre-commit; extra == 'dev'
Requires-Dist: ruff; extra == 'dev'
Provides-Extra: tests
Requires-Dist: pytest; extra == 'tests'
Description-Content-Type: text/markdown

<h1 align='center'>Cryo-electron microscopy ensemble optimization using individual particles and physical constraints</h1>


## Summary

CryoJAX ensemble optimization is a module of the [cryoJAX](https://github.com/michael-0brien/cryojax) library, a [JAX](https://github.com/jax-ml/jax) and [Equinox](https://docs.kidger.site/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](https://openmm.org/) library. Details and results are available in our [preprint](https://www.biorxiv.org/content/10.64898/2025.12.02.691891v1). Instructions for reproducing the paper results are provided below.


## Installation

Our library has been tested on the latest Ubuntu version. Availability for other platforms is dependent on the availability of OpenMM and JAX.

### CPU Installation

Our library can be installed to be used with a CPU via pip.

```bash
pip install git+https://github.com/flatironinstitute/cryojax-ensemble-optimization.git
```
We recommend using a freshly created virtual environment to install our library. A CPU installation is only recommended for dataset simulation, as OpenMM is built for GPU, and simulations will take a long time if run on CPU.

### GPU Installation

We recommend installing our library using conda (or one of its variants), as matching JAX's and OpenMM's CUDA versions can be difficult otherwise. Here we show an example of how to install our library with [mamba](https://github.com/mamba-org/mamba):

```bash
mamba create -n cryojax_eo_env python==3.11
mamba activate cryojax_eo_env
mamba install -c conda-forge openmm cuda-version==12.4 # Insert your cuda version!
pip install --upgrade "jax[cuda12]"
pip install git+https://github.com/flatironinstitute/cryojax-ensemble-optimization.git
```
To find your CUDA version, you can run `nvidia-smi` in a terminal. The CUDA version will appear in the top right corner of the output.


## Cryo-EM Ensemble Optimization Input

See the [input documentation](docs/ensemble_optimization.md)


## Cryo-EM heterogeneous dataset simulation

See the [input documentation](docs/dataset_simulation.md)


## Reproducing Paper Results

All the necessary data, atomic models, config files, and instructions to reproduce our results are available in [Zenodo](https://doi.org/10.5281/zenodo.17420048).

## Contact

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


## Contributing

If you are contributing to this project, please install the package with the following command:

```bash
pip install -e ".[dev]"
```


This will install the required dependencies for development, the most important being `Ruff` and `pre-commit`. After installation, activate your environment and install the `pre-commit` hooks by running

`pre-commit install`

Make sure that your code is formatted according to our guidelines by running:

```bash
pre-commit run --all-files
```

This will make sure the code is formatted correctly, fix whatever can be automatically fixed, and tell you if something else needs to be fixed.

## Acknowledgements
We thank Michael O'Brien, Miro Astore, Lars Dingeldein, Wai Shing Tang, Aaditya Rangan, and Sonya Hanson for helpful discussions.
