🍦 Icecream: High-Fidelity Equivariant Cryo-Electron Tomography

Official repo for (icecream).

Icecream is a self-supervised framework for cryo-ET reconstruction that integrates equivariance principles from modern imaging theory into a deep-learning architecture. Icecream provides a theoretically grounded and computationally efficient method that jointly performs denoising and missing-wedge correction.

🧊 Note

The codebase is under active development.

⚠️ Important: PyTorch 2.9 currently has a bug affecting half-precision (FP16) training. Please use PyTorch 2.8 or earlier for now. See this issue.

⚠️ Version Notice

v0.3 contained a bug that degraded reconstruction performance relative to v0.2.
If you downloaded the code between 05.11.2025–28.11.2025, please update to the latest version.

📄 Changelog

See the full list of updates in the CHANGELOG.md.

Installation

1. Create a Python environmnet

Create a conda environment or you can use other environment managers like pipenv, poetry, uv etc with Python 3.11 or above. We will use conda with Python 3.11 in this example:

conda create -n icecream python=3.11 -y
conda activate icecream

2. Install PyTorch (CUDA-enabled)

Install CUDA-enabled PyTorch from https://pytorch.org/get-started/locally/ based on your system configuration. For example, for Linux with CUDA 12.8

pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128

Note: PyTorch 2.9 currently has a bug affecting half-precision (FP16) training. Please use PyTorch 2.8 or earlier for now. See this issue.

3. Install Icecream and its dependencies:

You can install Icecream directly from the GitHub repository using pip or clone the repository and install it locally.

3.1 Direct installation (recommended):

pip install git+https://github.com/swing-research/icecream.git

To test the installation, run:

icecream --help

It should display the three main commands: train, predict, and split-tilt-series.

3.2 Local installation from cloned repository (for ones who want to modify the code):

Clone the repository using one of the following methods:

HTTPS

git clone https://github.com/swing-research/icecream.git
cd icecream

SSH (for users with SSH keys configured):

git clone git@github.com:swing-research/icecream.git
cd icecream

Install in editable mode:

pip install -e .

To test the installation, run:

icecream --help

It should display the three main commands: train, predict, and split-tilt-series.

Update

To get the latest version of ICECREAM, make sure to be on the main branch

git checkout main

and run

git pull

Use previous version

To use the first version, i.e. v0.1, you can select the appropriate Github branch:

git checkout v0.1

Usage

To train the model, use the train command. Note that the command also reconstructs the volume after training. However, you can use the predict command to reconstruct the volume from a trained model using intermediate checkpoints.

To train the model, you can use a config file or directly pass the parameters through the command line. This uses the default training parameters specified in src/icecream/default.yaml. You can override these by passing a config file of your own.

Now, we show how to train the model using default parameters.

icecream train \
  --tomo0 /path/to/tomogram_0.mrc \
  --tomo1 /path/to/tomogram_1.mrc \
  --angles /path/to/angles.tlt \
  --save-dir /path/to/save/dir \
  --batch-size 8
  --device 0 	# Optional (defaults to 0)

This will train the model using the two tomograms tomogram_0.mrc and tomogram_1.mrc with tilt angles specified in angles.tlt. The reconstructions will be saved in the directory /path/to/save/dir along with the 'config.json' file containing the training and model parameters. The actual model files will be saved in /path/to/save/dir/models. The training batch size is set to 8. You can change other training parameters like the number of epochs, eq_weight, etc.

Using a pre-trained model for initialization

You can optionally initialize the weights of the model using a pre-trained model. This can help in faster convergence. To do this, use the --pretrain-path option to specify the path to the pre-trained model file (.pt file). Note that the pre-trained model should be compatible with the current model architecture.

icecream train \
  --tomo0 /path/to/tomogram_0.mrc \
  --tomo1 /path/to/tomogram_1.mrc \
  --angles /path/to/angles.tlt \
  --save-dir /path/to/save/dir \
  --batch-size 8 \
  --pretrain-path /path/to/pretrained_model.pt

To train the model using a config file, create a YAML file with the contents similar to src/icecream/defaults.yaml and pass it to the --config option.

icecream train --config /path/to/config.yaml

In your config file, make sure to update the data section with the correct tomogram paths and angle information. For example:

data:
  save_dir: "runs/my_experiment/"   # directory to save outputs
  tomo0: 
    - "/path/to/tomogram_0.mrc"  # first tomogram
  tomo1: 
    - "/path/to/tomogram_1.mrc"  # second tomogram
  mask: null  # optional mask file

  # Option 1: specify tilt range
  # tilt_max: 57.1
  # tilt_min: -60.0

  # Option 2: specify an external angle file (recommended)
  angles: "/path/to/angles.tlt"

Note: The train command also reconstructs the volume after training. However, you can use the predict command to reconstruct the volume from a trained model using intermediate checkpoints.

To reconstruct the volume from a trained model, use the predict command. You can specify the epoch of the model to use for reconstruction. If not specified, it will use the latest epoch.

icecream predict --config /path/to/config.json \
  --tomo0 /path/to/tomogram_0.mrc \
  --tomo1 /path/to/tomogram_1.mrc \
  --angles /path/to/angles.tlt \
  --save_dir /path/to/save/dir \  # directory to save output if different from training (optional)
  --iter_load 50000  # iteration to use for reconstruction (optional)

Note that the config.json file is generated during training and contains the model and training parameters and saved in the training save directory. You can also use a YAML config file instead of JSON.

License

This project is licensed under the MIT License - see the LICENSE file for details