R5 parsing for Barrett

barrett.py

import ast

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

DO_VIS = True

particle_star = starfile.read('matching.star')
tomogram_star = starfile.read('matching_tomograms.star')

particle_df = particle_star['particles'].merge(particle_star['optics'], on='rlnOpticsGroup')

for key, particle_group_df in particle_df.groupby('rlnTomoName'):
    # get (n_particles, 3, 3) particle rotation_matrices
    particle_euler_angles = particle_group_df[['rlnAngleRot', 'rlnAngleTilt', 'rlnAnglePsi']].to_numpy()
    particle_rotation_matrices = R.from_euler(
        seq='ZYZ', angles=particle_euler_angles, degrees=True
    ).inv().as_matrix()

    # get table for specific table
    tomogram_df = tomogram_star[key]

    # get (n_tilts, 4) arrays for first three rows of projection matrices
    projection_matrices_x = np.stack(tomogram_df['rlnTomoProjX'].apply(lambda x: ast.literal_eval(x)))
    projection_matrices_y = np.stack(tomogram_df['rlnTomoProjY'].apply(lambda x: ast.literal_eval(x)))
    projection_matrices_z = np.stack(tomogram_df['rlnTomoProjZ'].apply(lambda x: ast.literal_eval(x)))

    # get (n_tilts, 3, 4) projection_matrices
    tilt_projection_matrices = einops.rearrange(
        [projection_matrices_x, projection_matrices_y, projection_matrices_z],
        pattern='projxyz b d -> b projxyz d'
    )

    # get (n_tilts, 3, 3) rotation matrices
    tilt_rotation_matrices = tilt_projection_matrices[:, :, :-1]

    # get (n_particles, n_tilts, 3, 3) per particle per tilt rotation matrices
    particle_rotation_matrices = einops.rearrange(particle_rotation_matrices, 'b i j -> b 1 i j')
    particle_tilt_rotation_matrices = particle_rotation_matrices @ tilt_rotation_matrices  # (n_particles, n_tilts, 3, 3)

    # get (n_particles * ntilts, 3) euler angles in RELION format
    particle_eulers = R.from_matrix(
        einops.rearrange(particle_tilt_rotation_matrices, 'nparticles ntilts i j -> (nparticles ntilts) i j')
    ).inv().as_euler(seq='ZYZ', degrees=True)
    
    # view first particle fourier planes rotated and unrotated
    if DO_VIS:
        unrotated = tilt_rotation_matrices
        rotated = particle_tilt_rotation_matrices[0]

        xy_plane = np.meshgrid(
            np.linspace(-1, 1, ),
            np.linspace(-1, 1, )
        )
        xy_plane = einops.rearrange(xy_plane, 'd h w -> (h w) d 1')  # (b, 2, 1)
        xy_plane = np.pad(xy_plane, ((0, 0), (0, 1), (0, 0)), mode='constant', constant_values=0)  # (b, 3, 1)

        unrotated = einops.rearrange(unrotated, 'ntilt i j -> ntilt 1 i j')
        unrotated_planes = unrotated @ xy_plane
        unrotated_planes = einops.rearrange(unrotated_planes, 'ntilt npoints d 1 -> (ntilt npoints) d')

        rotated = einops.rearrange(rotated, 'ntilt i j -> ntilt 1 i j')
        rotated_planes = rotated @ xy_plane
        rotated_planes = einops.rearrange(rotated_planes, 'ntilt npoints d 1 -> (ntilt npoints) d')

        import napari
        viewer = napari.Viewer(ndisplay=3)
        viewer.axes.visible = True
        viewer.add_points(unrotated_planes, size=0.01, face_color='grey')
        viewer.add_points(rotated_planes, size=0.01, face_color='purple')

        napari.run()