visualise projected positions from RELION tilt-series alignments

visualise_projected_points.py

import mrcfile
import numpy as np
import starfile
import napari

positions_xyz = np.loadtxt('picking/test_recon_noflipyz_noflipz.csv', delimiter=',') * 8
positions_xyz = np.c_[positions_xyz, np.ones(positions_xyz.shape[0])]  # add column of ones
positions_xyz = positions_xyz.reshape((-1, 1, 4, 1))
star = starfile.read('ImportTomo/job001/tomograms.star')
df = star['TS_01']

get_column = lambda x: [float(num) for num in x[1:-1].split(',')]

proj_x = np.stack(df['rlnTomoProjX'].apply(get_column))
proj_y = np.stack(df['rlnTomoProjY'].apply(get_column))
proj_z = np.stack(df['rlnTomoProjZ'].apply(get_column))
proj_w = np.stack(df['rlnTomoProjW'].apply(get_column))

proj = np.stack((proj_x, proj_y, proj_z, proj_w), axis=-1)
proj = np.transpose(proj, axes=(0, 2, 1))  # from documentation, star entries are rows of projection matrices
assert np.allclose(proj[:, 0, :], proj_x)

projected_positions = np.squeeze(proj @ positions_xyz)[:, :, :2]  # take xy coords
idx = np.arange(41)

visualised_positions = np.zeros((len(positions_xyz), 41, 3))
visualised_positions[:, :, 0] = idx
visualised_positions[:, :, 1:] = projected_positions[..., ::-1]  #yx for visualisation
visualised_positions = visualised_positions.reshape((-1, 3))

tilt_series = mrcfile.open('imod_alignments/my_cool_data_TS_01/my_cool_data_TS_01.mrc').data
viewer = napari.Viewer()
viewer.add_image(tilt_series)
viewer.add_points(visualised_positions, size=100)

napari.run()