pytme 0.3.1.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.2__cp311-cp311-macosx_15_0_arm64.whl

tme/filters/wedge.py

@@ -1,6 +1,5 @@
 """
-Implements class Wedge and WedgeReconstructed to create Fourier
-filter representations.
+Implements class Wedge and WedgeReconstructed.
 
 Copyright (c) 2024 European Molecular Biology Laboratory
 
@@ -16,16 +15,14 @@ from ..types import NDArray
 from ..backends import backend as be
 from .compose import ComposableFilter
 from ..matching_utils import center_slice
-from ..rotations import euler_to_rotationmatrix
 from ..parser import XMLParser, StarParser, MDOCParser
 from ._utils import (
     centered_grid,
     frequency_grid_at_angle,
     compute_tilt_shape,
-    crop_real_fourier,
     fftfreqn,
-    shift_fourier,
     create_reconstruction_filter,
+    shift_fourier,
 )
 
 __all__ = ["Wedge", "WedgeReconstructed"]
@@ -34,12 +31,10 @@ __all__ = ["Wedge", "WedgeReconstructed"]
 @dataclass
 class Wedge(ComposableFilter):
     """
-    Generate wedge mask for tomographic data.
+    Create per-tilt wedge mask for tomographic data.
     """
 
-    #: The shape of the reconstruction volume.
-    shape: Tuple[int] = None
-    #: The tilt angles.
+    #: Tilt angles in degrees.
     angles: Tuple[float] = None
     #: The weights corresponding to each tilt angle.
     weights: Tuple[float] = None
@@ -55,7 +50,7 @@ class Wedge(ComposableFilter):
     sampling_rate: Tuple[float] = 1
 
     @classmethod
-    def from_file(cls, filename: str) -> "Wedge":
+    def from_file(cls, filename: str, **kwargs) -> "Wedge":
         """
         Generate a :py:class:`Wedge` instance by reading tilt angles and weights.
         Supported extensions are:
@@ -100,180 +95,56 @@ class Wedge(ComposableFilter):
             raise ValueError("Length of weights and angles differ.")
 
         return cls(
-            shape=None,
             tilt_axis=0,
             opening_axis=2,
             angles=np.array(angles, dtype=np.float32),
             weights=np.array(weights, dtype=np.float32),
+            **kwargs,
         )
 
-    def __call__(self, **kwargs: Dict) -> NDArray:
-        """
-        Returns a Wedge stack of chosen parameters with DC component in the center.
-        """
-        func_args = vars(self).copy()
-        func_args.update(kwargs)
-
+    def _evaluate(self, shape: Tuple[int, ...], **kwargs: Dict) -> NDArray:
+        """Returns a Wedge stack of chosen parameters."""
         weight_types = {
-            None: self.weight_angle,
-            "angle": self.weight_angle,
-            "relion": self.weight_relion,
-            "grigorieff": self.weight_grigorieff,
+            None: weight_uniform,
+            "angle": weight_angle,
+            "relion": weight_relion,
+            "grigorieff": weight_grigorieff,
         }
 
-        weight_type = func_args.get("weight_type", None)
+        weight_type = kwargs.get("weight_type", None)
         if weight_type not in weight_types:
             raise ValueError(
                 f"Supported weight_types are {','.join(list(weight_types.keys()))}"
             )
 
         if weight_type == "angle":
-            func_args["weights"] = np.cos(np.radians(self.angles))
+            kwargs["weights"] = np.cos(np.radians(self.angles))
 
-        ret = weight_types[weight_type](**func_args)
+        ret = weight_types[weight_type](shape=shape, **kwargs)
 
-        frequency_cutoff = func_args.get("frequency_cutoff", None)
+        frequency_cutoff = kwargs.get("frequency_cutoff", None)
         if frequency_cutoff is not None:
-            for index, angle in enumerate(func_args["angles"]):
+            for index, angle in enumerate(kwargs["angles"]):
                 frequency_grid = frequency_grid_at_angle(
-                    shape=func_args["shape"],
-                    opening_axis=self.opening_axis,
-                    tilt_axis=self.tilt_axis,
+                    shape=shape,
+                    opening_axis=kwargs["opening_axis"],
+                    tilt_axis=kwargs["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 {
-            "data": ret,
-            "shape": func_args["shape"],
-            "return_real_fourier": func_args.get("return_real_fourier", False),
-            "is_multiplicative_filter": True,
-        }
-
-    @staticmethod
-    def weight_angle(
-        shape: Tuple[int],
-        weights: Tuple[float],
-        angles: Tuple[float],
-        opening_axis: int,
-        tilt_axis: int,
-        **kwargs,
-    ) -> NDArray:
-        """
-        Generate weighted wedges based on the cosine of the current angle.
-        """
-        tilt_shape = compute_tilt_shape(
-            shape=shape, opening_axis=opening_axis, reduce_dim=True
-        )
-        wedge, wedges = np.ones(tilt_shape), np.zeros((len(angles), *tilt_shape))
-        for index, angle in enumerate(angles):
-            wedge.fill(weights[index])
-            wedges[index] = wedge
-
-        return wedges
-
-    def weight_relion(
-        self,
-        shape: Tuple[int],
-        opening_axis: int,
-        tilt_axis: int,
-        sampling_rate: float = 1.0,
-        **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
-        -------
-        NDArray
-            Weighted wedges.
-        """
-        tilt_shape = compute_tilt_shape(
-            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=shape,
-                opening_axis=opening_axis,
-                tilt_axis=tilt_axis,
-                angle=angle,
-                sampling_rate=sampling_rate,
-            )
-            sigma = np.sqrt(self.weights[index] * 4 / (8 * np.pi**2))
-            sigma = -2 * np.pi**2 * sigma**2
-            frequency_grid = np.square(frequency_grid, out=frequency_grid)
-            frequency_grid = np.multiply(sigma, frequency_grid, out=frequency_grid)
-            frequency_grid = np.exp(frequency_grid, out=frequency_grid)
-            wedges[index] = np.multiply(frequency_grid, np.cos(np.radians(angle)))
-
-        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,
-        sampling_rate: float = 1.0,
-        **kwargs,
-    ) -> NDArray:
-        """
-        Generate weighted wedges based on the formalism introduced in [1]_.
-
-        Returns
-        -------
-        NDArray
-            Weighted wedges.
-
-        References
-        ----------
-        .. [1]  Timothy Grant, Nikolaus Grigorieff (2015), eLife 4:e06980.
-        """
-        tilt_shape = compute_tilt_shape(
-            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=shape,
-                opening_axis=opening_axis,
-                tilt_axis=tilt_axis,
-                angle=angle,
-                sampling_rate=sampling_rate,
-            )
-
-            with np.errstate(divide="ignore"):
-                np.power(frequency_grid, power, out=frequency_grid)
-                np.multiply(amplitude, frequency_grid, out=frequency_grid)
-                np.add(frequency_grid, offset, out=frequency_grid)
-                np.multiply(-2, frequency_grid, out=frequency_grid)
-                np.divide(
-                    self.weights[index],
-                    frequency_grid,
-                    out=frequency_grid,
-                )
-
-            wedges[index] = np.exp(frequency_grid)
-
-        return wedges
+        return {"data": ret, "shape": shape}
 
 
 @dataclass
-class WedgeReconstructed:
+class WedgeReconstructed(Wedge):
     """
-    Initialize :py:class:`WedgeReconstructed`.
+    Create wedge mask for tomographic reconstructions.
     """
 
-    #: The tilt angles, defaults to None.
+    #: Tilt angles in degrees.
     angles: Tuple[float] = None
     #: Weights to assign to individual wedge components. Not considered for continuous wedge
     weights: Tuple[float] = None
@@ -290,22 +161,14 @@ class WedgeReconstructed:
     #: Filter window applied during reconstruction.
     reconstruction_filter: str = None
 
-    def __post_init__(self):
-        if self.create_continuous_wedge:
-            self.angles = (min(self.angles), max(self.angles))
-
-    def __call__(
-        self, shape: Tuple[int], return_real_fourier: bool = False, **kwargs
-    ) -> Dict:
+    def _evaluate(self, shape: Tuple[int, ...], **kwargs) -> Dict:
         """
-        Generate the reconstructed wedge.
+        Generate a reconstructed wedge.
 
         Parameters
         ----------
         shape : tuple of int
-            The shape of the reconstruction volume.
-        return_real_fourier : tuple of int
-            Return a shape compliant with rfftn. Defaults to False.
+            The shape to build the filter for.
         **kwargs : dict
             Additional keyword arguments.
 
@@ -316,198 +179,306 @@ class WedgeReconstructed:
                 The filter mask.
             shape: tuple of ints
                 The requested filter shape
-            return_real_fourier: bool
-                Whether data is compliant with rfftn.
-            is_multiplicative_filter: bool
-                Whether the filter is multiplicative in Fourier space.
         """
-        func_args = vars(self).copy()
-        func_args.update(kwargs)
-
-        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)
+        func = step_wedge
+        angles = kwargs.pop("angles", (0,))
+        if kwargs.get("create_continuous_wedge", False):
+            func = continuous_wedge
+            if len(angles) != 2:
+                angles = (min(angles), max(angles))
+
+        weight_wedge = kwargs.get("weight_wedge", False)
+        if kwargs.get("wedge_weights") is None and weight_wedge:
+            kwargs["weights"] = np.cos(np.radians(be.to_numpy_array(angles)))
+        ret = func(shape=shape, angles=angles, **kwargs)
+
+        # Move DC component to origin
+        ret = shift_fourier(ret, shape_is_real_fourier=False)
+        frequency_cutoff = kwargs.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,
+            frequency_mask = (
+                fftfreqn(
+                    shape=shape,
+                    sampling_rate=1,
+                    compute_euclidean_norm=True,
+                    shape_is_real_fourier=False,
+                    fftshift=False,
+                )
+                <= frequency_cutoff
             )
-            ret = np.multiply(ret, frequency_mask <= frequency_cutoff, out=ret)
+            ret = np.multiply(ret, frequency_mask, out=ret)
 
         if not weight_wedge:
             ret = (ret > 0) * 1.0
-
         ret = be.astype(be.to_backend_array(ret), be._float_dtype)
+        return {"data": ret, "shape": shape}
 
-        ret = shift_fourier(data=ret, shape_is_real_fourier=False)
 
-        if return_real_fourier:
-            ret = crop_real_fourier(ret)
+def continuous_wedge(
+    shape: Tuple[int, ...],
+    angles: Tuple[float, float],
+    opening_axis: int,
+    tilt_axis: int,
+    **kwargs: Dict,
+) -> NDArray:
+    """
+    Generate a continous wedge mask with DC component at the center.
 
-        return {
-            "data": ret,
-            "shape": shape,
-            "return_real_fourier": return_real_fourier,
-            "is_multiplicative_filter": True,
-        }
+    Parameters
+    ----------
+    shape : tuple of int
+        The shape of the reconstruction volume.
+    angles : tuple of float
+        Start and stop tilt angle in degrees.
+    opening_axis : int
+        The axis around which the wedge is opened.
+    tilt_axis : int
+        The axis along which the tilt is applied.
 
-    @staticmethod
-    def continuous_wedge(
-        shape: Tuple[int],
-        angles: Tuple[float],
-        opening_axis: int,
-        tilt_axis: int,
-        **kwargs: Dict,
-    ) -> NDArray:
-        """
-        Generate a continous wedge mask with DC component at the center.
+    Returns
+    -------
+    NDArray
+        Wedge mask.
+    """
+    angles = np.abs(np.asarray(angles))
+    aspect_ratio = shape[opening_axis] / shape[tilt_axis]
+    angles = np.degrees(np.arctan(np.tan(np.radians(angles)) * aspect_ratio))
+
+    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],
+        )
 
-        Parameters
-        ----------
-        shape : tuple of int
-            The shape of the reconstruction volume.
-        angles : tuple of float
-            Start and stop tilt angle.
-        opening_axis : int
-            The axis around which the wedge is opened.
-        tilt_axis : int
-            The axis along which the tilt is applied.
+    wedge = np.logical_or(start_radians <= ratios, stop_radians >= ratios).astype(
+        np.float32
+    )
 
-        Returns
-        -------
-        NDArray
-            Wedge mask.
-        """
-        angles = np.abs(np.asarray(angles))
-        aspect_ratio = shape[opening_axis] / shape[tilt_axis]
-        angles = np.degrees(np.arctan(np.tan(np.radians(angles)) * aspect_ratio))
-
-        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],
-            )
+    return wedge
 
-        wedge = np.logical_or(start_radians <= ratios, stop_radians >= ratios).astype(
-            np.float32
-        )
 
-        return wedge
-
-    @staticmethod
-    def step_wedge(
-        shape: Tuple[int],
-        angles: Tuple[float],
-        opening_axis: int,
-        tilt_axis: int,
-        weights: Tuple[float] = None,
-        reconstruction_filter: str = None,
-        **kwargs: Dict,
-    ) -> NDArray:
-        """
-        Generate a per-angle wedge shape with DC component at the center.
+def step_wedge(
+    shape: Tuple[int, ...],
+    angles: Tuple[float, ...],
+    opening_axis: int,
+    tilt_axis: int,
+    weights: Tuple[float, ...] = None,
+    reconstruction_filter: str = None,
+    **kwargs: Dict,
+) -> NDArray:
+    """
+    Generate a per-angle wedge shape with DC component at the center.
 
-        Parameters
-        ----------
-        shape : tuple of int
-            The shape of the reconstruction volume.
-        angles : tuple of float
-            The tilt angles.
-        opening_axis : int
-            The axis around which the wedge is opened.
-        tilt_axis : int
-            The axis along which the tilt is applied.
-        reconstruction_filter : str
-            Filter used during reconstruction.
-        weights : tuple of float, optional
-            Weights to assign to individual tilts. Defaults to 1.
+    Parameters
+    ----------
+    shape : tuple of int
+        The shape of the reconstruction volume.
+    angles : tuple of float
+        The tilt angles in degrees.
+    opening_axis : int
+        The axis around which the wedge is opened.
+    tilt_axis : int
+        The axis along which the tilt is applied.
+    reconstruction_filter : str
+        Filter used during reconstruction.
+    weights : tuple of float, optional
+        Weights to assign to individual tilts. Defaults to 1.
 
-        Returns
-        -------
-        NDArray
-            Wege mask.
-        """
-        from ..backends import NumpyFFTWBackend
+    Returns
+    -------
+    NDArray
+        Wege mask.
+    """
+    from ..backends import NumpyFFTWBackend
+
+    angles = np.asarray(be.to_numpy_array(angles))
+
+    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,
+    )
+
+    aspect_ratio = plane.shape[0] / plane.shape[1]
+    angles = np.degrees(np.arctan(np.tan(np.radians(angles)) * aspect_ratio))
+
+    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
+
+        angle_rad = np.radians(angles[index])
+        rotation_matrix = np.array(
+            [
+                [np.cos(angle_rad), -np.sin(angle_rad)],
+                [np.sin(angle_rad), np.cos(angle_rad)],
+            ]
+        )
+        # We want a push rotation but rigid transform assumes pull
+        NumpyFFTWBackend().rigid_transform(
+            arr=plane * rec_filter,
+            rotation_matrix=rotation_matrix.T,
+            out=plane_rotated,
+            use_geometric_center=True,
+            order=1,
+        )
+        wedge_volume += plane_rotated * weights[index]
 
-        angles = np.asarray(be.to_numpy_array(angles))
+    subset = center_slice(wedge_volume.shape, (shape[opening_axis], shape[tilt_axis]))
+    wedge_volume = wedge_volume[subset]
 
-        if weights is None:
-            weights = np.ones(angles.size)
-        weights = np.asarray(weights)
+    np.fmin(wedge_volume, np.max(weights), wedge_volume)
 
-        shape = tuple(int(x) for x in shape)
-        opening_axis, tilt_axis = int(opening_axis), int(tilt_axis)
+    if opening_axis > tilt_axis:
+        wedge_volume = np.moveaxis(wedge_volume, 1, 0)
 
-        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,
-        )
+    reshape_dimensions = tuple(
+        x if i in (opening_axis, tilt_axis) else 1 for i, x in enumerate(shape)
+    )
 
-        aspect_ratio = plane.shape[0] / plane.shape[1]
-        angles = np.degrees(np.arctan(np.tan(np.radians(angles)) * aspect_ratio))
+    wedge_volume = wedge_volume.reshape(reshape_dimensions)
+    tile_dimensions = np.divide(shape, reshape_dimensions).astype(int)
+    return np.tile(wedge_volume, tile_dimensions)
 
-        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]
+def weight_uniform(angles: Tuple[float, ...], *args, **kwargs) -> NDArray:
+    """
+    Generate uniform weighted wedges.
+    """
+    return weight_angle(angles=np.zeros_like(angles), *args, **kwargs)
 
-        subset = center_slice(
-            wedge_volume.shape, (shape[opening_axis], shape[tilt_axis])
-        )
-        wedge_volume = wedge_volume[subset]
 
-        np.fmin(wedge_volume, np.max(weights), wedge_volume)
+def weight_angle(
+    shape: Tuple[int, ...],
+    angles: Tuple[float, ...],
+    opening_axis: int,
+    tilt_axis: int,
+    **kwargs,
+) -> NDArray:
+    """
+    Generate weighted wedges based on the cosine of the current angle.
+    """
+    tilt_shape = compute_tilt_shape(
+        shape=shape, opening_axis=opening_axis, reduce_dim=True
+    )
+    wedges = np.zeros((len(angles), *tilt_shape))
+    for index, angle in enumerate(angles):
+        wedges[index] = np.cos(np.radians(angle))
+    return wedges
+
+
+def weight_relion(
+    shape: Tuple[int, ...],
+    angles: Tuple[float, ...],
+    weights: Tuple[float, ...],
+    opening_axis: int,
+    tilt_axis: int,
+    sampling_rate: float = 1.0,
+    **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.
 
-        if opening_axis > tilt_axis:
-            wedge_volume = np.moveaxis(wedge_volume, 1, 0)
+    Returns
+    -------
+    NDArray
+        Weighted wedges.
+    """
+    tilt_shape = compute_tilt_shape(
+        shape=shape, opening_axis=opening_axis, reduce_dim=True
+    )
+    wedges = np.zeros((len(angles), *tilt_shape))
+    for index, angle in enumerate(angles):
+        frequency_grid = frequency_grid_at_angle(
+            shape=shape,
+            opening_axis=opening_axis,
+            tilt_axis=tilt_axis,
+            angle=angle,
+            sampling_rate=sampling_rate,
+            fftshift=False,
+        )
+        sigma = np.sqrt(weights[index] * 4 / (8 * np.pi**2))
+        sigma = -2 * np.pi**2 * sigma**2
+        frequency_grid = np.square(frequency_grid, out=frequency_grid)
+        frequency_grid = np.multiply(sigma, frequency_grid, out=frequency_grid)
+        frequency_grid = np.exp(frequency_grid, out=frequency_grid)
+        wedges[index] = np.multiply(frequency_grid, np.cos(np.radians(angle)))
+
+    return wedges
+
+
+def weight_grigorieff(
+    shape: Tuple[int, ...],
+    angles: Tuple[float, ...],
+    weights: Tuple[float, ...],
+    opening_axis: int,
+    tilt_axis: int,
+    amplitude: float = 0.245,
+    power: float = -1.665,
+    offset: float = 2.81,
+    sampling_rate: float = 1.0,
+    **kwargs,
+) -> NDArray:
+    """
+    Generate weighted wedges based on the formalism introduced in [1]_.
 
-        reshape_dimensions = tuple(
-            x if i in (opening_axis, tilt_axis) else 1 for i, x in enumerate(shape)
+    Returns
+    -------
+    NDArray
+        Weighted wedges.
+
+    References
+    ----------
+    .. [1]  Timothy Grant, Nikolaus Grigorieff (2015), eLife 4:e06980.
+    """
+    tilt_shape = compute_tilt_shape(
+        shape=shape, opening_axis=opening_axis, reduce_dim=True
+    )
+
+    wedges = np.zeros((len(angles), *tilt_shape), dtype=be._float_dtype)
+    for index, angle in enumerate(angles):
+        frequency_grid = frequency_grid_at_angle(
+            shape=shape,
+            opening_axis=opening_axis,
+            tilt_axis=tilt_axis,
+            angle=angle,
+            sampling_rate=sampling_rate,
+            fftshift=False,
         )
 
-        wedge_volume = wedge_volume.reshape(reshape_dimensions)
-        tile_dimensions = np.divide(shape, reshape_dimensions).astype(int)
-        wedge_volume = np.tile(wedge_volume, tile_dimensions)
+        with np.errstate(divide="ignore"):
+            np.power(frequency_grid, power, out=frequency_grid)
+            np.multiply(amplitude, frequency_grid, out=frequency_grid)
+            np.add(frequency_grid, offset, out=frequency_grid)
+            np.multiply(-2, frequency_grid, out=frequency_grid)
+            np.divide(weights[index], frequency_grid, out=frequency_grid)
+        wedges[index] = np.exp(frequency_grid)
 
-        return wedge_volume
+    return wedges
 
 
 def _from_xml(filename: str, **kwargs) -> Dict: