pytme 0.3b0.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1.dev20250731__cp311-cp311-macosx_15_0_arm64.whl

tme/backends/_numpyfftw_utils.py

@@ -0,0 +1,270 @@
+#!/usr/bin/env python
+#
+# Henry Gomersall
+# heng@kedevelopments.co.uk
+#
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its contributors
+# may be used to endorse or promote products derived from this software without
+# specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+#
+
+# This code has been adapted to add support for the out argument in rfftn, irfftn
+# to allow for reusing existing array buffers
+# Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+
+import threading
+
+import pyfftw
+import numpy as np
+import pyfftw.builders as builders
+from pyfftw.interfaces import cache
+from pyfftw.builders._utils import _norm_args, _default_effort, _default_threads
+
+
+def _Xfftn(
+    a,
+    s,
+    axes,
+    overwrite_input,
+    planner_effort,
+    threads,
+    auto_align_input,
+    auto_contiguous,
+    calling_func,
+    normalise_idft=True,
+    ortho=False,
+    real_direction_flag=None,
+    output_array=None,
+):
+
+    work_with_copy = False
+
+    a = np.asanyarray(a)
+
+    try:
+        s = tuple(s)
+    except TypeError:
+        pass
+
+    try:
+        axes = tuple(axes)
+    except TypeError:
+        pass
+
+    if calling_func in ("dct", "dst"):
+        # real-to-real transforms require passing an additional flag argument
+        avoid_copy = False
+        args = (
+            overwrite_input,
+            planner_effort,
+            threads,
+            auto_align_input,
+            auto_contiguous,
+            avoid_copy,
+            real_direction_flag,
+        )
+    elif calling_func in ("irfft2", "irfftn"):
+        # overwrite_input is not an argument to irfft2 or irfftn
+        args = (planner_effort, threads, auto_align_input, auto_contiguous)
+
+        if not overwrite_input:
+            # Only irfft2 and irfftn have overwriting the input
+            # as the default (and so require the input array to
+            # be reloaded).
+            work_with_copy = True
+    else:
+        args = (
+            overwrite_input,
+            planner_effort,
+            threads,
+            auto_align_input,
+            auto_contiguous,
+        )
+
+        if not a.flags.writeable:
+            # Special case of a locked array - always work with a
+            # copy.  See issue #92.
+            work_with_copy = True
+
+            if overwrite_input:
+                raise ValueError(
+                    "overwrite_input cannot be True when the "
+                    + "input array flags.writeable is False"
+                )
+
+    if work_with_copy:
+        # We make the copy before registering the key so that the
+        # copy's stride information will be cached since this will be
+        # used for planning.  Make sure the copy is byte aligned to
+        # prevent further copying
+        a_original = a
+        a = pyfftw.empty_aligned(shape=a.shape, dtype=a.dtype)
+        a[...] = a_original
+
+    if cache.is_enabled():
+        alignment = a.ctypes.data % pyfftw.simd_alignment
+
+        key = (
+            calling_func,
+            a.shape,
+            a.strides,
+            a.dtype,
+            s.__hash__(),
+            axes.__hash__(),
+            alignment,
+            args,
+            threading.get_ident(),
+        )
+
+        try:
+            if key in cache._fftw_cache:
+                FFTW_object = cache._fftw_cache.lookup(key)
+            else:
+                FFTW_object = None
+
+        except KeyError:
+            # This occurs if the object has fallen out of the cache between
+            # the check and the lookup
+            FFTW_object = None
+
+    if not cache.is_enabled() or FFTW_object is None:
+
+        # If we're going to create a new FFTW object and are not
+        # working with a copy, then we need to copy the input array to
+        # preserve it, otherwise we can't actually  take the transform
+        # of the input array! (in general, we have to assume that the
+        # input array will be destroyed during planning).
+        if not work_with_copy:
+            a_copy = a.copy()
+
+        planner_args = (a, s, axes) + args
+
+        FFTW_object = getattr(builders, calling_func)(*planner_args)
+
+        # Only copy if the input array is what was actually used
+        # (otherwise it shouldn't be overwritten)
+        if not work_with_copy and FFTW_object.input_array is a:
+            a[:] = a_copy
+
+        if cache.is_enabled():
+            cache._fftw_cache.insert(FFTW_object, key)
+
+        output_array = FFTW_object(normalise_idft=normalise_idft, ortho=ortho)
+
+    else:
+        orig_output_array = FFTW_object.output_array
+        output_shape = orig_output_array.shape
+        output_dtype = orig_output_array.dtype
+        output_alignment = FFTW_object.output_alignment
+
+        if output_array is None:
+            output_array = pyfftw.empty_aligned(
+                output_shape, output_dtype, n=output_alignment
+            )
+
+        FFTW_object(
+            input_array=a,
+            output_array=output_array,
+            normalise_idft=normalise_idft,
+            ortho=ortho,
+        )
+
+    return output_array
+
+
+def rfftn(
+    a,
+    s=None,
+    axes=None,
+    norm=None,
+    overwrite_input=False,
+    planner_effort=None,
+    threads=None,
+    auto_align_input=True,
+    auto_contiguous=True,
+    out=None,
+):
+    """Perform an n-D real FFT.
+
+    The first four arguments are as per :func:`numpy.fft.rfftn`;
+    the rest of the arguments are documented
+    in the :ref:`additional arguments docs<interfaces_additional_args>`.
+    """
+    calling_func = "rfftn"
+    planner_effort = _default_effort(planner_effort)
+    threads = _default_threads(threads)
+
+    return _Xfftn(
+        a,
+        s,
+        axes,
+        overwrite_input,
+        planner_effort,
+        threads,
+        auto_align_input,
+        auto_contiguous,
+        calling_func,
+        **_norm_args(norm),
+        output_array=out,
+    )
+
+
+def irfftn(
+    a,
+    s=None,
+    axes=None,
+    norm=None,
+    overwrite_input=False,
+    planner_effort=None,
+    threads=None,
+    auto_align_input=True,
+    auto_contiguous=True,
+    out=None,
+):
+    """Perform an n-D real inverse FFT.
+
+    The first four arguments are as per :func:`numpy.fft.rfftn`;
+    the rest of the arguments are documented
+    in the :ref:`additional arguments docs<interfaces_additional_args>`.
+    """
+    calling_func = "irfftn"
+    planner_effort = _default_effort(planner_effort)
+    threads = _default_threads(threads)
+
+    return _Xfftn(
+        a,
+        s,
+        axes,
+        overwrite_input,
+        planner_effort,
+        threads,
+        auto_align_input,
+        auto_contiguous,
+        calling_func,
+        **_norm_args(norm),
+        output_array=out,
+    )

tme/backends/cupy_backend.py

@@ -6,9 +6,9 @@ Copyright (c) 2023 European Molecular Biology Laboratory
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
+from typing import Tuple, List
 from importlib.util import find_spec
 from contextlib import contextmanager
-from typing import Tuple, Callable, List
 
 from .npfftw_backend import NumpyFFTWBackend
 from ..types import CupyArray, NDArray, shm_type
@@ -81,6 +81,17 @@ class CupyBackend(NumpyFFTWBackend):
             """,
             "norm_scores",
         )
+
+        # Sum of square computation similar to the demeaned variance in pytom
+        self.ssum = cp.ReductionKernel(
+            f"{ftype} arr",
+            f"{ftype} ret",
+            "arr * arr",
+            "a + b",
+            "ret = a",
+            "0",
+            f"ssum_{ftype}",
+        )
         self.texture_available = find_spec("voltools") is not None
 
     def to_backend_array(self, arr: NDArray) -> CupyArray:
@@ -111,53 +122,14 @@ class CupyBackend(NumpyFFTWBackend):
     def unravel_index(self, indices, shape):
         return self._array_backend.unravel_index(indices=indices, dims=shape)
 
-    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, overwrite_x=True)
-
-        def irfftn(
-            arr: CupyArray, out: CupyArray = None, s=irfft_shape, axes=inv_axes
-        ) -> CupyArray:
-            return self.irfftn(arr, s=s, axes=inv_axes, overwrite_x=True)
-
-        PLAN_CACHE[current_device] = [fwd_shape, inv_shape]
-        return rfftn, irfftn
+    def free_cache(self):
+        self._array_backend.fft.config.get_plan_cache().clear()
 
     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)
+        return self._cufft.irfftn(arr, **kwargs).astype(self._float_dtype)
 
     def compute_convolution_shapes(
         self, arr1_shape: Tuple[int], arr2_shape: Tuple[int]
@@ -178,17 +150,6 @@ class CupyBackend(NumpyFFTWBackend):
         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:
@@ -235,7 +196,7 @@ class CupyBackend(NumpyFFTWBackend):
             )
             return None
 
-        if data.ndim == 3 and cache and self.texture_available:
+        if data.ndim == 3 and cache and self.texture_available and not batched:
             # Device memory pool (should) come to rescue performance
             temp = self.zeros(data.shape, data.dtype)
             texture = self._get_texture(data, order=order, prefilter=prefilter)

tme/backends/jax_backend.py

@@ -7,7 +7,9 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from functools import wraps
-from typing import Tuple, List, Callable, Dict
+from typing import Tuple, List, Dict, Any
+
+import numpy as np
 
 from ..types import BackendArray
 from .npfftw_backend import NumpyFFTWBackend, shm_type
@@ -64,6 +66,10 @@ class JaxBackend(NumpyFFTWBackend):
         arr = arr.at[idx].set(value)
         return arr
 
+    def addat(self, arr, indices, values):
+        arr = arr.at[indices].add(values)
+        return arr
+
     def topleft_pad(
         self, arr: BackendArray, shape: Tuple[int], padval: int = 0
     ) -> BackendArray:
@@ -88,6 +94,7 @@ class JaxBackend(NumpyFFTWBackend):
             "sqrt",
             "maximum",
             "exp",
+            "mod",
         ]
         for ufunc in ufuncs:
             backend_method = emulate_out(getattr(self._array_backend, ufunc))
@@ -103,27 +110,6 @@ class JaxBackend(NumpyFFTWBackend):
             shape=arr.shape, dtype=arr.dtype, fill_value=value
         )
 
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int] = None,
-        inv_output_shape: Tuple[int] = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-        **kwargs,
-    ) -> Tuple[Callable, Callable]:
-        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, out=None, s=rfft_shape, axes=fwd_axes):
-            return self._array_backend.fft.rfftn(arr, s=s, axes=axes)
-
-        def irfftn(arr, out=None, s=irfft_shape, axes=inv_axes):
-            return self._array_backend.fft.irfftn(arr, s=s, axes=axes)
-
-        return rfftn, irfftn
-
     def rfftn(self, arr: BackendArray, *args, **kwargs) -> BackendArray:
         return self._array_backend.fft.rfftn(arr, **kwargs)
 
@@ -194,12 +180,37 @@ class JaxBackend(NumpyFFTWBackend):
 
         return convolution_shape, fast_shape, fast_ft_shape
 
+    def _to_hashable(self, obj: Any) -> Tuple[str, Tuple]:
+        if isinstance(obj, np.ndarray):
+            return ("array", (tuple(obj.flatten().tolist()), obj.shape))
+        return ("other", obj)
+
+    def _from_hashable(self, type_info: str, data: Any) -> Any:
+        if type_info == "array":
+            data, shape = data
+            return self.array(data).reshape(shape)
+        return data
+
+    def _dict_to_tuple(self, data: Dict) -> Tuple:
+        return tuple((k, self._to_hashable(v)) for k, v in data.items())
+
+    def _tuple_to_dict(self, data: Tuple) -> Dict:
+        return {x[0]: self._from_hashable(*x[1]) for x in data}
+
+    def _unbatch(self, data, target_ndim, index):
+        if not isinstance(data, type(self.zeros(1))):
+            return data
+        elif data.ndim <= target_ndim:
+            return data
+        return data[index]
+
     def scan(
         self,
         matching_data: type,
         splits: Tuple[Tuple[slice, slice]],
         n_jobs: int,
-        callback_class,
+        callback_class: object,
+        callback_class_args: Dict,
         rotate_mask: bool = False,
         **kwargs,
     ) -> List:
@@ -207,12 +218,14 @@ class JaxBackend(NumpyFFTWBackend):
         Emulates output of :py:meth:`tme.matching_exhaustive.scan` using
         :py:class:`tme.analyzer.MaxScoreOverRotations`.
         """
-        from ._jax_utils import scan as scan_inner
+        from ._jax_utils import setup_scan
+        from ..analyzer import MaxScoreOverRotations
 
         pad_target = True if len(splits) > 1 else False
         convolution_mode = "valid" if pad_target else "same"
         target_pad = matching_data.target_padding(pad_target=pad_target)
 
+        score_mask = self.full((1,), fill_value=1, dtype=bool)
         target_shape = tuple(
             (x.stop - x.start + p) for x, p in zip(splits[0][0], target_pad)
         )
@@ -220,16 +233,20 @@ class JaxBackend(NumpyFFTWBackend):
             target_shape=self.to_numpy_array(target_shape),
             template_shape=self.to_numpy_array(matching_data._template.shape),
             batch_mask=self.to_numpy_array(matching_data._batch_mask),
-            pad_target=pad_target,
         )
         analyzer_args = {
-            "convolution_mode": convolution_mode,
+            "shape": fast_shape,
             "fourier_shift": shift,
+            "fast_shape": fast_shape,
             "targetshape": target_shape,
             "templateshape": matching_data.template.shape,
             "convolution_shape": conv_shape,
+            "convolution_mode": convolution_mode,
+            "thread_safe": False,
+            "aggregate_axis": matching_data._batch_axis(matching_data._batch_mask),
+            "n_rotations": matching_data.rotations.shape[0],
+            "jax_mode": True,
         }
-
         create_target_filter = matching_data.target_filter is not None
         create_template_filter = matching_data.template_filter is not None
         create_filter = create_target_filter or create_template_filter
@@ -245,6 +262,9 @@ class JaxBackend(NumpyFFTWBackend):
             for i in range(matching_data.rotations.shape[0])
         }
         for split_start in range(0, len(splits), n_jobs):
+
+            analyzer_kwargs = []
+
             split_subset = splits[split_start : (split_start + n_jobs)]
             if not len(split_subset):
                 continue
@@ -256,8 +276,17 @@ class JaxBackend(NumpyFFTWBackend):
                     target_pad=target_pad,
                     template_slice=template_split,
                 )
+                cur_args = analyzer_args.copy()
+                cur_args["offset"] = translation_offset
+                cur_args.update(callback_class_args)
+                analyzer_kwargs.append(cur_args)
+
+                if pad_target:
+                    score_mask = base._score_mask(fast_shape, shift)
+
+                _target = self.astype(base._target, self._float_dtype)
                 translation_offsets.append(translation_offset)
-                targets.append(self.topleft_pad(base._target, fast_shape))
+                targets.append(self.topleft_pad(_target, fast_shape))
 
             if create_filter:
                 filter_args = {
@@ -279,24 +308,34 @@ class JaxBackend(NumpyFFTWBackend):
 
             create_filter, create_template_filter, create_target_filter = (False,) * 3
             base, targets = None, self._array_backend.stack(targets)
-            scores, rotations = scan_inner(
+
+            scan_inner = setup_scan(
+                analyzer_kwargs=analyzer_kwargs,
+                callback_class=callback_class,
+                fast_shape=fast_shape,
+                rotate_mask=rotate_mask,
+            )
+
+            states = scan_inner(
                 self.astype(targets, self._float_dtype),
-                matching_data.template,
-                matching_data.template_mask,
+                self.astype(matching_data.template, self._float_dtype),
+                self.astype(matching_data.template_mask, self._float_dtype),
                 matching_data.rotations,
                 template_filter,
                 target_filter,
-                fast_shape,
-                rotate_mask,
+                score_mask,
             )
 
-            for index in range(scores.shape[0]):
-                temp = callback_class(
-                    shape=scores[index].shape,
-                    offset=translation_offsets[index],
-                )
-                state = (scores[index], rotations[index], rotation_mapping)
-                ret.append(temp.result(state, **analyzer_args))
+            ndim = targets.ndim - 1
+            for index in range(targets.shape[0]):
+                kwargs = analyzer_kwargs[index]
+                analyzer = callback_class(**kwargs)
+                state = [self._unbatch(x, ndim, index) for x in states]
+
+                if isinstance(analyzer, MaxScoreOverRotations):
+                    state[2] = rotation_mapping
+
+                ret.append(analyzer.result(state, **kwargs))
         return ret
 
     def get_available_memory(self) -> int:

tme/backends/matching_backend.py

@@ -8,7 +8,7 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 
 from abc import ABC, abstractmethod
 from multiprocessing import shared_memory
-from typing import Tuple, Callable, List, Any, Union, Optional, Generator, Dict
+from typing import Tuple, Callable, List, Any, Union, Optional, Generator
 
 from ..types import BackendArray, NDArray, Scalar, shm_type
 
@@ -863,6 +863,17 @@ class MatchingBackend(ABC):
             Indices of ``k`` largest elements in ``arr``.
         """
 
+    @abstractmethod
+    def ssum(self, arr, *args, **kwargs) -> BackendArray:
+        """
+        Compute the sum of squares of ``arr``.
+
+        Returns
+        -------
+        BackendArray
+            Sum of squares with shape ().
+        """
+
     def indices(self, *args, **kwargs) -> BackendArray:
         """
         Creates an array representing the index grid of an input.
@@ -1087,57 +1098,12 @@ class MatchingBackend(ABC):
         """
 
     @abstractmethod
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int],
-        real_dtype: type,
-        cmpl_dtype: type,
-        inv_output_shape: Tuple[int] = None,
-        temp_fwd: NDArray = None,
-        temp_inv: NDArray = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-        fftargs: Dict = {},
-    ) -> Tuple[Callable, Callable]:
-        """
-        Build forward and inverse real fourier transform functions. The returned
-        callables have two parameters ``arr`` and ``out`` which correspond to the
-        input and output of the Fourier transform. The methods return the output
-        of the respective function call, regardless of ``out`` being provided or not,
-        analogous to most numpy functions.
+    def rfftn(self, **kwargs):
+        """Perform an n-D real FFT."""
 
-        Parameters
-        ----------
-        fwd_shape : tuple
-            Input shape for the forward Fourier transform.
-            (see `compute_convolution_shapes`).
-        inv_shape : tuple
-            Input shape for the inverse Fourier transform.
-        real_dtype : dtype
-            Data type of the forward Fourier transform.
-        complex_dtype : dtype
-            Data type of the inverse Fourier transform.
-        inv_output_shape : tuple, optional
-            Output shape of the inverse Fourier transform. By default fast_shape.
-        fftargs : dict, optional
-            Dictionary passed to pyFFTW builders.
-        temp_fwd : NDArray, optional
-            Temporary array to build the forward transform. Superseeds shape defined by
-            fwd_shape if provided.
-        temp_inv : NDArray, optional
-            Temporary array to build the inverse transform. Superseeds shape defined by
-            inv_shape if provided.
-        fwd_axes : tuple of int
-            Axes to perform the forward Fourier transform over.
-        inv_axes : tuple of int
-            Axes to perform the inverse Fourier transform over.
-
-        Returns
-        -------
-        tuple
-            Tuple of callables for forward and inverse real Fourier transform.
-        """
+    @abstractmethod
+    def irfftn(self, **kwargs):
+        """Perform an n-D real inverse FFT."""
 
     def extract_center(self, arr: BackendArray, newshape: Tuple[int]) -> BackendArray:
         """