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

tme/preprocessing/frequency_filters.py

@@ -8,12 +8,12 @@ from math import log, sqrt
 from typing import Tuple, Dict
 
 import numpy as np
-from numpy.typing import NDArray
 from scipy.ndimage import mean as ndimean
 from scipy.ndimage import map_coordinates
 
-from ._utils import fftfreqn, crop_real_fourier, shift_fourier
-from ..backends import backend
+from ..types import BackendArray
+from ..backends import backend as be
+from ._utils import fftfreqn, crop_real_fourier, shift_fourier, compute_fourier_shape
 
 
 class BandPassFilter:
@@ -29,7 +29,7 @@ class BandPassFilter:
     highpass : float, optional
         The highpass cutoff, defaults to None.
     sampling_rate : Tuple[float], optional
-        The sampling rate in Fourier space, defaults to 1.
+        The sampling r_position_to_molmapate in Fourier space, defaults to 1.
     use_gaussian : bool, optional
         Whether to use Gaussian bandpass filter, defaults to True.
     return_real_fourier : bool, optional
@@ -63,7 +63,7 @@ class BandPassFilter:
         return_real_fourier: bool = False,
         shape_is_real_fourier: bool = False,
         **kwargs,
-    ) -> NDArray:
+    ) -> BackendArray:
         """
         Generate a bandpass filter using discrete frequency cutoffs.
 
@@ -86,7 +86,7 @@ class BandPassFilter:
 
         Returns:
         --------
-        NDArray
+        BackendArray
             The bandpass filter in Fourier space.
         """
         if shape_is_real_fourier:
@@ -127,7 +127,7 @@ class BandPassFilter:
         return_real_fourier: bool = False,
         shape_is_real_fourier: bool = False,
         **kwargs,
-    ) -> NDArray:
+    ) -> BackendArray:
         """
         Generate a bandpass filter using Gaussian functions.
 
@@ -150,7 +150,7 @@ class BandPassFilter:
 
         Returns:
         --------
-        NDArray
+        BackendArray
             The bandpass filter in Fourier space.
         """
         if shape_is_real_fourier:
@@ -162,29 +162,32 @@ class BandPassFilter:
             shape_is_real_fourier=shape_is_real_fourier,
             compute_euclidean_norm=True,
         )
-        grid = -backend.square(grid)
+        grid = be.to_backend_array(grid)
+        grid = -be.square(grid)
 
         lowpass_filter, highpass_filter = 1, 1
         norm = float(sqrt(2 * log(2)))
-        upper_sampling = float(backend.max(backend.multiply(2, sampling_rate)))
+        upper_sampling = float(
+            be.max(be.multiply(2, be.to_backend_array(sampling_rate)))
+        )
 
         if lowpass is not None:
             lowpass = float(lowpass)
-            lowpass = backend.maximum(lowpass, backend.eps(backend._float_dtype))
+            lowpass = be.maximum(lowpass, be.eps(be._float_dtype))
         if highpass is not None:
             highpass = float(highpass)
-            highpass = backend.maximum(highpass, backend.eps(backend._float_dtype))
+            highpass = be.maximum(highpass, be.eps(be._float_dtype))
 
         if lowpass is not None:
             lowpass = upper_sampling / (lowpass * norm)
-            lowpass = backend.multiply(2, backend.square(lowpass))
-            lowpass_filter = backend.exp(backend.divide(grid, lowpass))
+            lowpass = be.multiply(2, be.square(lowpass))
+            lowpass_filter = be.exp(be.divide(grid, lowpass))
         if highpass is not None:
             highpass = upper_sampling / (highpass * norm)
-            highpass = backend.multiply(2, backend.square(highpass))
-            highpass_filter = 1 - backend.exp(backend.divide(grid, highpass))
+            highpass = be.multiply(2, be.square(highpass))
+            highpass_filter = 1 - be.exp(be.divide(grid, highpass))
 
-        bandpass_filter = backend.multiply(lowpass_filter, highpass_filter)
+        bandpass_filter = be.multiply(lowpass_filter, highpass_filter)
         bandpass_filter = shift_fourier(
             data=bandpass_filter, shape_is_real_fourier=shape_is_real_fourier
         )
@@ -205,7 +208,7 @@ class BandPassFilter:
         mask = func(**func_args)
 
         return {
-            "data": backend.to_backend_array(mask),
+            "data": be.to_backend_array(mask),
             "sampling_rate": func_args.get("sampling_rate", 1),
             "is_multiplicative_filter": True,
         }
@@ -237,14 +240,14 @@ class LinearWhiteningFilter:
 
     @staticmethod
     def _compute_spectrum(
-        data_rfft: NDArray, n_bins: int = None, batch_dimension: int = None
-    ) -> Tuple[NDArray, NDArray]:
+        data_rfft: BackendArray, n_bins: int = None, batch_dimension: int = None
+    ) -> Tuple[BackendArray, BackendArray]:
         """
         Compute the spectrum of the input data.
 
         Parameters:
         -----------
-        data_rfft : NDArray
+        data_rfft : BackendArray
             The Fourier transform of the input data.
         n_bins : int, optional
             The number of bins for computing the spectrum, defaults to None.
@@ -253,9 +256,9 @@ class LinearWhiteningFilter:
 
         Returns:
         --------
-        bins : NDArray
+        bins : BackendArray
             Array containing the bin indices for the spectrum.
-        radial_averages : NDArray
+        radial_averages : BackendArray
             Array containing the radial averages of the spectrum.
         """
         shape = tuple(x for i, x in enumerate(data_rfft.shape) if i != batch_dimension)
@@ -270,7 +273,7 @@ class LinearWhiteningFilter:
             shape_is_real_fourier=True,
             compute_euclidean_norm=True,
         )
-        bins = backend.to_numpy_array(bins)
+        bins = be.to_numpy_array(bins)
 
         # Implicit lowpass to nyquist
         bins = np.floor(bins * (n_bins - 1) + 0.5).astype(int)
@@ -292,11 +295,11 @@ class LinearWhiteningFilter:
 
     @staticmethod
     def _interpolate_spectrum(
-        spectrum: NDArray,
+        spectrum: BackendArray,
         shape: Tuple[int],
         shape_is_real_fourier: bool = True,
         order: int = 1,
-    ) -> NDArray:
+    ) -> BackendArray:
         """
         References
         ----------
@@ -306,19 +309,19 @@ class LinearWhiteningFilter:
         """
         grid = fftfreqn(
             shape=shape,
-            sampling_rate=.5,
+            sampling_rate=0.5,
             shape_is_real_fourier=shape_is_real_fourier,
             compute_euclidean_norm=True,
         )
-        grid = backend.to_numpy_array(grid)
-        np.multiply(grid, (spectrum.shape[0] - 1), out = grid) + 0.5
+        grid = be.to_numpy_array(grid)
+        np.multiply(grid, (spectrum.shape[0] - 1), out=grid) + 0.5
         spectrum = map_coordinates(spectrum, grid.reshape(1, -1), order=order)
         return spectrum.reshape(grid.shape)
 
     def __call__(
         self,
-        data: NDArray = None,
-        data_rfft: NDArray = None,
+        data: BackendArray = None,
+        data_rfft: BackendArray = None,
         n_bins: int = None,
         batch_dimension: int = None,
         order: int = 1,
@@ -329,9 +332,9 @@ class LinearWhiteningFilter:
 
         Parameters:
         -----------
-        data : NDArray, optional
+        data : BackendArray, optional
             The input data, defaults to None.
-        data_rfft : NDArray, optional
+        data_rfft : BackendArray, optional
             The Fourier transform of the input data, defaults to None.
         n_bins : int, optional
             The number of bins for computing the spectrum, defaults to None.
@@ -348,9 +351,9 @@ class LinearWhiteningFilter:
            Filter data and associated parameters.
         """
         if data_rfft is None:
-            data_rfft = np.fft.rfftn(backend.to_numpy_array(data))
+            data_rfft = np.fft.rfftn(be.to_numpy_array(data))
 
-        data_rfft = backend.to_numpy_array(data_rfft)
+        data_rfft = be.to_numpy_array(data_rfft)
 
         bins, radial_averages = self._compute_spectrum(
             data_rfft, n_bins, batch_dimension
@@ -358,21 +361,28 @@ class LinearWhiteningFilter:
 
         if order is None:
             cutoff = bins < radial_averages.size
-            filter_mask = np.zeros(data_rfft.shape, radial_averages.dtype)
+            filter_mask = np.zeros(bins.shape, radial_averages.dtype)
             filter_mask[cutoff] = radial_averages[bins[cutoff]]
         else:
+            shape = bins.shape
+            if kwargs.get("shape", False):
+                shape = compute_fourier_shape(
+                    shape=kwargs.get("shape"),
+                    shape_is_real_fourier=kwargs.get("shape_is_real_fourier", False),
+                )
+
             filter_mask = self._interpolate_spectrum(
                 spectrum=radial_averages,
-                shape=data_rfft.shape,
+                shape=shape,
                 shape_is_real_fourier=True,
             )
 
         filter_mask = np.fft.ifftshift(
             filter_mask,
-            axes=tuple(i for i in range(data_rfft.ndim - 1) if i != batch_dimension)
+            axes=tuple(i for i in range(data_rfft.ndim - 1) if i != batch_dimension),
         )
 
         return {
-            "data": backend.to_backend_array(filter_mask),
+            "data": be.to_backend_array(filter_mask),
             "is_multiplicative_filter": True,
         }

tme/preprocessing/tilt_series.py

@@ -9,17 +9,15 @@ from typing import Tuple, Dict
 from dataclasses import dataclass
 
 import numpy as np
-from numpy.typing import NDArray
 
 from .. import Preprocessor
-from ..backends import backend
+from ..types import NDArray
+from ..backends import backend as be
 from ..matching_utils import euler_to_rotationmatrix
-
 from ._utils import (
     frequency_grid_at_angle,
     compute_tilt_shape,
     crop_real_fourier,
-    centered_grid,
     fftfreqn,
     shift_fourier,
 )
@@ -51,7 +49,6 @@ def create_reconstruction_filter(
         +---------------+----------------------------------------------------+
         | hamming       | |w| * (.54 + .46 ( cos(|w| * pi))) [2]_            |
         +---------------+----------------------------------------------------+
-
     kwargs: Dict
         Keyword arguments for particular filter_types.
 
@@ -195,22 +192,20 @@ class ReconstructFromTilt:
         if data.shape == shape:
             return data
 
-        data = backend.to_backend_array(data)
-        volume_temp = backend.zeros(shape, dtype=backend._float_dtype)
-        volume_temp_rotated = backend.zeros(shape, dtype=backend._float_dtype)
-        volume = backend.zeros(shape, dtype=backend._float_dtype)
+        data = be.to_backend_array(data)
+        volume_temp = be.zeros(shape, dtype=be._float_dtype)
+        volume_temp_rotated = be.zeros(shape, dtype=be._float_dtype)
+        volume = be.zeros(shape, dtype=be._float_dtype)
 
-        slices = tuple(
-            slice(a, a + 1) for a in backend.astype(backend.divide(shape, 2), int)
-        )
+        slices = tuple(slice(a, a + 1) for a in be.astype(be.divide(shape, 2), int))
         subset = tuple(
             slice(None) if i != opening_axis else slices[opening_axis]
             for i in range(len(shape))
         )
-        angles_loop = backend.zeros(len(shape))
+        angles_loop = be.zeros(len(shape))
         wedge_dim = [x for x in data.shape]
         wedge_dim.insert(1 + opening_axis, 1)
-        wedges = backend.reshape(data, wedge_dim)
+        wedges = be.reshape(data, wedge_dim)
 
         rec_filter = 1
         if reconstruction_filter is not None:
@@ -226,31 +221,29 @@ class ReconstructFromTilt:
             if tilt_axis == 1 and opening_axis == 0:
                 rec_filter = rec_filter.T
 
-            rec_filter = backend.to_backend_array(rec_filter)
-            rec_filter = backend.reshape(rec_filter, wedges[0].shape)
+            rec_filter = be.to_backend_array(rec_filter)
+            rec_filter = be.reshape(rec_filter, wedges[0].shape)
 
         for index in range(len(angles)):
-            backend.fill(angles_loop, 0)
-            backend.fill(volume_temp, 0)
-            backend.fill(volume_temp_rotated, 0)
+            be.fill(angles_loop, 0)
+            be.fill(volume_temp, 0)
+            be.fill(volume_temp_rotated, 0)
 
             volume_temp[subset] = wedges[index] * rec_filter
 
             angles_loop[tilt_axis] = angles[index]
-            angles_loop = backend.roll(angles_loop, (opening_axis - 1,), axis=0)
-            rotation_matrix = euler_to_rotationmatrix(
-                backend.to_numpy_array(angles_loop)
-            )
-            rotation_matrix = backend.to_backend_array(rotation_matrix)
+            angles_loop = be.roll(angles_loop, (opening_axis - 1,), axis=0)
+            rotation_matrix = euler_to_rotationmatrix(be.to_numpy_array(angles_loop))
+            rotation_matrix = be.to_backend_array(rotation_matrix)
 
-            backend.rotate_array(
+            be.rigid_transform(
                 arr=volume_temp,
                 rotation_matrix=rotation_matrix,
                 out=volume_temp_rotated,
                 use_geometric_center=True,
                 order=interpolation_order,
             )
-            backend.add(volume, volume_temp_rotated, out=volume)
+            be.add(volume, volume_temp_rotated, out=volume)
 
         volume = shift_fourier(data=volume, shape_is_real_fourier=False)
 
@@ -387,7 +380,7 @@ class Wedge:
             func_args["weights"] = np.cos(np.radians(self.angles))
 
         ret = weight_types[weight_type](**func_args)
-        ret = backend.astype(backend.to_backend_array(ret), backend._float_dtype)
+        ret = be.astype(be.to_backend_array(ret), be._float_dtype)
 
         return {
             "data": ret,
@@ -483,7 +476,7 @@ class Wedge:
             reduce_dim=True,
         )
 
-        wedges = np.zeros((len(self.angles), *tilt_shape), dtype=backend._float_dtype)
+        wedges = np.zeros((len(self.angles), *tilt_shape), dtype=be._float_dtype)
         for index, angle in enumerate(self.angles):
             frequency_grid = frequency_grid_at_angle(
                 shape=self.shape,
@@ -573,7 +566,7 @@ class WedgeReconstructed:
             func = self.continuous_wedge
 
         ret = func(shape=shape, **func_args)
-        ret = backend.astype(backend.to_backend_array(ret), backend._float_dtype)
+        ret = be.astype(be.to_backend_array(ret), be._float_dtype)
 
         return {
             "data": ret,
@@ -664,7 +657,7 @@ class WedgeReconstructed:
         """
         preprocessor = Preprocessor()
 
-        angles = np.asarray(backend.to_numpy_array(angles))
+        angles = np.asarray(be.to_numpy_array(angles))
         weights = np.ones(angles.size)
         if weight_wedge:
             weights = np.cos(np.radians(angles))
@@ -858,7 +851,7 @@ class CTF:
             func_args["opening_axis"] = None
 
         ret = self.weight(**func_args)
-        ret = backend.astype(backend.to_backend_array(ret), backend._float_dtype)
+        ret = be.astype(be.to_backend_array(ret), be._float_dtype)
         return {
             "data": ret,
             "angles": func_args["angles"],
@@ -941,24 +934,26 @@ class CTF:
             shape=shape, opening_axis=opening_axis, reduce_dim=True
         )
         stack = np.zeros((len(angles), *tilt_shape))
-        electron_wavelength = self._compute_electron_wavelength() / sampling_rate
 
         correct_defocus_gradient &= len(shape) == 3
         correct_defocus_gradient &= tilt_axis is not None
         correct_defocus_gradient &= opening_axis is not None
 
         for index, angle in enumerate(angles):
-            grid = backend.to_numpy_array(centered_grid(shape=tilt_shape))
-            grid = np.divide(grid.T, sampling_rate).T
-
             defocus_x, defocus_y = defoci_x[index], defoci_y[index]
 
+            if correct_defocus_gradient or defocus_y is not None:
+                grid = fftfreqn(
+                    shape=shape,
+                    sampling_rate=be.divide(sampling_rate, shape),
+                    return_sparse_grid=True,
+                )
+
             # This should be done after defocus_x computation
             if correct_defocus_gradient:
                 angle_rad = np.radians(angle)
 
                 defocus_gradient = np.multiply(grid[1], np.sin(angle_rad))
-
                 remaining_axis = tuple(
                     i for i in range(len(shape)) if i not in (opening_axis, tilt_axis)
                 )[0]
@@ -983,7 +978,6 @@ class CTF:
                 angle=angle,
                 sampling_rate=1,
             )
-            frequency_grid *= frequency_grid <= 0.5
             np.square(frequency_grid, out=frequency_grid)
 
             electron_aberration = spherical_aberration * electron_wavelength**2
@@ -1001,13 +995,13 @@ class CTF:
             np.sin(-chi, out=chi)
             stack[index] = chi
 
+        # Avoid contrast inversion
+        np.negative(stack, out=stack)
         if flip_phase:
             np.abs(stack, out=stack)
 
-        np.negative(stack, out=stack)
         stack = np.squeeze(stack)
-
-        stack = backend.to_backend_array(stack)
+        stack = be.to_backend_array(stack)
 
         if len(angles) == 1:
             stack = shift_fourier(data=stack, shape_is_real_fourier=False)