pytme 0.2.9__cp311-cp311-macosx_15_0_arm64.whl

tme/backends/_jax_utils.py

@@ -0,0 +1,188 @@
+""" Utility functions for jax backend.
+
+    Copyright (c) 2023-2024 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Tuple
+from functools import partial
+
+import jax.numpy as jnp
+from jax import pmap, lax
+
+from ..types import BackendArray
+from ..backends import backend as be
+from ..matching_utils import normalize_template as _normalize_template
+
+
+def _correlate(template: BackendArray, ft_target: BackendArray) -> BackendArray:
+    """
+    Computes :py:meth:`tme.matching_exhaustive.cc_setup`.
+    """
+    template_ft = jnp.fft.rfftn(template, s=template.shape)
+    template_ft = template_ft.at[:].multiply(ft_target)
+    correlation = jnp.fft.irfftn(template_ft, s=template.shape)
+    return correlation
+
+
+def _flc_scoring(
+    template: BackendArray,
+    template_mask: BackendArray,
+    ft_target: BackendArray,
+    ft_target2: BackendArray,
+    n_observations: BackendArray,
+    eps: float,
+    **kwargs,
+) -> BackendArray:
+    """
+    Computes :py:meth:`tme.matching_exhaustive.flc_scoring`.
+    """
+    correlation = _correlate(template=template, ft_target=ft_target)
+    inv_denominator = _reciprocal_target_std(
+        ft_target=ft_target,
+        ft_target2=ft_target2,
+        template_mask=template_mask,
+        eps=eps,
+        n_observations=n_observations,
+    )
+    correlation = correlation.at[:].multiply(inv_denominator)
+    return correlation
+
+
+def _flcSphere_scoring(
+    template: BackendArray,
+    ft_target: BackendArray,
+    inv_denominator: BackendArray,
+    **kwargs,
+) -> BackendArray:
+    """
+    Computes :py:meth:`tme.matching_exhaustive.flc_scoring`.
+    """
+    correlation = _correlate(template=template, ft_target=ft_target)
+    correlation = correlation.at[:].multiply(inv_denominator)
+    return correlation
+
+
+def _reciprocal_target_std(
+    ft_target: BackendArray,
+    ft_target2: BackendArray,
+    template_mask: BackendArray,
+    n_observations: float,
+    eps: float,
+) -> BackendArray:
+    """
+    Computes reciprocal standard deviation of a target given a mask.
+
+    See Also
+    --------
+    :py:meth:`tme.matching_exhaustive.flc_scoring`.
+    """
+    ft_shape = template_mask.shape
+    ft_template_mask = jnp.fft.rfftn(template_mask, s=ft_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)
+
+    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 = sq_exp.at[:].power(2)
+
+    exp_sq = exp_sq.at[:].add(-sq_exp)
+    exp_sq = exp_sq.at[:].max(0)
+    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))
+    )
+    return exp_sq
+
+
+def _apply_fourier_filter(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
+    arr_ft = jnp.fft.rfftn(arr, s=arr.shape)
+    arr_ft = arr_ft.at[:].multiply(arr_filter)
+    return arr.at[:].set(jnp.fft.irfftn(arr_ft, s=arr.shape))
+
+
+def _identity(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
+    return arr
+
+
+@partial(
+    pmap,
+    in_axes=(0,) + (None,) * 6,
+    static_broadcasted_argnums=[6, 7],
+)
+def scan(
+    target: BackendArray,
+    template: BackendArray,
+    template_mask: BackendArray,
+    rotations: BackendArray,
+    template_filter: BackendArray,
+    target_filter: BackendArray,
+    fast_shape: Tuple[int],
+    rotate_mask: bool,
+) -> Tuple[BackendArray, BackendArray]:
+    eps = jnp.finfo(template.dtype).resolution
+
+    if hasattr(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
+    if not rotate_mask:
+        n_observations = 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),
+            eps=eps,
+            n_observations=n_observations,
+        )
+        ft_target2, scoring_func = None, _flcSphere_scoring
+
+    _template_filter_func = _identity
+    if template_filter.shape != ():
+        _template_filter_func = _apply_fourier_filter
+
+    def _sample_transform(ret, rotation_matrix):
+        max_scores, rotations, 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
+        )
+
+        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
+        )
+        template_rot = be.topleft_pad(template_rot, fast_shape)
+        template_mask_rot = be.topleft_pad(template_mask_rot, fast_shape)
+
+        scores = scoring_func(
+            template=template_rot,
+            template_mask=template_mask_rot,
+            ft_target=ft_target,
+            ft_target2=ft_target2,
+            inv_denominator=inv_denominator,
+            n_observations=n_observations,
+            eps=eps,
+        )
+        max_scores, rotations = be.max_score_over_rotations(
+            scores, max_scores, rotations, index
+        )
+        return (max_scores, rotations, index + 1), None
+
+    score_space = jnp.zeros(fast_shape)
+    rotation_space = jnp.full(shape=fast_shape, dtype=jnp.int32, fill_value=-1)
+    (score_space, rotation_space, _), _ = lax.scan(
+        _sample_transform, (score_space, rotation_space, 0), rotations
+    )
+
+    return score_space, rotation_space

tme/backends/cupy_backend.py

@@ -0,0 +1,294 @@
+""" Backend using cupy for template matching.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+import warnings
+from importlib.util import find_spec
+from contextlib import contextmanager
+from typing import Tuple, Callable, List
+
+import numpy as np
+
+from .npfftw_backend import NumpyFFTWBackend
+from ..types import CupyArray, NDArray, shm_type
+
+PLAN_CACHE = {}
+TEXTURE_CACHE = {}
+
+
+class CupyBackend(NumpyFFTWBackend):
+    """
+    A cupy-based matching backend.
+    """
+
+    def __init__(
+        self,
+        float_dtype: type = None,
+        complex_dtype: type = None,
+        int_dtype: type = None,
+        overflow_safe_dtype: type = None,
+        **kwargs,
+    ):
+        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
+        int_dtype = cp.int32 if int_dtype is None else int_dtype
+        if overflow_safe_dtype is None:
+            overflow_safe_dtype = cp.float32
+
+        super().__init__(
+            array_backend=cp,
+            float_dtype=float_dtype,
+            complex_dtype=complex_dtype,
+            int_dtype=int_dtype,
+            overflow_safe_dtype=overflow_safe_dtype,
+        )
+        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}"
+        self._max_score_over_rotations = self._array_backend.ElementwiseKernel(
+            f"{ftype} internal_scores, {ftype} scores, {itype} rot_index",
+            f"{ftype} out1, {itype} rotations",
+            "if (internal_scores < scores) {out1 = scores; rotations = rot_index;}",
+            "max_score_over_rotations",
+        )
+        self.norm_scores = cp.ElementwiseKernel(
+            f"{ftype} arr, {ftype} exp_sq, {ftype} sq_exp, {ftype} n_obs, {ftype} eps",
+            f"{ftype} out",
+            """
+            // tmp1 = E(X)^2; tmp2 = E(X^2)
+            float tmp1 = sq_exp / n_obs;
+            float tmp2 = exp_sq / n_obs;
+            tmp1 *= tmp1;
+
+            tmp2 = sqrt(max(tmp2 - tmp1, 0.0));
+            // out = (tmp2 < eps) ? 0.0 : arr / (tmp2 * n_obs);
+            tmp1 = arr;
+            if (tmp2 < eps){
+                tmp1 = 0;
+            }
+            tmp2 *= n_obs;
+            out = tmp1 / tmp2;
+            """,
+            "norm_scores",
+        )
+        self.texture_available = find_spec("voltools") is not None
+
+    def to_backend_array(self, arr: NDArray) -> CupyArray:
+        current_device = self._array_backend.cuda.device.get_device_id()
+        if (
+            isinstance(arr, self._array_backend.ndarray)
+            and arr.device.id == current_device
+        ):
+            return arr
+        return self._array_backend.asarray(arr)
+
+    def to_numpy_array(self, arr: CupyArray) -> NDArray:
+        return self._array_backend.asnumpy(arr)
+
+    def to_cpu_array(self, arr: NDArray) -> NDArray:
+        return self.to_numpy_array(arr)
+
+    def from_sharedarr(self, arr: CupyArray) -> CupyArray:
+        return arr
+
+    @staticmethod
+    def to_sharedarr(arr: CupyArray, shared_memory_handler: type = None) -> shm_type:
+        return arr
+
+    def zeros(self, *args, **kwargs):
+        return self._array_backend.zeros(*args, **kwargs)
+
+    def unravel_index(self, indices, shape):
+        return self._array_backend.unravel_index(indices=indices, dims=shape)
+
+    def unique(self, ar, axis=None, *args, **kwargs):
+        if axis is None:
+            return self._array_backend.unique(ar=ar, axis=axis, *args, **kwargs)
+
+        warnings.warn("Axis argument not yet supported in CupY, falling back to NumPy.")
+        ret = np.unique(ar=self.to_numpy_array(ar), axis=axis, *args, **kwargs)
+        if not isinstance(ret, tuple):
+            return self.to_backend_array(ret)
+        return tuple(self.to_backend_array(k) for k in ret)
+
+    def build_fft(
+        self,
+        fwd_shape: Tuple[int],
+        inv_shape: Tuple[int],
+        inv_output_shape: Tuple[int] = None,
+        fwd_axes: Tuple[int] = None,
+        inv_axes: Tuple[int] = None,
+        **kwargs,
+    ) -> Tuple[Callable, Callable]:
+        cache = self._array_backend.fft.config.get_plan_cache()
+        current_device = self._array_backend.cuda.device.get_device_id()
+
+        previous_transform = [fwd_shape, inv_shape]
+        if current_device in PLAN_CACHE:
+            previous_transform = PLAN_CACHE[current_device]
+
+        real_diff, cmplx_diff = True, True
+        if len(fwd_shape) == len(previous_transform[0]):
+            real_diff = fwd_shape == previous_transform[0]
+        if len(inv_shape) == len(previous_transform[1]):
+            cmplx_diff = inv_shape == previous_transform[1]
+
+        if real_diff or cmplx_diff:
+            cache.clear()
+
+        rfft_shape = self._format_fft_shape(fwd_shape, fwd_axes)
+        irfft_shape = fwd_shape if inv_output_shape is None else inv_output_shape
+        irfft_shape = self._format_fft_shape(irfft_shape, inv_axes)
+
+        def rfftn(
+            arr: CupyArray, out: CupyArray = None, s=rfft_shape, axes=fwd_axes
+        ) -> CupyArray:
+            return self.rfftn(arr, s=s, axes=fwd_axes)
+
+        def irfftn(
+            arr: CupyArray, out: CupyArray = None, s=irfft_shape, axes=inv_axes
+        ) -> CupyArray:
+            return self.irfftn(arr, s=s, axes=inv_axes)
+
+        PLAN_CACHE[current_device] = [fwd_shape, inv_shape]
+
+        return rfftn, irfftn
+
+    def rfftn(self, arr: CupyArray, out: CupyArray = None, **kwargs) -> CupyArray:
+        return self._cufft.rfftn(arr, **kwargs)
+
+    def irfftn(self, arr: CupyArray, out: CupyArray = None, **kwargs) -> CupyArray:
+        return self._cufft.irfftn(arr, **kwargs)
+
+    def compute_convolution_shapes(
+        self, arr1_shape: Tuple[int], arr2_shape: Tuple[int]
+    ) -> Tuple[List[int], List[int], List[int]]:
+        from cupyx.scipy.fft import next_fast_len
+
+        convolution_shape = [int(x + y - 1) for x, y in zip(arr1_shape, arr2_shape)]
+        fast_shape = [next_fast_len(x, real=True) for x in convolution_shape]
+        fast_ft_shape = list(fast_shape[:-1]) + [fast_shape[-1] // 2 + 1]
+
+        return convolution_shape, fast_shape, fast_ft_shape
+
+    def max_filter_coordinates(self, score_space, min_distance: Tuple[int]):
+        score_box = tuple(min_distance for _ in range(score_space.ndim))
+        max_filter = self.maximum_filter(score_space, size=score_box, mode="constant")
+        max_filter = max_filter == score_space
+
+        peaks = self._array_backend.array(self._array_backend.nonzero(max_filter)).T
+        return peaks
+
+    # The default methods in Cupy were oddly slow
+    def var(self, a, *args, **kwargs):
+        out = a - self._array_backend.mean(a, *args, **kwargs)
+        self._array_backend.square(out, out)
+        out = self._array_backend.mean(out, *args, **kwargs)
+        return out
+
+    def std(self, a, *args, **kwargs):
+        out = self.var(a, *args, **kwargs)
+        return self._array_backend.sqrt(out)
+
+    def _get_texture(self, arr: CupyArray, order: int = 3, prefilter: bool = False):
+        key = id(arr)
+        if key in TEXTURE_CACHE:
+            return TEXTURE_CACHE[key]
+
+        from voltools import StaticVolume
+
+        # Only keep template and potential corresponding mask in cache
+        if len(TEXTURE_CACHE) >= 2:
+            TEXTURE_CACHE.clear()
+
+        interpolation = "filt_bspline"
+        if order == 1:
+            interpolation = "linear"
+        elif order == 3 and not prefilter:
+            interpolation = "bspline"
+
+        current_device = self._array_backend.cuda.device.get_device_id()
+        TEXTURE_CACHE[key] = StaticVolume(
+            arr, interpolation=interpolation, device=f"gpu:{current_device}"
+        )
+
+        return TEXTURE_CACHE[key]
+
+    def _rigid_transform(
+        self,
+        data: CupyArray,
+        matrix: CupyArray,
+        output: CupyArray,
+        prefilter: bool,
+        order: int,
+        cache: bool = False,
+        batched: bool = False,
+    ) -> None:
+        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:
+        #     # Device memory pool (should) come to rescue performance
+        #     temp = self.zeros(data.shape, data.dtype)
+        #     texture = self._get_texture(data, order=order, prefilter=prefilter)
+        #     texture.affine(transform_m=matrix, profile=False, output=temp)
+        #     output[out_slice] = temp
+        #     return None
+
+        self.affine_transform(
+            input=data,
+            matrix=matrix,
+            mode="constant",
+            output=output[out_slice],
+            order=order,
+            prefilter=prefilter,
+        )
+
+    def get_available_memory(self) -> int:
+        with self._array_backend.cuda.Device():
+            free_memory, _ = self._array_backend.cuda.runtime.memGetInfo()
+        return free_memory
+
+    @contextmanager
+    def set_device(self, device_index: int):
+        with self._array_backend.cuda.Device(device_index):
+            yield
+
+    def device_count(self) -> int:
+        return self._array_backend.cuda.runtime.getDeviceCount()
+
+    def max_score_over_rotations(
+        self,
+        scores: CupyArray,
+        max_scores: CupyArray,
+        rotations: CupyArray,
+        rotation_index: int,
+    ) -> Tuple[CupyArray, CupyArray]:
+        return self._max_score_over_rotations(
+            max_scores,
+            scores,
+            rotation_index,
+            max_scores,
+            rotations,
+        )