pytme 0.2.1__cp311-cp311-macosx_14_0_arm64.whl → 0.2.2__cp311-cp311-macosx_14_0_arm64.whl

scripts/postprocess.py

@@ -8,9 +8,8 @@
 import argparse
 from sys import exit
 from os import getcwd
-from os.path import join, abspath
 from typing import List, Tuple
-from os.path import splitext
+from os.path import join, abspath, splitext
 
 import numpy as np
 from numpy.typing import NDArray
@@ -26,6 +25,7 @@ from tme.analyzer import (
 )
 from tme.matching_utils import (
     load_pickle,
+    centered_mask,
     euler_to_rotationmatrix,
     euler_from_rotationmatrix,
 )
@@ -41,9 +41,7 @@ PEAK_CALLERS = {
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Peak Calling for Template Matching Outputs"
-    )
+    parser = argparse.ArgumentParser(description="Analyze Template Matching Outputs")
 
     input_group = parser.add_argument_group("Input")
     output_group = parser.add_argument_group("Output")
@@ -56,6 +54,13 @@ def parse_args():
         nargs="+",
         help="Path to the output of match_template.py.",
     )
+    input_group.add_argument(
+        "--background_file",
+        required=False,
+        nargs="+",
+        help="Path to an output of match_template.py used for normalization. "
+        "For instance from --scramble_phases or a different template.",
+    )
     input_group.add_argument(
         "--target_mask",
         required=False,
@@ -87,7 +92,7 @@ def parse_args():
             "average",
         ],
         default="orientations",
-        help="Available output formats:"
+        help="Available output formats: "
         "orientations (translation, rotation, and score), "
         "alignment (aligned template to target based on orientations), "
         "extraction (extract regions around peaks from targets, i.e. subtomograms), "
@@ -206,6 +211,15 @@ def parse_args():
     elif args.number_of_peaks is None:
         args.number_of_peaks = 1000
 
+    if args.background_file is None:
+        args.background_file = [None]
+    if len(args.background_file) == 1:
+        args.background_file = args.background_file * len(args.input_file)
+    elif len(args.background_file) not in (0, len(args.input_file)):
+        raise ValueError(
+            "--background_file needs to be specified once or for each --input_file."
+        )
+
     return args
 
 
@@ -233,8 +247,8 @@ def load_template(
     return template, center, translation, template_is_density
 
 
-def merge_outputs(data, filepaths: List[str], args):
-    if len(filepaths) == 0:
+def merge_outputs(data, foreground_paths: List[str], background_paths: List[str], args):
+    if len(foreground_paths) == 0:
         return data, 1
 
     if data[0].ndim != data[2].ndim:
@@ -275,8 +289,11 @@ def merge_outputs(data, filepaths: List[str], args):
 
     entities = np.zeros_like(data[0])
     data[0] = _norm_scores(data=data, args=args)
-    for index, filepath in enumerate(filepaths):
-        new_scores = _norm_scores(data=load_pickle(filepath), args=args)
+    for index, filepath in enumerate(foreground_paths):
+        new_scores = _norm_scores(
+            data=load_match_template_output(filepath, background_paths[index]),
+            args=args,
+        )
         indices = new_scores > data[0]
         entities[indices] = index + 1
         data[0][indices] = new_scores[indices]
@@ -284,9 +301,18 @@ def merge_outputs(data, filepaths: List[str], args):
     return data, entities
 
 
+def load_match_template_output(foreground_path, background_path):
+    data = load_pickle(foreground_path)
+    if background_path is not None:
+        data_background = load_pickle(background_path)
+        data[0] = (data[0] - data_background[0]) / (1 - data_background[0])
+        np.fmax(data[0], 0, out=data[0])
+    return data
+
+
 def main():
     args = parse_args()
-    data = load_pickle(args.input_file[0])
+    data = load_match_template_output(args.input_file[0], args.background_file[0])
 
     target_origin, _, sampling_rate, cli_args = data[-1]
 
@@ -326,7 +352,12 @@ def main():
 
     entities = None
     if len(args.input_file) > 1:
-        data, entities = merge_outputs(data=data, filepaths=args.input_file, args=args)
+        data, entities = merge_outputs(
+            data=data,
+            foreground_paths=args.input_file,
+            background_paths=args.background_file,
+            args=args,
+        )
 
     orientations = args.orientations
     if orientations is None:
@@ -339,24 +370,27 @@ def main():
                 target_mask = Density.from_file(args.target_mask)
                 scores = scores * target_mask.data
 
-            if args.n_false_positives is not None:
-                args.n_false_positives = max(args.n_false_positives, 1)
-                cropped_shape = np.subtract(
-                    scores.shape, np.multiply(args.min_boundary_distance, 2)
-                ).astype(int)
+            cropped_shape = np.subtract(
+                scores.shape, np.multiply(args.min_boundary_distance, 2)
+            ).astype(int)
 
-                cropped_shape = tuple(
+            if args.min_boundary_distance > 0:
+                scores = centered_mask(scores, new_shape=cropped_shape)
+
+            if args.n_false_positives is not None:
+                # Rickgauer et al. 2017
+                cropped_slice = tuple(
                     slice(
                         int(args.min_boundary_distance),
                         int(x - args.min_boundary_distance),
                     )
                     for x in scores.shape
                 )
-                # Rickgauer et al. 2017
-                n_correlations = np.size(scores[cropped_shape]) * len(rotation_mapping)
+                args.n_false_positives = max(args.n_false_positives, 1)
+                n_correlations = np.size(scores[cropped_slice]) * len(rotation_mapping)
                 minimum_score = np.multiply(
                     erfcinv(2 * args.n_false_positives / n_correlations),
-                    np.sqrt(2) * np.std(scores[cropped_shape]),
+                    np.sqrt(2) * np.std(scores[cropped_slice]),
                 )
                 print(f"Determined minimum score cutoff: {minimum_score}.")
                 minimum_score = max(minimum_score, 0)
@@ -371,6 +405,8 @@ def main():
                 "min_distance": args.min_distance,
                 "min_boundary_distance": args.min_boundary_distance,
                 "batch_dims": args.batch_dims,
+                "minimum_score": args.minimum_score,
+                "maximum_score": args.maximum_score,
             }
 
             peak_caller = PEAK_CALLERS[args.peak_caller](**peak_caller_kwargs)
@@ -380,7 +416,6 @@ def main():
                 mask=template.data,
                 rotation_mapping=rotation_mapping,
                 rotation_array=rotation_array,
-                minimum_score=args.minimum_score,
             )
             candidates = peak_caller.merge(
                 candidates=[tuple(peak_caller)], **peak_caller_kwargs
@@ -388,10 +423,16 @@ def main():
             if len(candidates) == 0:
                 candidates = [[], [], [], []]
                 print("Found no peaks, consider changing peak calling parameters.")
-                exit(0)
+                exit(-1)
 
             for translation, _, score, detail in zip(*candidates):
-                rotations.append(rotation_mapping[rotation_array[tuple(translation)]])
+                rotation_index = rotation_array[tuple(translation)]
+                rotation = rotation_mapping.get(
+                    rotation_index, np.zeros(template.data.ndim, int)
+                )
+                if rotation.ndim == 2:
+                    rotation = euler_from_rotationmatrix(rotation)
+                rotations.append(rotation)
 
         else:
             candidates = data
@@ -430,7 +471,7 @@ def main():
             )
             exit(-1)
         orientations.translations = peak_caller.oversample_peaks(
-            score_space=data[0],
+            scores=data[0],
             peak_positions=orientations.translations,
             oversampling_factor=args.peak_oversampling,
         )
@@ -570,7 +611,7 @@ def main():
             return_orientations=True,
         )
         out = np.zeros_like(template.data)
-        out = np.zeros(np.multiply(template.shape, 2).astype(int))
+        # out = np.zeros(np.multiply(template.shape, 2).astype(int))
         for index in range(len(cand_slices)):
             from scipy.spatial.transform import Rotation
 

scripts/preprocessor_gui.py

@@ -1,7 +1,5 @@
 #!python3
-""" Simplify picking adequate filtering and masking parameters using a GUI.
-    Exposes tme.preprocessor.Preprocessor and tme.fitter_utils member functions
-    to achieve this aim.
+""" GUI for identifying adequate template matching filter and masks.
 
     Copyright (c) 2023 European Molecular Biology Laboratory
 
@@ -12,17 +10,20 @@ import argparse
 from typing import Tuple, Callable, List
 from typing_extensions import Annotated
 
+import napari
 import numpy as np
 import pandas as pd
-import napari
+from scipy.fft import next_fast_len
 from napari.layers import Image
 from napari.utils.events import EventedList
-
 from magicgui import widgets
 from qtpy.QtWidgets import QFileDialog
 from numpy.typing import NDArray
 
+from tme.backends import backend
 from tme import Preprocessor, Density
+from tme.preprocessing import BandPassFilter
+from tme.preprocessing.tilt_series import CTF
 from tme.matching_utils import create_mask, load_pickle
 
 preprocessor = Preprocessor()
@@ -35,19 +36,57 @@ def gaussian_filter(template: NDArray, sigma: float, **kwargs: dict) -> NDArray:
 
 def bandpass_filter(
     template: NDArray,
-    minimum_frequency: float,
-    maximum_frequency: float,
-    gaussian_sigma: float,
-    **kwargs: dict,
+    lowpass_angstrom: float = 30,
+    highpass_angstrom: float = 140,
+    hard_edges: bool = False,
+    sampling_rate=None,
 ) -> NDArray:
-    return preprocessor.bandpass_filter(
-        template=template,
-        minimum_frequency=minimum_frequency,
-        maximum_frequency=maximum_frequency,
-        sampling_rate=1,
-        gaussian_sigma=gaussian_sigma,
-        **kwargs,
+    bpf = BandPassFilter(
+        lowpass=lowpass_angstrom,
+        highpass=highpass_angstrom,
+        sampling_rate=np.max(sampling_rate),
+        use_gaussian=not hard_edges,
+        shape_is_real_fourier=True,
+        return_real_fourier=True,
     )
+    template_ft = np.fft.rfftn(template, s=template.shape)
+
+    mask = bpf(shape=template_ft.shape)["data"]
+    np.multiply(template_ft, mask, out=template_ft)
+    return np.fft.irfftn(template_ft, s=template.shape).real
+
+
+def ctf_filter(
+    template: NDArray,
+    defocus_angstrom: float = 30000,
+    acceleration_voltage: float = 300,
+    spherical_aberration: float = 2.7,
+    amplitude_contrast: float = 0.07,
+    phase_shift: float = 0,
+    defocus_angle: float = 0,
+    sampling_rate=None,
+    flip_phase: bool = False,
+) -> NDArray:
+    fast_shape = [next_fast_len(x) for x in np.multiply(template.shape, 2)]
+    template_pad = backend.topleft_pad(template, fast_shape)
+    template_ft = np.fft.rfftn(template_pad, s=template_pad.shape)
+    ctf = CTF(
+        angles=[0],
+        shape=fast_shape,
+        defocus_x=[defocus_angstrom],
+        acceleration_voltage=acceleration_voltage * 1e3,
+        spherical_aberration=spherical_aberration * 1e7,
+        amplitude_contrast=amplitude_contrast,
+        phase_shift=[phase_shift],
+        defocus_angle=[defocus_angle],
+        sampling_rate=np.max(sampling_rate),
+        return_real_fourier=True,
+        flip_phase=flip_phase,
+    )
+    np.multiply(template_ft, ctf()["data"], out=template_ft)
+    template_pad = np.fft.irfftn(template_ft, s=template_pad.shape).real
+    template = backend.topleft_pad(template_pad, template.shape)
+    return template
 
 
 def difference_of_gaussian_filter(
@@ -109,61 +148,6 @@ def mean(
     return preprocessor.mean_filter(template=template, width=width)
 
 
-def resolution_sphere(
-    template: NDArray,
-    cutoff_angstrom: float,
-    highpass: bool = False,
-    sampling_rate=None,
-) -> NDArray:
-    if cutoff_angstrom == 0:
-        return template
-
-    cutoff_frequency = np.max(2 * sampling_rate / cutoff_angstrom)
-
-    min_freq, max_freq = 0, cutoff_frequency
-    if highpass:
-        min_freq, max_freq = cutoff_frequency, 1e10
-
-    mask = preprocessor.bandpass_mask(
-        shape=template.shape,
-        minimum_frequency=min_freq,
-        maximum_frequency=max_freq,
-        omit_negative_frequencies=False,
-    )
-
-    template_ft = np.fft.fftn(template)
-    np.multiply(template_ft, mask, out=template_ft)
-    return np.fft.ifftn(template_ft).real
-
-
-def resolution_gaussian(
-    template: NDArray,
-    cutoff_angstrom: float,
-    highpass: bool = False,
-    sampling_rate=None,
-) -> NDArray:
-    if cutoff_angstrom == 0:
-        return template
-
-    grid = preprocessor.fftfreqn(
-        shape=template.shape, sampling_rate=sampling_rate / sampling_rate.max()
-    )
-
-    sigma_fourier = np.divide(
-        np.max(2 * sampling_rate / cutoff_angstrom), np.sqrt(2 * np.log(2))
-    )
-
-    mask = np.exp(-(grid**2) / (2 * sigma_fourier**2))
-    if highpass:
-        mask = 1 - mask
-
-    mask = np.fft.ifftshift(mask)
-
-    template_ft = np.fft.fftn(template)
-    np.multiply(template_ft, mask, out=template_ft)
-    return np.fft.ifftn(template_ft).real
-
-
 def wedge(
     template: NDArray,
     tilt_start: float,
@@ -274,8 +258,7 @@ WRAPPED_FUNCTIONS = {
     "mean_filter": mean,
     "wedge_filter": wedge,
     "power_spectrum": compute_power_spectrum,
-    "resolution_gaussian": resolution_gaussian,
-    "resolution_sphere": resolution_sphere,
+    "ctf": ctf_filter,
 }
 
 EXCLUDED_FUNCTIONS = [
@@ -421,6 +404,7 @@ def sphere_mask(
     center_y: float,
     center_z: float,
     radius: float,
+    sigma_decay: float = 0,
     **kwargs,
 ) -> NDArray:
     return create_mask(
@@ -428,6 +412,7 @@ def sphere_mask(
         shape=template.shape,
         center=(center_x, center_y, center_z),
         radius=radius,
+        sigma_decay=sigma_decay,
     )
 
 
@@ -439,6 +424,7 @@ def ellipsod_mask(
     radius_x: float,
     radius_y: float,
     radius_z: float,
+    sigma_decay: float = 0,
     **kwargs,
 ) -> NDArray:
     return create_mask(
@@ -446,6 +432,7 @@ def ellipsod_mask(
         shape=template.shape,
         center=(center_x, center_y, center_z),
         radius=(radius_x, radius_y, radius_z),
+        sigma_decay=sigma_decay,
     )
 
 
@@ -457,6 +444,7 @@ def box_mask(
     height_x: int,
     height_y: int,
     height_z: int,
+    sigma_decay: float = 0,
     **kwargs,
 ) -> NDArray:
     return create_mask(
@@ -464,6 +452,7 @@ def box_mask(
         shape=template.shape,
         center=(center_x, center_y, center_z),
         height=(height_x, height_y, height_z),
+        sigma_decay=sigma_decay,
     )
 
 
@@ -476,6 +465,7 @@ def tube_mask(
     inner_radius: float,
     outer_radius: float,
     height: int,
+    sigma_decay: float = 0,
     **kwargs,
 ) -> NDArray:
     return create_mask(
@@ -486,6 +476,7 @@ def tube_mask(
         inner_radius=inner_radius,
         outer_radius=outer_radius,
         height=height,
+        sigma_decay=sigma_decay,
     )
 
 
@@ -533,13 +524,23 @@ def wedge_mask(
 
 
 def threshold_mask(
-    template: NDArray, standard_deviation: float = 5.0, invert: bool = False, **kwargs
+    template: NDArray,
+    invert: bool = False,
+    standard_deviation: float = 5.0,
+    sigma: float = 0.0,
+    **kwargs,
 ) -> NDArray:
     template_mean = template.mean()
     template_deviation = standard_deviation * template.std()
     upper = template_mean + template_deviation
     lower = template_mean - template_deviation
-    mask = np.logical_and(template > lower, template < upper)
+    mask = np.logical_or(template <= lower, template >= upper)
+
+    if sigma != 0:
+        mask_filter = preprocessor.gaussian_filter(template=mask * 1.0, sigma=sigma)
+        mask = np.add(mask, (1 - mask) * mask_filter)
+        mask[mask < np.exp(-np.square(sigma))] = 0
+
     if invert:
         np.invert(mask, out=mask)
 
@@ -890,6 +891,7 @@ class PointCloudWidget(widgets.Container):
 
         self.viewer = viewer
         self.dataframes = {}
+        self.selected_category = -1
 
         self.import_button = widgets.PushButton(
             name="Import", text="Import Point Cloud"
@@ -902,10 +904,98 @@ class PointCloudWidget(widgets.Container):
         self.export_button.clicked.connect(self._export_point_cloud)
         self.export_button.enabled = False
 
+        self.annotation_container = widgets.Container(name="Label", layout="horizontal")
+        self.positive_button = widgets.PushButton(name="Positive", text="Positive")
+        self.negative_button = widgets.PushButton(name="Negative", text="Negative")
+        self.positive_button.clicked.connect(self._set_positive)
+        self.negative_button.clicked.connect(self._set_negative)
+        self.annotation_container.append(self.positive_button)
+        self.annotation_container.append(self.negative_button)
+
+        self.face_color_select = widgets.ComboBox(
+            name="Color", choices=["Label", "Score"], value=None, nullable=True
+        )
+        self.face_color_select.changed.connect(self._update_face_color_mode)
+
         self.append(self.import_button)
         self.append(self.export_button)
+        self.append(self.annotation_container)
+        self.append(self.face_color_select)
+
         self.viewer.layers.selection.events.changed.connect(self._update_buttons)
 
+        self.viewer.layers.events.inserted.connect(self._initialize_points_layer)
+
+    def _update_face_color_mode(self, event: str = None):
+        for layer in self.viewer.layers:
+            if not isinstance(layer, napari.layers.Points):
+                continue
+
+            layer.face_color = "white"
+            if event == "Label":
+                if len(layer.properties.get("detail", ())) == 0:
+                    continue
+                layer.face_color = "detail"
+                layer.face_color_cycle = {
+                    -1: "grey",
+                    0: "red",
+                    1: "green",
+                }
+                layer.face_color_mode = "cycle"
+            elif event == "Score":
+                if len(layer.properties.get("score_scaled", ())) == 0:
+                    continue
+                layer.face_color = "score_scaled"
+                layer.face_colormap = "turbo"
+                layer.face_color_mode = "colormap"
+
+            layer.refresh_colors()
+
+        return None
+
+    def _set_positive(self, event):
+        self.selected_category = 1 if self.selected_category != 1 else -1
+        self._update_annotation_buttons()
+
+    def _set_negative(self, event):
+        self.selected_category = 0 if self.selected_category != 0 else -1
+        self._update_annotation_buttons()
+
+    def _update_annotation_buttons(self):
+        selected_style = "background-color: darkgrey"
+        default_style = "background-color: none"
+
+        self.positive_button.native.setStyleSheet(
+            selected_style if self.selected_category == 1 else default_style
+        )
+        self.negative_button.native.setStyleSheet(
+            selected_style if self.selected_category == 0 else default_style
+        )
+
+    def _initialize_points_layer(self, event):
+        layer = event.value
+        if not isinstance(layer, napari.layers.Points):
+            return
+        if len(layer.properties) == 0:
+            layer.properties = {"detail": [-1]}
+
+        if "detail" not in layer.properties:
+            layer["detail"] = [-1]
+
+        layer.mouse_drag_callbacks.append(self._on_click)
+        return None
+
+    def _on_click(self, layer, event):
+        if layer.mode == "add":
+            layer.current_properties["detail"][-1] = self.selected_category
+        elif layer.mode == "select":
+            for index in layer.selected_data:
+                layer.properties["detail"][index] = self.selected_category
+
+        # TODO: Check whether current face color is the desired one already
+        self._update_face_color_mode(self.face_color_select.value)
+        layer.refresh_colors()
+
     def _update_buttons(self, event):
         is_pointcloud = isinstance(
             self.viewer.layers.selection.active, napari.layers.Points
@@ -951,9 +1041,7 @@ class PointCloudWidget(widgets.Container):
 
             if "score" in merged_data.columns:
                 merged_data["score"] = merged_data["score"].fillna(1)
-            if "detail" in merged_data.columns:
-                merged_data["detail"] = merged_data["detail"].fillna(2)
-
+            merged_data["detail"] = layer.properties["detail"]
             merged_data.to_csv(filename, sep="\t", index=False)
 
     def _get_load_path(self, event):
@@ -977,7 +1065,7 @@ class PointCloudWidget(widgets.Container):
             dataframe["score"] = 1
 
         if "detail" not in dataframe.columns:
-            dataframe["detail"] = -2
+            dataframe["detail"] = -1
 
         point_properties = {
             "score": np.array(dataframe["score"].values),
@@ -991,8 +1079,6 @@ class PointCloudWidget(widgets.Container):
             points,
             size=10,
             properties=point_properties,
-            face_color="score_scaled",
-            face_colormap="turbo",
             name=layer_name,
         )
         self.dataframes[layer_name] = dataframe
@@ -1025,9 +1111,14 @@ class MatchingWidget(widgets.Container):
     def _load_data(self, filename):
         data = load_pickle(filename)
 
-        _ = self.viewer.add_image(data=data[2], name="Rotations", colormap="orange")
+        metadata = {"origin": data[-1][1], "sampling_rate": data[-1][2]}
+        _ = self.viewer.add_image(
+            data=data[2], name="Rotations", colormap="orange", metadata=metadata
+        )
 
-        _ = self.viewer.add_image(data=data[0], name="Scores", colormap="turbo")
+        _ = self.viewer.add_image(
+            data=data[0], name="Scores", colormap="turbo", metadata=metadata
+        )
 
 
 def main():
@@ -1045,11 +1136,10 @@ def main():
         widget=alignment_widget, name="Alignment", area="right"
     )
     viewer.window.add_dock_widget(widget=mask_widget, name="Mask", area="right")
+    viewer.window.add_dock_widget(widget=export_widget, name="Export", area="right")
     viewer.window.add_dock_widget(widget=point_cloud, name="PointCloud", area="left")
     viewer.window.add_dock_widget(widget=matching_widget, name="Matching", area="left")
 
-    viewer.window.add_dock_widget(widget=export_widget, name="Export", area="right")
-
     napari.run()