transform classes in a rln3.1 STAR file based on manipulations in ChimeraX

realign_classes.py

from typing import Tuple

import starfile
import numpy as np
import einops

from scipy.spatial.transform import Rotation as R

CHIMERAX_VIEW_MATRIX_OUTPUT_FILE = 'chimerax/view_matrix_output_origin_at_rotation_center.txt'
PARTICLE_STAR_FILE = 'class3d_dir/run_it025_data.star'
OUTPUT_STAR_FILE = 'output_realigned_data.star'
VOLUME_VOXEL_SPACING = 3.3
MODEL_TO_CLASS = {
    2: 1,
    3: 2,
    4: 3,
    5: 4,
    6: 5,
    7: 6,
}


def read_chimerax_matrices(output_text_file):
    """read 4x4 matrices from cx 'view matrix' output."""
    with open(output_text_file) as f:
        [matrix_line] = [
            line.strip()
            for line in f.readlines()
            if line.startswith('model positions:')
        ]
    matrix_line = matrix_line.strip('model positions: ')
    line_per_model = matrix_line.split('#')

    # submodels always identity transform so skip
    # assumes less than 9 models present
    model_to_matrix = {}
    for entry in line_per_model:
        if entry == '' or entry[1] == '.':
            continue
        model_number = int(entry[0])
        matrix_data = entry[2:].split(',')
        matrix_data = np.array([float(char) for char in matrix_data if char != ''])
        model_to_matrix[model_number] = matrix_data.reshape((3, 4))

    return model_to_matrix


def Rx(theta: np.ndarray) -> np.ndarray:
    """Generate 4x4 matrices which rotate vectors around the X-axis by the angle `theta`.
    The rotation angle `theta` is expected to be in degrees.
    Application of these matrices is performed by left-multiplying xyzw column vectors.
    Parameters
    ----------
    theta: np.ndarray
        An `(n, )` array of rotation angles in degrees.
    Returns
    -------
    matrices: np.ndarray
        An `(n, 4, 4)` array of matrices which left-multiply xyzw column vectors.
    """
    theta = np.asarray(theta).reshape(-1)
    c = np.cos(np.deg2rad(theta))
    s = np.sin(np.deg2rad(theta))
    matrices = einops.repeat(
        np.eye(4), 'i j -> n i j', n=len(theta)
    )
    matrices[:, 1, 1] = c
    matrices[:, 1, 2] = -s
    matrices[:, 2, 1] = s
    matrices[:, 2, 2] = c
    return matrices


def Ry(theta: np.ndarray) -> np.ndarray:
    """Generate 4x4 matrices which rotate vectors around the Y-axis by the angle `theta`.
    The rotation angle `theta` is expected to be in degrees.
    Application of these matrices is performed by left-multiplying xyzw column vectors.
    Parameters
    ----------
    theta: np.ndarray
        An `(n, )` array of rotation angles in degrees.
    Returns
    -------
    matrices: np.ndarray
        An `(n, 4, 4)` array of matrices which left-multiply xyzw column vectors."""
    theta = np.asarray(theta).reshape(-1)
    c = np.cos(np.deg2rad(theta))
    s = np.sin(np.deg2rad(theta))
    matrices = einops.repeat(np.eye(4), 'i j -> n i j', n=len(theta))
    matrices[:, 0, 0] = c
    matrices[:, 0, 2] = s
    matrices[:, 2, 0] = -s
    matrices[:, 2, 2] = c
    return matrices


def Rz(theta: float) -> np.ndarray:
    """Generate 4x4 matrices which rotate vectors around the Z-axis by the angle `theta`.
    The rotation angle `theta` is expected to be in degrees.
    Application of these matrices is performed by left-multiplying xyzw column vectors.
    Parameters
    ----------
    theta: np.ndarray
        An `(n, )` array of rotation angles in degrees.
    Returns
    -------
    matrices: np.ndarray
        An `(n, 4, 4)` array of matrices which left-multiply xyzw column vectors."""
    theta = np.asarray(theta).reshape(-1)
    c = np.cos(np.deg2rad(theta))
    s = np.sin(np.deg2rad(theta))
    matrices = einops.repeat(
        np.eye(4), 'i j -> n i j', n=len(theta)
    )
    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:
    """Generate 4x4 matrices for 2D (xy) or 3D (xyz) shifts.
    Application of these matrices is performed by left-multiplying xyzw column vectors.
    Parameters
    ----------
    shifts: np.ndarray
        An `(n, 2)` or `(n, 3)` array of xy(z) shifts.
    Returns
    -------
    matrices: np.ndarray
        An `(n, 4, 4)` array of
    """
    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(len(shifts))]
    return np.squeeze(shifts)


def decompose_chimerax_matrix_for_relion(view_matrix):
    """Convert a chimerax matrix to a set of shifts and euler angles for particles."""
    # we need to decompose the chimerax matrix in a way that we can then use to construct RELION
    # transformations (shift, then rot around new origin, origin is at particle center)

    # get chimerax matrix and embed in 4x4
    CX = np.eye(4)
    CX[:3, :] = view_matrix  # (3, 4)
    # T[:3, 3] = [0, 0, 0]  # TEST KNOWN SHIFTS

    # okay, now construct the thing you need to do to each particle in RELION...
    T = np.linalg.pinv(CX)

    # now decompose transformation into matrices for the rotation, then the shift
    TR = np.eye(4)
    TR[:3, :3] = T[:3, :3]
    TS = S(T[:3, 3])
    assert np.allclose(T, TS @ TR)

    # in relion, we need to first apply shifts, then rotations
    TS = np.linalg.pinv(TR) @ TS
    TS[:3, :3] = np.eye(3)
    assert np.allclose(T, TR @ TS)

    particle_shifts_ang = TS[:3, 3]  # (3, )
    particle_rotation = TR[:3, :3]  # (3, 3)
    assert particle_shifts_ang.ndim == 1

    particle_shifts_ang = einops.rearrange(particle_shifts_ang, 'xyz -> xyz 1')
    particle_shifts_ang = particle_rotation @ particle_shifts_ang
    particle_shifts_ang = einops.rearrange(particle_shifts_ang, 'xyz 1 -> xyz')
    return particle_shifts_ang, particle_rotation


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


def get_extraction_xyz(df, idx):
    headings = (
        'rlnCoordinateX',
        'rlnCoordinateY',
        'rlnCoordinateZ'
    )
    return df.loc[idx, headings].to_numpy()


def get_relion_shifts(df, idx):
    headings = (
        'rlnOriginXAngst',
        'rlnOriginYAngst',
        'rlnOriginZAngst'
    )
    return df.loc[idx, headings].to_numpy()


def get_eulers(df, idx):
    headings = (
        'rlnAngleRot',
        'rlnAngleTilt',
        'rlnAnglePsi'
    )
    return df.loc[idx, headings].to_numpy()


def get_active_rotation_matrices(df, idx):
    zyz_intrinsic_euler_angles = get_eulers(df, idx)
    rotations = R.from_euler(
        seq='ZYZ',
        angles=zyz_intrinsic_euler_angles,
        degrees=True
    ).inv()  # passive -> active
    return rotations.as_matrix()


def get_particle_pose(df, idx):
    particle_extraction_positions_px = get_extraction_xyz(df, idx)
    particle_shifts_ang = get_relion_shifts(df, idx)
    particle_shifts_px = particle_shifts_ang / VOLUME_VOXEL_SPACING
    particle_positions_px = particle_extraction_positions_px - particle_shifts_px
    particle_rotation_matrices = get_active_rotation_matrices(df, idx)
    return particle_positions_px, particle_rotation_matrices


def active_rotation_matrix_to_relion_euler_angles(rotation_matrix):
    rotation = R.from_matrix(rotation_matrix).inv()
    return rotation.as_euler(seq='ZYZ', degrees=True)


if __name__ == '__main__':
    star = starfile.read(PARTICLE_STAR_FILE)
    df = star['particles']
    chimerax_view_matrices = read_chimerax_matrices(CHIMERAX_VIEW_MATRIX_OUTPUT_FILE)
    for model_number, chimerax_view_matrix in chimerax_view_matrices.items():
        class_number = MODEL_TO_CLASS[model_number]
        idx = df['rlnClassNumber'] == class_number

        # get particle info for class
        particle_extraction_positions_px = get_extraction_xyz(df, idx)
        particle_positions_px, particle_rotation_matrices = get_particle_pose(df, idx)

        # decompose chimerax matrix and apply shift and rotation
        class_shift_ang, class_rotation = decompose_chimerax_matrix_for_relion(chimerax_view_matrix)
        class_shift_px = class_shift_ang / VOLUME_VOXEL_SPACING

        class_positions_px, class_rotation_matrices = shift_then_rotate_particles(
            particle_positions=particle_positions_px,
            particle_orientations=particle_rotation_matrices,
            shift=class_shift_px,
            rotation=class_rotation
        )
        class_euler_angles = active_rotation_matrix_to_relion_euler_angles(class_rotation_matrices)

        # we don't want to reextract particles
        # instead, find the shifts relative to the original extraction position
        class_shifts_px = class_positions_px - particle_extraction_positions_px
        class_shifts_ang = class_shifts_px * VOLUME_VOXEL_SPACING

        # update dataframe
        df.loc[idx, ('rlnOriginXAngst', 'rlnOriginYAngst', 'rlnOriginZAngst')] = -class_shifts_ang
        df.loc[idx, ('rlnAngleRot', 'rlnAngleTilt', 'rlnAnglePsi')] = class_euler_angles
        starfile.write({'optics': star['optics'], 'particles': df.loc[idx, :]},
                       f'output_class{class_number}_data.star', overwrite=True)
    starfile.write({'optics': star['optics'], 'particles': df}, OUTPUT_STAR_FILE, overwrite=True)

transform_classes_in_chimerax.md

notes

1. open all maps in chimerax
2. set origin of all maps to sidelength / 2
3. reorient any misaligned classes (shift and rotate)
4. save the output of view matrix in ChiemraX into a text file
5. run the above python script, pointing it at the file containing the chimerax output and the particle star file