recenter particles r3.1+ for bryce

recenter.py

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


from pathlib import Path
from typing import Tuple

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


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)
    df = star['particles'].merge(star['optics'], on='rlnOpticsGroup')

    # get positions and shifts
    xy = df[['rlnCoordinateX', 'rlnCoordinateY']].to_numpy()
    shifts = df[['rlnOriginXAngst', 'rlnOriginYAngst']].to_numpy()
    pixel_spacing = df['rlnImagePixelSize'].to_numpy()
    shifts = shifts / pixel_spacing.reshape((-1, 1))
    xy -= shifts

    # get particle rotation matrices
    euler_angles = df[['rlnAngleRot', 'rlnAngleTilt', 'rlnAnglePsi']].to_numpy()
    rotation_matrices = R.from_euler(angles=euler_angles, seq='ZYZ', degrees=True).inv().as_matrix()

    # make pseudo-3d positions
    xyz = np.pad(xy, ((0, 0), (0, 1)), mode='constant')

    # recenter particles, we don't care about orientations
    new_positions, _ = shift_then_rotate_particles(
        particle_positions=xyz,
        particle_orientations=rotation_matrices,
        shift=np.asarray(shift),
        rotation=np.eye(3),
    )

    new_xy = new_positions[:, :2]

    # express new positions relative to old positions in star file
    new_shifts_xy_px = new_xy - df[['rlnCoordinateX', 'rlnCoordinateY']].to_numpy()
    new_shifts_xy_ang = -1 * new_shifts_xy_px * pixel_spacing.reshape((-1, 1))
    star['particles'][['rlnOriginXAngst', 'rlnOriginYAngst']] = new_shifts_xy_ang

    # write output
    starfile.write(star, 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()