pytme 0.3b0.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1.dev20250731__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_analyzer.py

@@ -165,7 +165,6 @@ class TestMaxScoreOverRotations:
         assert res[0].dtype == be._float_dtype
         assert res[1].size == self.data.ndim
         assert np.allclose(res[2].shape, self.data.shape)
-        assert len(res) == 4
 
     @pytest.mark.parametrize("use_memmap", [False, True])
     @pytest.mark.parametrize("score_threshold", [0, 1e10, -1e10])
@@ -181,7 +180,7 @@ class TestMaxScoreOverRotations:
 
         data2 = self.data * 2
         score_analyzer(state, data2, rotation_matrix=self.rotation_matrix)
-        scores, translation_offset, rotations, mapping = score_analyzer.result(state)
+        scores, offset, rotations, mapping, *_ = score_analyzer.result(state)
 
         assert np.all(scores >= score_threshold)
         max_scores = np.maximum(self.data, data2)
@@ -214,7 +213,7 @@ class TestMaxScoreOverRotations:
         ret = MaxScoreOverRotations.merge(
             results=states, use_memmap=use_memmap, score_threshold=score_threshold
         )
-        scores, translation, rotations, mapping = ret
+        scores, translation, rotations, mapping, *_ = ret
         assert np.all(scores >= score_threshold)
         max_scores = np.maximum(self.data, data2)
         max_scores = np.maximum(max_scores, score_threshold)

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_extensions.py

@@ -53,7 +53,6 @@ class TestExtensions:
     @pytest.mark.parametrize("min_distance", [0, 5, 10])
     def test_find_candidate_indices(self, dimension, dtype, min_distance):
         coordinates = COORDINATES[dimension].astype(dtype)
-        print(coordinates.shape)
 
         min_distance = np.array([min_distance]).astype(dtype)[0]
 

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_orientations.py

@@ -95,18 +95,6 @@ class TestDensity:
             self.orientations.rotations, orientations_new.rotations, atol=1e-3
         )
 
-    @pytest.mark.parametrize("input_format", ("text", "star", "tbl"))
-    @pytest.mark.parametrize("output_format", ("text", "star", "tbl"))
-    def test_file_format_io(self, input_format: str, output_format: str):
-        _, output_file = mkstemp(suffix=f".{input_format}")
-        _, output_file2 = mkstemp(suffix=f".{output_format}")
-
-        self.orientations.to_file(output_file)
-        orientations_new = Orientations.from_file(output_file)
-        orientations_new.to_file(output_file2)
-
-        assert True
-
     @pytest.mark.parametrize("drop_oob", (True, False))
     @pytest.mark.parametrize("shape", (10, 40, 80))
     @pytest.mark.parametrize("odd", (True, False))

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

@@ -132,12 +132,14 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             - scores : BackendArray of shape `self._shape` filled with `score_threshold`.
             - rotations : BackendArray of shape `self._shape` filled with -1.
             - rotation_mapping : dict, empty mapping from rotation bytes to indices.
+            - ssum : BackendArray, accumulator for sum of squared scores.
         """
         scores = be.full(
             shape=self._shape, dtype=be._float_dtype, fill_value=self._score_threshold
         )
         rotations = be.full(self._shape, dtype=be._int_dtype, fill_value=-1)
-        return scores, rotations, {}
+        ssum = be.full((1), dtype=be._float_dtype, fill_value=0)
+        return scores, rotations, {}, ssum
 
     def __call__(
         self,
@@ -156,6 +158,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             - scores : BackendArray, current maximum scores.
             - rotations : BackendArray, current rotation indices.
             - rotation_mapping : dict, mapping from rotation bytes to indices.
+            - ssum : BackendArray, accumulator for sum of squared scores.
         scores : BackendArray
             Array of new scores to update analyzer with.
         rotation_matrix : BackendArray
@@ -168,7 +171,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         # be.tobytes behaviour caused overhead for certain GPU/CUDA combinations
         # If the analyzer is not shared and each rotation is unique, we can
         # use index to rotation mapping and invert prior to merging.
-        prev_scores, rotations, rotation_mapping = state
+        prev_scores, rotations, rotation_mapping, ssum = state
 
         rotation_index = len(rotation_mapping)
         rotation_matrix = be.astype(rotation_matrix, be._float_dtype)
@@ -180,13 +183,14 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             rotation = be.tobytes(rotation_matrix)
             rotation_index = rotation_mapping.setdefault(rotation, rotation_index)
 
+        ssum = be.add(ssum, be.ssum(scores), out=ssum)
         scores, rotations = be.max_score_over_rotations(
             scores=scores,
             max_scores=prev_scores,
             rotations=rotations,
             rotation_index=rotation_index,
         )
-        return scores, rotations, rotation_mapping
+        return scores, rotations, rotation_mapping, ssum
 
     @staticmethod
     def _invert_rmap(rotation_mapping: dict) -> dict:
@@ -224,6 +228,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             - scores : BackendArray, current maximum scores.
             - rotations : BackendArray, current rotation indices.
             - rotation_mapping : dict, mapping from rotation indices to matrices.
+            - ssum : BackendArray, accumulator for sum of squared scores.
         targetshape : Tuple[int], optional
             Shape of the target for convolution mode correction.
         templateshape : Tuple[int], optional
@@ -240,9 +245,9 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         Returns
         -------
         tuple
-            Final result tuple (scores, offset, rotations, rotation_mapping).
+            Final result tuple (scores, offset, rotations, rotation_mapping, ssum).
         """
-        scores, rotations, rotation_mapping = state
+        scores, rotations, rotation_mapping, ssum = state
 
         # Apply postprocessing if parameters are provided
         if fourier_shift is not None:
@@ -269,11 +274,13 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         if self._inversion_mapping:
             rotation_mapping = {be.tobytes(v): k for k, v in rotation_mapping.items()}
 
+        n_rotations = max(len(rotation_mapping), 1)
         return (
             scores,
             be.to_numpy_array(self._offset),
             rotations,
             self._invert_rmap(rotation_mapping),
+            be.to_numpy_array(ssum) / (scores.size * n_rotations),
         )
 
     def _harmonize_states(states: List[Tuple]):
@@ -287,18 +294,18 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             if states[i] is None:
                 continue
 
-            scores, offset, rotations, rotation_mapping = states[i]
+            scores, offset, rotations, rotation_mapping, ssum = states[i]
             if out_shape is None:
                 out_shape = np.zeros(scores.ndim, int)
             out_shape = np.maximum(out_shape, np.add(offset, scores.shape))
 
             new_param = {}
             for key, value in rotation_mapping.items():
-                rotation_bytes = be.tobytes(value)
+                rotation_bytes = np.asarray(value).tobytes()
                 new_param[rotation_bytes] = key
                 if rotation_bytes not in new_rotation_mapping:
                     new_rotation_mapping[rotation_bytes] = len(new_rotation_mapping)
-            states[i] = (scores, offset, rotations, new_param)
+            states[i] = (scores, offset, rotations, new_param, ssum)
         out_shape = tuple(int(x) for x in out_shape)
         return new_rotation_mapping, out_shape, states
 
@@ -329,11 +336,10 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         if len(results) == 1:
             ret = results[0]
             if use_memmap:
-                scores, offset, rotations, rotation_mapping = ret
+                scores, offset, rotations, rotation_mapping, ssum = ret
                 scores = array_to_memmap(scores)
                 rotations = array_to_memmap(rotations)
-                ret = (scores, offset, rotations, rotation_mapping)
-
+                ret = (scores, offset, rotations, rotation_mapping, ssum)
             return ret
 
         # Determine output array shape and create consistent rotation map
@@ -368,6 +374,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             )
             rotations_out = np.full(out_shape, fill_value=-1, dtype=rotations_dtype)
 
+        total_ssum = 0
         for i in range(len(results)):
             if results[i] is None:
                 continue
@@ -385,7 +392,9 @@ class MaxScoreOverRotations(AbstractAnalyzer):
                     shape=out_shape,
                     dtype=rotations_dtype,
                 )
-            scores, offset, rotations, rotation_mapping = results[i]
+            scores, offset, rotations, rotation_mapping, ssum = results[i]
+
+            total_ssum = np.add(total_ssum, ssum)
             stops = np.add(offset, scores.shape).astype(int)
             indices = tuple(slice(*pos) for pos in zip(offset, stops))
 
@@ -428,6 +437,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             np.zeros(scores_out.ndim, dtype=int),
             rotations_out,
             cls._invert_rmap(master_rotation_mapping),
+            total_ssum / len(results),
         )
 
 
@@ -545,13 +555,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 +652,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

@@ -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, vmap, jit
 
 from ..types import BackendArray
 from ..backends import backend as be
 from ..matching_utils import normalize_template as _normalize_template
 
 
-__all__ = ["scan"]
+__all__ = ["scan", "setup_scan"]
 
 
 def _correlate(template: BackendArray, ft_target: BackendArray) -> BackendArray:
@@ -112,11 +112,56 @@ def _identity(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
     return arr
 
 
-@partial(
-    pmap,
-    in_axes=(0,) + (None,) * 6,
-    static_broadcasted_argnums=[6, 7],
-)
+def _mask_scores(arr, mask):
+    return arr.at[:].multiply(mask)
+
+
+def _select_config(analyzer_kwargs, device_idx):
+    return analyzer_kwargs[device_idx]
+
+
+def setup_scan(analyzer_kwargs, callback_class, fast_shape, rotate_mask):
+    """Create separate scan function with initialized analyzer for each device"""
+    device_scans = [
+        partial(
+            scan,
+            fast_shape=fast_shape,
+            rotate_mask=rotate_mask,
+            analyzer=callback_class(**device_config),
+        )
+        for device_config in analyzer_kwargs
+    ]
+
+    @partial(
+        pmap,
+        in_axes=(0,) + (None,) * 6,
+        axis_name="batch",
+    )
+    def scan_combined(
+        target,
+        template,
+        template_mask,
+        rotations,
+        template_filter,
+        target_filter,
+        score_mask,
+    ):
+        return lax.switch(
+            lax.axis_index("batch"),
+            device_scans,
+            target,
+            template,
+            template_mask,
+            rotations,
+            template_filter,
+            target_filter,
+            score_mask,
+        )
+
+    return scan_combined
+
+
+@partial(jit, static_argnums=(7, 8, 9))
 def scan(
     target: BackendArray,
     template: BackendArray,
@@ -124,8 +169,10 @@ def scan(
     rotations: BackendArray,
     template_filter: BackendArray,
     target_filter: BackendArray,
+    score_mask: BackendArray,
     fast_shape: Tuple[int],
     rotate_mask: bool,
+    analyzer: object,
 ) -> Tuple[BackendArray, BackendArray]:
     eps = jnp.finfo(template.dtype).resolution
 
@@ -150,8 +197,12 @@ def scan(
     if template_filter.shape != ():
         _template_filter_func = _apply_fourier_filter
 
+    _score_mask_func = _identity
+    if score_mask.shape != ():
+        _score_mask_func = _mask_scores
+
     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 +227,10 @@ 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
+        scores = _score_mask_func(scores, score_mask)
 
-    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