pytme 0.3b0__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1__cp311-cp311-macosx_15_0_arm64.whl

tme/backends/_jax_utils.py

@@ -10,16 +10,19 @@ from typing import Tuple
 from functools import partial
 
 import jax.numpy as jnp
-from jax import pmap, lax
+from jax import pmap, lax, vmap
 
 from ..types import BackendArray
 from ..backends import backend as be
 from ..matching_utils import normalize_template as _normalize_template
 
 
+__all__ = ["scan"]
+
+
 def _correlate(template: BackendArray, ft_target: BackendArray) -> BackendArray:
     """
-    Computes :py:meth:`tme.matching_exhaustive.cc_setup`.
+    Computes :py:meth:`tme.matching_scores.cc_setup`.
     """
     template_ft = jnp.fft.rfftn(template, s=template.shape)
     template_ft = template_ft.at[:].multiply(ft_target)
@@ -28,18 +31,17 @@ def _correlate(template: BackendArray, ft_target: BackendArray) -> BackendArray:
 
 
 def _flc_scoring(
-    template: BackendArray,
-    template_mask: BackendArray,
     ft_target: BackendArray,
     ft_target2: BackendArray,
+    template: BackendArray,
+    template_mask: BackendArray,
     n_observations: BackendArray,
     eps: float,
     **kwargs,
 ) -> BackendArray:
     """
-    Computes :py:meth:`tme.matching_exhaustive.flc_scoring`.
+    Computes :py:meth:`tme.matching_scores.flc_scoring`.
     """
-    correlation = _correlate(template=template, ft_target=ft_target)
     inv_denominator = _reciprocal_target_std(
         ft_target=ft_target,
         ft_target2=ft_target2,
@@ -47,18 +49,17 @@ def _flc_scoring(
         eps=eps,
         n_observations=n_observations,
     )
-    correlation = correlation.at[:].multiply(inv_denominator)
-    return correlation
+    return _flcSphere_scoring(ft_target, template, inv_denominator)
 
 
 def _flcSphere_scoring(
-    template: BackendArray,
     ft_target: BackendArray,
+    template: BackendArray,
     inv_denominator: BackendArray,
     **kwargs,
 ) -> BackendArray:
     """
-    Computes :py:meth:`tme.matching_exhaustive.flc_scoring`.
+    Computes :py:meth:`tme.matching_scores.corr_scoring`.
     """
     correlation = _correlate(template=template, ft_target=ft_target)
     correlation = correlation.at[:].multiply(inv_denominator)
@@ -77,7 +78,7 @@ def _reciprocal_target_std(
 
     See Also
     --------
-    :py:meth:`tme.matching_exhaustive.flc_scoring`.
+    :py:meth:`tme.matching_scores.flc_scoring`.
     """
     ft_shape = template_mask.shape
     ft_template_mask = jnp.fft.rfftn(template_mask, s=ft_shape)
@@ -114,7 +115,8 @@ def _identity(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
 @partial(
     pmap,
     in_axes=(0,) + (None,) * 6,
-    static_broadcasted_argnums=[6, 7],
+    static_broadcasted_argnums=[6, 7, 8, 9],
+    axis_name="batch",
 )
 def scan(
     target: BackendArray,
@@ -125,9 +127,17 @@ def scan(
     target_filter: BackendArray,
     fast_shape: Tuple[int],
     rotate_mask: bool,
+    analyzer_class: object,
+    analyzer_kwargs: Tuple[Tuple],
 ) -> Tuple[BackendArray, BackendArray]:
     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"):
         target = _apply_fourier_filter(target, target_filter)
 
@@ -150,7 +160,7 @@ def scan(
         _template_filter_func = _apply_fourier_filter
 
     def _sample_transform(ret, rotation_matrix):
-        max_scores, rotations, index = ret
+        state, index = ret
         template_rot, template_mask_rot = be.rigid_transform(
             arr=template,
             arr_mask=template_mask,
@@ -163,27 +173,20 @@ def scan(
         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)
+        rot_pad = be.topleft_pad(template_rot, fast_shape)
+        mask_rot_pad = be.topleft_pad(template_mask_rot, fast_shape)
 
         scores = scoring_func(
-            template=template_rot,
-            template_mask=template_mask_rot,
+            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,
         )
-        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
-    )
+        state = analyzer(state, scores, rotation_matrix, rotation_index=index)
+        return (state, index + 1), None
 
-    return score_space, rotation_space
+    (state, _), _ = lax.scan(_sample_transform, (analyzer.init_state(), 0), rotations)
+    return state

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,12 +6,9 @@ Copyright (c) 2023 European Molecular Biology Laboratory
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-import warnings
+from typing import Tuple, List
 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
@@ -114,54 +111,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)
-
-        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 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]
@@ -239,13 +196,13 @@ class CupyBackend(NumpyFFTWBackend):
             )
             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
+        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)
+            texture.affine(transform_m=matrix, profile=False, output=temp)
+            output[out_slice] = temp
+            return None
 
         self.affine_transform(
             input=data,