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

tme/matching_exhaustive.py

@@ -308,6 +308,49 @@ def cam_setup(**kwargs):
     return corr_setup(**kwargs)
 
 
+def _normalize_under_mask(template: NDArray, mask: NDArray, mask_intensity) -> None:
+    """
+    Standardizes the values in in template by subtracting the mean and dividing by the
+    standard deviation based on the elements in mask. Subsequently, the template is
+    multiplied by the mask.
+
+    Parameters
+    ----------
+    template : NDArray
+        The data array to be normalized. This array is modified in-place.
+    mask : NDArray
+        A boolean array of the same shape as `template`. True values indicate the positions in `template`
+        to consider for normalization.
+    mask_intensity : float
+        Mask intensity used to compute expectations.
+
+    References
+    ----------
+    .. [1]  T. Hrabe, Y. Chen, S. Pfeffer, L. Kuhn Cuellar, A.-V. Mangold,
+            and F. Förster, J. Struct. Biol. 178, 177 (2012).
+    .. [2]  M. L. Chaillet, G. van der Schot, I. Gubins, S. Roet,
+            R. C. Veltkamp, and F. Förster, Int. J. Mol. Sci. 24,
+            13375 (2023)
+
+    Returns
+    -------
+    None
+        This function modifies `template` in-place and does not return any value.
+    """
+    masked_mean = backend.sum(backend.multiply(template, mask))
+    masked_mean = backend.divide(masked_mean, mask_intensity)
+    masked_std = backend.sum(backend.multiply(backend.square(template), mask))
+    masked_std = backend.subtract(
+        masked_std / mask_intensity, backend.square(masked_mean)
+    )
+    masked_std = backend.sqrt(backend.maximum(masked_std, 0))
+
+    backend.subtract(template, masked_mean, out=template)
+    backend.divide(template, masked_std, out=template)
+    backend.multiply(template, mask, out=template)
+    return None
+
+
 def flc_setup(
     rfftn: Callable,
     irfftn: Callable,
@@ -359,7 +402,8 @@ def flc_setup(
     ft_target = backend.preallocate_array(fast_ft_shape, complex_dtype)
     ft_target2 = backend.preallocate_array(fast_ft_shape, complex_dtype)
     rfftn(target_pad, ft_target)
-    rfftn(backend.square(target_pad), ft_target2)
+    backend.square(target_pad, out=target_pad)
+    rfftn(target_pad, ft_target2)
 
     # Convert arrays used in subsequent fitting to SharedMemory objects
     ft_target = backend.arr_to_sharedarr(
@@ -369,13 +413,9 @@ def flc_setup(
         arr=ft_target2, shared_memory_handler=shared_memory_handler
     )
 
-    template_mask = template_mask > 0
-    template_mean = backend.mean(template[template_mask])
-    template_std = backend.std(template[template_mask])
-    template_mask = backend.astype(template_mask, real_dtype)
-
-    backend.divide(template - template_mean, template_std, out=template)
-    backend.multiply(template, template_mask, out=template)
+    _normalize_under_mask(
+        template=template, mask=template_mask, mask_intensity=backend.sum(template_mask)
+    )
 
     template_buffer = backend.arr_to_sharedarr(
         arr=template, shared_memory_handler=shared_memory_handler
@@ -455,7 +495,6 @@ def flcSphericalMask_setup(
     :py:meth:`corr_scoring`
     """
     target_pad = backend.topleft_pad(target, fast_shape)
-    template_mask_pad = backend.topleft_pad(template_mask, fast_shape)
 
     # Target and squared target window sums
     ft_target = backend.preallocate_array(fast_ft_shape, complex_dtype)
@@ -467,16 +506,17 @@ def flcSphericalMask_setup(
     numerator2 = backend.preallocate_array(1, real_dtype)
 
     eps = backend.eps(real_dtype)
-    n_observations = backend.sum(template_mask_pad > np.exp(-2))
+    n_observations = backend.sum(template_mask)
+
+    template_mask_pad = backend.topleft_pad(template_mask, fast_shape)
     rfftn(template_mask_pad, ft_template_mask)
 
-    # Variance part denominator
+    # Denominator E(X^2) - E(X)^2
     rfftn(backend.square(target_pad), ft_target)
     backend.multiply(ft_target, ft_template_mask, out=ft_temp)
     irfftn(ft_temp, temp2)
     backend.divide(temp2, n_observations, out=temp2)
 
-    # Mean part denominator
     rfftn(target_pad, ft_target)
     backend.multiply(ft_target, ft_template_mask, out=ft_temp)
     irfftn(ft_temp, temp)
@@ -495,14 +535,10 @@ def flcSphericalMask_setup(
     backend.fill(temp2, 0)
     temp2[nonzero_indices] = 1 / temp[nonzero_indices]
 
-    template_mask = template_mask > np.exp(-2)
-    template_mean = backend.mean(template[template_mask])
-    template_std = backend.std(template[template_mask])
-
-    template = backend.divide(backend.subtract(template, template_mean), template_std)
-    backend.multiply(template, template_mask, out=template)
+    _normalize_under_mask(
+        template=template, mask=template_mask, mask_intensity=backend.sum(template_mask)
+    )
 
-    # Convert arrays used in subsequent fitting to SharedMemory objects
     template_buffer = backend.arr_to_sharedarr(
         arr=template, shared_memory_handler=shared_memory_handler
     )
@@ -773,6 +809,7 @@ def corr_scoring(
     fourier_shift = callback_class_args.get("fourier_shift", backend.zeros(arr.ndim))
     fourier_shift_scores = backend.sum(fourier_shift != 0) != 0
 
+    template_sum = backend.sum(template)
     for index in range(rotations.shape[0]):
         rotation = rotations[index]
         backend.fill(arr, 0)
@@ -783,6 +820,9 @@ def corr_scoring(
             use_geometric_center=False,
             order=interpolation_order,
         )
+        rotation_norm = template_sum / backend.sum(arr)
+        backend.multiply(arr, rotation_norm, out=arr)
+
         rfftn(arr, ft_temp)
         template_filter_func(ft_temp, template_filter, out=ft_temp)
 
@@ -905,6 +945,7 @@ def flc_scoring(
     fourier_shift = callback_class_args.get("fourier_shift", backend.zeros(arr.ndim))
     fourier_shift_scores = backend.sum(fourier_shift != 0) != 0
 
+    unpadded_slice = tuple(slice(0, stop) for stop in template.shape)
     for index in range(rotations.shape[0]):
         rotation = rotations[index]
         backend.fill(arr, 0)
@@ -918,8 +959,15 @@ def flc_scoring(
             use_geometric_center=False,
             order=interpolation_order,
         )
+        # Given the amount of FFTs, might aswell normalize properly
         n_observations = backend.sum(temp)
 
+        _normalize_under_mask(
+            template=arr[unpadded_slice],
+            mask=temp[unpadded_slice],
+            mask_intensity=n_observations,
+        )
+
         rfftn(temp, ft_temp)
 
         backend.multiply(ft_target, ft_temp, out=ft_denom)

tme/matching_utils.py

@@ -763,10 +763,14 @@ def euler_to_rotationmatrix(angles: Tuple[float]) -> NDArray:
     NDArray
         The generated rotation matrix.
     """
-    if len(angles) == 1:
+    n_angles = len(angles)
+    angle_convention = "zyx"[:n_angles]
+    if n_angles == 1:
         angles = (angles, 0, 0)
     rotation_matrix = (
-        Rotation.from_euler("zyx", angles, degrees=True).as_matrix().astype(np.float32)
+        Rotation.from_euler(angle_convention, angles, degrees=True)
+        .as_matrix()
+        .astype(np.float32)
     )
     return rotation_matrix
 
@@ -1052,7 +1056,7 @@ def tube_mask(
     symmetry_axis : int
         The axis of symmetry for the tube.
     base_center : tuple
-        Center of the base circle of the tube.
+        Center of the tube.
     inner_radius : float
         Inner radius of the tube.
     outer_radius : float
@@ -1068,8 +1072,9 @@ def tube_mask(
     Raises
     ------
     ValueError
-        If the inner radius is larger than the outer radius. Or height is larger
-        than the symmetry axis shape.
+        If the inner radius is larger than the outer radius, height is larger
+        than the symmetry axis shape, or if base_center and shape do not have the
+        same length.
     """
     if inner_radius > outer_radius:
         raise ValueError("inner_radius should be smaller than outer_radius.")
@@ -1080,8 +1085,11 @@ def tube_mask(
     if symmetry_axis > len(shape):
         raise ValueError(f"symmetry_axis can be not larger than {len(shape)}.")
 
+    if len(base_center) != len(shape):
+        raise ValueError("shape and base_center need to have the same length.")
+
     circle_shape = tuple(b for ix, b in enumerate(shape) if ix != symmetry_axis)
-    base_center = tuple(b for ix, b in enumerate(base_center) if ix != symmetry_axis)
+    circle_center = tuple(b for ix, b in enumerate(base_center) if ix != symmetry_axis)
 
     inner_circle = np.zeros(circle_shape)
     outer_circle = np.zeros_like(inner_circle)
@@ -1090,34 +1098,39 @@ def tube_mask(
             mask_type="ellipse",
             shape=circle_shape,
             radius=inner_radius,
-            center=base_center,
+            center=circle_center,
         )
     if outer_radius > 0:
         outer_circle = create_mask(
             mask_type="ellipse",
             shape=circle_shape,
             radius=outer_radius,
-            center=base_center,
+            center=circle_center,
         )
     circle = outer_circle - inner_circle
     circle = np.expand_dims(circle, axis=symmetry_axis)
 
-    center = shape[symmetry_axis] // 2
-    start_idx = center - height // 2
-    stop_idx = center + height // 2 + height % 2
+    center = base_center[symmetry_axis]
+    start_idx = int(center - height // 2)
+    stop_idx = int(center + height // 2 + height % 2)
+
+    start_idx, stop_idx = max(start_idx, 0), min(stop_idx, shape[symmetry_axis])
 
     slice_indices = tuple(
         slice(None) if i != symmetry_axis else slice(start_idx, stop_idx)
         for i in range(len(shape))
     )
     tube = np.zeros(shape)
-    tube[slice_indices] = np.repeat(circle, height, axis=symmetry_axis)
+    tube[slice_indices] = circle
 
     return tube
 
 
 def scramble_phases(
-    arr: NDArray, noise_proportion: float = 0.5, seed: int = 42
+    arr: NDArray,
+    noise_proportion: float = 0.5,
+    seed: int = 42,
+    normalize_power: bool = True,
 ) -> NDArray:
     """
     Applies random phase scrambling to a given array.
@@ -1135,6 +1148,8 @@ def scramble_phases(
         The proportion of noise in the phase scrambling, by default 0.5.
     seed : int, optional
         The seed for the random phase scrambling, by default 42.
+    normalize_power : bool, optional
+        Whether the returned template should have the same sum of squares as arr.
 
     Returns
     -------
@@ -1158,6 +1173,17 @@ def scramble_phases(
     ph_noise = np.random.permutation(ph)
     ph_new = ph * (1 - noise_proportion) + ph_noise * noise_proportion
     ret = np.real(np.fft.ifftn(amp * np.exp(1j * ph_new)))
+
+    if normalize_power:
+        np.divide(
+            np.subtract(ret, ret.min()), np.subtract(ret.max(), ret.min()), out=ret
+        )
+        np.multiply(ret, np.subtract(arr.max(), arr.min()), out=ret)
+        np.add(ret, arr.min(), out=ret)
+
+        scaling = np.divide(np.abs(arr).sum(), np.abs(ret).sum())
+        np.multiply(ret, scaling, out=ret)
+
     return ret
 
 

tme/preprocessor.py

@@ -1055,8 +1055,8 @@ class Preprocessor:
         be equivalent to the following
 
         >>> wedge = Preprocessor().continuous_wedge_mask(
-        >>>    shape = (50,50,50),
-        >>>    start_tilt = 50,
+        >>>    shape=(50,50,50),
+        >>>    start_tilt=50,
         >>>    stop_tilt=55,
         >>>    tilt_axis=1,
         >>>    omit_negative_frequencies=False,
@@ -1074,6 +1074,7 @@ class Preprocessor:
 
         See Also
         --------
+        :py:meth:`Preprocessor.step_wedge_mask`
         :py:meth:`Preprocessor.continuous_wedge_mask`
         """
         opening_axes = np.asarray(opening_axes)
@@ -1124,6 +1125,112 @@ class Preprocessor:
 
         return wedge_volume
 
+    def step_wedge_mask(
+        self,
+        start_tilt: float,
+        stop_tilt: float,
+        tilt_step: float,
+        shape: Tuple[int],
+        opening_axis: int = 0,
+        tilt_axis: int = 2,
+        sigma: float = 0,
+        omit_negative_frequencies: bool = True,
+    ) -> NDArray:
+        """
+        Create a wedge mask with the same shape as template by rotating a
+        plane according to tilt angles. The DC component of the filter is at the origin.
+
+        Parameters
+        ----------
+        start_tilt : float
+            Starting tilt angle in degrees, e.g. a stage tilt of 70 degrees
+            would yield a start_tilt value of 70.
+        stop_tilt : float
+            Ending tilt angle in degrees, , e.g. a stage tilt of -70 degrees
+            would yield a stop_tilt value of 70.
+        tilt_step : float
+            Angle between the different tilt planes.
+        shape : Tuple of ints
+            Shape of the output wedge array.
+        tilt_axis : int, optional
+            Axis that the plane is tilted over.
+            - 0 for Z-axis
+            - 1 for Y-axis
+            - 2 for X-axis
+        opening_axis : int, optional
+            Axis running through the void defined by the wedge.
+            - 0 for Z-axis
+            - 1 for Y-axis
+            - 2 for X-axis
+        sigma : float, optional
+            Standard deviation for Gaussian kernel used for smoothing the wedge.
+        omit_negative_frequencies : bool, optional
+            Whether the wedge mask should omit negative frequencies, i.e. be
+            applicable to symmetric Fourier transforms (see :obj:`numpy.fft.fftn`)
+
+        Returns
+        -------
+        NDArray
+            A numpy array containing the wedge mask.
+
+        Notes
+        -----
+        This function is equivalent to :py:meth:`Preprocessor.wedge_mask`, but much faster
+        for large shapes because it only considers a single tilt angle rather than the rotation
+        of an N-1 dimensional hyperplane in N dimensions.
+
+        See Also
+        --------
+        :py:meth:`Preprocessor.wedge_mask`
+        :py:meth:`Preprocessor.continuous_wedge_mask`
+        """
+        tilt_angles = np.arange(-start_tilt, stop_tilt + tilt_step, tilt_step)
+        plane = np.zeros((shape[opening_axis], shape[tilt_axis]), dtype=np.float32)
+        subset = tuple(
+            slice(None) if i != 0 else slice(x // 2, x // 2 + 1)
+            for i, x in enumerate(plane.shape)
+        )
+        plane[subset] = 1
+        plane_rotated, wedge_volume = np.zeros_like(plane), np.zeros_like(plane)
+        for index in range(tilt_angles.shape[0]):
+            plane_rotated.fill(0)
+
+            rotation_matrix = euler_to_rotationmatrix((tilt_angles[index], 0))
+            rotation_matrix = rotation_matrix[np.ix_((0, 1), (0, 1))]
+
+            Density.rotate_array(
+                arr=plane,
+                rotation_matrix=rotation_matrix,
+                out=plane_rotated,
+                use_geometric_center=True,
+                order=1,
+            )
+            wedge_volume += plane_rotated
+
+        wedge_volume = self.gaussian_filter(
+            template=wedge_volume, sigma=sigma, fourier=False
+        )
+        wedge_volume = np.where(wedge_volume > np.exp(-2), 1, 0)
+
+        if opening_axis > tilt_axis:
+            wedge_volume = np.moveaxis(wedge_volume, 1, 0)
+
+        reshape_dimensions = tuple(
+            x if i in (opening_axis, tilt_axis) else 1 for i, x in enumerate(shape)
+        )
+
+        wedge_volume = wedge_volume.reshape(reshape_dimensions)
+        tile_dimensions = np.divide(shape, reshape_dimensions).astype(int)
+        wedge_volume = np.tile(wedge_volume, tile_dimensions)
+
+        wedge_volume = np.fft.ifftshift(wedge_volume)
+
+        if omit_negative_frequencies:
+            stop = 1 + (wedge_volume.shape[-1] // 2)
+            wedge_volume = wedge_volume[..., :stop]
+
+        return wedge_volume
+
     def continuous_wedge_mask(
         self,
         start_tilt: float,
@@ -1192,6 +1299,7 @@ class Preprocessor:
         See Also
         --------
         :py:meth:`Preprocessor.wedge_mask`
+        :py:meth:`Preprocessor.step_wedge_mask`
         """
         shape_center = np.divide(shape, 2).astype(int)
 
@@ -1348,7 +1456,7 @@ class LinearWhiteningFilter:
     def filter(
         self, template: NDArray, n_bins: int = None
     ) -> Tuple[NDArray, NDArray, NDArray]:
-        max_bins = int(np.linalg.norm(template.shape) // 2 + 1)
+        max_bins = np.max(template.shape) // 2 + 1
         n_bins = max_bins if n_bins is None else n_bins
         n_bins = int(min(n_bins, max_bins))
 
@@ -1371,7 +1479,7 @@ class LinearWhiteningFilter:
         np.multiply(fourier_transform, radial_averages[bins], out=fourier_transform)
 
         ret = np.fft.irfftn(
-            np.fft.ifftshift(fourier_transform, axes=fft_shift_axes)
+            np.fft.ifftshift(fourier_transform, axes=fft_shift_axes), s=template.shape
         ).real
         return ret, bin_edges, radial_averages
 
@@ -1389,7 +1497,7 @@ class LinearWhiteningFilter:
         bins = np.digitize(frequency_grid, bins=bin_edges, right=True)
         np.multiply(fourier_transform, radial_averages[bins], out=fourier_transform)
         ret = np.fft.irfftn(
-            np.fft.ifftshift(fourier_transform, axes=fft_shift_axes)
+            np.fft.ifftshift(fourier_transform, axes=fft_shift_axes), s=template.shape
         ).real
 
         return ret