pytme 0.3.1.post2__cp311-cp311-macosx_15_0_arm64.whl → 0.3.2.dev0__cp311-cp311-macosx_15_0_arm64.whl

tme/density.py

@@ -16,7 +16,6 @@ from os.path import splitext, basename
 import h5py
 import mrcfile
 import numpy as np
-import skimage.io as skio
 
 from scipy.ndimage import (
     zoom,
@@ -26,7 +25,6 @@ from scipy.ndimage import (
     binary_erosion,
     generic_gradient_magnitude,
 )
-from scipy.spatial import ConvexHull
 
 from .types import NDArray
 from .rotations import align_to_axis
@@ -571,6 +569,8 @@ class Density:
         --------
         :py:meth:`Density.from_file`
         """
+        import skimage.io as skio
+
         swap = filename
         if is_gzipped(filename):
             with gzip_open(filename, "rb") as infile:
@@ -938,6 +938,8 @@ class Density:
         ----------
         .. [1] https://scikit-image.org/docs/stable/api/skimage.io.html
         """
+        import skimage.io as skio
+
         swap, kwargs = filename, {}
         if gzip:
             swap = BytesIO()
@@ -1403,8 +1405,7 @@ class Density:
         cutoff : float
             Above this value arr elements are considered. Defaults to 0.
         use_geometric_center : bool, optional
-            Whether the box should accommodate the geometric or the coordinate
-            center. Defaults to False.
+            Accommodate the geometric instead of the mass center.
 
         Returns
         -------
@@ -1416,24 +1417,25 @@ class Density:
         :py:meth:`Density.adjust_box`
         :py:meth:`tme.matching_utils.minimum_enclosing_box`
         """
-        coordinates = self.to_pointcloud(threshold=cutoff)
-        starts, stops = coordinates.min(axis=1), coordinates.max(axis=1)
+        if cutoff is None:
+            cutoff = self.data.min() - 1
 
+        coordinates = self.to_pointcloud(threshold=cutoff)
         shape = minimum_enclosing_box(
             coordinates=coordinates,
             use_geometric_center=use_geometric_center,
         )
-        difference = np.maximum(np.subtract(shape, np.subtract(stops, starts)), 0)
 
-        shift_start = np.divide(difference, 2).astype(int)
-        shift_stop = shift_start + np.mod(difference, 2)
+        starts, stops = coordinates.min(axis=1), coordinates.max(axis=1)
+        diff = np.maximum(np.subtract(shape, np.subtract(stops, starts)), 0)
+
+        shift_start = np.divide(diff, 2).astype(int)
+        shift_stop = shift_start + np.mod(diff, 2)
 
+        # These are purposefully negative to indicate left and right pad
         starts = (starts - shift_start).astype(int)
         stops = (stops + shift_stop).astype(int)
-
-        enclosing_box = tuple(slice(start, stop) for start, stop in zip(starts, stops))
-
-        return tuple(enclosing_box)
+        return tuple(slice(start, stop) for start, stop in zip(starts, stops))
 
     def pad(
         self, new_shape: Tuple[int], center: bool = True, padding_value: float = 0
@@ -1500,7 +1502,9 @@ class Density:
 
         self.adjust_box(new_box, pad_kwargs={"constant_values": padding_value})
 
-    def centered(self, cutoff: float = 0) -> Tuple["Density", NDArray]:
+    def centered(
+        self, cutoff: float = 0.0, order: int = 1
+    ) -> Tuple["Density", NDArray]:
         """
         Shifts the data center of mass to the center of the data array using linear
         interpolation. The box size of the returned :py:class:`Density` object is at
@@ -1509,8 +1513,10 @@ class Density:
         Parameters
         ----------
         cutoff : float, optional
-            Only elements in data larger than cutoff will be considered for
-            computing the new box. By default considers only positive elements.
+            Consider all elements larger than cutoff for computing the new box.
+            Default is 0.0.
+        order : int, optional
+            Interpolation order, defaults to 1.
 
         Notes
         -----
@@ -1534,16 +1540,15 @@ class Density:
         Examples
         --------
         :py:meth:`Density.centered` returns a tuple containing a centered version
-        of the current :py:class:`Density` instance, as well as an array with
-        translations. The translation corresponds to the shift between the original and
-        current center of mass.
+        of the current :py:class:`Density` instance. Centering is achieved via padding
+        and rigid-transform of the internal :py:attr:`Density.data` attribute.
+        `centered_dens` is sufficiently large to represent all rotations of the
+        :py:attr:`Density.data` attribute.
 
         >>> import numpy as np
         >>> from tme import Density
         >>> dens = Density(np.ones((5,5,5)))
-        >>> centered_dens, translation = dens.centered(0)
-        >>> translation
-        array([0., 0., 0.])
+        >>> centered_dens = dens.centered(0)
 
         :py:meth:`Density.centered` extended the :py:attr:`Density.data` attribute
         of the current :py:class:`Density` instance and modified
@@ -1551,22 +1556,6 @@ class Density:
 
         >>> centered_dens
         Origin: (-2.0, -2.0, -2.0), sampling_rate: (1, 1, 1), Shape: (9, 9, 9)
-
-        :py:meth:`Density.centered` achieves centering via zero-padding and
-        rigid-transform of the internal :py:attr:`Density.data` attribute.
-        `centered_dens` is sufficiently large to represent all rotations of the
-        :py:attr:`Density.data` attribute, such as ones obtained from
-        :py:meth:`tme.matching_utils.get_rotation_matrices`.
-
-        >>> from tme.matching_utils import get_rotation_matrices
-        >>> rotation_matrix = get_rotation_matrices(dim = 3 ,angular_sampling = 10)[0]
-        >>> rotated_centered_dens = centered_dens.rigid_transform(
-        >>>     rotation_matrix = rotation_matrix,
-        >>>     order = None
-        >>> )
-        >>> print(centered_dens.data.sum(), rotated_centered_dens.data.sum())
-        125.0 125.0
-
         """
         ret = self.copy()
 
@@ -1583,12 +1572,11 @@ class Density:
         ret = ret.rigid_transform(
             translation=shift,
             rotation_matrix=np.eye(ret.data.ndim),
-            use_geometric_center=False,
-            order=1,
+            use_geometric_center=True,
+            order=order,
         )
-
-        shift = np.subtract(center, self.center_of_mass(ret.data, cutoff))
-        return ret, shift
+        ret.origin = np.subtract(ret.origin, np.multiply(shift, ret.sampling_rate))
+        return ret
 
     def rigid_transform(
         self,
@@ -1895,6 +1883,8 @@ class Density:
 
         lower_bound, upper_bound = density_boundaries
         if method == "ConvexHull":
+            from scipy.spatial import ConvexHull
+
             binary = np.transpose(np.where(self.data > lower_bound))
             hull = ConvexHull(binary)
             surface_points = binary[hull.vertices[:]]
@@ -2032,27 +2022,24 @@ class Density:
             eroded_mask = binary_erosion(eroded_mask)
         return core_indices
 
-    @staticmethod
-    def center_of_mass(arr: NDArray, cutoff: float = None) -> NDArray:
+    def center_of_mass(self, arr: NDArray = None, cutoff: float = None) -> NDArray:
         """
-        Computes the center of mass of a numpy ndarray instance using all available
-        elements. For template matching it typically makes sense to only input
-        positive densities.
+        Computes the center of mass of a numpy ndarray instance.
 
         Parameters
         ----------
-        arr : NDArray
-            Array to compute the center of mass of.
+        arr : NDArray, optional
+            Array to compute the center of mass of, default is :py:attr:`Density.data`.
         cutoff : float, optional
-            Densities less than or equal to cutoff are nullified for center
-            of mass computation. By default considers all values.
+            Density cutoff for calculation. Defaults to None
 
         Returns
         -------
         NDArray
             Center of mass with shape (arr.ndim).
         """
-        return NumpyFFTWBackend().center_of_mass(arr**2, cutoff)
+        arr = self.data if arr is None else arr
+        return NumpyFFTWBackend().center_of_mass(arr, cutoff)
 
     @classmethod
     def match_densities(
@@ -2101,13 +2088,13 @@ class Density:
         -----
         No densities below cutoff_template are present in the returned Density object.
         """
-        from .matching_utils import normalize_template
+        from .matching_utils import standardize
         from .matching_optimization import optimize_match, create_score_object
 
         template_mask = template.empty
         template_mask.data.fill(1)
 
-        normalize_template(
+        template.data = standardize(
             template=template.data,
             mask=template_mask.data,
             n_observations=template_mask.data.sum(),
@@ -2141,8 +2128,8 @@ class Density:
         template_coordinates = template_coordinates * template_scaling[:, None]
 
         mass_center_difference = np.subtract(
-            cls.center_of_mass(target.data, cutoff_target),
-            cls.center_of_mass(template.data, cutoff_template),
+            target.center_of_mass(target.data, cutoff_target),
+            target.center_of_mass(template.data, cutoff_template),
         ).astype(int)
         template_coordinates += mass_center_difference[:, None]
 

tme/filters/_utils.py

@@ -1,5 +1,5 @@
 """
-Utilities for the generation of frequency grids.
+Utilities for the creation of composable filters.
 
 Copyright (c) 2024 European Molecular Biology Laboratory
 
@@ -79,6 +79,7 @@ def frequency_grid_at_angle(
     sampling_rate: Tuple[float],
     opening_axis: int = None,
     tilt_axis: int = None,
+    fftshift: bool = False,
 ) -> NDArray:
     """
     Generate a frequency grid from 0 to 1/(2 * sampling_rate) in each axis.
@@ -94,13 +95,16 @@ def frequency_grid_at_angle(
     shape : tuple of int
         The shape of the grid.
     angle : float
-        The angle at which to generate the grid.
+        The angle at which to generate the grid in degrees.
     sampling_rate : tuple of float
         The sampling rate for each dimension.
     opening_axis : int, optional
         The projection axis, defaults to None.
     tilt_axis : int, optional
         The axis along which the grid is tilted, defaults to None.
+    fftshift : bool, optional
+        Whether to return a grid centered at shape // 2. Default is grid centered around
+        origin, which is compliant with the rfftn definitions used in this project
 
     Returns
     -------
@@ -114,14 +118,17 @@ def frequency_grid_at_angle(
         shape=shape, opening_axis=opening_axis, reduce_dim=False
     )
 
-    if angle == 0:
+    missing_axes = opening_axis is None or tilt_axis is None
+    if angle == 0 or missing_axes or len(set(shape)) == 1:
+        # Crop the sampling rate to tilt shape
         sampling_rate = compute_tilt_shape(
             shape=sampling_rate, opening_axis=opening_axis, reduce_dim=True
         )
-        index_grid = fftfreqn(
+        return fftfreqn(
             tuple(x for x in tilt_shape if x != 1),
             sampling_rate=sampling_rate,
             compute_euclidean_norm=True,
+            fftshift=fftshift,
         )
 
     if angle != 0:
@@ -130,9 +137,11 @@ def frequency_grid_at_angle(
 
         angles = np.zeros(len(shape))
         angles[tilt_axis] = angle
-        rotation_matrix = euler_to_rotationmatrix(np.roll(angles, opening_axis - 1))
+        rotation_matrix = euler_to_rotationmatrix(
+            np.roll(angles, opening_axis - 1), seq="zyz"
+        )
 
-        index_grid = fftfreqn(tilt_shape, sampling_rate=None)
+        index_grid = fftfreqn(tilt_shape, sampling_rate=None, fftshift=fftshift)
         index_grid = np.einsum("ij,j...->i...", rotation_matrix, index_grid)
         norm = np.multiply(sampling_rate, shape).astype(int)
 
@@ -146,19 +155,25 @@ def frequency_grid_at_angle(
 def fftfreqn(
     shape: Tuple[int],
     sampling_rate: Tuple[float],
+    fftshift: bool = False,
     compute_euclidean_norm: bool = False,
     shape_is_real_fourier: bool = False,
     return_sparse_grid: bool = False,
 ) -> NDArray:
     """
-    Generate the n-dimensional discrete Fourier transform sample frequencies.
+    Generate n-dimensional (frequency) grids.
 
     Parameters
     ----------
     shape : Tuple[int]
         The shape of the data.
     sampling_rate : float or Tuple[float]
-        The sampling rate.
+        Sets the maximum value along each axis in shape to x=1/(2*sampling_rate), e.g.,
+        a sampling_rate of 1 yields a grid from -n/x * 1/n to (n)/x -1 * 1/n. A sampling
+        rate of None returns a grid from -n/2 to n/2 - 1
+    fftshift : bool, optional
+        Whether to return a grid centered at shape // 2. Default is grid centered around
+        origin, which is compliant with the rfftn definitions used in this project.
     compute_euclidean_norm : bool, optional
         Whether to compute the Euclidean norm, defaults to False.
     shape_is_real_fourier : bool, optional
@@ -181,10 +196,18 @@ def fftfreqn(
         if sampling_rate is not None:
             norm[-1] = (shape[-1] - 1) * 2 * sampling_rate
 
-    grids = []
+    ndim, grids = len(shape), []
     for i, x in enumerate(shape):
         baseline_dims = tuple(1 if i != t else x for t in range(len(shape)))
         grid = (np_be.arange(x, dtype=np_be._int_dtype) - center[i]) / norm[i]
+
+        # We have to invert because we build the grid centered around shape // 2
+        if not fftshift:
+            if shape_is_real_fourier and i == (ndim - 1):
+                pass
+            else:
+                grid = np.fft.ifftshift(grid)
+
         grid = np_be.astype(grid, np_be._float_dtype)
         grids.append(np_be.reshape(grid, baseline_dims))
 
@@ -197,9 +220,7 @@ def fftfreqn(
         return grids
 
     grid_flesh = np_be.full(shape, fill_value=1, dtype=np_be._float_dtype)
-    grids = np_be.stack(tuple(grid * grid_flesh for grid in grids))
-
-    return grids
+    return np_be.stack(tuple(grid * grid_flesh for grid in grids))
 
 
 def crop_real_fourier(data: BackendArray) -> BackendArray:
@@ -231,27 +252,28 @@ def compute_fourier_shape(
 
 
 def shift_fourier(
-    data: BackendArray, shape_is_real_fourier: bool = False
+    data: BackendArray, shape_is_real_fourier: bool = False, ifftshift: bool = True
 ) -> BackendArray:
     comp = be
     if isinstance(data, np.ndarray):
         comp = NumpyFFTWBackend()
+
     shape = comp.to_backend_array(data.shape)
-    shift = comp.add(comp.divide(shape, 2), comp.mod(shape, 2))
+    shift = comp.divide(shape, 2)
+    if ifftshift:
+        shift = comp.add(shift, comp.mod(shape, 2))
+
     shift = [int(x) for x in shift]
     if shape_is_real_fourier:
         shift[-1] = 0
-
-    data = comp.roll(data, shift, tuple(i for i in range(len(shift))))
-    return data
+    return comp.roll(data, shift, tuple(i for i in range(len(shift))))
 
 
 def create_reconstruction_filter(
-    filter_shape: Tuple[int], filter_type: str, **kwargs: Dict
+    filter_shape: Tuple[int], filter_type: str, fftshift: bool = True, **kwargs: Dict
 ):
     """
-    Create a reconstruction filter of given filter_type. The DC component of
-    the filter will be located in the array center.
+    Create a reconstruction filter of given filter_type.
 
     Parameters
     ----------
@@ -274,6 +296,8 @@ def create_reconstruction_filter(
         +---------------+----------------------------------------------------+
         | hamming       | |w| * (.54 + .46 ( cos(|w| * pi))) [2]_            |
         +---------------+----------------------------------------------------+
+    fftshift : bool, optional
+        Should the DC component be located at the center, default is True.
     kwargs: Dict
         Keyword arguments for particular filter_types.
 
@@ -288,7 +312,9 @@ def create_reconstruction_filter(
     .. [2]  https://odlgroup.github.io/odl/index.html
     """
     filter_type = str(filter_type).lower()
-    freq = fftfreqn(filter_shape, sampling_rate=0.5, compute_euclidean_norm=True)
+    freq = fftfreqn(
+        filter_shape, sampling_rate=0.5, compute_euclidean_norm=True, fftshift=fftshift
+    )
 
     if filter_type == "ram-lak":
         ret = np.copy(freq)
@@ -297,8 +323,8 @@ def create_reconstruction_filter(
         for dim, size in enumerate(filter_shape):
             n = np.concatenate(
                 (
-                    np.arange(1, size / 2 + 1, 2, dtype=int),
-                    np.arange(size / 2 - 1, 0, -2, dtype=int),
+                    np.arange(1, size // 2 + 1, 2, dtype=int),
+                    np.arange(size // 2 - 1, 0, -2, dtype=int),
                 )
             )
             ret1d = np.zeros(size)
@@ -316,7 +342,9 @@ def create_reconstruction_filter(
         if tilt_angles is False:
             raise ValueError("'ramp' filter requires specifying tilt angles.")
         size = filter_shape[0]
-        ret = fftfreqn((size,), sampling_rate=1, compute_euclidean_norm=True)
+        ret = fftfreqn(
+            (size,), sampling_rate=1, compute_euclidean_norm=True, fftshift=fftshift
+        )
         min_increment = np.radians(np.min(np.abs(np.diff(np.sort(tilt_angles)))))
         ret *= min_increment * size
         ret = np.fmin(ret, 1, out=ret)