pytme 0.3.1.post2__cp311-cp311-macosx_15_0_arm64.whl → 0.3.2.dev0__cp311-cp311-macosx_15_0_arm64.whl

tme/analyzer/peaks.py

@@ -18,7 +18,6 @@ from .base import AbstractAnalyzer
 from ._utils import score_to_cart
 from ..backends import backend as be
 from ..types import BackendArray, NDArray
-from ..rotations import euler_to_rotationmatrix
 from ..matching_utils import split_shape, compute_extraction_box
 
 __all__ = [
@@ -182,6 +181,7 @@ class PeakCaller(AbstractAnalyzer):
         min_score: float = None,
         max_score: float = None,
         batch_dims: Tuple[int] = None,
+        projection_dims: Tuple[int] = None,
         shm_handler: object = None,
         **kwargs,
     ):
@@ -197,9 +197,13 @@ class PeakCaller(AbstractAnalyzer):
         self.min_distance = int(min_distance)
         self.min_boundary_distance = int(min_boundary_distance)
 
-        self.batch_dims = batch_dims
+        self.batch_dims = ()
         if batch_dims is not None:
-            self.batch_dims = tuple(int(x) for x in self.batch_dims)
+            self.batch_dims = tuple(int(x) for x in batch_dims)
+
+        self.projection_dims = ()
+        if projection_dims is not None:
+            self.projection_dims = tuple(int(x) for x in projection_dims)
 
         self.min_score, self.max_score = min_score, max_score
 
@@ -231,7 +235,7 @@ class PeakCaller(AbstractAnalyzer):
 
         rdim = len(self.shape)
         if self.batch_dims:
-            rdim -= len(self.batch_dims)
+            rdim = rdim - len(self.batch_dims) + len(self.projection_dims)
 
         rotations = be.full(
             (self.num_peaks, rdim, rdim), fill_value=0, dtype=be._float_dtype
@@ -388,6 +392,20 @@ class PeakCaller(AbstractAnalyzer):
 
         return state
 
+    def correct_background(self, state, mean, inv_std=1, **kwargs):
+        arr_type = type(be.zeros((1,), be._float_dtype))
+        translations, rotations, scores, details = state
+
+        if isinstance(mean, arr_type):
+            mean = mean[tuple(be.astype(translations.T, int))]
+        scores = be.subtract(scores, mean, out=scores)
+
+        if isinstance(inv_std, arr_type):
+            inv_std = inv_std[tuple(be.astype(translations.T, int))]
+        scores = be.multiply(scores, inv_std, out=scores)
+
+        return translations, rotations, scores, details
+
     @classmethod
     def merge(cls, results=List[Tuple], **kwargs) -> Tuple:
         """
@@ -778,6 +796,9 @@ class PeakCallerRecursiveMasking(PeakCaller):
             mask = be.to_backend_array(mask)
             mask_buffer = be.zeros(mask.shape, dtype=mask.dtype)
 
+        if min_score is None:
+            min_score = self.min_score
+
         if min_score is None:
             min_score = be.min(scores) - 1
 
@@ -849,15 +870,7 @@ class PeakCallerRecursiveMasking(PeakCaller):
         """
         if rotation_space is None or rotation_mapping is None:
             return rotation_matrix
-
-        rotation = rotation_mapping[rotation_space[tuple(peak)]]
-
-        # 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))
-            )
-        return rotation
+        return rotation_mapping[rotation_space[tuple(peak)]]
 
 
 class PeakCallerScipy(PeakCaller):

tme/analyzer/proxy.py

@@ -85,6 +85,16 @@ class StatelessSharedAnalyzerProxy:
             final_state = tuple(self._shared_to_object(x) for x in final_state)
         return self._analyzer.result(final_state, **kwargs)
 
+    def correct_background(self, state, *args, **kwargs):
+        if self._shared:
+            # Copy to not correct the internal score array across processes
+            backend_arr = type(be.zeros((1), dtype=be._float_dtype))
+            state = tuple(self._shared_to_object(x) for x in state)
+            state = tuple(
+                be.copy(x) if isinstance(x, backend_arr) else x for x in state
+            )
+        return self._analyzer.correct_background(state, *args, **kwargs)
+
     def merge(self, *args, **kwargs):
         return self._analyzer.merge(*args, **kwargs)
 
@@ -121,3 +131,7 @@ class SharedAnalyzerProxy(StatelessSharedAnalyzerProxy):
     def result(self, **kwargs):
         """Extract final result"""
         return super().result(self._state, **kwargs)
+
+    def correct_background(self, *args, **kwargs):
+        # We always assign to state as this operation can not be shared
+        self._state = super().correct_background(self._state, *args, **kwargs)

tme/backends/_jax_utils.py

@@ -10,11 +10,11 @@ from typing import Tuple
 from functools import partial
 
 import jax.numpy as jnp
-from jax import pmap, lax, vmap
+from jax import pmap, lax, jit
 
 from ..types import BackendArray
 from ..backends import backend as be
-from ..matching_utils import normalize_template as _normalize_template
+from ..matching_utils import standardize, to_padded
 
 
 __all__ = ["scan", "setup_scan"]
@@ -62,15 +62,14 @@ def _flcSphere_scoring(
     Computes :py:meth:`tme.matching_scores.corr_scoring`.
     """
     correlation = _correlate(template=template, ft_target=ft_target)
-    correlation = correlation.at[:].multiply(inv_denominator)
-    return correlation
+    return correlation.at[:].multiply(inv_denominator)
 
 
 def _reciprocal_target_std(
     ft_target: BackendArray,
     ft_target2: BackendArray,
     template_mask: BackendArray,
-    n_observations: float,
+    n_obs: float,
     eps: float,
 ) -> BackendArray:
     """
@@ -80,16 +79,16 @@ def _reciprocal_target_std(
     --------
     :py:meth:`tme.matching_scores.flc_scoring`.
     """
-    ft_shape = template_mask.shape
-    ft_template_mask = jnp.fft.rfftn(template_mask, s=ft_shape)
+    shape = template_mask.shape
+    ft_template_mask = jnp.fft.rfftn(template_mask, s=shape)
 
     # E(X^2)- E(X)^2
-    exp_sq = jnp.fft.irfftn(ft_target2 * ft_template_mask, s=ft_shape)
-    exp_sq = exp_sq.at[:].divide(n_observations)
+    exp_sq = jnp.fft.irfftn(ft_target2 * ft_template_mask, s=shape)
+    exp_sq = exp_sq.at[:].divide(n_obs)
 
     ft_template_mask = ft_template_mask.at[:].multiply(ft_target)
-    sq_exp = jnp.fft.irfftn(ft_template_mask, s=ft_shape)
-    sq_exp = sq_exp.at[:].divide(n_observations)
+    sq_exp = jnp.fft.irfftn(ft_template_mask, s=shape)
+    sq_exp = sq_exp.at[:].divide(n_obs)
     sq_exp = sq_exp.at[:].power(2)
 
     exp_sq = exp_sq.at[:].add(-sq_exp)
@@ -97,7 +96,7 @@ def _reciprocal_target_std(
     exp_sq = exp_sq.at[:].power(0.5)
 
     exp_sq = exp_sq.at[:].set(
-        jnp.where(exp_sq <= eps, 0, jnp.reciprocal(exp_sq * n_observations))
+        jnp.where(exp_sq <= eps, 0, jnp.reciprocal(exp_sq * n_obs))
     )
     return exp_sq
 
@@ -108,32 +107,21 @@ def _apply_fourier_filter(arr: BackendArray, arr_filter: BackendArray) -> Backen
     return arr.at[:].set(jnp.fft.irfftn(arr_ft, s=arr.shape))
 
 
-def _identity(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
-    return arr
-
-
-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):
+def setup_scan(analyzer_kwargs, analyzer, fast_shape, rotate_mask, match_projection):
     """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
+            analyzer=analyzer(**device_config),
+        )
+        for device_config in analyzer_kwargs
     ]
 
     @partial(
         pmap,
-        in_axes=(0,) + (None,) * 6,
+        in_axes=(0,) + (None,) * 7,
         axis_name="batch",
     )
     def scan_combined(
@@ -144,6 +132,7 @@ def setup_scan(analyzer_kwargs, callback_class, fast_shape, rotate_mask):
         template_filter,
         target_filter,
         score_mask,
+        background_template,
     ):
         return lax.switch(
             lax.axis_index("batch"),
@@ -155,10 +144,13 @@ def setup_scan(analyzer_kwargs, callback_class, fast_shape, rotate_mask):
             template_filter,
             target_filter,
             score_mask,
+            background_template,
         )
+
     return scan_combined
 
 
+@partial(jit, static_argnums=(8, 9, 10))
 def scan(
     target: BackendArray,
     template: BackendArray,
@@ -167,67 +159,98 @@ def scan(
     template_filter: BackendArray,
     target_filter: BackendArray,
     score_mask: BackendArray,
+    background_template: BackendArray,
     fast_shape: Tuple[int],
     rotate_mask: bool,
     analyzer: object,
-) -> Tuple[BackendArray, BackendArray]:
+) -> Tuple:
     eps = jnp.finfo(template.dtype).resolution
 
-    if hasattr(target_filter, "shape"):
+    if target_filter.shape != ():
         target = _apply_fourier_filter(target, target_filter)
 
     ft_target = jnp.fft.rfftn(target, s=fast_shape)
     ft_target2 = jnp.fft.rfftn(jnp.square(target), s=fast_shape)
-    inv_denominator, target, scoring_func = None, None, _flc_scoring
+    _n_obs, _inv_denominator, target = None, None, None
+
+    unpadded_slice = tuple(slice(0, x) for x in template.shape)
+    rot_buffer, mask_rot_buffer = jnp.zeros(fast_shape), jnp.zeros(fast_shape)
     if not rotate_mask:
-        n_observations = jnp.sum(template_mask)
-        inv_denominator = _reciprocal_target_std(
+        _n_obs = jnp.sum(template_mask)
+        _inv_denominator = _reciprocal_target_std(
             ft_target=ft_target,
             ft_target2=ft_target2,
-            template_mask=be.topleft_pad(template_mask, fast_shape),
+            template_mask=to_padded(mask_rot_buffer, template_mask, unpadded_slice),
             eps=eps,
-            n_observations=n_observations,
+            n_obs=_n_obs,
         )
-        ft_target2, scoring_func = None, _flcSphere_scoring
+        ft_target2 = None
 
-    _template_filter_func = _identity
-    if template_filter.shape != ():
-        _template_filter_func = _apply_fourier_filter
+    mask_scores = score_mask.shape != ()
+    filter_template = template_filter.shape != ()
+    bg_correction = background_template.shape != ()
+    bg_scores = jnp.zeros(fast_shape) if bg_correction else 0
 
-    _score_mask_func = _identity
-    if score_mask.shape != ():
-        _score_mask_func = _mask_scores
+    _template_mask_rot = template_mask
+    template_indices = be._index_grid(template.shape)
+    center = be.divide(be.to_backend_array(template.shape) - 1, 2)
 
     def _sample_transform(ret, rotation_matrix):
-        state, index = ret
-        template_rot, template_mask_rot = be.rigid_transform(
-            arr=template,
-            arr_mask=template_mask,
-            rotation_matrix=rotation_matrix,
-            order=1,  # thats all we get for now
+        matrix = be._build_transform_matrix(
+            rotation_matrix=rotation_matrix, center=center
         )
+        indices = be._transform_indices(template_indices, matrix)
+
+        template_rot = be._interpolate(template, indices, order=1)
+        n_obs, template_mask_rot = _n_obs, _template_mask_rot
+        if rotate_mask:
+            template_mask_rot = be._interpolate(template_mask, indices, order=1)
+            n_obs = jnp.sum(template_mask_rot)
+
+        if filter_template:
+            template_rot = _apply_fourier_filter(template_rot, template_filter)
+        template_rot = standardize(template_rot, template_mask_rot, n_obs)
+
+        rot_pad = to_padded(rot_buffer, template_rot, unpadded_slice)
+
+        inv_denominator = _inv_denominator
+        if rotate_mask:
+            mask_rot_pad = to_padded(mask_rot_buffer, template_mask_rot, unpadded_slice)
+            inv_denominator = _reciprocal_target_std(
+                ft_target=ft_target,
+                ft_target2=ft_target2,
+                template_mask=mask_rot_pad,
+                n_obs=n_obs,
+                eps=eps,
+            )
+
+        scores = _flcSphere_scoring(ft_target, rot_pad, inv_denominator)
+        if mask_scores:
+            scores = scores.at[:].multiply(score_mask)
+
+        state, bg_scores, index = ret
+        state = analyzer(state, scores, rotation_matrix, rotation_index=index)
 
-        n_observations = jnp.sum(template_mask_rot)
-        template_rot = _template_filter_func(template_rot, template_filter)
-        template_rot = _normalize_template(
-            template_rot, template_mask_rot, n_observations
-        )
-        rot_pad = be.topleft_pad(template_rot, fast_shape)
-        mask_rot_pad = be.topleft_pad(template_mask_rot, fast_shape)
+        if bg_correction:
+            template_rot = be._interpolate(background_template, indices, order=1)
+            if filter_template:
+                template_rot = _apply_fourier_filter(template_rot, template_filter)
+            template_rot = standardize(template_rot, template_mask_rot, n_obs)
 
-        scores = scoring_func(
-            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,
-        )
-        scores = _score_mask_func(scores, score_mask)
+            rot_pad = to_padded(rot_buffer, template_rot, unpadded_slice)
+            scores = _flcSphere_scoring(ft_target, rot_pad, inv_denominator)
+            bg_scores = jnp.maximum(bg_scores, scores)
 
-        state = analyzer(state, scores, rotation_matrix, rotation_index=index)
-        return (state, index + 1), None
+        return (state, bg_scores, index + 1), None
+
+    (state, bg_scores, _), _ = lax.scan(
+        _sample_transform, (analyzer.init_state(), bg_scores, 0), rotations
+    )
+
+    if bg_correction:
+        if mask_scores:
+            bg_scores = bg_scores.at[:].multiply(score_mask)
+        bg_scores = bg_scores.at[:].add(-be.mean(bg_scores))
+        state = analyzer.correct_background(state, bg_scores)
 
-    (state, _), _ = lax.scan(_sample_transform, (analyzer.init_state(), 0), rotations)
     return state

tme/backends/cupy_backend.py

@@ -33,7 +33,6 @@ class CupyBackend(NumpyFFTWBackend):
         import cupy as cp
         import cupyx.scipy.fft as cufft
         from cupyx.scipy.ndimage import affine_transform, maximum_filter
-        from ._cupy_utils import affine_transform_batch
 
         float_dtype = cp.float32 if float_dtype is None else float_dtype
         complex_dtype = cp.complex64 if complex_dtype is None else complex_dtype
@@ -51,7 +50,6 @@ class CupyBackend(NumpyFFTWBackend):
         self._cufft = cufft
         self.maximum_filter = maximum_filter
         self.affine_transform = affine_transform
-        self.affine_transform_batch = affine_transform_batch
 
         itype = f"int{self.datatype_bytes(int_dtype) * 8}"
         ftype = f"float{self.datatype_bytes(float_dtype) * 8}"
@@ -157,8 +155,8 @@ class CupyBackend(NumpyFFTWBackend):
 
         from voltools import StaticVolume
 
-        # Only keep template and potential corresponding mask in cache
-        if len(TEXTURE_CACHE) >= 2:
+        # Only keep template, mask and noise template in cache
+        if len(TEXTURE_CACHE) >= 3:
             TEXTURE_CACHE.clear()
 
         interpolation = "filt_bspline"
@@ -174,7 +172,7 @@ class CupyBackend(NumpyFFTWBackend):
 
         return TEXTURE_CACHE[key]
 
-    def _rigid_transform(
+    def _transform(
         self,
         data: CupyArray,
         matrix: CupyArray,
@@ -182,21 +180,10 @@ class CupyBackend(NumpyFFTWBackend):
         prefilter: bool,
         order: int,
         cache: bool = False,
-        batched: bool = False,
-    ) -> None:
+    ) -> CupyArray:
         out_slice = tuple(slice(0, stop) for stop in data.shape)
-        if batched:
-            self.affine_transform_batch(
-                input=data,
-                matrix=matrix,
-                mode="constant",
-                output=output[out_slice],
-                order=order,
-                prefilter=prefilter,
-            )
-            return None
 
-        if data.ndim == 3 and cache and self.texture_available and not batched:
+        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)
@@ -204,7 +191,7 @@ class CupyBackend(NumpyFFTWBackend):
             output[out_slice] = temp
             return None
 
-        self.affine_transform(
+        return self.affine_transform(
             input=data,
             matrix=matrix,
             mode="constant",