Interpret/visualise composition of prior and pose in RELION

eric_bogdan_thing.py

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

df = starfile.read('particles_warp.star')

xyz = df[['rlnCoordinateX', 'rlnCoordinateY', 'rlnCoordinateZ']].to_numpy()
eulers = df[['rlnAngleTiltPrior', 'rlnAnglePsiPrior']].to_numpy()
rotation_matrices = R.from_euler(seq='YZ', degrees=True, angles=eulers).inv().as_matrix()
print(rotation_matrices.shape)


vec_shape = (xyz.shape[0], 2, 3)

z_vectors_napari = np.empty(vec_shape)
z_vectors_napari[:, 0, :] = xyz
z_vectors_napari[:, 1, :] = rotation_matrices[:, :, 2]

y_vectors_napari = np.empty(vec_shape)
y_vectors_napari[:, 0, :] = xyz
y_vectors_napari[:, 1, :] = rotation_matrices[:, :, 1]

# debug vis, check everything looks ~correct
# viewer = napari.Viewer(ndisplay=3)
# viewer.add_vectors(z_vectors_napari, length=10, edge_color='red')
# viewer.add_vectors(y_vectors_napari, length=10, edge_color='blue')
# napari.run()

# turn into 'new weird RELION rotated basis for filaments thing'
rotated_basis = R.from_euler('y', angles=90, degrees=True).as_matrix()
rotated_orientations = rotation_matrices @ rotated_basis
rotated_eulers = R.from_matrix(rotated_orientations).inv().as_euler(seq='ZYZ', degrees=True)

rot_prior, tilt_prior, psi_prior = R.from_matrix(rotated_basis).as_euler(seq='ZYZ', degrees=True)
df['rlnAngleRot'] = [rot_prior] * len(df)
df['rlnAngleTilt'] = [tilt_prior] * len(df)
df['rlnAnglePsi'] = [psi_prior] * len(df)
df['rlnAngleTiltPrior'] = [tilt_prior] * len(df)
df['rlnAnglePsiPrior'] = [psi_prior] * len(df)
df['rlnTomoSubtomogramRot'] = rotated_eulers[:, 0]
df['rlnTomoSubtomogramTilt'] = rotated_eulers[:, 1]
df['rlnTomoSubtomogramPsi'] = rotated_eulers[:, 2]

# this is how we compose subtomogramRot/Tilt/Psi with Rot/Tilt/PsiPrior
composed_rotations = rotated_orientations @ np.linalg.inv(rotated_basis)

z_vectors_napari = np.empty(vec_shape)
z_vectors_napari[:, 0, :] = xyz
z_vectors_napari[:, 1, :] = composed_rotations[:, :, 2]

y_vectors_napari = np.empty(vec_shape)
y_vectors_napari[:, 0, :] = xyz
y_vectors_napari[:, 1, :] = composed_rotations[:, :, 1]

x_vectors_napari = np.empty(vec_shape)
x_vectors_napari[:, 0, :] = xyz
x_vectors_napari[:, 1, :] = composed_rotations[:, :, 0]


# debug vis, check everything looks ~correct
# these should be correctly oriented in the tomogram, particle z along filament axis
viewer = napari.Viewer(ndisplay=3)
viewer.add_vectors(z_vectors_napari, length=10, edge_color='red')
viewer.add_vectors(y_vectors_napari, length=4, edge_color='blue')
viewer.add_vectors(x_vectors_napari, length=4, edge_color='green')
napari.run()


starfile.write(df, 'output.star')