torch-tiltxcorr 0.0.8

Cross correlation with cosine stretching for cryo-EM data in PyTorch

torch-tiltxcorr

Cross correlation with image stretching for coarse alignment of cryo-EM tilt series data in PyTorch.

Overview

torch-tiltxcorr reimplements the IMOD program tiltxcorr in PyTorch.

Installation

pip install torch-tiltxcorr

Usage

import torch
from torch_fourier_shift import fourier_shift_image_2d
from torch_tiltxcorr import tiltxcorr

# Load or create your tilt series
# tilt_series shape: (batch, height, width) - batch is number of tilt images
# Example: tilt_series with shape (61, 512, 512) - 61 tilt images of 512x512 pixels
tilt_series = torch.randn(61, 512, 512)

# Define tilt angles (in degrees)
# Shape: (batch,) - one angle per tilt image
tilt_angles = torch.linspace(-60, 60, steps=61)

# Define tilt axis angle (in degrees)
tilt_axis_angle = 45

# Run tiltxcorr
shifts = tiltxcorr(
    tilt_series=tilt_series,
    tilt_angles=tilt_angles,
    tilt_axis_angle=tilt_axis_angle,
    low_pass_cutoff=.5,
)
# shifts shape: (batch, 2) - (dy, dx) shifts which center each tilt image

# Apply shifts to align the tilt series
aligned_tilt_series = fourier_shift_image_2d(tilt_series, shifts=shifts)
# aligned_tilt_series shape: (batch, height, width)

Use uv to run an example with simulated data and visualize the results.

uv run examples/tiltxcorr_example_simulated_data.py

How It Works

torch-tiltxcorr performs coarse tilt series alignment by:

1. Sorting images by tilt angle
2. Dividing the series into groups of positive and negative tilt angles
3. For each adjacent pair of images in each group:
   • Applying a stretch perpendicular to the tilt axis on the image with the larger tilt angle
   • Calculating cross-correlation between the images
   • Extracting the shift from the position of the correlation peak
   • Transforming the shift to account for the stretch applied to the image
4. Accumulating shifts to align the entire series

With Pretilt Offset Search

For samples with significant pretilt, use tiltxcorr_with_pretilt_offset to automatically find the optimal pretilt offset:

import torch
from torch_fourier_shift import fourier_shift_image_2d
from torch_tiltxcorr import tiltxcorr_with_pretilt_offset

# Load or create your tilt series
tilt_series = torch.randn(61, 512, 512)
tilt_angles = torch.linspace(-60, 60, steps=61)
tilt_axis_angle = 45

# Run tiltxcorr with pretilt offset search
shifts, optimal_pretilt = tiltxcorr_with_pretilt_offset(
    tilt_series=tilt_series,
    tilt_angles=tilt_angles,
    tilt_axis_angle=tilt_axis_angle,
    low_pass_cutoff=0.5, # cycles/px
    pretilt_range=(-30.0, 30.0),  # search range in degrees
)
# shifts shape: (batch, 2) tensor of (dy, dx) shifts which center each tilt image
# optimal_pretilt: float value optimal pretilt offset in degrees

# Apply shifts to align the tilt series
aligned_tilt_series = fourier_shift_image_2d(tilt_series, shifts=shifts)

License

This package is distributed under the BSD 3-Clause License.