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

tme/analyzer/aggregation.py

@@ -192,6 +192,15 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         )
         return scores, rotations, rotation_mapping, ssum
 
+    def correct_background(self, state, mean=0, inv_std=1, **kwargs):
+        scores, rotations, rotation_mapping, ssum = state
+
+        scores = be.subtract(scores, mean, out=scores)
+        scores = be.multiply(scores, inv_std, out=scores)
+
+        scores = be.maximum(scores, self._score_threshold, out=scores)
+        return scores, rotations, rotation_mapping, ssum
+
     @staticmethod
     def _invert_rmap(rotation_mapping: dict) -> dict:
         """
@@ -201,7 +210,12 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         new_map, ndim = {}, None
         for k, v in rotation_mapping.items():
             nbytes = be.datatype_bytes(be._float_dtype)
-            dtype = np.float32 if nbytes == 4 else np.float16
+            if nbytes == 8:
+                dtype = np.float64
+            elif nbytes == 4:
+                dtype = np.float32
+            else:
+                np.float16
             rmat = np.frombuffer(k, dtype=dtype)
             if ndim is None:
                 ndim = int(np.sqrt(rmat.size))
@@ -451,7 +465,7 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
         Maximum accepted rotational deviation in degrees.
     positions : BackendArray
         Array of shape (n, d) with n seed point translations.
-    positions : BackendArray
+    rotations : BackendArray
         Array of shape (n, d, d) with n seed point rotation matrices.
     reference : BackendArray
         Reference orientation of the template, wlog defaults to (0,0,1).
@@ -489,6 +503,7 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
             be.reshape(be.to_backend_array(reference), (-1,)), be._float_dtype
         )
         positions = be.astype(be.to_backend_array(positions), be._int_dtype)
+        rotations = be.astype(be.to_backend_array(rotations), be._float_dtype)
 
         ndim = positions.shape[1]
         rotate_mask = len(set(acceptance_radius)) != 1
@@ -515,7 +530,13 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
         )
 
         self._positions = positions[valid_positions]
-        rotations = be.to_backend_array(rotations)[valid_positions]
+        rotations = rotations[valid_positions]
+
+        # Convert to pull matrix to remain consistent with rotation convention
+        rotations = be.concatenate(
+            [rotations[i].T[None] for i in range(rotations.shape[0])]
+        )
+
         ex = be.astype(be.to_backend_array((1, 0, 0)), be._float_dtype)
         ey = be.astype(be.to_backend_array((0, 1, 0)), be._float_dtype)
         ez = be.astype(be.to_backend_array((0, 0, 1)), be._float_dtype)
@@ -524,6 +545,15 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
         self._normals_y = (rotations @ ey[..., None])[..., 0]
         self._normals_z = (rotations @ ez[..., None])[..., 0]
 
+        # All scores will be rejected in this case. We should think about a
+        # unified interface for checking analyzer validity to skip such runs
+        if self._positions.shape[0] == 0:
+
+            def _get_score_mask(*args, **kwargs):
+                return 0
+
+            self._get_score_mask = _get_score_mask
+
         # Periodic wrapping could be avoided by padding the target
         shape = be.to_backend_array(self._shape)
         starts = be.subtract(self._positions, extend)
@@ -539,9 +569,7 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
         self._mask_shape = tuple(1 if i != 0 else -1 for i in range(1 + ndim))
 
         if rotate_mask:
-            self._score_mask = be.zeros(
-                (rotations.shape[0], *self._score_mask.shape), dtype=be._float_dtype
-            )
+            self._score_mask = []
             for i in range(rotations.shape[0]):
                 mask = create_mask(
                     mask_type="ellipse",
@@ -550,9 +578,10 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
                     center=tuple(extend for _ in range(ndim)),
                     orientation=be.to_numpy_array(rotations[i]),
                 )
-                self._score_mask[i] = be.astype(
-                    be.to_backend_array(mask), be._float_dtype
+                self._score_mask.append(
+                    be.astype(be.to_backend_array(mask), be._float_dtype)[None]
                 )
+            self._score_mask = be.concatenate(self._score_mask)
 
     def __call__(
         self,
@@ -573,7 +602,7 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
         """
         Determine whether the angle between projection of reference w.r.t to
         a given rotation matrix and a set of rotations fall within the set
-        cone_angle cutoff.
+        cone angle cutoff.
 
         Parameters
         ----------
@@ -585,7 +614,7 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
         BackerndArray
             Boolean mask of shape (n, )
         """
-        template_rot = rotation_matrix @ self._reference
+        template_rot = rotation_matrix.T @ self._reference
 
         x = be.sum(be.multiply(self._normals_x, template_rot), axis=1)
         y = be.sum(be.multiply(self._normals_y, template_rot), axis=1)
@@ -596,10 +625,9 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
     def _get_score_mask(self, mask: BackendArray, scores: BackendArray, **kwargs):
         score_mask = be.zeros(scores.shape, scores.dtype)
 
-        if be.sum(mask) == 0:
-            return score_mask
+        # The indexing could be improved to avoid expanding the mask to
+        # the number of seed points
         mask = be.reshape(mask, self._mask_shape)
-
         score_mask = be.addat(score_mask, self._index_grid, self._score_mask * mask)
         return score_mask > 0
 
@@ -663,13 +691,16 @@ class MaxScoreOverTranslations(MaxScoreOverRotations):
         rotation_index = len(rotation_mapping)
         if self._inversion_mapping:
             rotation_mapping[rotation_index] = rotation_matrix
+        elif self._jax_mode:
+            rotation_index = kwargs.get("rotation_index", 0)
         else:
             rotation = be.tobytes(rotation_matrix)
             rotation_index = rotation_mapping.setdefault(rotation, rotation_index)
-        max_score = be.max(scores, axis=self._aggregate_axis)
 
-        update = prev_scores[rotation_index]
-        update = be.maximum(max_score, update, out=update)
+        scores = be.max(scores, axis=self._aggregate_axis)
+        scores = be.maximum(scores, prev_scores[rotation_index])
+        prev_scores = be.at(prev_scores, rotation_index, scores)
+
         return prev_scores, rotations, rotation_mapping
 
     @classmethod

tme/analyzer/base.py

@@ -73,6 +73,40 @@ class AbstractAnalyzer(ABC):
             Updated analyzer state incorporating the new data.
         """
 
+    @abstractmethod
+    def correct_background(self, state, mean=0, inv_std=1, **kwargs):
+        """
+        Applies flat-fielding correction to scores f as
+
+        .. math::
+
+            f' = (f - \\text{mean}) \\cdot \\text{inv_std},
+
+        transforming raw correlations to SNR-like scores.
+
+        Parameters
+        ----------
+        state : tuple
+            Current analyzer state as returned :py:meth:`AbstractAnalyzer.init_state`
+            or previous invocations of :py:meth:`AbstractAnalyzer.__call__`.
+        mean : BackendArray, optional
+            Background mean (or equivalent), defaults to 0.
+        inv_std : BackendArray, optional
+            Reciprocal background standard deviation (or equivalent), defaults to 1.
+
+        Notes
+        -----
+        This method should be called after all rotations have been processed
+        but before calling :py:meth:`result`. The correction helps distinguish genuine
+        template matches from systematic background artifacts that may arise from
+        template edges, interpolation artifacts, or structured noise in the target.
+
+        Returns
+        -------
+        tuple
+            Updated analyzer state incorporating the new data.
+        """
+
     @abstractmethod
     def result(self, state: Tuple, **kwargs) -> Tuple:
         """

tme/analyzer/peaks.py

@@ -18,7 +18,6 @@ from .base import AbstractAnalyzer
 from ._utils import score_to_cart
 from ..backends import backend as be
 from ..types import BackendArray, NDArray
-from ..rotations import euler_to_rotationmatrix
 from ..matching_utils import split_shape, compute_extraction_box
 
 __all__ = [
@@ -182,6 +181,7 @@ class PeakCaller(AbstractAnalyzer):
         min_score: float = None,
         max_score: float = None,
         batch_dims: Tuple[int] = None,
+        projection_dims: Tuple[int] = None,
         shm_handler: object = None,
         **kwargs,
     ):
@@ -197,9 +197,13 @@ class PeakCaller(AbstractAnalyzer):
         self.min_distance = int(min_distance)
         self.min_boundary_distance = int(min_boundary_distance)
 
-        self.batch_dims = batch_dims
+        self.batch_dims = ()
         if batch_dims is not None:
-            self.batch_dims = tuple(int(x) for x in self.batch_dims)
+            self.batch_dims = tuple(int(x) for x in batch_dims)
+
+        self.projection_dims = ()
+        if projection_dims is not None:
+            self.projection_dims = tuple(int(x) for x in projection_dims)
 
         self.min_score, self.max_score = min_score, max_score
 
@@ -231,7 +235,7 @@ class PeakCaller(AbstractAnalyzer):
 
         rdim = len(self.shape)
         if self.batch_dims:
-            rdim -= len(self.batch_dims)
+            rdim = rdim - len(self.batch_dims) + len(self.projection_dims)
 
         rotations = be.full(
             (self.num_peaks, rdim, rdim), fill_value=0, dtype=be._float_dtype
@@ -388,6 +392,20 @@ class PeakCaller(AbstractAnalyzer):
 
         return state
 
+    def correct_background(self, state, mean, inv_std=1, **kwargs):
+        arr_type = type(be.zeros((1,), be._float_dtype))
+        translations, rotations, scores, details = state
+
+        if isinstance(mean, arr_type):
+            mean = mean[tuple(be.astype(translations.T, int))]
+        scores = be.subtract(scores, mean, out=scores)
+
+        if isinstance(inv_std, arr_type):
+            inv_std = inv_std[tuple(be.astype(translations.T, int))]
+        scores = be.multiply(scores, inv_std, out=scores)
+
+        return translations, rotations, scores, details
+
     @classmethod
     def merge(cls, results=List[Tuple], **kwargs) -> Tuple:
         """
@@ -778,6 +796,9 @@ class PeakCallerRecursiveMasking(PeakCaller):
             mask = be.to_backend_array(mask)
             mask_buffer = be.zeros(mask.shape, dtype=mask.dtype)
 
+        if min_score is None:
+            min_score = self.min_score
+
         if min_score is None:
             min_score = be.min(scores) - 1
 
@@ -849,15 +870,7 @@ class PeakCallerRecursiveMasking(PeakCaller):
         """
         if rotation_space is None or rotation_mapping is None:
             return rotation_matrix
-
-        rotation = rotation_mapping[rotation_space[tuple(peak)]]
-
-        # Old versions of rotation mapping contained Euler angles
-        if rotation.ndim != 2:
-            rotation = be.to_backend_array(
-                euler_to_rotationmatrix(be.to_numpy_array(rotation))
-            )
-        return rotation
+        return rotation_mapping[rotation_space[tuple(peak)]]
 
 
 class PeakCallerScipy(PeakCaller):

tme/analyzer/proxy.py

@@ -85,6 +85,16 @@ class StatelessSharedAnalyzerProxy:
             final_state = tuple(self._shared_to_object(x) for x in final_state)
         return self._analyzer.result(final_state, **kwargs)
 
+    def correct_background(self, state, *args, **kwargs):
+        if self._shared:
+            # Copy to not correct the internal score array across processes
+            backend_arr = type(be.zeros((1), dtype=be._float_dtype))
+            state = tuple(self._shared_to_object(x) for x in state)
+            state = tuple(
+                be.copy(x) if isinstance(x, backend_arr) else x for x in state
+            )
+        return self._analyzer.correct_background(state, *args, **kwargs)
+
     def merge(self, *args, **kwargs):
         return self._analyzer.merge(*args, **kwargs)
 
@@ -121,3 +131,7 @@ class SharedAnalyzerProxy(StatelessSharedAnalyzerProxy):
     def result(self, **kwargs):
         """Extract final result"""
         return super().result(self._state, **kwargs)
+
+    def correct_background(self, *args, **kwargs):
+        # We always assign to state as this operation can not be shared
+        self._state = super().correct_background(self._state, *args, **kwargs)

tme/backends/_jax_utils.py

@@ -10,14 +10,14 @@ from typing import Tuple
 from functools import partial
 
 import jax.numpy as jnp
-from jax import pmap, lax, vmap
+from jax import pmap, lax, jit
 
 from ..types import BackendArray
 from ..backends import backend as be
-from ..matching_utils import normalize_template as _normalize_template
+from ..matching_utils import standardize, to_padded
 
 
-__all__ = ["scan"]
+__all__ = ["scan", "setup_scan"]
 
 
 def _correlate(template: BackendArray, ft_target: BackendArray) -> BackendArray:
@@ -62,15 +62,14 @@ def _flcSphere_scoring(
     Computes :py:meth:`tme.matching_scores.corr_scoring`.
     """
     correlation = _correlate(template=template, ft_target=ft_target)
-    correlation = correlation.at[:].multiply(inv_denominator)
-    return correlation
+    return correlation.at[:].multiply(inv_denominator)
 
 
 def _reciprocal_target_std(
     ft_target: BackendArray,
     ft_target2: BackendArray,
     template_mask: BackendArray,
-    n_observations: float,
+    n_obs: float,
     eps: float,
 ) -> BackendArray:
     """
@@ -80,16 +79,16 @@ def _reciprocal_target_std(
     --------
     :py:meth:`tme.matching_scores.flc_scoring`.
     """
-    ft_shape = template_mask.shape
-    ft_template_mask = jnp.fft.rfftn(template_mask, s=ft_shape)
+    shape = template_mask.shape
+    ft_template_mask = jnp.fft.rfftn(template_mask, s=shape)
 
     # E(X^2)- E(X)^2
-    exp_sq = jnp.fft.irfftn(ft_target2 * ft_template_mask, s=ft_shape)
-    exp_sq = exp_sq.at[:].divide(n_observations)
+    exp_sq = jnp.fft.irfftn(ft_target2 * ft_template_mask, s=shape)
+    exp_sq = exp_sq.at[:].divide(n_obs)
 
     ft_template_mask = ft_template_mask.at[:].multiply(ft_target)
-    sq_exp = jnp.fft.irfftn(ft_template_mask, s=ft_shape)
-    sq_exp = sq_exp.at[:].divide(n_observations)
+    sq_exp = jnp.fft.irfftn(ft_template_mask, s=shape)
+    sq_exp = sq_exp.at[:].divide(n_obs)
     sq_exp = sq_exp.at[:].power(2)
 
     exp_sq = exp_sq.at[:].add(-sq_exp)
@@ -97,7 +96,7 @@ def _reciprocal_target_std(
     exp_sq = exp_sq.at[:].power(0.5)
 
     exp_sq = exp_sq.at[:].set(
-        jnp.where(exp_sq <= eps, 0, jnp.reciprocal(exp_sq * n_observations))
+        jnp.where(exp_sq <= eps, 0, jnp.reciprocal(exp_sq * n_obs))
     )
     return exp_sq
 
@@ -108,20 +107,50 @@ def _apply_fourier_filter(arr: BackendArray, arr_filter: BackendArray) -> Backen
     return arr.at[:].set(jnp.fft.irfftn(arr_ft, s=arr.shape))
 
 
-def _identity(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
-    return arr
+def setup_scan(analyzer_kwargs, analyzer, fast_shape, rotate_mask, match_projection):
+    """Create separate scan function with initialized analyzer for each device"""
+    device_scans = [
+        partial(
+            scan,
+            fast_shape=fast_shape,
+            rotate_mask=rotate_mask,
+            analyzer=analyzer(**device_config),
+        )
+        for device_config in analyzer_kwargs
+    ]
 
+    @partial(
+        pmap,
+        in_axes=(0,) + (None,) * 7,
+        axis_name="batch",
+    )
+    def scan_combined(
+        target,
+        template,
+        template_mask,
+        rotations,
+        template_filter,
+        target_filter,
+        score_mask,
+        background_template,
+    ):
+        return lax.switch(
+            lax.axis_index("batch"),
+            device_scans,
+            target,
+            template,
+            template_mask,
+            rotations,
+            template_filter,
+            target_filter,
+            score_mask,
+            background_template,
+        )
 
-def _mask_scores(arr, mask):
-    return arr.at[:].multiply(mask)
+    return scan_combined
 
 
-@partial(
-    pmap,
-    in_axes=(0,) + (None,) * 7,
-    static_broadcasted_argnums=[7, 8, 9, 10],
-    axis_name="batch",
-)
+@partial(jit, static_argnums=(8, 9, 10))
 def scan(
     target: BackendArray,
     template: BackendArray,
@@ -130,74 +159,98 @@ def scan(
     template_filter: BackendArray,
     target_filter: BackendArray,
     score_mask: BackendArray,
+    background_template: BackendArray,
     fast_shape: Tuple[int],
     rotate_mask: bool,
-    analyzer_class: object,
-    analyzer_kwargs: Tuple[Tuple],
-) -> Tuple[BackendArray, BackendArray]:
+    analyzer: object,
+) -> Tuple:
     eps = jnp.finfo(template.dtype).resolution
 
-    kwargs = lax.switch(
-        lax.axis_index("batch"),
-        [lambda: analyzer_kwargs[i] for i in range(len(analyzer_kwargs))],
-    )
-    analyzer = analyzer_class(**be._tuple_to_dict(kwargs))
-
-    if hasattr(target_filter, "shape"):
+    if target_filter.shape != ():
         target = _apply_fourier_filter(target, target_filter)
 
     ft_target = jnp.fft.rfftn(target, s=fast_shape)
     ft_target2 = jnp.fft.rfftn(jnp.square(target), s=fast_shape)
-    inv_denominator, target, scoring_func = None, None, _flc_scoring
+    _n_obs, _inv_denominator, target = None, None, None
+
+    unpadded_slice = tuple(slice(0, x) for x in template.shape)
+    rot_buffer, mask_rot_buffer = jnp.zeros(fast_shape), jnp.zeros(fast_shape)
     if not rotate_mask:
-        n_observations = jnp.sum(template_mask)
-        inv_denominator = _reciprocal_target_std(
+        _n_obs = jnp.sum(template_mask)
+        _inv_denominator = _reciprocal_target_std(
             ft_target=ft_target,
             ft_target2=ft_target2,
-            template_mask=be.topleft_pad(template_mask, fast_shape),
+            template_mask=to_padded(mask_rot_buffer, template_mask, unpadded_slice),
             eps=eps,
-            n_observations=n_observations,
+            n_obs=_n_obs,
         )
-        ft_target2, scoring_func = None, _flcSphere_scoring
+        ft_target2 = None
 
-    _template_filter_func = _identity
-    if template_filter.shape != ():
-        _template_filter_func = _apply_fourier_filter
+    mask_scores = score_mask.shape != ()
+    filter_template = template_filter.shape != ()
+    bg_correction = background_template.shape != ()
+    bg_scores = jnp.zeros(fast_shape) if bg_correction else 0
 
-    _score_mask_func = _identity
-    if score_mask.shape != ():
-        _score_mask_func = _mask_scores
+    _template_mask_rot = template_mask
+    template_indices = be._index_grid(template.shape)
+    center = be.divide(be.to_backend_array(template.shape) - 1, 2)
 
     def _sample_transform(ret, rotation_matrix):
-        state, index = ret
-        template_rot, template_mask_rot = be.rigid_transform(
-            arr=template,
-            arr_mask=template_mask,
-            rotation_matrix=rotation_matrix,
-            order=1,  # thats all we get for now
+        matrix = be._build_transform_matrix(
+            rotation_matrix=rotation_matrix, center=center
         )
+        indices = be._transform_indices(template_indices, matrix)
+
+        template_rot = be._interpolate(template, indices, order=1)
+        n_obs, template_mask_rot = _n_obs, _template_mask_rot
+        if rotate_mask:
+            template_mask_rot = be._interpolate(template_mask, indices, order=1)
+            n_obs = jnp.sum(template_mask_rot)
+
+        if filter_template:
+            template_rot = _apply_fourier_filter(template_rot, template_filter)
+        template_rot = standardize(template_rot, template_mask_rot, n_obs)
+
+        rot_pad = to_padded(rot_buffer, template_rot, unpadded_slice)
+
+        inv_denominator = _inv_denominator
+        if rotate_mask:
+            mask_rot_pad = to_padded(mask_rot_buffer, template_mask_rot, unpadded_slice)
+            inv_denominator = _reciprocal_target_std(
+                ft_target=ft_target,
+                ft_target2=ft_target2,
+                template_mask=mask_rot_pad,
+                n_obs=n_obs,
+                eps=eps,
+            )
+
+        scores = _flcSphere_scoring(ft_target, rot_pad, inv_denominator)
+        if mask_scores:
+            scores = scores.at[:].multiply(score_mask)
+
+        state, bg_scores, index = ret
+        state = analyzer(state, scores, rotation_matrix, rotation_index=index)
 
-        n_observations = jnp.sum(template_mask_rot)
-        template_rot = _template_filter_func(template_rot, template_filter)
-        template_rot = _normalize_template(
-            template_rot, template_mask_rot, n_observations
-        )
-        rot_pad = be.topleft_pad(template_rot, fast_shape)
-        mask_rot_pad = be.topleft_pad(template_mask_rot, fast_shape)
+        if bg_correction:
+            template_rot = be._interpolate(background_template, indices, order=1)
+            if filter_template:
+                template_rot = _apply_fourier_filter(template_rot, template_filter)
+            template_rot = standardize(template_rot, template_mask_rot, n_obs)
 
-        scores = scoring_func(
-            template=rot_pad,
-            template_mask=mask_rot_pad,
-            ft_target=ft_target,
-            ft_target2=ft_target2,
-            inv_denominator=inv_denominator,
-            n_observations=n_observations,
-            eps=eps,
-        )
-        scores = _score_mask_func(scores, score_mask)
+            rot_pad = to_padded(rot_buffer, template_rot, unpadded_slice)
+            scores = _flcSphere_scoring(ft_target, rot_pad, inv_denominator)
+            bg_scores = jnp.maximum(bg_scores, scores)
 
-        state = analyzer(state, scores, rotation_matrix, rotation_index=index)
-        return (state, index + 1), None
+        return (state, bg_scores, index + 1), None
+
+    (state, bg_scores, _), _ = lax.scan(
+        _sample_transform, (analyzer.init_state(), bg_scores, 0), rotations
+    )
+
+    if bg_correction:
+        if mask_scores:
+            bg_scores = bg_scores.at[:].multiply(score_mask)
+        bg_scores = bg_scores.at[:].add(-be.mean(bg_scores))
+        state = analyzer.correct_background(state, bg_scores)
 
-    (state, _), _ = lax.scan(_sample_transform, (analyzer.init_state(), 0), rotations)
     return state

tme/backends/cupy_backend.py

@@ -33,7 +33,6 @@ class CupyBackend(NumpyFFTWBackend):
         import cupy as cp
         import cupyx.scipy.fft as cufft
         from cupyx.scipy.ndimage import affine_transform, maximum_filter
-        from ._cupy_utils import affine_transform_batch
 
         float_dtype = cp.float32 if float_dtype is None else float_dtype
         complex_dtype = cp.complex64 if complex_dtype is None else complex_dtype
@@ -51,7 +50,6 @@ class CupyBackend(NumpyFFTWBackend):
         self._cufft = cufft
         self.maximum_filter = maximum_filter
         self.affine_transform = affine_transform
-        self.affine_transform_batch = affine_transform_batch
 
         itype = f"int{self.datatype_bytes(int_dtype) * 8}"
         ftype = f"float{self.datatype_bytes(float_dtype) * 8}"
@@ -157,8 +155,8 @@ class CupyBackend(NumpyFFTWBackend):
 
         from voltools import StaticVolume
 
-        # Only keep template and potential corresponding mask in cache
-        if len(TEXTURE_CACHE) >= 2:
+        # Only keep template, mask and noise template in cache
+        if len(TEXTURE_CACHE) >= 3:
             TEXTURE_CACHE.clear()
 
         interpolation = "filt_bspline"
@@ -174,7 +172,7 @@ class CupyBackend(NumpyFFTWBackend):
 
         return TEXTURE_CACHE[key]
 
-    def _rigid_transform(
+    def _transform(
         self,
         data: CupyArray,
         matrix: CupyArray,
@@ -182,21 +180,10 @@ class CupyBackend(NumpyFFTWBackend):
         prefilter: bool,
         order: int,
         cache: bool = False,
-        batched: bool = False,
-    ) -> None:
+    ) -> CupyArray:
         out_slice = tuple(slice(0, stop) for stop in data.shape)
-        if batched:
-            self.affine_transform_batch(
-                input=data,
-                matrix=matrix,
-                mode="constant",
-                output=output[out_slice],
-                order=order,
-                prefilter=prefilter,
-            )
-            return None
 
-        if data.ndim == 3 and cache and self.texture_available and not batched:
+        if data.ndim == 3 and cache and self.texture_available:
             # Device memory pool (should) come to rescue performance
             temp = self.zeros(data.shape, data.dtype)
             texture = self._get_texture(data, order=order, prefilter=prefilter)
@@ -204,7 +191,7 @@ class CupyBackend(NumpyFFTWBackend):
             output[out_slice] = temp
             return None
 
-        self.affine_transform(
+        return self.affine_transform(
             input=data,
             matrix=matrix,
             mode="constant",