tomo_vis_napari_demo.py

tomo_vis_napari_demo.py

import numpy as np
import mrcfile
import starfile
from eulerangles import euler2matrix, invert_rotation_matrices
import napari
from fuzzywuzzy import fuzz, process

# Read in data
particle_file = 'particles_10.00Apx.star'
volume_file = 'TS_01.mrc_10.00Apx.mrc'

all_particles = starfile.read(particle_file)
volume = mrcfile.open(volume_file).data

## Isolate the subset of particles which correspond to the volume
# get all volume names from the STAR file
choices = list(all_particles['rlnMicrographName'].unique())

# fuzzy match the strings, find the closest match
matched_volume, score = process.extractOne(volume_file, choices, scorer=fuzz.token_set_ratio)

# subset the dataframe
idx = all_particles['rlnMicrographName'] == matched_volume
particles = all_particles[idx]


## Get info from dataframe
# get euler angles
euler_headings = [f'rlnAngle{angle}' for angle in ('Rot', 'Tilt', 'Psi')]
eulers = particles[euler_headings].to_numpy()

# get coords
coord_headings = [f'rlnCoordinate{axis}' for axis in 'XYZ']
coords = particles[coord_headings].to_numpy()

# Calculate rotation matrices from euler angles
rotation_matrices = euler2matrix(eulers, axes='zyz', intrinsic=True, right_handed_rotation=True)
rotation_matrices = invert_rotation_matrices(rotation_matrices)

# calculate rotated unit z vectors (must be column vector)
unit_z = np.array([0, 0, 1]).reshape((3, 1))
rotated_unit_z = rotation_matrices @ unit_z

## Set up napari visualisation

# imaging data axis order is zyx, our coords are xyz, reorder coords and vectors
coords_napari = coords[:, ::-1]
rotated_unit_z_napari = rotated_unit_z.squeeze()[:, ::-1]

# set up vectors layer - https://napari.org/tutorials/fundamentals/vectors.html
vec_shape = (coords.shape[0], 2, 3)
vectors_napari = np.empty(vec_shape)

vectors_napari[:, 0, :] = coords_napari
vectors_napari[:, 1, :] = rotated_unit_z_napari

with napari.gui_qt():
    viewer = napari.Viewer(ndisplay=3)

    vol_layer = viewer.add_image(volume)

    pos_layer = viewer.add_points(coords_napari, size=4.5, face_color='k')
    pos_layer.opacity = 0.5

    vec_layer = viewer.add_vectors(vectors_napari, length=10, edge_color='cornflowerblue')
    vec_layer.opacity = 0.5