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

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,40 +1085,52 @@ def tube_mask(
     if symmetry_axis > len(shape):
         raise ValueError(f"symmetry_axis can be not larger than {len(shape)}.")
 
-    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)
+    if len(base_center) != len(shape):
+        raise ValueError("shape and base_center need to have the same length.")
 
-    inner_circle = create_mask(
-        mask_type="ellipse",
-        shape=circle_shape,
-        radius=inner_radius,
-        center=base_center,
-    )
-    outer_circle = create_mask(
-        mask_type="ellipse",
-        shape=circle_shape,
-        radius=outer_radius,
-        center=base_center,
-    )
+    circle_shape = tuple(b for ix, b in enumerate(shape) 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)
+    if inner_radius > 0:
+        inner_circle = create_mask(
+            mask_type="ellipse",
+            shape=circle_shape,
+            radius=inner_radius,
+            center=circle_center,
+        )
+    if outer_radius > 0:
+        outer_circle = create_mask(
+            mask_type="ellipse",
+            shape=circle_shape,
+            radius=outer_radius,
+            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.
@@ -1131,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
     -------
@@ -1154,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/parser.py

@@ -229,7 +229,7 @@ class PDBParser(Parser):
 class MMCIFParser(Parser):
     """
     A Parser subclass for converting MMCIF file data into a dictionary representation.
-    This implementation heavily relies on the atomium library:
+    This implementation heavily relies on the atomium library [1]_.
 
     References
     ----------
@@ -308,7 +308,7 @@ class MMCIFParser(Parser):
         -------
         list of dict
             A list of dictionaries where each dictionary represents a block
-             of data from the MMCIF file.
+            of data from the MMCIF file.
         """
         category = None
         block, blocks = [], []

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)
@@ -1095,7 +1096,7 @@ class Preprocessor:
 
             opening_axis = opening_axes[index]
             rotation_matrix = euler_to_rotationmatrix(
-                np.roll(tilt_angles[:, index], 1 - opening_axes)
+                np.roll(tilt_angles[:, index], opening_axis - 1)
             )
 
             subset = tuple(
@@ -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)
 
@@ -1218,7 +1326,7 @@ class Preprocessor:
             distances = np.linalg.norm(grid, axis=0)
 
         if not infinite_plane:
-            np.multiply(wedge, distances <= shape[opening_axis] // 2, out=wedge)
+            np.multiply(wedge, distances <= shape[tilt_axis] // 2, out=wedge)
 
         wedge = self.gaussian_filter(template=wedge, sigma=sigma, fourier=False)
         wedge = np.fft.ifftshift(wedge > np.exp(-2))
@@ -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))
 
@@ -1360,31 +1468,36 @@ class LinearWhiteningFilter:
         _, bin_edges = np.histogram(frequency_grid, bins=n_bins - 1)
         bins = np.digitize(frequency_grid, bins=bin_edges, right=True)
 
-        fourier_transform = np.fft.fftshift(np.fft.rfftn(template))
+        fft_shift_axes = tuple(range(template.ndim - 1))
+        fourier_transform = np.fft.fftshift(np.fft.rfftn(template), axes=fft_shift_axes)
         fourier_spectrum = np.abs(fourier_transform)
 
         radial_averages = ndimean(fourier_spectrum, labels=bins, index=np.unique(bins))
         np.reciprocal(radial_averages, out=radial_averages)
         np.divide(radial_averages, radial_averages.max(), out=radial_averages)
 
-        bins = radial_averages[bins]
-        np.multiply(fourier_transform, bins, out=fourier_transform)
-        center_indices = tuple([dim_size // 2 for dim_size in fourier_transform.shape])
-        fourier_transform[center_indices] = 0
-        ret = np.fft.irfftn(np.fft.ifftshift(fourier_transform)).real
+        np.multiply(fourier_transform, radial_averages[bins], out=fourier_transform)
 
+        ret = np.fft.irfftn(
+            np.fft.ifftshift(fourier_transform, axes=fft_shift_axes), s=template.shape
+        ).real
         return ret, bin_edges, radial_averages
 
     def apply(
         self, template: NDArray, bin_edges: NDArray, radial_averages: NDArray
     ) -> NDArray:
-        fourier_transform = np.fft.fftshift(np.fft.fftn(template))
-
-        grid = self._fftfreqn(shape=fourier_transform.shape, sampling_rate=1)
+        grid = self._fftfreqn(
+            shape=template.shape, sampling_rate=1, omit_negative_frequencies=True
+        )
         frequency_grid = np.linalg.norm(grid, axis=0)
 
+        fft_shift_axes = tuple(range(template.ndim - 1))
+        fourier_transform = np.fft.fftshift(np.fft.rfftn(template), axes=fft_shift_axes)
+
         bins = np.digitize(frequency_grid, bins=bin_edges, right=True)
         np.multiply(fourier_transform, radial_averages[bins], out=fourier_transform)
-        ret = np.fft.ifftn(np.fft.ifftshift(fourier_transform)).real
+        ret = np.fft.irfftn(
+            np.fft.ifftshift(fourier_transform, axes=fft_shift_axes), s=template.shape
+        ).real
 
         return ret

tme/types.py

@@ -3,6 +3,7 @@ from typing import Union, TypeVar
 NDArray = TypeVar("numpy.ndarray")
 CupyArray = TypeVar("cupy.ndarray")
 TorchTensor = TypeVar("torch.Tensor")
+MlxArray = TypeVar("mlx.core.array")
 
 Scalar = Union[int, float, complex]