warp template matching visualisation

warp_tm_vis.py

from enum import Enum
from pathlib import Path

import starfile
import mrcfile
import napari
import numpy as np
from magicgui import magicgui, widgets

RECONSTRUCTION_DIR = '/Users/burta2/Downloads/10499/reconstruction'
MATCHING_DIR = '/Users/burta2/Downloads/10499/matching'
MATCHING_PATTERN = '*_flipx.star'
CORRELATION_VOLUME_PATTERN = '*_flipx_corr.mrc'
LOAD_TOMOGRAMS = True

particle_files = list(Path(MATCHING_DIR).glob(MATCHING_PATTERN))
tomogram_files = list(Path(RECONSTRUCTION_DIR).glob('*.mrc'))
correlation_volume_files = list(Path(MATCHING_DIR).glob(CORRELATION_VOLUME_PATTERN))


# match tomogram files to particles
def find_particles_file(tomogram_file: Path) -> Path | None:
    matching_files = [
        f for f in particle_files
        if f.name.startswith(tomogram_file.stem)
    ]
    if len(matching_files) == 1:
        return matching_files[0]
    else:
        raise RuntimeError('found no files or too many files')


def find_correlation_volume_file(tomogram_file: Path) -> Path | None:
    matching_files = [
        f for f in correlation_volume_files
        if f.name.startswith(tomogram_file.stem)
    ]
    if len(matching_files) == 1:
        return matching_files[0]
    else:
        raise RuntimeError('found no files or too many files')


if LOAD_TOMOGRAMS is True:
    name_to_tomo_file = {file.name: file for file in tomogram_files}
    tomogram_names = list(name_to_tomo_file.keys())
    name_to_particle_file = {file.name: find_particles_file(file) for file in tomogram_files}
    name_to_correlation_volume_file = {file.name: find_correlation_volume_file(file) for file in tomogram_files}

tomogram_files = [str(path) for path in tomogram_files]
Tomogram = Enum('Tomogram', ' '.join(tomogram_files))


def get_particle_positions_and_cc(tomogram_file: str) -> np.ndarray:
    tomogram_name = Path(tomogram_file).name
    df = starfile.read(name_to_particle_file[tomogram_name])
    zyx = df[['rlnCoordinateZ', 'rlnCoordinateY', 'rlnCoordinateX']].to_numpy()
    cc = df['rlnAutopickFigureOfMerit'].to_numpy()
    with mrcfile.open(tomogram_file, header_only=True) as mrc:
        nz, ny, nx = mrc.header.nz, mrc.header.ny, mrc.header.nx
    return zyx * np.array([nz, ny, nx]), cc


viewer = napari.Viewer(ndisplay=3)


@magicgui(auto_call=True)
def add_tomogram(tomogram: Tomogram):
    if LOAD_TOMOGRAMS is True:
        volume = mrcfile.read(tomogram.name)
        correlation_volume_file = find_correlation_volume_file(Path(tomogram.name))
        correlation_volume = mrcfile.read(correlation_volume_file)
    zyx, cc = get_particle_positions_and_cc(tomogram.name)

    if LOAD_TOMOGRAMS is True:
        if 'tomogram' not in viewer.layers:
            viewer.add_image(data=volume, name='tomogram', colormap='gray_r')
        else:
            viewer.layers['tomogram'].data = volume
        if 'correlation_volume' not in viewer.layers:
            viewer.add_image(data=correlation_volume, name='correlation_volume', colormap='inferno', contrast_limits=[2, 10], blending='additive')
        else:
            viewer.layers['correlation_volume'].data = correlation_volume

    if 'particles' not in viewer.layers:
        viewer.add_points(
            zyx,
            name='particles',
            size=4,
            metadata={'positions': zyx, 'cc': cc},
            face_color='orange',
            opacity=0.5,
        )
    else:
        viewer.layers['particles'].data = zyx
        viewer.layers['particles'].metadata['positions'] = zyx
        viewer.layers['particles'].metadata['cc'] = cc


@magicgui(auto_call=True)
def subset_particles(min_cc: float):
    cc = viewer.layers['particles'].metadata['cc']
    zyx = viewer.layers['particles'].metadata['positions']
    zyx = zyx[cc >= min_cc]
    viewer.layers['particles'].data = zyx

viewer.window.add_dock_widget(add_tomogram, area='bottom')
viewer.window.add_dock_widget(subset_particles, area='bottom')
add_tomogram()
napari.run()