recenter 3D particles in RELION 3.1+ style STAR files

recenter_3d.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "pandas",
#     "scipy",
#     "starfile",
#     "typer",
#     "einops",
#     "rich",
# ]
# [tool.uv]
# exclude-newer = "2025-01-01T00:00:00Z"
# ///


from pathlib import Path
from typing import Tuple

import einops
import numpy as np
import rich
import starfile
import typer
from scipy.spatial.transform import Rotation as R

console = rich.console.Console()


def cli(
    input_star_file: Path = typer.Option(..., '--input', '-i', help="input star file"),
    shift: tuple[float, float, float] = typer.Option(..., '--shift', '-s', help="shift x, y and z"),
    output_star_file: Path = typer.Option(..., '--output', '-o', help="output star file"),
):
    star = starfile.read(input_star_file, always_dict=True)
    console.log(f"{input_star_file} read")

    if not all(key in star for key in ('particles', 'optics')):
        console.log("expected RELION 3.1+ style STAR file containing particles and optics blocks", style="bold red")

    df = star['particles'].merge(star['optics'], on='rlnOpticsGroup')
    console.log("optics table merged")
    console.log(f"{len(df)} particles found")

    # get relevant info from star file as numpy arrays
    console.log('grabbing relevant info...')

    xyz = df[['rlnCoordinateX', 'rlnCoordinateY', 'rlnCoordinateZ']].to_numpy()
    console.log("got shifts from 'rlnCoordinateX', 'rlnCoordinateY', 'rlnCoordinateZ'")

    pixel_spacing = df['rlnImagePixelSize'].to_numpy()
    console.log("got pixel spacing from 'rlnImagePixelSize'")

    euler_angles = df[['rlnAngleRot', 'rlnAngleTilt', 'rlnAnglePsi']].to_numpy()
    console.log("got euler angles from 'rlnAngleRot', 'rlnAngleTilt', 'rlnAnglePsi'")

    if all(col in df.columns for col in ['rlnOriginXAngst', 'rlnOriginYAngst', 'rlnOriginZAngst']):
        shifts = df[['rlnOriginXAngst', 'rlnOriginYAngst', 'rlnOriginZAngst']].to_numpy()
        console.log("got shifts from 'rlnOriginXAngst', 'rlnOriginYAngst', 'rlnOriginZAngst'")
    else:
        shifts = np.zeros(shape=(3,))
        console.log("no shifts found in 'rlnOriginXAngst', 'rlnOriginYAngst', 'rlnOriginZAngst', setting to 0")

    # convert shifts to angstrom then apply shifts to calculate absolute particle position
    pixel_spacing = einops.rearrange(pixel_spacing, 'b -> b 1')
    shifts = shifts / pixel_spacing
    console.log("converted shifts to angstroms")
    xyz -= shifts
    console.log("applied shifts to particle positions")

    # get particle rotation matrices (column vectors are particle x/y/z in tomogram)
    rotation_matrices = R.from_euler(angles=euler_angles, seq='ZYZ', degrees=True).inv().as_matrix()
    console.log("calculated rotation matrices from euler angles")

    # recenter particles, we don't care about orientations so apply identity rotation
    new_xyz, _ = shift_then_rotate_particles(
        particle_positions=xyz,
        particle_orientations=rotation_matrices,
        shift=np.asarray(shift),
        rotation=np.eye(3),
    )
    console.log('calculated shifted particle positions')

    # express new positions relative to old positions in star file
    new_shifts = new_xyz - df[['rlnCoordinateX', 'rlnCoordinateY', 'rlnCoordinateZ']].to_numpy()
    new_shifts = -1 * new_shifts * pixel_spacing
    console.log("calculated new shifts from shifted particle positions")

    star['particles'][['rlnOriginXAngst', 'rlnOriginYAngst', 'rlnOriginZAngst']] = new_shifts
    console.log("updated shift values in 'rlnOriginXAngst', 'rlnOriginYAngst', 'rlnOriginZAngst'")

    # write output
    with console.status(f"writing output STAR file {output_star_file}", spinner="arc"):
        starfile.write(star, output_star_file)
    console.log(f"Output with updated shifts written to {output_star_file}")


def shift_then_rotate_particles(
    particle_positions,  # (n, 3)
    particle_orientations,  # (n, 3, 3)
    shift,  # (3, )
    rotation,  # (3, 3)
) -> Tuple[np.ndarray, np.ndarray]:  # positions, orientations
    # goal: apply transformations in the local coordinate
    # system of each particle

    # transform the shifts into the local particle reference frame
    shift = einops.rearrange(shift, 'xyz -> xyz 1')
    local_shifts = particle_orientations @ shift
    local_shifts = einops.rearrange(local_shifts, 'b xyz 1 -> b xyz')

    # apply the shifts
    updated_particle_positions = particle_positions + local_shifts

    # transform the reference rotation to find the new particle orientation
    particle_orientations = particle_orientations @ rotation

    return updated_particle_positions, particle_orientations


if __name__ == "__main__":
    app = typer.Typer(add_completion=False)
    app.command(no_args_is_help=True)(cli)
    app()