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

scripts/postprocess.py

@@ -15,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 (
@@ -28,6 +29,7 @@ from tme.matching_utils import (
     load_pickle,
     euler_to_rotationmatrix,
     euler_from_rotationmatrix,
+    centered_mask,
 )
 
 PEAK_CALLERS = {
@@ -75,6 +77,13 @@ def parse_args():
         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",
         type=str,
@@ -403,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(
@@ -419,6 +442,17 @@ 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,
@@ -458,19 +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 = Density.from_structure(template, sampling_rate=sampling_rate)
+    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))

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,22 +94,30 @@ 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):
+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
@@ -123,34 +131,39 @@ def resolution_sphere(template: NDArray, cutoff_angstrom: float, highpass : bool
         omit_negative_frequencies=False,
     )
 
-    mask = np.fft.ifftshift(mask)
     template_ft = np.fft.fftn(template)
-    np.multiply(template_ft, mask, out = template_ft)
-
+    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):
+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()
+        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))
+        np.max(2 * sampling_rate / cutoff_angstrom), np.sqrt(2 * np.log(2))
     )
 
-    mask = np.exp(-grid ** 2 / (2 * sigma_fourier ** 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)
 
+    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,
@@ -162,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:
@@ -181,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)
@@ -197,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))))
 
 
@@ -262,8 +274,8 @@ WRAPPED_FUNCTIONS = {
     "mean_filter": mean,
     "wedge_filter": wedge,
     "power_spectrum": compute_power_spectrum,
-    "resolution_gaussian" : resolution_gaussian,
-    "resolution_sphere" : resolution_sphere,
+    "resolution_gaussian": resolution_gaussian,
+    "resolution_sphere": resolution_sphere,
 }
 
 EXCLUDED_FUNCTIONS = [
@@ -405,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,
@@ -422,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,
@@ -439,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,
@@ -457,7 +469,7 @@ def tube_mask(
     inner_radius: float,
     outer_radius: float,
     height: int,
-):
+) -> NDArray:
     return create_mask(
         mask_type="tube",
         shape=template.shape,
@@ -480,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,
@@ -496,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
@@ -526,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")
@@ -543,6 +563,7 @@ class MaskWidget(widgets.Container):
             "Box": box_mask,
             "Wedge": wedge_mask,
             "Threshold": threshold_mask,
+            "Lowpass": lowpass_mask,
         }
 
         self.method_dropdown = widgets.ComboBox(
@@ -606,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
@@ -636,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:

tme/__version__.py

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

tme/backends/npfftw_backend.py

@@ -613,13 +613,6 @@ class NumpyFFTWBackend(MatchingBackend):
         rotate_mask = arr_mask is not None and mask_coordinates is not None
         return_type = (out is None) + 2 * rotate_mask * (out_mask is None)
 
-        # Otherwise array might be slightly shifted by centering
-        if np.allclose(
-            rotation_matrix,
-            np.eye(rotation_matrix.shape[0], dtype=rotation_matrix.dtype),
-        ):
-            center_rotation = False
-
         coordinates_rotated = np.empty(coordinates.shape, dtype=rotation_matrix.dtype)
         mask_rotated = (
             np.empty(mask_coordinates.shape, dtype=rotation_matrix.dtype)

tme/density.py

@@ -199,16 +199,16 @@ class Density:
             func = cls._load_mrc
             if filename.endswith(".em") or filename.endswith(".em.gz"):
                 func = cls._load_em
-            data, origin, sampling_rate = func(
+            data, origin, sampling_rate, meta = func(
                 filename=filename, subset=subset, use_memmap=use_memmap
             )
         except ValueError:
-            data, origin, sampling_rate = cls._load_skio(filename=filename)
+            data, origin, sampling_rate, meta = cls._load_skio(filename=filename)
             if subset is not None:
                 cls._validate_slices(slices=subset, shape=data.shape)
                 data = data[subset].copy()
 
-        return cls(data=data, origin=origin, sampling_rate=sampling_rate)
+        return cls(data=data, origin=origin, sampling_rate=sampling_rate, metadata=meta)
 
     @classmethod
     def _load_mrc(
@@ -307,6 +307,13 @@ class Density:
 
             extended_header = mrc.header.nsymbt
 
+            metadata = {
+                "min": float(mrc.header.dmin),
+                "max": float(mrc.header.dmax),
+                "mean": float(mrc.header.dmean),
+                "std": float(mrc.header.rms),
+            }
+
         if is_gzipped(filename):
             if use_memmap:
                 warnings.warn(
@@ -329,7 +336,7 @@ class Density:
                 dtype=data_type,
                 header_size=1024 + extended_header,
             )
-            return data, origin, sampling_rate
+            return data, origin, sampling_rate, metadata
 
         if not use_memmap:
             with mrcfile.open(filename, header_only=False) as mrc:
@@ -343,7 +350,7 @@ class Density:
             data = np.transpose(data, crs_index)
             start = np.take(start, crs_index)
 
-        return data, origin, sampling_rate
+        return data, origin, sampling_rate, metadata
 
     @classmethod
     def _load_em(
@@ -448,7 +455,7 @@ class Density:
             pixel_size = 1
         sampling_rate = np.repeat(pixel_size, data.ndim).astype(data.dtype)
 
-        return data, origin, sampling_rate
+        return data, origin, sampling_rate, {}
 
     @staticmethod
     def _validate_slices(slices: Tuple[slice], shape: Tuple[int]):
@@ -592,7 +599,7 @@ class Density:
         warnings.warn(
             "origin and sampling_rate are not yet extracted from non CCP4/MRC files."
         )
-        return data, np.zeros(data.ndim), np.ones(data.ndim)
+        return data, np.zeros(data.ndim), np.ones(data.ndim), {}
 
     @classmethod
     def from_structure(
@@ -909,7 +916,8 @@ class Density:
         Returns a copy of the current class instance with all elements in
         :py:attr:`Density.data` set to zero. :py:attr:`Density.origin` and
         :py:attr:`Density.sampling_rate` will be copied, while
-        :py:attr:`Density.metadata` will be initialized to an empty dictionary.
+        :py:attr:`Density.metadata` will be initialized to contain min, max,
+        mean and standard deviation of :py:attr:`Density.data`.
 
         Examples
         --------
@@ -924,6 +932,7 @@ class Density:
             data=np.zeros_like(self.data),
             origin=deepcopy(self.origin),
             sampling_rate=deepcopy(self.sampling_rate),
+            metadata={"min": 0, "max": 0, "mean": 0, "std": 0},
         )
 
     def copy(self) -> "Density":

tme/matching_data.py

@@ -4,7 +4,7 @@
 
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
-
+import warnings
 from typing import Tuple, List
 
 import numpy as np
@@ -81,7 +81,11 @@ class MatchingData:
         return arr
 
     def subset_array(
-        self, arr: NDArray, arr_slice: Tuple[slice], padding: NDArray
+        self,
+        arr: NDArray,
+        arr_slice: Tuple[slice],
+        padding: NDArray,
+        invert: bool = False,
     ) -> NDArray:
         """
         Extract a subset of the input array according to the given slice and
@@ -96,8 +100,11 @@ class MatchingData:
         padding : NDArray
             Padding values for each dimension. If the padding exceeds the array
             dimensions, the extra regions are filled with the mean of the array
-            values, otherwise, the
-            values in ``arr`` are used.
+            values, otherwise, the values in ``arr`` are used.
+        invert : bool, optional
+            Whether the returned array should be inverted and normalized to the interval
+            [0, 1]. If available, uses the metadata information of the Density object,
+            otherwise computes min and max on the extracted subset.
 
         Returns
         -------
@@ -109,7 +116,6 @@ class MatchingData:
 
         slice_start = np.array([x.start for x in arr_slice], dtype=int)
         slice_stop = np.array([x.stop for x in arr_slice], dtype=int)
-        slice_shape = np.subtract(slice_stop, slice_start)
 
         padding = np.add(padding, np.mod(padding, 2))
         left_pad = right_pad = np.divide(padding, 2).astype(int)
@@ -119,20 +125,18 @@ class MatchingData:
             np.subtract(arr.shape, slice_stop), right_pad
         ).astype(int)
 
-        ret_shape = np.add(slice_shape, padding)
         arr_start = np.subtract(slice_start, data_voxels_left)
         arr_stop = np.add(slice_stop, data_voxels_right)
         arr_slice = tuple(slice(*pos) for pos in zip(arr_start, arr_stop))
         arr_mesh = self._slice_to_mesh(arr_slice, arr.shape)
 
-        subset_start = np.subtract(left_pad, data_voxels_left)
-        subset_stop = np.add(subset_start, np.subtract(arr_stop, arr_start))
-        subset_slice = tuple(slice(*prod) for prod in zip(subset_start, subset_stop))
-        subset_mesh = self._slice_to_mesh(subset_slice, ret_shape)
-
+        arr_min, arr_max = None, None
         if type(arr) == Density:
             if type(arr.data) == np.memmap:
-                arr = Density.from_file(arr.data.filename, subset=arr_slice).data
+                dens = Density.from_file(arr.data.filename, subset=arr_slice)
+                arr = dens.data
+                arr_min = dens.metadata.get("min", None)
+                arr_max = dens.metadata.get("max", None)
             else:
                 arr = np.asarray(arr.data[*arr_mesh])
         else:
@@ -141,10 +145,38 @@ class MatchingData:
                     arr.filename, mode="r", shape=arr.shape, dtype=arr.dtype
                 )
             arr = np.asarray(arr[*arr_mesh])
-        ret = np.full(
-            shape=np.add(slice_shape, padding), fill_value=arr.mean(), dtype=arr.dtype
+
+        def _warn_on_mismatch(
+            expectation: float, computation: float, name: str
+        ) -> float:
+            expectation, computation = float(expectation), float(computation)
+            if expectation is None:
+                expectation = computation
+
+            if abs(computation) > abs(expectation):
+                warnings.warn(
+                    f"Computed {name} value is more extreme than value specified in file"
+                    f" (|{computation}| > |{expectation}|). This may lead to issues"
+                    " with padding and contrast inversion."
+                )
+
+            return expectation
+
+        padding = tuple(
+            (left, right)
+            for left, right in zip(
+                np.subtract(left_pad, data_voxels_left),
+                np.subtract(right_pad, data_voxels_right),
+            )
         )
-        ret[*subset_mesh] = arr
+        ret = np.pad(arr, padding, mode="reflect")
+
+        if invert:
+            arr_min = _warn_on_mismatch(arr_min, arr.min(), "min")
+            arr_max = _warn_on_mismatch(arr_max, arr.max(), "max")
+
+            np.subtract(-ret, arr_min, out=ret)
+            np.divide(ret, arr_max - arr_min, out=ret)
 
         return ret
 
@@ -198,12 +230,11 @@ class MatchingData:
         )
 
         target_subset = self.subset_array(
-            arr=self._target, arr_slice=target_slice, padding=target_pad
+            arr=self._target,
+            arr_slice=target_slice,
+            padding=target_pad,
+            invert=self._invert_target,
         )
-        if self._invert_target:
-            target_subset *= -1
-            target_min, target_max = target_subset.min(), target_subset.max()
-            target_subset = (target_subset - target_min) / (target_max - target_min)
 
         template_subset = self.subset_array(
             arr=self._template,
@@ -488,14 +519,14 @@ class MatchingData:
     @property
     def target(self):
         """Returns the target NDArray."""
-        if type(self._target) == Density:
+        if isinstance(self._target, Density):
             return self._target.data
         return self._target
 
     @property
     def template(self):
         """Returns the reversed template NDArray."""
-        if type(self._template) == Density:
+        if isinstance(self._template, Density):
             return backend.reverse(self._template.data)
         return backend.reverse(self._template)
 
@@ -524,7 +555,7 @@ class MatchingData:
     @property
     def target_mask(self):
         """Returns the target mask NDArray."""
-        if type(self._target_mask) == Density:
+        if isinstance(self._target_mask, Density):
             return self._target_mask.data
         return self._target_mask
 
@@ -553,7 +584,7 @@ class MatchingData:
         template : NDArray
             Array to set as the template.
         """
-        if type(self._template_mask) == Density:
+        if isinstance(self._template_mask, Density):
             return backend.reverse(self._template_mask.data)
         return backend.reverse(self._template_mask)