particle viewer/selector for Euan/Giulia

viewer.py

from pathlib import Path
from typing import Dict
from enum import Enum

import napari
import numpy as np
import mrcfile
import pandas as pd
import starfile
import eulerangles

from magicgui import magicgui

viewer = napari.Viewer(ndisplay=3)
viewer.text_overlay.visible = True

placeholder_image = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]])
image_layer = viewer.add_image(
    data=placeholder_image,
    metadata={'tomogram_id': 'placeholder'}
)
points_layer = viewer.add_points([], ndim=3)

image_layer.depiction = 'plane'
image_layer.plane.thickness = 10
image_layer.rendering = 'minip'


def reset_view():
    viewer.reset_view()
    image_layer.reset_contrast_limits()
    image_layer.reset_contrast_limits_range()


@magicgui(auto_call=True)
def load_tomogram(file: Path):
    with mrcfile.open(file) as mrc:
        data = mrc.data
    image_layer.data = data
    image_layer.metadata['tomogram_id'] = file.stem
    image_layer.plane.position = np.array(data.shape) // 2
    viewer.text_overlay.text = file.stem
    reset_view()
    load_particles()


def get_positions(star: Dict[str, pd.DataFrame], tomogram_id: str) -> np.ndarray:
    pixel_size = float(star['optics']['rlnImagePixelSize'])
    df = star['particles']
    subset = df[df['rlnTomoName'] == tomogram_id]
    xyz = subset[[f'rlnCoordinate{ax}' for ax in 'XYZ']].to_numpy()
    if 'rlnOriginXAngst' in star['particles'].columns:
        shifts = subset[[f'rlnOrigin{ax}Angst' for ax in 'XYZ']].to_numpy()
        xyz -= shifts / pixel_size
    return xyz


def get_rotation_matrices(star: Dict[str, pd.DataFrame], tomogram_id: str) -> np.ndarray:
    df = star['particles']
    subset = df[df['rlnTomoName'] == tomogram_id]
    euler_angles = subset[[f'rlnAngle{e}' for e in ('Rot', 'Tilt', 'Psi')]].to_numpy()
    rotation_matrices = eulerangles.euler2matrix(
        euler_angles=euler_angles,
        axes='ZYZ',
        intrinsic=True,
        right_handed_rotation=True,
    )
    rotation_matrices = eulerangles.invert_rotation_matrices(rotation_matrices)
    return rotation_matrices


def get_features(star: Dict[str, pd.DataFrame], tomogram_id: str) -> pd.DataFrame:
    df = star['particles']
    subset = df[df['rlnTomoName'] == tomogram_id]
    return subset


@magicgui(auto_call=True)
def load_particles(file: Path, division_factor: float = 1):
    if file.is_dir():
        return
    tomogram_id = image_layer.metadata['tomogram_id']
    star = starfile.read(file)
    xyz = get_positions(star, tomogram_id=tomogram_id)
    rotation_matrices = get_rotation_matrices(star, tomogram_id=tomogram_id)
    features = get_features(star, tomogram_id=tomogram_id)
    zyx = xyz[:, ::-1]
    points_layer.data = zyx / division_factor
    points_layer.features = features


@magicgui(auto_call=True)
def color_by_column(column_name: str):
    if column_name not in points_layer.features.columns:
        return
    points_layer.face_color = column_name


class Operator(Enum):
    equals = np.equal
    greater_than = np.greater
    less_than = np.less


@magicgui(
    auto_call=True,
    value={'min': -1e10, 'max': 1e10}
          )
def show_subset(column_name: str, operator: Operator, value: float = 1):
    if column_name not in points_layer.features.columns:
        return
    idx = operator.value(points_layer.features[column_name], value)
    points_layer.shown = idx

viewer.window.add_dock_widget(load_tomogram)
viewer.window.add_dock_widget(load_particles)
viewer.window.add_dock_widget(color_by_column)
viewer.window.add_dock_widget(show_subset)
napari.run()