X-RAI: Scalable 3D Reconstruction From Single Particle X-Ray Diffraction Images Based on Online Machine Learning

This repository is the official implementation of X-RAI: Scalable 3D Reconstruction From Single Particle X-Ray Diffraction Images Based on Online Machine Learning

Requirements

We use Docker to containerize all the dependencies. So far, our system has been tested on Linux Red Hat Enterprise Linux 8.6 (Ootpa) equipped with a single 40 GB A100 GPU.

To pull the latest image from DockerHub and convert it to a Singularity image, do the following:

singularity pull xrai_latest.sif docker://jayshenoy/cryoai:spi

Downloading and installing the image should take a few hours.

Running a Job

All jobs can be run by specifying a config file (.ini). The config file must at least specify the experiment_name and experiment_type. To run a job with a Singularity container, use the following command:

singularity exec --nv ABSOLUTE_PATH_TO_CONTAINER python -m src.reconstruct.main -c RELATVE_PATH_TO_CONFIG_FILE

If using a shared computational resource managed with slurm, the job can be launched from the script run_from_config.sh

sbatch run_from_config.sh -c RELATVE_PATH_TO_CONFIG_FILE --sif ABSOLUTE_PATH_TO_CONTAINER

Generating and Reconstructing a Synthetic Dataset

We simulated synthetic datasets using Skopi, a modified copy of which is provided in this repository.

We provide a configuration file in simulation/configfiles/6J5I_50K_dataset_sim.ini to generate one of the synthetic datasets shown in the paper, namely 50,000 images of the 6J5I protein. Open the file and insert the desired output directory for parent_dir:

beam_file = input/amo86615.beam
pdb_file = input/6j5i.pdb
det_info = [128, 0.1, 0.2]
num_images_train = 50000
num_images_test = 0
reconstruction_dataset_sizes = [50000]
parent_dir = [ABSOLUTE PATH TO OUTPUT DIRECTORY]
train_epochs = [100]
add_poisson_noise = 1
fluence_level = [100]

Open simulation/simulate.sh and fill in the required Slurm configuration parameters, and specify the mount directory after the -B flag in the singularity command. Generate the dataset with the following command:

cd simulation
sbatch simulate.sh -c configfiles/6J5I_50K_dataset_sim.ini --sif ABSOLUTE_PATH_TO_CONTAINER

This script will create a new dataset directory within simulation/datasets. Within the dataset directory, there is a configuration file inside xrai_configs/unknown_poses/reconstruction_1000_fluence_100.0.ini that can be used to reconstruct the 6J5I protein.

Be sure to edit run_from_config.sh with the correct Slurm configuration parameters and specify the mount directory after the -B flag in the singularity command. To launch a reconstruction job, issue the command:

cd [TOP-LEVEL DIRECTORY]
sbatch run_from_config.sh -c [ABSOLUTE PATH TO reconstruction_1000_fluence_100.0.ini] --sif ABSOLUTE_PATH_TO_CONTAINER

On a single A100 GPU with 40 GB of memory, reconstruction should take between 8 and 9 hours. The reconstruction results will be saved in a directory called xrai within the dataset directory. Each reconstruction run is stored in a separate directory, with the 3D intensity file written to intensity.mrc. Performing phase retrieval on the intensity volumes must be done with separate softare, after which you can view the density volumes.

Monitoring Training

By default, log files containing tensorboard summaries, mrc files (reconstructed volumes) and config files are stored in the logs/ directory. You can monitor your model with the following command:

tensorboard --logdir logs --port 8888 --bind_all

Citing this work

Please cite this work as

@article{shenoy2023xrai,
  author    = {Shenoy, Jay and Levy, Axel and Ayyer, Kartik and Poitevin, Frédéric and Wetzstein, Gordon},
  title     = {Scalable 3D Reconstruction for X-Ray Single Particle Imaging With Online Machine Learning},
  journal   = {arXiv:preprint},
  year      = {2023},
}

Acknowledgments

This codebase builds off of the CryoAI project, which introduced amortized pose estimation for reconstruction in cryo-EM.

We use Skopi to simulate X-ray single particle images, a modified copy of which is provided in this repository.

List of Authors

J. Shenoy, A. Levy, K. Ayyer, F. Poitevin, G. Wetzstein

Stanford Computational Imaging Group

SLAC

MPSD Hamburg