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

tests/preprocessing/test_utils.py

@@ -51,6 +51,24 @@ class TestPreprocessUtils:
         assert fgrid.shape == tuple(tilt_shape)
         assert fgrid.max() <= np.sqrt(1 / sampling_rate * len(shape))
 
+    @pytest.mark.parametrize("shape", ((15, 15, 15), (31, 31, 31), (64, 64, 64)))
+    @pytest.mark.parametrize("sampling_rate", (0.5, 1, 2))
+    @pytest.mark.parametrize("angle", (-5, 0, 5))
+    def test_freqgrid_comparison(self, shape, sampling_rate, angle):
+        grid = frequency_grid_at_angle(
+            shape=shape,
+            angle=angle,
+            sampling_rate=sampling_rate,
+            opening_axis=2,
+            tilt_axis=0,
+        )
+        grid2 = fftfreqn(
+            shape=shape[1:], sampling_rate=sampling_rate, compute_euclidean_norm=True
+        )
+
+        # These should be equal for cubical input shapes
+        assert np.allclose(grid, grid2)
+
     @pytest.mark.parametrize("n", [10, 100, 1000])
     @pytest.mark.parametrize("sampling_rate", range(1, 4))
     def test_fftfreqn(self, n, sampling_rate):

tests/test_backends.py

@@ -292,16 +292,10 @@ class TestBackends:
 
     @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
     @pytest.mark.parametrize("fast_shape", ((10, 15, 100), (55, 23, 17)))
-    def test_build_fft(self, backend, fast_shape):
+    def test_fft(self, backend, fast_shape):
         _, fast_shape, fast_ft_shape = backend.compute_convolution_shapes(
             fast_shape, (1 for _ in range(len(fast_shape)))
         )
-        rfftn, irfftn = backend.build_fft(
-            fwd_shape=fast_shape,
-            inv_shape=fast_ft_shape,
-            real_dtype=backend._float_dtype,
-            cmpl_dtype=backend._complex_dtype,
-        )
         arr = np.random.rand(*fast_shape)
         out = np.zeros(fast_ft_shape)
 
@@ -310,11 +304,11 @@ class TestBackends:
             backend.to_backend_array(out), backend._complex_dtype
         )
 
-        rfftn(
+        backend.rfftn(
             backend.astype(backend.to_backend_array(arr), backend._float_dtype),
             complex_arr,
         )
-        irfftn(complex_arr, real_arr)
+        backend.irfftn(complex_arr, real_arr)
         assert np.allclose(arr, backend.to_numpy_array(real_arr), rtol=0.3)
 
     @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)

tests/test_density.py

@@ -98,7 +98,6 @@ class TestDensity:
         assert np.allclose(density.data, self.density.data)
         assert np.allclose(density.sampling_rate, self.density.sampling_rate)
         assert np.allclose(density.origin, self.density.origin)
-        assert density.metadata == self.density.metadata
 
     def test_from_file_baseline(self):
         self.test_to_file(gzip=False)

tests/test_matching_utils.py

@@ -10,9 +10,6 @@ from tme.backends import backend as be
 from tme.memory import MATCHING_MEMORY_REGISTRY
 from tme.matching_utils import (
     compute_parallelization_schedule,
-    elliptical_mask,
-    box_mask,
-    tube_mask,
     create_mask,
     scramble_phases,
     apply_convolution_mode,
@@ -50,73 +47,26 @@ class TestMatchingUtils:
             max_splits=256,
         )
 
-    def test_create_mask(self):
+    @pytest.mark.parametrize("mask_type", ["ellipse", "box", "tube", "membrane"])
+    def test_create_mask(self, mask_type: str):
         create_mask(
-            mask_type="ellipse",
+            mask_type=mask_type,
             shape=self.density.shape,
             radius=5,
             center=np.divide(self.density.shape, 2),
+            height=np.max(self.density.shape) // 2,
+            size=np.divide(self.density.shape, 2).astype(int),
+            thickness=2,
+            separation=2,
+            symmetry_axis=1,
+            inner_radius=5,
+            outer_radius=10,
         )
 
     def test_create_mask_error(self):
         with pytest.raises(ValueError):
             create_mask(mask_type=None)
 
-    def test_elliptical_mask(self):
-        elliptical_mask(
-            shape=self.density.shape,
-            radius=5,
-            center=np.divide(self.density.shape, 2),
-        )
-
-    def test_box_mask(self):
-        box_mask(
-            shape=self.density.shape,
-            height=[5, 10, 20],
-            center=np.divide(self.density.shape, 2),
-        )
-
-    def test_tube_mask(self):
-        tube_mask(
-            shape=self.density.shape,
-            outer_radius=10,
-            inner_radius=5,
-            height=5,
-            base_center=np.divide(self.density.shape, 2),
-            symmetry_axis=1,
-        )
-
-    def test_tube_mask_error(self):
-        with pytest.raises(ValueError):
-            tube_mask(
-                shape=self.density.shape,
-                outer_radius=5,
-                inner_radius=10,
-                height=5,
-                base_center=np.divide(self.density.shape, 2),
-                symmetry_axis=1,
-            )
-
-        with pytest.raises(ValueError):
-            tube_mask(
-                shape=self.density.shape,
-                outer_radius=5,
-                inner_radius=10,
-                height=10 * np.max(self.density.shape),
-                base_center=np.divide(self.density.shape, 2),
-                symmetry_axis=1,
-            )
-
-        with pytest.raises(ValueError):
-            tube_mask(
-                shape=self.density.shape,
-                outer_radius=5,
-                inner_radius=10,
-                height=10 * np.max(self.density.shape),
-                base_center=np.divide(self.density.shape, 2),
-                symmetry_axis=len(self.density.shape) + 1,
-            )
-
     def test_scramble_phases(self):
         scramble_phases(arr=self.density.data, noise_proportion=0.5)
 

tests/test_rotations.py

@@ -8,6 +8,7 @@ from tme import Density
 from scipy.spatial.transform import Rotation
 from scipy.signal import correlate
 
+from tme.mask import elliptical_mask
 from tme.rotations import (
     euler_from_rotationmatrix,
     euler_to_rotationmatrix,
@@ -16,7 +17,6 @@ from tme.rotations import (
     get_rotation_matrices,
 )
 from tme.matching_utils import (
-    elliptical_mask,
     split_shape,
     compute_full_convolution_index,
 )

tme/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.b0.post1"
+__version__ = "0.3.1"

tme/analyzer/_utils.py

@@ -93,22 +93,22 @@ def cart_to_score(
     templateshape = be.to_backend_array(templateshape)
     convolution_shape = be.to_backend_array(convolution_shape)
 
-    # Compute removed padding
     output_shape = _convmode_to_shape(
         convolution_mode=convolution_mode,
         targetshape=targetshape,
         templateshape=templateshape,
         convolution_shape=convolution_shape,
     )
-    valid_positions = be.multiply(positions >= 0, positions < output_shape)
-    valid_positions = be.sum(valid_positions, axis=1) == positions.shape[1]
 
+    # Offset from padding the target
     starts = be.astype(
         be.divide(be.subtract(convolution_shape, output_shape), 2),
         be._int_dtype,
     )
-
     positions = be.add(positions, starts)
+
+    valid_positions = be.multiply(positions >= 0, positions < fast_shape)
+    valid_positions = be.sum(valid_positions, axis=1) == positions.shape[1]
     if fourier_shift is not None:
         fourier_shift = be.to_backend_array(fourier_shift)
         positions = be.subtract(positions, fourier_shift)

tme/analyzer/aggregation.py

@@ -545,13 +545,19 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
                 )
 
     def __call__(
-        self, state: Tuple, scores: BackendArray, rotation_matrix: BackendArray
+        self,
+        state: Tuple,
+        scores: BackendArray,
+        rotation_matrix: BackendArray,
+        **kwargs,
     ) -> Tuple:
         mask = self._get_constraint(rotation_matrix)
         mask = self._get_score_mask(mask=mask, scores=scores)
 
         scores = be.multiply(scores, mask, out=scores)
-        return super().__call__(state, scores=scores, rotation_matrix=rotation_matrix)
+        return super().__call__(
+            state, scores=scores, rotation_matrix=rotation_matrix, **kwargs
+        )
 
     def _get_constraint(self, rotation_matrix: BackendArray) -> BackendArray:
         """
@@ -636,7 +642,11 @@ class MaxScoreOverTranslations(MaxScoreOverRotations):
         return scores, rotations, {}
 
     def __call__(
-        self, state, scores: BackendArray, rotation_matrix: BackendArray
+        self,
+        state,
+        scores: BackendArray,
+        rotation_matrix: BackendArray,
+        **kwargs,
     ) -> Tuple:
         prev_scores, rotations, rotation_mapping = state
 

tme/analyzer/peaks.py

@@ -228,10 +228,15 @@ class PeakCaller(AbstractAnalyzer):
         translations = be.full(
             (self.num_peaks, ndim), fill_value=-1, dtype=be._int_dtype
         )
+
+        rdim = len(self.shape)
+        if self.batch_dims:
+            rdim -= len(self.batch_dims)
+
         rotations = be.full(
-            (self.num_peaks, ndim, ndim), fill_value=0, dtype=be._float_dtype
+            (self.num_peaks, rdim, rdim), fill_value=0, dtype=be._float_dtype
         )
-        for i in range(ndim):
+        for i in range(rdim):
             rotations[:, i, i] = 1.0
 
         scores = be.full((self.num_peaks,), fill_value=-1, dtype=be._float_dtype)
@@ -750,8 +755,7 @@ class PeakCallerRecursiveMasking(PeakCaller):
             Dictionary mapping values in rotations to Euler angles.
             By default None
         min_score : float
-            Minimum score value to consider. If provided, superseeds limit given
-            by :py:attr:`PeakCaller.num_peaks`.
+            Minimum score value to consider.
 
         Returns
         -------
@@ -774,10 +778,7 @@ class PeakCallerRecursiveMasking(PeakCaller):
             mask = be.to_backend_array(mask)
             mask_buffer = be.zeros(mask.shape, dtype=mask.dtype)
 
-        peak_limit = self.num_peaks
-        if min_score is not None:
-            peak_limit = be.size(scores)
-        else:
+        if min_score is None:
             min_score = be.min(scores) - 1
 
         _scores = be.zeros(scores.shape, dtype=scores.dtype)
@@ -815,7 +816,7 @@ class PeakCallerRecursiveMasking(PeakCaller):
                 score_mask = mask_buffer[tmpl_slice] <= 0.1
 
             _scores[score_slice] = be.multiply(_scores[score_slice], score_mask)
-            if len(peaks) >= peak_limit:
+            if len(peaks) >= self.num_peaks:
                 break
 
         return be.to_backend_array(peaks), None
@@ -851,7 +852,7 @@ class PeakCallerRecursiveMasking(PeakCaller):
 
         rotation = rotation_mapping[rotation_space[tuple(peak)]]
 
-        # TODO: Newer versions of rotation mapping contain rotation matrices not angles
+        # 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))

tme/backends/_jax_utils.py

@@ -115,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,
@@ -126,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)
 
@@ -151,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,
@@ -176,15 +185,8 @@ def scan(
             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,
+    )