particle vis for paris - to run do "uv run https://gist.githubusercontent.com/alisterburt/8f8858dbbb4bc11f62b21df34e3cbeee/raw/d78ff3730e97aae2662bc14d032fc4addb7c43d8/warp_particle_vis.py"

warp_particle_vis.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "numpy",
#     "scipy",
#     "starfile",
#     "napari[pyqt5]",
#     "magicgui",
#     "typer",
# ]
# [tool.uv]
# exclude-newer = "2025-01-01T00:00:00Z"
# ///

from enum import Enum
from pathlib import Path

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

def main(star: Path = typer.Option(..., '--input-star-file', '-i')):
    # read file
    df = starfile.read(star)

    # construct enum containing tomogram IDs
    tomogram_ids = df['rlnMicrographName'].unique().tolist()
    TomogramEnum = Enum('TomogramEnum', {Path(id).stem: id for id in tomogram_ids})

    # construct viewer
    viewer = napari.Viewer(ndisplay=3)

    # write function to load particles to gui when tomogram ID changes
    @magicgui(auto_call=True)
    def _add_particle_data(tomogram: TomogramEnum):
        # subset df for just this tomogram
        df_subset = df[df['rlnMicrographName'] == tomogram.value]

        # extract info from dataframe
        xyz = df_subset[['rlnCoordinateX', 'rlnCoordinateY', 'rlnCoordinateZ']].to_numpy()
        eulers = df_subset[['rlnAngleRot', 'rlnAngleTilt', 'rlnAnglePsi']].to_numpy()

        # convert eulers to rotation matrix
        rotation_matrices = R.from_euler(angles=eulers, seq='ZYZ', degrees=True).inv().as_matrix()

        # grab vectors from columns of rotation matrix
        z_vectors = rotation_matrices[:, :, 2]
        y_vectors = rotation_matrices[:, :, 1]
        x_vectors = rotation_matrices[:, :, 0]

        # construct napari layers
        points_layer = napari.layers.Points(xyz, size=1, face_color='cornflowerblue', name="points")
        z_vectors = napari.layers.Vectors(np.stack([xyz, z_vectors], axis=-2), length=3, edge_width=0.2, name="z vectors",
                                          edge_color='cornflowerblue')
        y_vectors = napari.layers.Vectors(np.stack([xyz, y_vectors], axis=-2), length=0.4, edge_width=0.2, name="y vectors",
                                          edge_color='red')
        x_vectors = napari.layers.Vectors(np.stack([xyz, x_vectors], axis=-2), length=0.4, edge_width=0.2, name="x vectors",
                                          edge_color='green')

        # add layers to viewer
        if "points" in viewer.layers:
            viewer.layers["points"].data = points_layer.data
        else:
            viewer.add_layer(points_layer)

        if "z vectors" in viewer.layers:
            viewer.layers["z vectors"].data = z_vectors.data
        else:
            viewer.add_layer(z_vectors)

        if "y vectors" in viewer.layers:
            viewer.layers["y vectors"].data = y_vectors.data
        else:
            viewer.add_layer(y_vectors)

        if "x vectors" in viewer.layers:
            viewer.layers["x vectors"].data = x_vectors.data
        else:
            viewer.add_layer(x_vectors)

        # reset camera
        viewer.reset_view()

        return

    # add widget to GUI
    viewer.window.add_dock_widget(_add_particle_data, area='left')

    # call func once to populate viewer with first tomogram particles on load
    _add_particle_data()

    # launch napari
    napari.run()

if __name__ == "__main__":
    typer.run(main)