projection matrix generation from AreTomo aln files for Zhen

zhen.py

import pandas as pd
import einops
import numpy as np


def Rx(angles_degrees: np.ndarray) -> np.ndarray:
    """Affine matrix for a rotation around the X-axis."""
    angles_degrees = np.asarray(angles_degrees).reshape(-1)
    c = np.cos(np.deg2rad(angles_degrees))
    s = np.sin(np.deg2rad(angles_degrees))
    matrices = einops.repeat(
        np.eye(4), 'i j -> n i j', n=len(angles_degrees)
    )
    matrices[:, 1, 1] = c
    matrices[:, 1, 2] = -s
    matrices[:, 2, 1] = s
    matrices[:, 2, 2] = c
    return np.squeeze(matrices)


def Ry(angles_degrees: np.ndarray) -> np.ndarray:
    """Affine matrix for a rotation around the Y-axis."""
    angles_degrees = np.asarray(angles_degrees).reshape(-1)
    c = np.cos(np.deg2rad(angles_degrees))
    s = np.sin(np.deg2rad(angles_degrees))
    matrices = einops.repeat(
        np.eye(4), 'i j -> n i j', n=len(angles_degrees)
    )
    matrices[:, 0, 0] = c
    matrices[:, 0, 2] = s
    matrices[:, 2, 0] = -s
    matrices[:, 2, 2] = c
    return np.squeeze(matrices)


def Rz(angles_degrees: float) -> np.ndarray:
    """Affine matrix for a rotation around the Z-axis."""
    angle_degrees = np.asarray(angles_degrees).reshape(-1)
    c = np.cos(np.deg2rad(angle_degrees))
    s = np.sin(np.deg2rad(angle_degrees))
    matrices = einops.repeat(
        np.eye(4), 'i j -> n i j', n=len(angle_degrees)
    )
    matrices[:, 0, 0] = c
    matrices[:, 0, 1] = -s
    matrices[:, 1, 0] = s
    matrices[:, 1, 1] = c
    return np.squeeze(matrices)


def S(shifts: np.ndarray) -> np.ndarray:
    """Affine matrices for shifts.
    Shifts supplied can be 2D or 3D.
    """
    shifts = np.asarray(shifts, dtype=float)
    if shifts.shape[-1] == 2:
        shifts = promote_2d_to_3d(shifts)
    shifts = np.array(shifts).reshape((-1, 3))
    matrices = einops.repeat(np.eye(4), 'i j -> n i j', n=shifts.shape[0])
    matrices[:, 0:3, 3] = shifts
    return np.squeeze(matrices)


def promote_2d_to_3d(shifts: np.ndarray) -> np.ndarray:
    """Promote 2D vectors to 3D with zeros in the last dimension."""
    shifts = np.asarray(shifts).reshape(-1, 2)
    shifts = np.c_[shifts, np.zeros(shifts.shape[0])]
    return np.squeeze(shifts)
def read_aln(filename: os.PathLike) -> pd.DataFrame:
    """Read an AreTomo .aln file"""
    df = pd.read_csv(filename, header='infer', skiprows=2, delimiter=r'\s+')

    # '#' character in header line is parsed as a column name
    # drop empty column on the far right and shift column names to the left by 1
    column_names = list(df.columns)
    df.drop(df.columns[len(df.columns) - 1], axis=1, inplace=True)
    df.columns = column_names[1:]
    return df

def get_specimen_shifts(aln_file: Path) -> np.ndarray:
    """Get specimen shifts from AreTomo alignments file."""
    df = read_aln(aln_file)
    return np.array(df[['TX', 'TY']])


def get_xyz_extrinsic_euler_angles(aln_file: Path) -> np.ndarray:
    """Get XYZ-extrinsic Euler angles froom AreTomo alignments file."""
    df = read_aln(aln_file)
    n_images = len(df)
    euler_angles = np.empty(shape=(n_images, 3))
    euler_angles[:, 0] = 0
    euler_angles[:, 1] = df['TILT']
    euler_angles[:, 2] = df['ROT']
    return euler_angles

def tilt_series_alignment_parameters_to_relion_projection_matrices(
        specimen_shifts: pd.DataFrame,
        euler_angles: pd.DataFrame,
        tilt_image_dimensions: np.ndarray,
        tomogram_dimensions: np.ndarray,
):
    """Generate affine matrices transforming points in 3D to 2D in tilt-images.
    Projection model:
    3D specimen is rotated about its center then translated such that the projection
    of points onto the XY-plane gives their position in a tilt-image.
    More specifically
    - 3D specimen is rotated about its center by
        - shifting the origin to the specimen center
        - rotated extrinsically about the Y-axis by the tilt angle
        - rotated extrinsically about the Z-axis by the in plane rotation angle
    - 3D specimen is translated in the camera to align coordinate system with tilt-image
        - move center-of-rotation of specimen to center of tilt-image
        - move center-of-rotation of specimen to rotation center in tilt-image
    Parameters
    ----------
    specimen_shifts: XY-shifts which align the projected specimen with tilt-images
    euler_angles: YZX intrinsic Euler angles which transform the specimen
    tilt_image_dimensions: XY-dimensions of tilt-series.
    tomogram_dimensions: size of tomogram in XYZ
    """
    tilt_image_center = tilt_image_dimensions / 2
    specimen_center = tomogram_dimensions / 2

    # Transformations, defined in order of application
    s0 = S(-specimen_center)  # put specimen center-of-rotation at the origin
    r0 = Rx(euler_angles['rlnTomoXTilt'])  # rotate specimen around X-axis
    r1 = Ry(euler_angles['rlnTomoYTilt'])  # rotate specimen around Y-axis
    r2 = Rz(euler_angles['rlnTomoZRot'])  # rotate specimen around Z-axis
    s1 = S(specimen_shifts)  # shift projected specimen in xy (camera) plane
    s2 = S(tilt_image_center)  # move specimen back into tilt-image coordinate system

    # compose matrices
    transformations = s2 @ s1 @ r2 @ r1 @ r0 @ s0
    return np.squeeze(transformations)