pytme 0.3.1__tar.gz → 0.3.1.dev20250731__tar.gz

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytme
-Version: 0.3.1
+Version: 0.3.1.dev20250731
 Summary: Python Template Matching Engine
 Author: Valentin Maurer
 Author-email: Valentin Maurer <valentin.maurer@embl-hamburg.de>
@@ -26,9 +26,7 @@ Requires-Dist: importlib_resources
 Requires-Dist: joblib
 Provides-Extra: cupy
 Requires-Dist: cupy-cuda12x>13.0.0; extra == "cupy"
-Provides-Extra: cupy-voltools
-Requires-Dist: cupy-cuda12x>13.0.0; extra == "cupy-voltools"
-Requires-Dist: voltools; extra == "cupy-voltools"
+Requires-Dist: voltools; extra == "cupy"
 Provides-Extra: pytorch
 Requires-Dist: torch; extra == "pytorch"
 Requires-Dist: torchvision; extra == "pytorch"

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/pyproject.toml

@@ -7,7 +7,7 @@ name="pytme"
 authors = [
   { name = "Valentin Maurer", email = "valentin.maurer@embl-hamburg.de" },
 ]
-version="0.3.1"
+version="0.3.1.dev20250731"
 description="Python Template Matching Engine"
 readme="README.md"
 requires-python = ">=3.11"
@@ -33,8 +33,7 @@ classifiers = [
 ]
 
 [project.optional-dependencies]
-cupy = ["cupy-cuda12x>13.0.0"]
-cupy_voltools = ["cupy-cuda12x>13.0.0", "voltools"]
+cupy = ["cupy-cuda12x>13.0.0", "voltools"]
 pytorch = ["torch", "torchvision"]
 jax = ["jax[cuda12]", "jaxlib"]
 jax_cpu = ["jax", "jaxlib"]

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/scripts/match_template.py

@@ -359,8 +359,8 @@ def parse_args():
         "--invert-target-contrast",
         action="store_true",
         default=False,
-        help="Invert the target's contrast for cases where templates to-be-matched have "
-        "negative values, e.g. tomograms.",
+        help="Invert the target contrast. Useful for matching on tomograms if the "
+        "template has not been inverted.",
     )
     io_group.add_argument(
         "--scramble-phases",

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/scripts/postprocess.py

@@ -188,7 +188,7 @@ def parse_args():
     )
     additional_group.add_argument(
         "--n-false-positives",
-        type=int,
+        type=float,
         default=None,
         required=False,
         help="Number of accepted false-positives picks to determine minimum score.",
@@ -318,11 +318,7 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
         data = load_matching_output(foreground)
         scores, _, rotations, rotation_mapping, *_ = data
 
-        # We could normalize to unit sdev, but that might lead to unexpected
-        # results for flat background distributions
-        # scores -= scores.mean()
         indices = tuple(slice(0, x) for x in scores.shape)
-
         indices_update = scores > scores_out[indices]
         scores_out[indices][indices_update] = scores[indices_update]
 
@@ -369,9 +365,7 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
     scores_norm = np.full(out_shape_norm, fill_value=0, dtype=np.float32)
     for background in backgrounds:
         data_norm = load_matching_output(background)
-
-        scores = data_norm[0]
-        # scores -= scores.mean()
+        scores, _, rotations, rotation_mapping, *_ = data_norm
 
         indices = tuple(slice(0, x) for x in scores.shape)
         indices_update = scores > scores_norm[indices]
@@ -381,9 +375,10 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
     update = tuple(slice(0, int(x)) for x in np.minimum(out_shape, scores.shape))
     scores_out = np.full(out_shape, fill_value=0, dtype=np.float32)
     scores_out[update] = data[0][update] - scores_norm[update]
+    scores_out = np.fmax(scores_out, 0, out=scores_out)
+    scores_out[update] += scores_norm[update].mean()
 
     # scores_out[update] = np.divide(scores_out[update], 1 - scores_norm[update])
-    scores_out = np.fmax(scores_out, 0, out=scores_out)
     data[0] = scores_out
 
     fg, bg = simple_stats(data[0]), simple_stats(scores_norm)
@@ -485,8 +480,11 @@ def main():
     if orientations is None:
         translations, rotations, scores, details = [], [], [], []
 
-        # Data processed by normalize_input is guaranteed to have this shape
-        scores, offset, rotation_array, rotation_mapping, meta = data
+        var = None
+        # Data processed by normalize_input is guaranteed to have this shape)
+        scores, _, rotation_array, rotation_mapping, *_ = data
+        if len(data) == 6:
+            scores, _, rotation_array, rotation_mapping, var, *_ = data
 
         cropped_shape = np.subtract(
             scores.shape, np.multiply(args.min_boundary_distance, 2)
@@ -509,13 +507,16 @@ def main():
             )
             args.n_false_positives = max(args.n_false_positives, 1)
             n_correlations = np.size(scores[cropped_slice]) * len(rotation_mapping)
+            std = np.std(scores[cropped_slice])
+            if var is not None:
+                std = np.asarray(np.sqrt(var)).reshape(())
+
             minimum_score = np.multiply(
                 erfcinv(2 * args.n_false_positives / n_correlations),
-                np.sqrt(2) * np.std(scores[cropped_slice]),
+                np.sqrt(2) * std,
             )
-            print(f"Determined minimum score cutoff: {minimum_score}.")
-            minimum_score = max(minimum_score, 0)
-            args.min_score = minimum_score
+            print(f"Determined cutoff --min-score {minimum_score}.")
+            args.min_score = max(minimum_score, 0)
 
         args.batch_dims = None
         if hasattr(cli_args, "batch_dims"):

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/scripts/preprocessor_gui.py

@@ -474,6 +474,7 @@ def membrane_mask(
     **kwargs,
 ) -> NDArray:
     return create_mask(
+        center=(center_x, center_y, center_z),
         mask_type="membrane",
         shape=template.shape,
         radius=radius,

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/scripts/refine_matches.py

@@ -10,11 +10,9 @@ import subprocess
 from sys import exit
 from os import unlink
 from time import time
-from os.path import join
 from typing import Tuple, List, Dict
 
 import numpy as np
-from scipy import optimize
 from sklearn.metrics import roc_auc_score
 
 from tme import Orientations, Density
@@ -66,7 +64,6 @@ def parse_args():
     matching_group.add_argument(
         "-i",
         "--template",
-        dest="template",
         type=str,
         required=True,
         help="Path to a template in PDB/MMCIF or other supported formats (see target).",
@@ -102,7 +99,7 @@ def parse_args():
     )
     matching_group.add_argument(
         "-s",
-        dest="score",
+        "--score",
         type=str,
         default="batchFLCSphericalMask",
         choices=list(MATCHING_EXHAUSTIVE_REGISTER.keys()),
@@ -197,6 +194,7 @@ def create_matching_argdict(args) -> Dict:
         "-n": args.cores,
         "--ctf-file": args.ctf_file,
         "--invert-target-contrast": args.invert_target_contrast,
+        "--backend" : args.backend,
     }
     return arg_dict
 
@@ -252,7 +250,7 @@ class DeepMatcher:
         if args.lowpass_range:
             self.filter_parameters["--lowpass"] = 0
         if args.highpass_range:
-            self.filter_parameters["--highpass"] = 200
+            self.filter_parameters["--highpass"] = 0
 
         self.postprocess_args = create_postprocessing_argdict(args)
         self.log_file = f"{args.output_prefix}_optimization_log.txt"
@@ -309,14 +307,14 @@ class DeepMatcher:
 
         match_template = argdict_to_command(
             self.match_template_args,
-            executable="match_template.py",
+            executable="match_template",
         )
         run_command(match_template)
 
         # Assume we get a new peak for each input in the same order
         postprocess = argdict_to_command(
             self.postprocess_args,
-            executable="postprocess.py",
+            executable="postprocess",
         )
         run_command(postprocess)
 

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/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)

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/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]
 

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/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))

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/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),
         )
 
 

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/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,12 +112,56 @@ def _identity(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
     return arr
 
 
-@partial(
-    pmap,
-    in_axes=(0,) + (None,) * 6,
-    static_broadcasted_argnums=[6, 7, 8, 9],
-    axis_name="batch",
-)
+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,
@@ -125,19 +169,13 @@ def scan(
     rotations: BackendArray,
     template_filter: BackendArray,
     target_filter: BackendArray,
+    score_mask: BackendArray,
     fast_shape: Tuple[int],
     rotate_mask: bool,
-    analyzer_class: object,
-    analyzer_kwargs: Tuple[Tuple],
+    analyzer: object,
 ) -> 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)
 
@@ -159,6 +197,10 @@ 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):
         state, index = ret
         template_rot, template_mask_rot = be.rigid_transform(
@@ -185,6 +227,8 @@ def scan(
             n_observations=n_observations,
             eps=eps,
         )
+        scores = _score_mask_func(scores, score_mask)
+
         state = analyzer(state, scores, rotation_matrix, rotation_index=index)
         return (state, index + 1), None
 

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/tme/backends/cupy_backend.py

@@ -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:
@@ -139,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:

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/tme/backends/jax_backend.py

@@ -197,6 +197,13 @@ class JaxBackend(NumpyFFTWBackend):
     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,
@@ -211,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)
         )
@@ -270,8 +279,10 @@ class JaxBackend(NumpyFFTWBackend):
                 cur_args = analyzer_args.copy()
                 cur_args["offset"] = translation_offset
                 cur_args.update(callback_class_args)
+                analyzer_kwargs.append(cur_args)
 
-                analyzer_kwargs.append(self._dict_to_tuple(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)
@@ -298,7 +309,13 @@ class JaxBackend(NumpyFFTWBackend):
             create_filter, create_template_filter, create_target_filter = (False,) * 3
             base, targets = None, self._array_backend.stack(targets)
 
-            analyzer_kwargs = tuple(analyzer_kwargs)
+            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),
                 self.astype(matching_data.template, self._float_dtype),
@@ -306,17 +323,18 @@ class JaxBackend(NumpyFFTWBackend):
                 matching_data.rotations,
                 template_filter,
                 target_filter,
-                fast_shape,
-                rotate_mask,
-                callback_class,
-                analyzer_kwargs,
+                score_mask,
             )
 
+            ndim = targets.ndim - 1
             for index in range(targets.shape[0]):
-                kwargs = self._tuple_to_dict(analyzer_kwargs[index])
+                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
 
-                state = (states[0][index], states[1][index], rotation_mapping)
                 ret.append(analyzer.result(state, **kwargs))
         return ret
 

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/tme/backends/matching_backend.py

@@ -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.

{pytme-0.3.1 → pytme-0.3.1.dev20250731}/tme/backends/npfftw_backend.py

@@ -201,6 +201,9 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         sorted_indices = self.unravel_index(indices=sorted_indices, shape=arr.shape)
         return sorted_indices
 
+    def ssum(self, arr, *args, **kwargs):
+        return self.sum(self.square(arr), *args, **kwargs)
+
     def indices(self, *args, **kwargs) -> NDArray:
         return self._array_backend.indices(*args, **kwargs)