Skip to content

Instantly share code, notes, and snippets.

@alexlib
Forked from alisterburt/viewer.py
Created May 12, 2022 18:12
Show Gist options
  • Save alexlib/01c389ee6629b2b6f80f9bba572cfdca to your computer and use it in GitHub Desktop.
Save alexlib/01c389ee6629b2b6f80f9bba572cfdca to your computer and use it in GitHub Desktop.
particle viewer/selector for Euan/Giulia
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment