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

{pytme-0.2.1.data → pytme-0.2.2.data}/scripts/preprocessor_gui.py

@@ -1,7 +1,5 @@
 #!python
-""" 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()
 
 

pytme-0.2.2.dist-info/METADATA

@@ -0,0 +1,91 @@
+Metadata-Version: 2.1
+Name: pytme
+Version: 0.2.2
+Summary: Python Template Matching Engine
+Author: Valentin Maurer
+Author-email: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+License: Proprietary
+Project-URL: Homepage, https://github.com/KosinskiLab/pyTME
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: mrcfile >=1.4.3
+Requires-Dist: numpy >=1.22.2
+Requires-Dist: pyfftw >=0.13.1
+Requires-Dist: pytest >=6.2.5
+Requires-Dist: PyYAML >=6.0
+Requires-Dist: scikit-image >=0.19.0
+Requires-Dist: scikit-learn >=1.2.1
+Requires-Dist: scipy >=1.9.1
+Requires-Dist: pybind11
+Requires-Dist: psutil
+Requires-Dist: tifffile
+Requires-Dist: h5py
+Provides-Extra: all
+Requires-Dist: cupy ; extra == 'all'
+Requires-Dist: voltools ==0.6.0 ; extra == 'all'
+Requires-Dist: torch ; extra == 'all'
+Requires-Dist: jax ; extra == 'all'
+Requires-Dist: jaxlib ; extra == 'all'
+Provides-Extra: cupy
+Requires-Dist: cupy ; extra == 'cupy'
+Requires-Dist: voltools ==0.6.0 ; extra == 'cupy'
+Provides-Extra: jax
+Requires-Dist: jax ; extra == 'jax'
+Requires-Dist: jaxlib ; extra == 'jax'
+Provides-Extra: pytorch
+Requires-Dist: torch ; extra == 'pytorch'
+
+# Python Template Matching Engine (PyTME)
+
+[![Build Status](https://img.shields.io/github/actions/workflow/status/KosinskiLab/pyTME/main.yml?label=CI)](https://github.com/KosinskiLab/pyTME/actions)
+[![PyPI](https://img.shields.io/pypi/v/pytme.svg)](https://pypi.org/project/pytme/)
+
+**[Documentation](https://kosinskilab.github.io/pyTME/)** | **[Installation](https://kosinskilab.github.io/pyTME/quickstart/installation.html)** | **[API](https://kosinskilab.github.io/pyTME/reference/index.html)**
+
+PyTME is a Python library for data-intensive n-dimensional template matching using CPUs and GPUs.
+
+With its [backend-agnostic design](https://kosinskilab.github.io/pyTME/reference/backends.html), the same code can be run on diverse hardware platforms using a best-of-breed approach. The underyling abstract backend specification allows for adding new backends to benefit from gains in performance and capabilities without modifying the library's core routines. The implementation of template matching scores is modular and provides developers with a flexible framework for rapid prototyping. Furthermore, pyTME supports a unique callback capability through [analyzers](https://kosinskilab.github.io/pyTME/reference/analyzer.html), which allows for injection of custom code, enabling real-time processing and manipulation of results
+
+PyTME includes a [graphical user interface](https://kosinskilab.github.io/pyTME/quickstart/preprocessing.html#practical-example) that provides simplified mask creation, interactive filter exploration, result visualization, and manual refinement capabilities. This GUI serves as an accessible entry point to the library's core functionalities, allowing users to efficiently interact with and analyze their data.
+
+Finally, pyTME offers specialized tools for cryogenic electron microscopy data, such as wedge masks, CTF correction, as well as [means for handling structural data](https://kosinskilab.github.io/pyTME/reference/data_structure.html). Through dedicated [integrations](https://kosinskilab.github.io/pyTME/quickstart/integrations.html), the output of pyTME seamlessly integrates with commonly used cryogenic electron microscopy software such as RELION, Dynamo and IMOD.
+
+Running into bugs or missing a feature? Help us improve the project by opening an [issue](https://github.com/KosinskiLab/pyTME/issues).
+
+## Installation
+
+We recommend installation using one of the following methods
+
+| Method   | Command                                                 |
+|----------|---------------------------------------------------------|
+| PyPi     | `pip install pytme`                                     |
+| Source   | `pip install git+https://github.com/KosinskiLab/pyTME`  |
+| Docker   | `docker build -t pytme -f docker/Dockerfile_GPU .`      |
+
+You can find alternative installation methods in the [documentation](https://kosinskilab.github.io/pyTME/quickstart/installation.html).
+
+
+## Quickstart
+
+Learn how to get started with
+
+- [Preprocessing:](https://kosinskilab.github.io/pyTME/quickstart/preprocessing.html) Picking the right mask and filter for template matching.
+- [Template matching:](https://kosinskilab.github.io/pyTME/quickstart/match_template.html) Find your template of interest.
+- [Postprocessing](https://kosinskilab.github.io/pyTME/quickstart/postprocessing.html) Analyze template matching results and downstream integrations.
+
+## How to Cite
+
+If PyTME contributed significantly to your research, please cite the corresponding publication on [SoftwareX](https://www.sciencedirect.com/science/article/pii/S2352711024000074).
+
+```bibtex
+@article{Maurer:2024aa,
+    author = {Maurer, Valentin J. and Siggel, Marc and Kosinski, Jan},
+    journal = {SoftwareX},
+    pages = {101636},
+    title = {PyTME (Python Template Matching Engine): A fast, flexible, and multi-purpose template matching library for cryogenic electron microscopy data},
+    volume = {25},
+    year = {2024}}
+```

pytme-0.2.2.dist-info/RECORD

@@ -0,0 +1,74 @@
+pytme-0.2.2.data/scripts/estimate_ram_usage.py,sha256=R1NDpFajcF-MonJ4a43SfDlA-nxBYwK7D2quzCdsVFM,2767
+pytme-0.2.2.data/scripts/match_template.py,sha256=FXAzPUnKMM2_thOtTZUKtIDV7GCmhFmLXm2yJDCI6G0,41456
+pytme-0.2.2.data/scripts/postprocess.py,sha256=K3tP8licTsCW1LKy8hQ8rMA7UrJ0AakM3kffx_1d4n4,24780
+pytme-0.2.2.data/scripts/preprocess.py,sha256=zog-l2Je-GeouJ6SnamOMuHgTn7fFPiGnO5X03y5qSY,2527
+pytme-0.2.2.data/scripts/preprocessor_gui.py,sha256=3eGC5RqnnZ3aXyfKEVstddLvKQP46kvyclj4-q01vAA,39275
+scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+scripts/estimate_ram_usage.py,sha256=rN7haobnHg3YcgGJIp81FNiCzy8-saJGeEurQlmQmNQ,2768
+scripts/extract_candidates.py,sha256=DAfNyuauogWvSdRWIbtH44tsk9buLn13JrL1zJjJGLE,8373
+scripts/match_template.py,sha256=tEHfYXvfoBZB7bFjaBrHG3qH_e-U96fRfruB4n1uH9U,41457
+scripts/match_template_filters.py,sha256=Gj4a1b_S5NWp_dfFEPFn0D7jGf-qYgBbnTvZZ4bwqOQ,42036
+scripts/postprocess.py,sha256=TngTI7R8VaotxyOR9o5Hgmppjn_ILJ6ek1uwyHjkuco,24781
+scripts/preprocess.py,sha256=ebJVLxbRlB6TI5YHNr0VavZ4lmaRdf8QVafyiDhh_oU,2528
+scripts/preprocessor_gui.py,sha256=BUm8dsbe5LEo0FyjKrK1KdbeM75aYrS1yiCp2GfI19M,39276
+scripts/refine_matches.py,sha256=Y17Ku_t0W9vglPNF2oU5EFrqoedJIm3lCGl-hXRHvjc,21920
+tme/__init__.py,sha256=liARhKZ565F2Y02iWRbFzfrTCkL5You3YaUSiqUkees,202
+tme/__version__.py,sha256=m6kyaNpwBcP1XYcqrelX2oS3PJuOnElOcRdBa9pEb8c,22
+tme/analyzer.py,sha256=WfL6RbJDFpjwF1YiIa9NGOxEL8NddlekQIuG1r6B5fA,50325
+tme/density.py,sha256=JAC2TKzD2OWBa45MCEwT6WbL1XcTdcbiErMylwIwnH0,83837
+tme/extensions.cpython-311-darwin.so,sha256=e48Rshv8cWPSiBD-xqq5L8ekQJn84MqZse05fndMnkc,411984
+tme/matching_data.py,sha256=QvPozAnlT9GvEUaXwqmN05j-axQ5EKuoPrNBEELMxvg,25255
+tme/matching_exhaustive.py,sha256=QlFOORb2a2uDy6-wk4YU3X04y3XW4sfzwk3gC9gL27I,18843
+tme/matching_memory.py,sha256=bmCAUYyXWEet-1XXhldtc0irio2ytMSsAzWYyFI5LNM,11273
+tme/matching_optimization.py,sha256=Y8HfecXiOvAHXM1viBaQ_aXljqqTnGwlOlFe0MJpDRQ,45082
+tme/matching_scores.py,sha256=3XxvwRNztdC-4hKdloyCfwjzt4GYzA0T-WF6YUnOcvs,30867
+tme/matching_utils.py,sha256=VdQZPad9uVCQ9-32Vulh5r7bhBYylPvqv2X2DQNk91M,39736
+tme/memory.py,sha256=6xeIMAncQkgYDi6w-PIYgFEWRTUPu0_OTCeRO0p9r9Q,11029
+tme/orientations.py,sha256=Ul-1g2ci3QdRLfD447ZkcIo6Rv6f7gQH83JZcRkAru0,25384
+tme/parser.py,sha256=fNiCAdsWI4ql92F1Ob4suiVzpjUOBlh2lad1iNY_FP8,13772
+tme/preprocessor.py,sha256=ZhlP-8b4NzAzRogWOnk5DLNh_tqCMOpc4msQIJEKBTU,48566
+tme/structure.py,sha256=k5kgFNHkM2mUvXS7VAwaevltLC1mUxEiDhMW0TqTNjo,65834
+tme/types.py,sha256=NAY7C4qxE6yz-DXVtClMvFfoOV-spWGLNfpLATZ1LcU,442
+tme/backends/__init__.py,sha256=Y69k5_Vx7VpcBNV8W0pQrebzo42lp3WBfEGAl1s4Vd4,5187
+tme/backends/_jax_utils.py,sha256=v9T36Jv4ztyC38IhrbNU5tf80VKAdAXYX5ZxCWe02qU,5757
+tme/backends/cupy_backend.py,sha256=rK3pjNNB4D46MqzY2DrB_GSObFJlG5GQDOt-eOwhwTg,9429
+tme/backends/jax_backend.py,sha256=CUEJynVw9jqWid_vARMpNrhORgUR1SlLjz4ZTXkLuvI,9730
+tme/backends/matching_backend.py,sha256=KfCOKD_rA9el3Y7BeH17KJ1apCUIIhhvn-vmbkb3CB0,33750
+tme/backends/mlx_backend.py,sha256=FJhqmCzgjXAjWGX1HhHFrCy_We4YwQQBkKFNG05ctzM,7788
+tme/backends/npfftw_backend.py,sha256=QAQxr3OSq9bBN_Aq0qps-nz5vvUsx0_fFn9TocI_QCw,16988
+tme/backends/pytorch_backend.py,sha256=fzLucJZH5VYhrUtcX9osirNvLMsyMNPt_JNfS2l-kqg,14940
+tme/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tme/data/c48n309.npy,sha256=NwH64mOEbm3tStq5c98o81fY1vMOoq4nvXDAh7Z7iZg,296768
+tme/data/c48n527.npy,sha256=saSUMTa1R0MisPvgFL02a7IHQSwEZ-mJu0v3qJjg5AU,506048
+tme/data/c48n9.npy,sha256=bDVLV6mWjZHSQfeDc-MOCKKarfc1jaNeVvpoe2xMUy4,8768
+tme/data/c48u1.npy,sha256=JeXMFzFITs2ezdc3x5lp3jo1cHHHHVADSA1Tpf77kXs,1088
+tme/data/c48u1153.npy,sha256=ECiEximtYDWtIux3Fwe_EJlyn08gUqP85DN9gjkT9_k,1107008
+tme/data/c48u1201.npy,sha256=aceC_Jeienz_81X4520nPpZcg5tnRhbW795EqbpWkrg,1153088
+tme/data/c48u1641.npy,sha256=p4LwW3LzdTjrUUpA7H53RfNWxYfPX0XjeSwZ39Ac78Q,1575488
+tme/data/c48u181.npy,sha256=mLYXrv1YHLH6DsBp5MkxHkxlxgMnj1mw_KKI0udH-FY,173888
+tme/data/c48u2219.npy,sha256=p8TQeX8YHu4pdxnwJjEAlQWAPa66W7kpK96iZKZr9JE,2130368
+tme/data/c48u27.npy,sha256=k03ZNEsoPwBKCy8IeIa5G0WRZqjGZMtX6Ibu7EpJHvU,26048
+tme/data/c48u2947.npy,sha256=icI97ED6ct66y7FIaJAugmjzrIWk7CINCxtO3wDTnrU,2829248
+tme/data/c48u3733.npy,sha256=tla-__Pf-hpN6h04vtFIfkkFdCLple11VO06kr1dXkM,3583808
+tme/data/c48u4749.npy,sha256=tItOA4oV7SiqCCREwz3fyEpZoxM0lCq_jfEo5_-fp2s,4559168
+tme/data/c48u5879.npy,sha256=bFk89MllIFCX_sLXTYWFquSyN1NuahH4wwnEsPJLxzA,5643968
+tme/data/c48u7111.npy,sha256=CMy9kI2edH-q9eTIVdgUtXurplYNI7Uqp4dXfkkVdf8,6826688
+tme/data/c48u815.npy,sha256=bCuJxLtm0Sjg3GGxtyjGzRYZ1G0Gz79XHI-71GvqQnI,782528
+tme/data/c48u83.npy,sha256=7ODJYnsiuDjGbgd9GFopsyIW2IjrYI0J2X2f-cK868U,79808
+tme/data/c48u8649.npy,sha256=-IPlpR4zrPQZWhhSPu4zEulFdrCEVgTMFffCB5d-huE,8303168
+tme/data/c600v.npy,sha256=JqSu3ALoL1A9iguehc0YGUMFPsh2fprHHp76VXeFXIw,2528
+tme/data/c600vc.npy,sha256=Yht-GFXDSjjGvsjFBvyxxEZAI-ODADPd5gEgFNZQVTA,14528
+tme/data/metadata.yaml,sha256=fAgX-mEzB0QMHTEtYDG4cSMbJhYxBbDJH3sdvJvL7a8,750
+tme/data/quat_to_numpy.py,sha256=-gkDZb10fKBxwfYrSLCUWvMB76TzZWELCeKsYProwws,1333
+tme/preprocessing/__init__.py,sha256=7O3vDzJcIfxovJkf7avWSPtzaIVlTbmsW7egQFukC_s,98
+tme/preprocessing/_utils.py,sha256=hC8Qfh_a_fk4o7v82c4-TZYYRFfo5XttcghvpkzBo0A,6193
+tme/preprocessing/composable_filter.py,sha256=sp3bN8JeFB2r384cEIgvN6yXjC53GCoPhBmGI0S0kbI,781
+tme/preprocessing/compose.py,sha256=tilp14UhBh98SJVJhC_STbgi7i8HeU4FB7B7KUOD8gc,1458
+tme/preprocessing/frequency_filters.py,sha256=x900ntX0QI46lDdycPISO-7K0XgM0015ksCOFlm-yM4,12964
+tme/preprocessing/tilt_series.py,sha256=Yb3dk7TZf7OfPWZPwVuFTirOZK4x6mjLbxYX9KtRBdg,34545
+pytme-0.2.2.dist-info/LICENSE,sha256=K1IUNSVAz8BXbpH5EA8y5FpaHdvFXnAF2zeK95Lr2bY,18467
+pytme-0.2.2.dist-info/METADATA,sha256=axenBTwu62Z-5jY-4kScnqv7bbWWwDU7S-VatPzSz0k,5206
+pytme-0.2.2.dist-info/WHEEL,sha256=IHsX_ZtCj0uMf8jtN3eK3zfXfu1SfKRsLbDhsprnhR0,109
+pytme-0.2.2.dist-info/entry_points.txt,sha256=ff3LQL3FCWfCYOwFiP9zatm7laUbnwCkuPELkQVyUO4,241
+pytme-0.2.2.dist-info/top_level.txt,sha256=J8FUkazOb2fZ0n_KexnqCGyNOtie2bwisFSUBiM5-0w,12
+pytme-0.2.2.dist-info/RECORD,,

{pytme-0.2.1.dist-info → pytme-0.2.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: setuptools (72.1.0)
 Root-Is-Purelib: false
 Tag: cp311-cp311-macosx_14_0_arm64
 

scripts/extract_candidates.py

@@ -108,6 +108,12 @@ def parse_args():
         type=int,
         help="Box size for extraction, defaults to two times the template.",
     )
+    extraction_group.add_argument(
+        "--translation_uncertainty",
+        required=False,
+        type=int,
+        help="Sets box size for extraction to template box plus this value.",
+    )
     extraction_group.add_argument(
         "--keep_out_of_box",
         action="store_true",
@@ -150,6 +156,9 @@ def main():
     box_size = args.box_size
     if box_size is None:
         box_size = np.multiply(template.shape, 2)
+    if args.translation_uncertainty is not None:
+        box_size = np.add(template.shape, args.translation_uncertainty)
+
     box_size = np.array(box_size)
     box_size = np.repeat(box_size, template.data.ndim // box_size.size).astype(int)
 
@@ -195,8 +204,6 @@ def main():
     )
     dens.data[:] = target.metadata["mean"]
 
-    print(target.data.shape)
-    # There appears to be an isseu with the stack creation. Trace this further
     data_subset = np.zeros(extraction_shape, dtype = target.data.dtype)
     pbar = ProgressBar(message = "Orientation ", nchars = 80, total = len(obs_slices))
     for index, (obs_slice, cand_slice) in enumerate(zip(obs_slices, cand_slices)):
@@ -219,13 +226,13 @@ def main():
             )
         target_subset.pad(box_size, center=True)
 
-        target_value = target.data[tuple(orientations.translations[index].astype(int))]
-        center = np.divide(target_subset.data.shape, 2).astype(int ) + np.mod(target_subset.shape, 2)
-        print(np.where(target_subset.data == target_value), center)
-        print(target_subset.data[tuple(center.astype(int))],
-            target_value,
-            target_subset.data[tuple(center.astype(int))] == target_value
-        )
+        # target_value = target.data[tuple(orientations.translations[index].astype(int))]
+        # center = np.divide(target_subset.data.shape, 2).astype(int ) + np.mod(target_subset.shape, 2)
+        # print(np.where(target_subset.data == target_value), center)
+        # print(target_subset.data[tuple(center.astype(int))],
+            # target_value,
+            # target_subset.data[tuple(center.astype(int))] == target_value
+        # )
 
         dens.data[index] = target_subset.data
     print("")
@@ -241,10 +248,10 @@ def main():
         f"{splitext(args.output_file)[0]}_aligned.tsv",
         file_format="text"
     )
-    orientations.to_file(
-        f"{splitext(args.output_file)[0]}_aligned.star",
-        file_format="relion"
-    )
+    # orientations.to_file(
+    #     f"{splitext(args.output_file)[0]}_aligned.star",
+    #     file_format="relion"
+    # )
 
 if __name__ == "__main__":
     main()