pytme 0.1.6__cp311-cp311-macosx_14_0_arm64.whl → 0.1.8__cp311-cp311-macosx_14_0_arm64.whl

scripts/postprocess.py

@@ -5,6 +5,8 @@
 
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
+from os import getcwd
+from os.path import join
 import argparse
 from sys import exit
 from typing import List, Tuple
@@ -13,6 +15,7 @@ from dataclasses import dataclass
 
 import numpy as np
 from scipy.spatial.transform import Rotation
+from numpy.typing import NDArray
 
 from tme import Density, Structure
 from tme.analyzer import (
@@ -26,6 +29,7 @@ from tme.matching_utils import (
     load_pickle,
     euler_to_rotationmatrix,
     euler_from_rotationmatrix,
+    centered_mask,
 )
 
 PEAK_CALLERS = {
@@ -52,16 +56,33 @@ def parse_args():
         help="Prefix for the output file name. Extension depends on output_format.",
     )
     parser.add_argument(
-        "--number_of_peaks", type=int, default=1000, help="Number of peaks to consider."
+        "--number_of_peaks",
+        type=int,
+        default=1000,
+        help="Number of peaks to consider. Note, this is the number of called peaks "
+        ", subject to min_distance and min_boundary_distance filtering. Therefore, the "
+        "returned number of peaks will be at most equal to number_of_peaks. "
+        "Ignored when --orientations is provided.",
     )
     parser.add_argument(
-        "--min_distance", type=int, default=5, help="Minimum distance between peaks."
+        "--min_distance",
+        type=int,
+        default=5,
+        help="Minimum distance between peaks. Ignored when --orientations is provided.",
     )
     parser.add_argument(
         "--min_boundary_distance",
         type=int,
         default=0,
-        help="Minimum distance from target boundaries.",
+        help="Minimum distance from target boundaries. Ignored when --orientations "
+        "is provided.",
+    )
+    parser.add_argument(
+        "--mask_edges",
+        action="store_true",
+        default=False,
+        help="Whether to mask edges of the input score array according to the template shape."
+        "Uses twice the value of --min_boundary_distance if boht are provided.",
     )
     parser.add_argument(
         "--wedge_mask",
@@ -73,7 +94,8 @@ def parse_args():
         "--peak_caller",
         choices=list(PEAK_CALLERS.keys()),
         default="PeakCallerScipy",
-        help="Peak caller to use for analysis. Ignored if input_file contains peaks.",
+        help="Peak caller to use for analysis. Ignored if input_file contains peaks or when "
+        "--orientations is provided.",
     )
     parser.add_argument(
         "--orientations",
@@ -205,8 +227,12 @@ class Orientations:
         return None
 
     def _to_relion_star(
-        self, filename: str, name_prefix: str = None, ctf_image: str = None,
-        sampling_rate : float = None, subtomogram_size : int = None
+        self,
+        filename: str,
+        name_prefix: str = None,
+        ctf_image: str = None,
+        sampling_rate: float = 1.0,
+        subtomogram_size: int = 0,
     ) -> None:
         """
         Save orientations in RELION's STAR file format.
@@ -249,12 +275,11 @@ class Orientations:
             "300.000000",
             str(int(subtomogram_size)),
             "3",
-            str(float(sampling_rate))
+            str(float(sampling_rate)),
         ]
         optics_header = "\n".join(optics_header)
         optics_data = "\t".join(optics_data)
 
-
         header = [
             "data_particles",
             "",
@@ -283,13 +308,14 @@ class Orientations:
             _ = ofile.write("\n# version 30001\n")
             _ = ofile.write(f"{header}\n")
 
+            # pyTME uses a zyx data layout
             for index, (translation, rotation, score, detail) in enumerate(self):
                 rotation = Rotation.from_euler("zyx", rotation, degrees=True)
-                rotation = rotation.as_euler(seq="zyz", degrees=True)
+                rotation = rotation.as_euler(seq="xyx", degrees=True)
 
                 translation_string = "\t".join([str(x) for x in translation][::-1])
-                angle_string = "\t".join([str(x) for x in rotation[::-1]])
-                name = f"{name_prefix}{index}.mrc"
+                angle_string = "\t".join([str(x) for x in rotation])
+                name = f"{name_prefix}_{index}.mrc"
                 _ = ofile.write(
                     f"{translation_string}\t{name}\t{angle_string}\t1{ctf_image}\n"
                 )
@@ -386,11 +412,25 @@ class Orientations:
         return translation, rotation, score, detail
 
 
+def load_template(filepath: str, sampling_rate: NDArray) -> "Density":
+    try:
+        template = Density.from_file(filepath)
+        template, _ = template.centered(0)
+        center_of_mass = template.center_of_mass(template.data)
+    except ValueError:
+        template = Structure.from_file(filepath)
+        center_of_mass = template.center_of_mass()[::-1]
+        template = Density.from_structure(template, sampling_rate=sampling_rate)
+
+    return template, center_of_mass
+
+
 def main():
     args = parse_args()
     data = load_pickle(args.input_file)
 
     meta = data[-1]
+    target_origin, _, sampling_rate, cli_args = meta
 
     if args.orientations is not None:
         orientations = Orientations.from_file(
@@ -402,13 +442,29 @@ def main():
         # Output is MaxScoreOverRotations
         if data[0].ndim == data[2].ndim:
             scores, offset, rotation_array, rotation_mapping, meta = data
+            if args.mask_edges:
+                template, center_of_mass = load_template(
+                    cli_args.template, sampling_rate=sampling_rate
+                )
+                if not cli_args.no_centering:
+                    template, *_ = template.centered(0)
+                mask_size = template.shape
+                if args.min_boundary_distance > 0:
+                    mask_size = 2 * args.min_boundary_distance
+                scores = centered_mask(scores, np.subtract(scores.shape, mask_size) + 1)
+
             peak_caller = PEAK_CALLERS[args.peak_caller](
                 number_of_peaks=args.number_of_peaks,
                 min_distance=args.min_distance,
                 min_boundary_distance=args.min_boundary_distance,
             )
             peak_caller(scores, rotation_matrix=np.eye(3))
-            candidates = peak_caller.merge([tuple(peak_caller)])
+            candidates = peak_caller.merge(
+                candidates=[tuple(peak_caller)],
+                number_of_peaks=args.number_of_peaks,
+                min_distance=args.min_distance,
+                min_boundary_distance=args.min_boundary_distance,
+            )
             if len(candidates) == 0:
                 exit(
                     "Found no peaks. Try reducing min_distance or min_boundary_distance."
@@ -436,18 +492,11 @@ def main():
         orientations.to_file(filename=f"{args.output_prefix}.tsv", file_format="text")
         exit(0)
 
-    target_origin, _, sampling_rate, cli_args = meta
-
-    template_is_density, index = True, 0
     _, template_extension = splitext(cli_args.template)
-    try:
-        template = Density.from_file(cli_args.template)
-        template, _ = template.centered(0)
-        center_of_mass = template.center_of_mass(template.data)
-    except ValueError:
-        template_is_density = False
-        template = Structure.from_file(cli_args.template)
-        center_of_mass = template.center_of_mass()[::-1]
+    template, center_of_mass = load_template(
+        filepath=cli_args.template, sampling_rate=sampling_rate
+    )
+    template_is_density, index = isinstance(template, Density), 0
 
     if args.output_format == "relion":
         new_shape = np.add(template.shape, np.mod(template.shape, 2))
@@ -457,10 +506,6 @@ def main():
 
     if args.output_format in ("extraction", "relion"):
         target = Density.from_file(cli_args.target)
-        if isinstance(template, Structure):
-            template = Density.from_structure(
-                template, sampling_rate=target.sampling_rate
-            )
 
         if not np.all(np.divide(target.shape, template.shape) > 2):
             print(
@@ -489,14 +534,15 @@ def main():
         )
 
         orientations = orientations[keep_peaks]
+        working_directory = getcwd()
         if args.output_format == "relion":
             orientations.to_file(
                 filename=f"{args.output_prefix}.star",
                 file_format="relion",
-                name_prefix=args.output_prefix,
+                name_prefix=join(working_directory, args.output_prefix),
                 ctf_image=args.wedge_mask,
-                sampling_rate = target.sampling_rate.max(),
-                subtomogram_size = template.shape[0]
+                sampling_rate=target.sampling_rate.max(),
+                subtomogram_size=template.shape[0],
             )
 
         peaks = peaks[keep_peaks,]
@@ -543,7 +589,9 @@ def main():
                 origin=candidate_starts[index] * sampling_rate,
             )
             # out_density.data = out_density.data * template_mask.data
-            out_density.to_file(f"{args.output_prefix}{index}.mrc")
+            out_density.to_file(
+                join(working_directory, f"{args.output_prefix}_{index}.mrc")
+            )
 
         exit(0)
 
@@ -566,7 +614,10 @@ def main():
                 translation=translation[::-1],
                 rotation_matrix=rotation_matrix[::-1, ::-1],
             )
-        transformed_template.to_file(f"{args.output_prefix}{index}{template_extension}")
+        # template_extension should contain the extension '.'
+        transformed_template.to_file(
+            f"{args.output_prefix}_{index}{template_extension}"
+        )
         index += 1
 
 

scripts/preprocessor_gui.py

@@ -29,17 +29,17 @@ preprocessor = Preprocessor()
 SLIDER_MIN, SLIDER_MAX = 0, 25
 
 
-def gaussian_filter(template, sigma: float, **kwargs: dict):
+def gaussian_filter(template: NDArray, sigma: float, **kwargs: dict) -> NDArray:
     return preprocessor.gaussian_filter(template=template, sigma=sigma, **kwargs)
 
 
 def bandpass_filter(
-    template,
+    template: NDArray,
     minimum_frequency: float,
     maximum_frequency: float,
     gaussian_sigma: float,
     **kwargs: dict,
-):
+) -> NDArray:
     return preprocessor.bandpass_filter(
         template=template,
         minimum_frequency=minimum_frequency,
@@ -51,8 +51,8 @@ def bandpass_filter(
 
 
 def difference_of_gaussian_filter(
-    template, sigmas: Tuple[float, float], **kwargs: dict
-):
+    template: NDArray, sigmas: Tuple[float, float], **kwargs: dict
+) -> NDArray:
     low_sigma, high_sigma = sigmas
     return preprocessor.difference_of_gaussian_filter(
         template=template, low_sigma=low_sigma, high_sigma=high_sigma, **kwargs
@@ -60,7 +60,7 @@ def difference_of_gaussian_filter(
 
 
 def edge_gaussian_filter(
-    template,
+    template: NDArray,
     sigma: float,
     edge_algorithm: Annotated[
         str,
@@ -68,7 +68,7 @@ def edge_gaussian_filter(
     ],
     reverse: bool = False,
     **kwargs: dict,
-):
+) -> NDArray:
     return preprocessor.edge_gaussian_filter(
         template=template,
         sigma=sigma,
@@ -78,13 +78,13 @@ def edge_gaussian_filter(
 
 
 def local_gaussian_filter(
-    template,
+    template: NDArray,
     lbd: float,
     sigma_range: Tuple[float, float],
     gaussian_sigma: float,
     reverse: bool = False,
     **kwargs: dict,
-):
+) -> NDArray:
     return preprocessor.local_gaussian_filter(
         template=template,
         lbd=lbd,
@@ -94,21 +94,76 @@ def local_gaussian_filter(
 
 
 def ntree(
-    template,
+    template: NDArray,
     sigma_range: Tuple[float, float],
     **kwargs: dict,
-):
+) -> NDArray:
     return preprocessor.ntree_filter(template=template, sigma_range=sigma_range)
 
 
 def mean(
-    template,
+    template: NDArray,
     width: int,
     **kwargs: dict,
-):
+) -> NDArray:
     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,
@@ -120,7 +175,7 @@ def wedge(
     omit_negative_frequencies: bool = True,
     extrude_plane: bool = True,
     infinite_plane: bool = True,
-):
+) -> NDArray:
     template_ft = np.fft.rfftn(template)
 
     if tilt_step <= 0:
@@ -139,15 +194,14 @@ def wedge(
         template = np.real(np.fft.irfftn(template_ft))
         return template
 
-    tilt_angles = np.arange(-tilt_start, tilt_stop, tilt_step)
-    angles = np.zeros((template.ndim, tilt_angles.size))
-    angles[tilt_axis, :] = tilt_angles
-
-    wedge_mask = preprocessor.wedge_mask(
-        tilt_angles=angles,
+    wedge_mask = preprocessor.step_wedge_mask(
+        start_tilt=tilt_start,
+        stop_tilt=tilt_stop,
+        tilt_axis=tilt_axis,
+        tilt_step=tilt_step,
+        opening_axis=opening_axis,
         shape=template.shape,
         sigma=gaussian_sigma,
-        opening_axes=opening_axis,
         omit_negative_frequencies=omit_negative_frequencies,
     )
     np.multiply(template_ft, wedge_mask, out=template_ft)
@@ -155,7 +209,7 @@ def wedge(
     return template
 
 
-def compute_power_spectrum(template: NDArray):
+def compute_power_spectrum(template: NDArray) -> NDArray:
     return np.fft.fftshift(np.log(np.abs(np.fft.fftn(template))))
 
 
@@ -220,6 +274,8 @@ WRAPPED_FUNCTIONS = {
     "mean_filter": mean,
     "wedge_filter": wedge,
     "power_spectrum": compute_power_spectrum,
+    "resolution_gaussian": resolution_gaussian,
+    "resolution_sphere": resolution_sphere,
 }
 
 EXCLUDED_FUNCTIONS = [
@@ -332,6 +388,9 @@ class FilterWidget(widgets.Container):
         function_name = self.name_mapping.get(self.method_dropdown.value)
         function = self._get_function(function_name)
 
+        if "sampling_rate" in inspect.getfullargspec(function).args:
+            kwargs["sampling_rate"] = selected_layer_metadata["sampling_rate"]
+
         processed_data = function(selected_layer.data, **kwargs)
 
         new_layer_name = f"{selected_layer.name} ({self.method_dropdown.value})"
@@ -358,7 +417,7 @@ class FilterWidget(widgets.Container):
 
 def sphere_mask(
     template: NDArray, center_x: float, center_y: float, center_z: float, radius: float
-):
+) -> NDArray:
     return create_mask(
         mask_type="ellipse",
         shape=template.shape,
@@ -375,7 +434,7 @@ def ellipsod_mask(
     radius_x: float,
     radius_y: float,
     radius_z: float,
-):
+) -> NDArray:
     return create_mask(
         mask_type="ellipse",
         shape=template.shape,
@@ -392,7 +451,7 @@ def box_mask(
     height_x: int,
     height_y: int,
     height_z: int,
-):
+) -> NDArray:
     return create_mask(
         mask_type="box",
         shape=template.shape,
@@ -410,7 +469,7 @@ def tube_mask(
     inner_radius: float,
     outer_radius: float,
     height: int,
-):
+) -> NDArray:
     return create_mask(
         mask_type="tube",
         shape=template.shape,
@@ -433,7 +492,7 @@ def wedge_mask(
     omit_negative_frequencies: bool = False,
     extrude_plane: bool = True,
     infinite_plane: bool = True,
-):
+) -> NDArray:
     if tilt_step <= 0:
         wedge_mask = preprocessor.continuous_wedge_mask(
             start_tilt=tilt_start,
@@ -449,25 +508,24 @@ def wedge_mask(
         wedge_mask = np.fft.fftshift(wedge_mask)
         return wedge_mask
 
-    tilt_angles = np.arange(-tilt_start, tilt_stop + tilt_step, tilt_step)
-    angles = np.zeros((template.ndim, tilt_angles.size))
-
-    angles[tilt_axis, :] = tilt_angles
-
-    wedge_mask = preprocessor.wedge_mask(
-        tilt_angles=angles,
+    wedge_mask = preprocessor.step_wedge_mask(
+        start_tilt=tilt_start,
+        stop_tilt=tilt_stop,
+        tilt_axis=tilt_axis,
+        tilt_step=tilt_step,
+        opening_axis=opening_axis,
         shape=template.shape,
         sigma=gaussian_sigma,
-        opening_axes=opening_axis,
         omit_negative_frequencies=omit_negative_frequencies,
     )
+
     wedge_mask = np.fft.fftshift(wedge_mask)
     return wedge_mask
 
 
 def threshold_mask(
     template: NDArray, standard_deviation: float = 5.0, invert: bool = False
-):
+) -> NDArray:
     template_mean = template.mean()
     template_deviation = standard_deviation * template.std()
     upper = template_mean + template_deviation
@@ -479,6 +537,15 @@ def threshold_mask(
     return mask
 
 
+def lowpass_mask(template: NDArray, sigma: float = 1.0):
+    template = template / template.max()
+    template = (template > np.exp(-2)) * 128.0
+    template = preprocessor.gaussian_filter(template=template, sigma=sigma)
+    mask = template > np.exp(-2)
+
+    return mask
+
+
 class MaskWidget(widgets.Container):
     def __init__(self, viewer):
         super().__init__(layout="vertical")
@@ -496,6 +563,7 @@ class MaskWidget(widgets.Container):
             "Box": box_mask,
             "Wedge": wedge_mask,
             "Threshold": threshold_mask,
+            "Lowpass": lowpass_mask,
         }
 
         self.method_dropdown = widgets.ComboBox(
@@ -509,7 +577,6 @@ class MaskWidget(widgets.Container):
 
         self.adapt_button = widgets.PushButton(text="Adapt to layer", enabled=False)
         self.adapt_button.changed.connect(self._update_initial_values)
-
         self.viewer.layers.selection.events.active.connect(
             self._update_action_button_state
         )
@@ -520,8 +587,9 @@ class MaskWidget(widgets.Container):
         # self.density_field.value = f"Positive Density in Mask: {0:.2f}%"
 
         self.append(self.method_dropdown)
-        self.append(self.percentile_range_edit)
         self.append(self.adapt_button)
+        self.append(self.percentile_range_edit)
+
         self.append(self.align_button)
         self.append(self.action_button)
         self.append(self.density_field)
@@ -559,6 +627,7 @@ class MaskWidget(widgets.Container):
             arr=active_layer.data,
             rotation_matrix=rotation_matrix,
             use_geometric_center=False,
+            order=1,
         )
         eps = np.finfo(rotated_data.dtype).eps
         rotated_data[rotated_data < eps] = 0
@@ -589,10 +658,10 @@ class MaskWidget(widgets.Container):
             dict(zip(["height_x", "height_y", "height_z"], coordinates_heights))
         )
 
-        defaults["radius"] = np.min(coordinate_radius)
+        defaults["radius"] = np.max(coordinate_radius)
         defaults["inner_radius"] = np.min(coordinate_radius)
         defaults["outer_radius"] = np.max(coordinate_radius)
-        defaults["height"] = defaults["radius"]
+        defaults["height"] = np.max(coordinates_heights)
 
         for widget in self.action_widgets:
             if widget.name in defaults:
@@ -804,6 +873,7 @@ def parse_args():
     args = parser.parse_args()
     return args
 
+
 if __name__ == "__main__":
     parse_args()
     main()

tme/__version__.py

@@ -1 +1 @@
-__version__ = "0.1.6"
+__version__ = "0.1.8"

tme/analyzer.py

@@ -136,13 +136,16 @@ class PeakCaller(ABC):
             upper_limit = backend.subtract(
                 score_space.shape, self.min_boundary_distance
             )
-            valid_peaks = backend.sum(
-                backend.multiply(
-                    peak_positions < upper_limit,
-                    peak_positions >= self.min_boundary_distance,
-                ),
-                axis=1,
-            ) == peak_positions.shape[1]
+            valid_peaks = (
+                backend.sum(
+                    backend.multiply(
+                        peak_positions < upper_limit,
+                        peak_positions >= self.min_boundary_distance,
+                    ),
+                    axis=1,
+                )
+                == peak_positions.shape[1]
+            )
             if backend.sum(valid_peaks) == 0:
                 return None
 
@@ -261,7 +264,9 @@ class PeakCaller(ABC):
         backend.add(peak_positions, translation_offset, out=peak_positions)
         if not len(self.peak_list):
             self.peak_list = [peak_positions, rotations, peak_scores, peak_details]
-            peak_scores, peak_details, dim = (), (), peak_positions.shape[1]
+            dim = peak_positions.shape[1]
+            peak_scores = backend.zeros((0,), peak_scores.dtype)
+            peak_details = backend.zeros((0,), peak_details.dtype)
             rotations = backend.zeros((0, dim, dim), rotations.dtype)
             peak_positions = backend.zeros((0, dim), peak_positions.dtype)
 

tme/backends/__init__.py

@@ -11,6 +11,7 @@ from .matching_backend import MatchingBackend
 from .npfftw_backend import NumpyFFTWBackend
 from .pytorch_backend import PytorchBackend
 from .cupy_backend import CupyBackend
+from .mlx_backend import MLXBackend
 
 
 class BackendManager:

tme/backends/cupy_backend.py

@@ -45,7 +45,11 @@ class CupyBackend(NumpyFFTWBackend):
         self.maximum_filter = maximum_filter
 
     def to_backend_array(self, arr: NDArray) -> CupyArray:
-        if isinstance(arr, self._array_backend.ndarray):
+        current_device = self._array_backend.cuda.device.get_device_id()
+        if (
+            isinstance(arr, self._array_backend.ndarray)
+            and arr.device.id == current_device
+        ):
             return arr
         return self._array_backend.asarray(arr)
 

tme/backends/matching_backend.py

@@ -373,7 +373,7 @@ class MatchingBackend(ABC):
         """
 
     @abstractmethod
-    def std(self, arr: ArrayLike) -> Scalar:
+    def std(self, arr: ArrayLike, axis: Scalar) -> Scalar:
         """
         Compute the standad deviation of array elements.
 
@@ -381,6 +381,8 @@ class MatchingBackend(ABC):
         ----------
         arr : Scalar
             The array whose standard deviation should be computed.
+        axis : Scalar
+            Axis to perform the operation on.
 
         Returns
         -------