pytme 0.2.3__cp311-cp311-macosx_14_0_arm64.whl → 0.2.5__cp311-cp311-macosx_14_0_arm64.whl

tme/preprocessing/tilt_series.py

@@ -10,11 +10,11 @@ from dataclasses import dataclass
 
 import numpy as np
 
-from .. import Preprocessor
 from ..types import NDArray
 from ..backends import backend as be
-from ..matching_utils import euler_to_rotationmatrix
+from ..matching_utils import euler_to_rotationmatrix, centered
 from ._utils import (
+    centered_grid,
     frequency_grid_at_angle,
     compute_tilt_shape,
     crop_real_fourier,
@@ -90,12 +90,8 @@ def create_reconstruction_filter(
         tilt_angles = kwargs.get("tilt_angles", False)
         if tilt_angles is False:
             raise ValueError("'ramp' filter requires specifying tilt angles.")
-        size, odd = filter_shape[0], filter_shape[0] % 2
-        ret = np.arange(-size // 2 + odd, size // 2 + odd, 1, dtype=np.float32)
-        ret /= size // 2
-        ret *= 0.5
-        np.abs(ret, out=ret)
-
+        size = filter_shape[0]
+        ret = fftfreqn((size,), sampling_rate=1, compute_euclidean_norm=True)
         min_increment = np.radians(np.min(np.abs(np.diff(np.sort(tilt_angles)))))
         ret *= min_increment * size
         np.fmin(ret, 1, out=ret)
@@ -164,8 +160,8 @@ class ReconstructFromTilt:
         """
         Reconstruct a volume from a tilt series.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         data : NDArray
             The tilt series data.
         shape : tuple of int
@@ -177,15 +173,15 @@ class ReconstructFromTilt:
         tilt_axis : int
             Axis the plane is tilted over.
         interpolation_order : int, optional
-            Interpolation order used for rotation, defautls to 1.
+            Interpolation order used for rotation, defaults to 1.
         return_real_fourier : bool, optional
            Whether to return a shape compliant with rfftn, defaults to True.
         reconstruction_filter : bool, optional
            Filter window applied during reconstruction.
            See :py:meth:`create_reconstruction_filter` for available options.
 
-        Returns:
-        --------
+        Returns
+        -------
         NDArray
             The reconstructed volume.
         """
@@ -197,7 +193,7 @@ class ReconstructFromTilt:
         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 be.astype(be.divide(shape, 2), int))
+        slices = tuple(slice(a // 2, (a // 2) + 1) for a in shape)
         subset = tuple(
             slice(None) if i != opening_axis else slices[opening_axis]
             for i in range(len(shape))
@@ -224,6 +220,7 @@ class ReconstructFromTilt:
             rec_filter = be.to_backend_array(rec_filter)
             rec_filter = be.reshape(rec_filter, wedges[0].shape)
 
+        angles = be.to_backend_array(angles)
         for index in range(len(angles)):
             be.fill(angles_loop, 0)
             be.fill(volume_temp, 0)
@@ -257,8 +254,8 @@ class Wedge:
     """
     Generate wedge mask for tomographic data.
 
-    Parameters:
-    -----------
+    Parameters
+    ----------
     shape : tuple of int
         The shape of the reconstruction volume.
     tilt_axis : int
@@ -272,10 +269,10 @@ class Wedge:
     weight_type : str, optional
         The type of weighting to apply, defaults to None.
     frequency_cutoff : float, optional
-        The frequency cutoff value, defaults to 0.5.
+        Frequency cutoff for created mask. Nyquist 0.5 by default.
 
-    Returns:
-    --------
+    Returns
+    -------
     Dict
         A dictionary containing weighted wedges and related information.
     """
@@ -303,13 +300,13 @@ class Wedge:
         Generate a :py:class:`Wedge` instance by reading tilt angles and weights
         from a tab-separated text file.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         filename : str
             The path to the file containing tilt angles and weights.
 
-        Returns:
-        --------
+        Returns
+        -------
         :py:class:`Wedge`
            Class instance instance initialized with angles and weights from the file.
         """
@@ -338,15 +335,15 @@ class Wedge:
         """
         Read column data from a text file.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         filename : str
             The path to the text file.
         delimiter : str, optional
             The delimiter used in the file, defaults to '\t'.
 
-        Returns:
-        --------
+        Returns
+        -------
         Dict
             A dictionary with one key for each column.
         """
@@ -380,6 +377,19 @@ class Wedge:
             func_args["weights"] = np.cos(np.radians(self.angles))
 
         ret = weight_types[weight_type](**func_args)
+
+        frequency_cutoff = func_args.get("frequency_cutoff", None)
+        if frequency_cutoff is not None:
+            for index, angle in enumerate(self.angles):
+                frequency_grid = frequency_grid_at_angle(
+                    shape=func_args["shape"],
+                    opening_axis=self.opening_axis,
+                    tilt_axis=self.tilt_axis,
+                    angle=angle,
+                    sampling_rate=1,
+                )
+                ret[index] = np.multiply(ret[index], frequency_grid <= frequency_cutoff)
+
         ret = be.astype(be.to_backend_array(ret), be._float_dtype)
 
         return {
@@ -413,46 +423,47 @@ class Wedge:
 
         return wedges
 
-    def weight_relion(self, **kwargs) -> NDArray:
+    def weight_relion(
+        self, shape: Tuple[int], opening_axis: int, tilt_axis: int, **kwargs
+    ) -> NDArray:
         """
         Generate weighted wedges based on the RELION 1.4 formalism, weighting each
         angle using the cosine of the angle and a Gaussian lowpass filter computed
         with respect to the exposure per angstrom.
 
-        Returns:
-        --------
+        Returns
+        -------
         NDArray
             Weighted wedges.
         """
         tilt_shape = compute_tilt_shape(
-            shape=self.shape, opening_axis=self.opening_axis, reduce_dim=True
+            shape=shape, opening_axis=opening_axis, reduce_dim=True
         )
 
         wedges = np.zeros((len(self.angles), *tilt_shape))
         for index, angle in enumerate(self.angles):
             frequency_grid = frequency_grid_at_angle(
-                shape=self.shape,
-                opening_axis=self.opening_axis,
-                tilt_axis=self.tilt_axis,
+                shape=shape,
+                opening_axis=opening_axis,
+                tilt_axis=tilt_axis,
                 angle=angle,
                 sampling_rate=1,
             )
-            # frequency_mask = frequency_grid <= self.frequency_cutoff
-
             sigma = np.sqrt(self.weights[index] * 4 / (8 * np.pi**2))
             sigma = -2 * np.pi**2 * sigma**2
             np.square(frequency_grid, out=frequency_grid)
             np.multiply(sigma, frequency_grid, out=frequency_grid)
             np.exp(frequency_grid, out=frequency_grid)
             np.multiply(frequency_grid, np.cos(np.radians(angle)), out=frequency_grid)
-            # np.multiply(frequency_grid, frequency_mask, out=frequency_grid)
-
             wedges[index] = frequency_grid
 
         return wedges
 
     def weight_grigorieff(
         self,
+        shape: Tuple[int],
+        opening_axis: int,
+        tilt_axis: int,
         amplitude: float = 0.245,
         power: float = -1.665,
         offset: float = 2.81,
@@ -461,31 +472,28 @@ class Wedge:
         """
         Generate weighted wedges based on the formalism introduced in [1]_.
 
-        Returns:
-        --------
+        Returns
+        -------
         NDArray
             Weighted wedges.
 
         References
         ----------
-        .. [1]  Timothy GrantNikolaus Grigorieff (2015), eLife 4:e06980.
+        .. [1]  Timothy Grant, Nikolaus Grigorieff (2015), eLife 4:e06980.
         """
         tilt_shape = compute_tilt_shape(
-            shape=self.shape,
-            opening_axis=self.opening_axis,
-            reduce_dim=True,
+            shape=shape, opening_axis=opening_axis, reduce_dim=True
         )
 
         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,
-                opening_axis=self.opening_axis,
-                tilt_axis=self.tilt_axis,
+                shape=shape,
+                opening_axis=opening_axis,
+                tilt_axis=tilt_axis,
                 angle=angle,
                 sampling_rate=1,
             )
-            # frequency_mask = frequency_grid <= self.frequency_cutoff
 
             with np.errstate(divide="ignore"):
                 np.power(frequency_grid, power, out=frequency_grid)
@@ -498,10 +506,7 @@ class Wedge:
                     out=frequency_grid,
                 )
 
-            np.exp(frequency_grid, out=frequency_grid)
-            # np.multiply(frequency_grid, frequency_mask, out=frequency_grid)
-
-            wedges[index] = frequency_grid
+            wedges[index] = np.exp(frequency_grid)
 
         return wedges
 
@@ -510,14 +515,24 @@ class WedgeReconstructed:
     """
     Initialize :py:class:`WedgeReconstructed`.
 
-    Parameters:
-    -----------
-    angles : Tuple[float], optional
+    Parameters
+    ----------
+    angles :tuple of float, optional
         The tilt angles, defaults to None.
     opening_axis : int, optional
         The axis around which the wedge is opened, defaults to 0.
     tilt_axis : int, optional
         The axis along which the tilt is applied, defaults to 2.
+    weights : tuple of float, optional
+        Weights to assign to individual wedge components.
+    weight_wedge : bool, optional
+        Whether individual wedge components should be weighted. If True and weights
+        is None, uses the cosine of the angle otherwise weights.
+    create_continuous_wedge: bool, optional
+        Whether to create a continous wedge or a per-component wedge. Weights are only
+        considered for non-continuous wedges.
+    frequency_cutoff : float, optional
+        Filter window applied during reconstruction.
     **kwargs : Dict
         Additional keyword arguments.
     """
@@ -525,33 +540,37 @@ class WedgeReconstructed:
     def __init__(
         self,
         angles: Tuple[float] = None,
-        start_tilt: float = None,
-        stop_tilt: float = None,
         opening_axis: int = 0,
         tilt_axis: int = 2,
+        weights: Tuple[float] = None,
         weight_wedge: bool = False,
         create_continuous_wedge: bool = False,
+        frequency_cutoff: float = 0.5,
+        reconstruction_filter: str = None,
         **kwargs: Dict,
     ):
         self.angles = angles
         self.opening_axis = opening_axis
         self.tilt_axis = tilt_axis
+        self.weights = weights
         self.weight_wedge = weight_wedge
+        self.reconstruction_filter = reconstruction_filter
         self.create_continuous_wedge = create_continuous_wedge
+        self.frequency_cutoff = frequency_cutoff
 
     def __call__(self, shape: Tuple[int], **kwargs: Dict) -> Dict:
         """
         Generate the reconstructed wedge.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         shape : tuple of int
             The shape of the reconstruction volume.
         **kwargs : Dict
             Additional keyword arguments.
 
-        Returns:
-        --------
+        Returns
+        -------
         Dict
             A dictionary containing the reconstructed wedge and related information.
         """
@@ -559,15 +578,39 @@ class WedgeReconstructed:
         func_args.update(kwargs)
 
         if kwargs.get("is_fourier_shape", False):
-            print("Cannot create continuous wedge mask basde on real fourier shape.")
+            print("Cannot create continuous wedge mask based on real fourier shape.")
 
         func = self.step_wedge
         if func_args.get("create_continuous_wedge", False):
             func = self.continuous_wedge
 
+        weight_wedge = func_args.get("weight_wedge", False)
+        if func_args.get("wedge_weights") is None and weight_wedge:
+            func_args["weights"] = np.cos(
+                np.radians(be.to_numpy_array(func_args.get("angles", (0,))))
+            )
+
         ret = func(shape=shape, **func_args)
+
+        frequency_cutoff = func_args.get("frequency_cutoff", None)
+        if frequency_cutoff is not None:
+            frequency_mask = fftfreqn(
+                shape=shape,
+                sampling_rate=1,
+                compute_euclidean_norm=True,
+                shape_is_real_fourier=False,
+            )
+            ret = np.multiply(ret, frequency_mask <= frequency_cutoff, out=ret)
+
+        if not weight_wedge:
+            ret = (ret > 0) * 1.0
+
         ret = be.astype(be.to_backend_array(ret), be._float_dtype)
 
+        ret = shift_fourier(data=ret, shape_is_real_fourier=False)
+        if func_args.get("return_real_fourier", False):
+            ret = crop_real_fourier(ret)
+
         return {
             "data": ret,
             "shape_is_real_fourier": func_args["return_real_fourier"],
@@ -584,14 +627,13 @@ class WedgeReconstructed:
         angles: Tuple[float],
         opening_axis: int,
         tilt_axis: int,
-        return_real_fourier: bool,
         **kwargs: Dict,
     ) -> NDArray:
         """
-        Generate a continuous reconstructed wedge.
+        Generate a continous wedge mask with DC component at the center.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         shape : tuple of int
             The shape of the reconstruction volume.
         angles : tuple of float
@@ -600,27 +642,28 @@ class WedgeReconstructed:
             The axis around which the wedge is opened.
         tilt_axis : int
             The axis along which the tilt is applied.
-        return_real_fourier : bool
-            Whether to return the real part of the Fourier transform.
 
-        Returns:
-        --------
+        Returns
+        -------
         NDArray
-            The reconstructed wedge.
+            Wedge mask.
         """
-        preprocessor = Preprocessor()
-        start_tilt, stop_tilt = angles
-        ret = preprocessor.continuous_wedge_mask(
-            start_tilt=start_tilt,
-            stop_tilt=stop_tilt,
-            shape=shape,
-            opening_axis=opening_axis,
-            tilt_axis=tilt_axis,
-            omit_negative_frequencies=return_real_fourier,
-            infinite_plane=False,
+        start_radians = np.tan(np.radians(90 - angles[0]))
+        stop_radians = np.tan(np.radians(-1 * (90 - angles[1])))
+
+        grid = centered_grid(shape)
+        with np.errstate(divide="ignore", invalid="ignore"):
+            ratios = np.where(
+                grid[opening_axis] == 0,
+                np.tan(np.radians(90)) + 1,
+                grid[tilt_axis] / grid[opening_axis],
+            )
+
+        wedge = np.logical_or(start_radians <= ratios, stop_radians >= ratios).astype(
+            np.float32
         )
 
-        return ret
+        return wedge
 
     @staticmethod
     def step_wedge(
@@ -628,15 +671,15 @@ class WedgeReconstructed:
         angles: Tuple[float],
         opening_axis: int,
         tilt_axis: int,
-        return_real_fourier: bool,
-        weight_wedge: bool = False,
+        weights: Tuple[float] = None,
+        reconstruction_filter: str = None,
         **kwargs: Dict,
     ) -> NDArray:
         """
-        Generate a step-wise reconstructed wedge.
+        Generate a per-angle wedge shape with DC component at the center.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         shape : tuple of int
             The shape of the reconstruction volume.
         angles : tuple of float
@@ -645,38 +688,76 @@ class WedgeReconstructed:
             The axis around which the wedge is opened.
         tilt_axis : int
             The axis along which the tilt is applied.
-        weight_wedge : bool, optional
-            Whether to weight the wedge by the cosine of the angle.
-        return_real_fourier : bool
-            Whether to return the real part of the Fourier transform.
+        reconstruction_filter : str
+            Filter used during reconstruction.
+        weights : tuple of float, optional
+            Weights to assign to individual tilts. Defaults to 1.
 
-        Returns:
-        --------
+        Returns
+        -------
         NDArray
-            The reconstructed wedge.
+            Wege mask.
         """
-        preprocessor = Preprocessor()
+        from ..backends import NumpyFFTWBackend
 
         angles = np.asarray(be.to_numpy_array(angles))
-        weights = np.ones(angles.size)
-        if weight_wedge:
-            weights = np.cos(np.radians(angles))
-        ret = preprocessor.step_wedge_mask(
-            tilt_angles=angles,
-            weights=weights,
-            start_tilt=None,
-            stop_tilt=None,
-            tilt_step=None,
-            shape=shape,
-            opening_axis=opening_axis,
-            tilt_axis=tilt_axis,
-            omit_negative_frequencies=return_real_fourier,
+
+        if weights is None:
+            weights = np.ones(angles.size)
+        weights = np.asarray(weights)
+
+        shape = tuple(int(x) for x in shape)
+        opening_axis, tilt_axis = int(opening_axis), int(tilt_axis)
+
+        weights = np.repeat(weights, angles.size // weights.size)
+        plane = np.zeros(
+            (shape[opening_axis], shape[tilt_axis] + (1 - shape[tilt_axis] % 2)),
+            dtype=np.float32,
         )
 
-        if not weight_wedge:
-            ret = (ret > 0) * 1.0
+        # plane = np.zeros((shape[opening_axis], int(2 * np.max(shape)) + 1), dtype=np.float32)
 
-        return ret
+        rec_filter = 1
+        if reconstruction_filter is not None:
+            rec_filter = create_reconstruction_filter(
+                plane.shape[::-1], filter_type=reconstruction_filter, tilt_angles=angles
+            ).T
+
+        subset = tuple(
+            slice(None) if i != 0 else slice(x // 2, x // 2 + 1)
+            for i, x in enumerate(plane.shape)
+        )
+        plane_rotated, wedge_volume = np.zeros_like(plane), np.zeros_like(plane)
+        for index in range(angles.shape[0]):
+            plane_rotated.fill(0)
+            plane[subset] = 1
+            rotation_matrix = euler_to_rotationmatrix((angles[index], 0))
+            rotation_matrix = rotation_matrix[np.ix_((0, 1), (0, 1))]
+
+            NumpyFFTWBackend().rigid_transform(
+                arr=plane * rec_filter,
+                rotation_matrix=rotation_matrix,
+                out=plane_rotated,
+                use_geometric_center=True,
+                order=1,
+            )
+            wedge_volume += plane_rotated * weights[index]
+
+        wedge_volume = centered(wedge_volume, (shape[opening_axis], shape[tilt_axis]))
+        np.fmin(wedge_volume, np.max(weights), wedge_volume)
+
+        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)
+
+        return wedge_volume
 
 
 @dataclass
@@ -726,8 +807,8 @@ class CTF:
         """
         Initialize :py:class:`CTF` from file.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         filename : str
             The path to a file with ctf parameters. Supports the following formats:
             - CTFFIND4
@@ -857,8 +938,8 @@ class CTF:
         """
         Compute the CTF weight tilt stack.
 
-        Parameters:
-        -----------
+        Parameters
+        ----------
         shape : tuple of int
             The shape of the CTF.
         defocus_x : tuple of float
@@ -890,8 +971,8 @@ class CTF:
         **kwargs : Dict
             Additional keyword arguments.
 
-        Returns:
-        --------
+        Returns
+        -------
         NDArray
             A stack containing the CTF weight.
         """