rln41 particle visualisation for ricardo

ricardo.py

from enum import Enum

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

PARTICLE_STAR_FILE = 'ricardo_data.star'
UNBINNED_PIXEL_SIZE = 1.7

viewer = napari.Viewer(ndisplay=3)

star = starfile.read(PARTICLE_STAR_FILE)
df = star['particles']


tomogram_names = df['rlnTomoName'].unique()
Tomogram = Enum('Tomogram', ' '.join(tomogram_names))


def get_napari_points(df: pd.DataFrame, unbinned_pixel_size: float):
    xyz = df[['rlnCoordinateX', 'rlnCoordinateY', 'rlnCoordinateZ']].to_numpy()
    if 'rlnOriginXAngst' in df.columns:
        shifts = df[['rlnOriginXAngst', 'rlnOriginYAngst', 'rlnOriginZAngst']].to_numpy()
        xyz -= shifts / unbinned_pixel_size
    return xyz[:, ::-1]


def get_napari_vectors(
    df: pd.DataFrame, initial_vector: np.ndarray, unbinned_pixel_size: float
):
    points = get_napari_points(df, unbinned_pixel_size)
    euler_angles = df[['rlnAngleRot', 'rlnAngleTilt', 'rlnAnglePsi']].to_numpy()
    rotations = R.from_euler(seq='ZYZ', angles=euler_angles, degrees=True).inv()
    rotated_vectors = rotations.apply(initial_vector)
    vectors = np.zeros((rotated_vectors.shape[0], 2, 3))
    vectors[:, 0, :] = points
    vectors[:, 1, :] = rotated_vectors[:, ::-1]
    return vectors

@magicgui(auto_call=True)
def select_tomogram(tomogram: Tomogram):
    idx = df['rlnTomoName'] == tomogram.name
    subset = df.loc[idx]
    points = get_napari_points(subset, UNBINNED_PIXEL_SIZE)
    z_vectors = get_napari_vectors(subset, [0, 0, 1], UNBINNED_PIXEL_SIZE)
    x_vectors = get_napari_vectors(subset, [1, 0, 0], UNBINNED_PIXEL_SIZE)
    viewer.layers.clear()
    viewer.add_points(points, size=20)
    viewer.add_vectors(z_vectors,  length=80, edge_width=5, edge_color='cornflowerblue')
    viewer.add_vectors(x_vectors, length=80, edge_width=5, edge_color='orange')


viewer.window.add_dock_widget(select_tomogram, area='left')
select_tomogram()
napari.run()