pytme 0.3b0__cp311-cp311-macosx_15_0_arm64.whl → 0.3b0.post1__cp311-cp311-macosx_15_0_arm64.whl

tme/analyzer/base.py

@@ -38,16 +38,16 @@ class AbstractAnalyzer(ABC):
 
         Returns
         -------
-        state
-            Initial state tuple containing the analyzer's internal data
-            structures. The exact structure depends on the specific
-            implementation.
+        state : tuple
+            Initial state tuple of the analyzer instance. The exact structure
+            depends on the specific implementation.
 
         Notes
         -----
         This method creates the initial state that will be passed to
-        subsequent calls to __call__. The state should contain all
-        necessary data structures for accumulating analysis results.
+        :py:meth:`AbstractAnalyzer.__call__` and finally to
+        :py:meth:`AbstractAnalyzer.result`. The state should contain all necessary
+        data structures for accumulating analysis results.
         """
 
     @abstractmethod
@@ -57,49 +57,39 @@ class AbstractAnalyzer(ABC):
 
         Parameters
         ----------
-        state : object
-            Current analyzer state as returned by init_state() or
-            previous calls to __call__.
+        state : tuple
+            Current analyzer state as returned :py:meth:`AbstractAnalyzer.init_state`
+            or previous invocations of :py:meth:`AbstractAnalyzer.__call__`.
         scores : BackendArray
-            Array of scores computed for the current rotation.
+            Array of new scores with dimensionality d.
         rotation_matrix : BackendArray
-            Rotation matrix used to generate the scores.
+            Rotation matrix used to generate scores with shape (d,d).
         **kwargs : dict
-            Additional keyword arguments specific to the analyzer
-            implementation.
+            Keyword arguments used by specific implementations.
 
         Returns
         -------
-        state
-            Updated analyzer state with the new scoring data incorporated.
-
-        Notes
-        -----
-        This method should be pure functional - it should not modify
-        the input state but return a new state with the updates applied.
-        The exact signature may vary between implementations.
+        tuple
+            Updated analyzer state incorporating the new data.
         """
-        pass
 
     @abstractmethod
     def result(self, state: Tuple, **kwargs) -> Tuple:
         """
-        Finalize the analysis and produce the final result.
+        Finalize the analysis by performing potential post processing.
 
         Parameters
         ----------
         state : tuple
-            Final analyzer state containing all accumulated data.
+            Analyzer state containing accumulated data.
         **kwargs : dict
-            Additional keyword arguments for result processing,
-            such as postprocessing parameters.
+            Keyword arguments used by specific implementations.
 
         Returns
         -------
         result
-            Final analysis result. The exact format depends on the
-            analyzer implementation but typically includes processed
-            scores, rotation information, and metadata.
+            Final analysis result. The exact struccture depends on the
+            analyzer implementation.
 
         Notes
         -----
@@ -108,25 +98,24 @@ class AbstractAnalyzer(ABC):
         It may apply postprocessing operations like convolution mode
         correction or coordinate transformations.
         """
-        pass
 
     @classmethod
     @abstractmethod
     def merge(cls, results: List[Tuple], **kwargs) -> Tuple:
         """
-        Merge results from multiple analyzer instances.
+        Merge multiple analyzer results.
 
         Parameters
         ----------
-        results : list
-            List of result objects as returned by the result() method
-            from multiple analyzer instances.
+        results : list of tuple
+            List of tuple objects returned by :py:meth:`AbstractAnalyzer.result`
+            from different instances of the same analyzer class.
         **kwargs : dict
-            Additional keyword arguments for merge configuration.
+            Keyword arguments used by specific implementations.
 
         Returns
         -------
-        merged_result
+        tuple
             Single result object combining all input results.
 
         Notes

tme/analyzer/peaks.py

@@ -17,9 +17,9 @@ from skimage.registration._phase_cross_correlation import _upsampled_dft
 from .base import AbstractAnalyzer
 from ._utils import score_to_cart
 from ..backends import backend as be
-from ..matching_utils import split_shape
 from ..types import BackendArray, NDArray
 from ..rotations import euler_to_rotationmatrix
+from ..matching_utils import split_shape, compute_extraction_box
 
 __all__ = [
     "PeakCaller",
@@ -765,14 +765,14 @@ class PeakCallerRecursiveMasking(PeakCaller):
         values. If rotations and rotation_mapping is provided, the respective
         rotation will be applied to the mask, otherwise rotation_matrix is used.
         """
-        coordinates, masking_function = [], self._mask_scores_rotate
+        peaks = []
+        box = tuple(self.min_distance for _ in range(scores.ndim))
 
-        if mask is None:
-            masking_function = self._mask_scores_box
-            shape = tuple(self.min_distance for _ in range(scores.ndim))
-            mask = be.zeros(shape, dtype=be._float_dtype)
-
-        rotated_template = be.zeros(mask.shape, dtype=mask.dtype)
+        scores = be.to_backend_array(scores)
+        if mask is not None:
+            box = mask.shape
+            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:
@@ -780,39 +780,45 @@ class PeakCallerRecursiveMasking(PeakCaller):
         else:
             min_score = be.min(scores) - 1
 
-        scores_copy = be.zeros(scores.shape, dtype=scores.dtype)
-        scores_copy[:] = scores
-
+        _scores = be.zeros(scores.shape, dtype=scores.dtype)
+        _scores[:] = scores[:]
         while True:
-            be.argmax(scores_copy)
-            peak = be.unravel_index(
-                indices=be.argmax(scores_copy), shape=scores_copy.shape
-            )
-            if scores_copy[tuple(peak)] < min_score:
+            peak = be.unravel_index(indices=be.argmax(_scores), shape=_scores.shape)
+            if _scores[tuple(peak)] < min_score:
                 break
+            peaks.append(peak)
 
-            coordinates.append(peak)
-
-            current_rotation_matrix = self._get_rotation_matrix(
-                peak=peak,
-                rotation_space=rotations,
-                rotation_mapping=rotation_mapping,
-                rotation_matrix=rotation_matrix,
+            score_beg, score_end, tmpl_beg, tmpl_end, _ = compute_extraction_box(
+                centers=be.to_backend_array(peak)[None],
+                extraction_shape=box,
+                original_shape=scores.shape,
             )
-
-            masking_function(
-                scores=scores_copy,
-                rotation_matrix=current_rotation_matrix,
-                peak=peak,
-                mask=mask,
-                rotated_template=rotated_template,
+            score_slice = tuple(
+                slice(int(x), int(y)) for x, y in zip(score_beg[0], score_end[0])
+            )
+            tmpl_slice = tuple(
+                slice(int(x), int(y)) for x, y in zip(tmpl_beg[0], tmpl_end[0])
             )
 
-            if len(coordinates) >= peak_limit:
+            score_mask = 0
+            if mask is not None:
+                mask_buffer.fill(0)
+                rmat = self._get_rotation_matrix(
+                    peak=peak,
+                    rotation_space=rotations,
+                    rotation_mapping=rotation_mapping,
+                    rotation_matrix=rotation_matrix,
+                )
+                be.rigid_transform(
+                    arr=mask, rotation_matrix=rmat, order=1, out=mask_buffer
+                )
+                score_mask = mask_buffer[tmpl_slice] <= 0.1
+
+            _scores[score_slice] = be.multiply(_scores[score_slice], score_mask)
+            if len(peaks) >= peak_limit:
                 break
 
-        peaks = be.to_backend_array(coordinates)
-        return peaks, None
+        return be.to_backend_array(peaks), None
 
     @staticmethod
     def _get_rotation_matrix(
@@ -852,86 +858,6 @@ class PeakCallerRecursiveMasking(PeakCaller):
             )
         return rotation
 
-    @staticmethod
-    def _mask_scores_box(
-        scores: BackendArray, peak: BackendArray, mask: BackendArray, **kwargs: Dict
-    ) -> None:
-        """
-        Mask scores in a box around a peak.
-
-        Parameters
-        ----------
-        scores : BackendArray
-            Data array of scores.
-        peak : BackendArray
-            Peak coordinates.
-        mask : BackendArray
-            Mask array.
-        """
-        start = be.maximum(be.subtract(peak, mask.shape), 0)
-        stop = be.minimum(be.add(peak, mask.shape), scores.shape)
-        start, stop = be.astype(start, int), be.astype(stop, int)
-        coords = tuple(slice(*pos) for pos in zip(start, stop))
-        scores[coords] = 0
-        return None
-
-    @staticmethod
-    def _mask_scores_rotate(
-        scores: BackendArray,
-        peak: BackendArray,
-        mask: BackendArray,
-        rotated_template: BackendArray,
-        rotation_matrix: BackendArray,
-        **kwargs: Dict,
-    ) -> None:
-        """
-        Mask scores using mask rotation around a peak.
-
-        Parameters
-        ----------
-        scores : BackendArray
-            Data array of scores.
-        peak : BackendArray
-            Peak coordinates.
-        mask : BackendArray
-            Mask array.
-        rotated_template : BackendArray
-            Empty array to write mask rotations to.
-        rotation_matrix : BackendArray
-            Rotation matrix.
-        """
-        left_pad = be.divide(mask.shape, 2).astype(int)
-        right_pad = be.add(left_pad, be.mod(mask.shape, 2).astype(int))
-
-        score_start = be.subtract(peak, left_pad)
-        score_stop = be.add(peak, right_pad)
-
-        template_start = be.subtract(be.maximum(score_start, 0), score_start)
-        template_stop = be.subtract(score_stop, be.minimum(score_stop, scores.shape))
-        template_stop = be.subtract(mask.shape, template_stop)
-
-        score_start = be.maximum(score_start, 0)
-        score_stop = be.minimum(score_stop, scores.shape)
-        score_start = be.astype(score_start, int)
-        score_stop = be.astype(score_stop, int)
-
-        template_start = be.astype(template_start, int)
-        template_stop = be.astype(template_stop, int)
-        coords_score = tuple(slice(*pos) for pos in zip(score_start, score_stop))
-        coords_template = tuple(
-            slice(*pos) for pos in zip(template_start, template_stop)
-        )
-
-        rotated_template.fill(0)
-        be.rigid_transform(
-            arr=mask, rotation_matrix=rotation_matrix, order=1, out=rotated_template
-        )
-
-        scores[coords_score] = be.multiply(
-            scores[coords_score], (rotated_template[coords_template] <= 0.1)
-        )
-        return None
-
 
 class PeakCallerScipy(PeakCaller):
     """

tme/analyzer/proxy.py

@@ -18,6 +18,7 @@ from ..backends import backend as be
 
 __all__ = ["StatelessSharedAnalyzerProxy", "SharedAnalyzerProxy"]
 
+
 class StatelessSharedAnalyzerProxy:
     """
     Proxy that wraps functional analyzers for concurrent access via shared memory.

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)
@@ -163,12 +164,12 @@ 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,

tme/backends/cupy_backend.py

@@ -6,13 +6,10 @@ Copyright (c) 2023 European Molecular Biology Laboratory
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-import warnings
 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
 
@@ -146,15 +143,14 @@ class CupyBackend(NumpyFFTWBackend):
         def rfftn(
             arr: CupyArray, out: CupyArray = None, s=rfft_shape, axes=fwd_axes
         ) -> CupyArray:
-            return self.rfftn(arr, s=s, axes=fwd_axes)
+            return self.rfftn(arr, s=s, axes=fwd_axes, overwrite_x=True)
 
         def irfftn(
             arr: CupyArray, out: CupyArray = None, s=irfft_shape, axes=inv_axes
         ) -> CupyArray:
-            return self.irfftn(arr, s=s, axes=inv_axes)
+            return self.irfftn(arr, s=s, axes=inv_axes, overwrite_x=True)
 
         PLAN_CACHE[current_device] = [fwd_shape, inv_shape]
-
         return rfftn, irfftn
 
     def rfftn(self, arr: CupyArray, out: CupyArray = None, **kwargs) -> CupyArray:
@@ -239,13 +235,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:
+            # 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,

tme/backends/jax_backend.py

@@ -7,7 +7,7 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from functools import wraps
-from typing import Tuple, List, Callable
+from typing import Tuple, List, Callable, Dict
 
 from ..types import BackendArray
 from .npfftw_backend import NumpyFFTWBackend, shm_type
@@ -51,12 +51,6 @@ class JaxBackend(NumpyFFTWBackend):
         )
         self.scipy = jsp
         self._create_ufuncs()
-        try:
-            from ._jax_utils import scan as _
-
-            self.scan = self._scan
-        except Exception:
-            pass
 
     def from_sharedarr(self, arr: BackendArray) -> BackendArray:
         return arr
@@ -189,7 +183,18 @@ class JaxBackend(NumpyFFTWBackend):
         rotations = rotations.at[:].set(self.where(update, rotations, rotation_index))
         return max_scores, rotations
 
-    def _scan(
+    def compute_convolution_shapes(
+        self, arr1_shape: Tuple[int], arr2_shape: Tuple[int]
+    ) -> Tuple[List[int], List[int], List[int]]:
+        from scipy.fft import next_fast_len
+
+        convolution_shape = [int(x + y - 1) for x, y in zip(arr1_shape, arr2_shape)]
+        fast_shape = [next_fast_len(x, real=True) for x in convolution_shape]
+        fast_ft_shape = list(fast_shape[:-1]) + [fast_shape[-1] // 2 + 1]
+
+        return convolution_shape, fast_shape, fast_ft_shape
+
+    def scan(
         self,
         matching_data: type,
         splits: Tuple[Tuple[slice, slice]],
@@ -214,9 +219,9 @@ class JaxBackend(NumpyFFTWBackend):
         conv_shape, fast_shape, fast_ft_shape, shift = matching_data._fourier_padding(
             target_shape=self.to_numpy_array(target_shape),
             template_shape=self.to_numpy_array(matching_data._template.shape),
-            pad_fourier=False,
+            batch_mask=self.to_numpy_array(matching_data._batch_mask),
+            pad_target=pad_target,
         )
-
         analyzer_args = {
             "convolution_mode": convolution_mode,
             "fourier_shift": shift,
@@ -246,19 +251,18 @@ class JaxBackend(NumpyFFTWBackend):
 
             targets, translation_offsets = [], []
             for target_split, template_split in split_subset:
-                base = matching_data.subset_by_slice(
+                base, translation_offset = matching_data.subset_by_slice(
                     target_slice=target_split,
                     target_pad=target_pad,
                     template_slice=template_split,
                 )
-                translation_offsets.append(base._translation_offset)
+                translation_offsets.append(translation_offset)
                 targets.append(self.topleft_pad(base._target, fast_shape))
 
             if create_filter:
                 filter_args = {
                     "data_rfft": self.fft.rfftn(targets[0]),
                     "return_real_fourier": True,
-                    "shape_is_real_fourier": False,
                 }
 
             if create_template_filter:
@@ -288,12 +292,11 @@ class JaxBackend(NumpyFFTWBackend):
 
             for index in range(scores.shape[0]):
                 temp = callback_class(
-                    shape=scores.shape,
+                    shape=scores[index].shape,
                     offset=translation_offsets[index],
                 )
-                state = (scores, rotations, rotation_mapping)
+                state = (scores[index], rotations[index], rotation_mapping)
                 ret.append(temp.result(state, **analyzer_args))
-
         return ret
 
     def get_available_memory(self) -> int:

tme/backends/npfftw_backend.py

@@ -398,33 +398,33 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         out_mask: NDArray = None,
         order: int = 3,
         cache: bool = False,
+        batched: bool = False,
     ) -> Tuple[NDArray, NDArray]:
-        out = self.zeros_like(arr) if out is None else out
-        batched = arr.ndim != rotation_matrix.shape[0]
-
-        center = self.divide(self.to_backend_array(arr.shape) - 1, 2)
-        if not use_geometric_center:
-            center = self.center_of_mass(arr, cutoff=0)
-
-        offset = int(arr.ndim - rotation_matrix.shape[0])
-        center = center[offset:]
-        translation = self.zeros(center.size) if translation is None else translation
-        matrix = self._rigid_transform_matrix(
-            rotation_matrix=rotation_matrix,
-            translation=translation,
-            center=center,
-        )
-
-        subset = tuple(slice(None) for _ in range(arr.ndim))
-        if offset > 1:
-            subset = tuple(
-                0 if i < (offset - 1) else slice(None) for i in range(arr.ndim)
+        if out is None:
+            out = self.zeros_like(arr)
+
+        # Check whether rotation_matrix is already a rigid transform matrix
+        matrix = rotation_matrix
+        if matrix.shape[-1] == (arr.ndim - int(batched)):
+            center = self.divide(self.to_backend_array(arr.shape) - 1, 2)
+            if not use_geometric_center:
+                center = self.center_of_mass(arr, cutoff=0)
+
+            offset = int(arr.ndim - rotation_matrix.shape[0])
+            center = center[offset:]
+            translation = (
+                self.zeros(center.size) if translation is None else translation
+            )
+            matrix = self._rigid_transform_matrix(
+                rotation_matrix=rotation_matrix,
+                translation=translation,
+                center=center,
             )
 
         self._rigid_transform(
-            data=arr[subset],
+            data=arr,
             matrix=matrix,
-            output=out[subset],
+            output=out,
             order=order,
             prefilter=True,
             cache=cache,
@@ -433,11 +433,13 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
 
         # Applying the prefilter leads to artifacts in the mask.
         if arr_mask is not None:
-            out_mask = self.zeros_like(arr_mask) if out_mask is None else out_mask
+            if out_mask is None:
+                out_mask = self.zeros_like(arr_mask)
+
             self._rigid_transform(
-                data=arr_mask[subset],
+                data=arr_mask,
                 matrix=matrix,
-                output=out_mask[subset],
+                output=out_mask,
                 order=order,
                 prefilter=False,
                 cache=cache,

tme/backends/pytorch_backend.py

@@ -306,6 +306,9 @@ class PytorchBackend(NumpyFFTWBackend):
         kwargs["dim"] = kwargs.pop("axes", None)
         return self._array_backend.fft.irfftn(arr, **kwargs)
 
+    def _rigid_transform_matrix(self, rotation_matrix, *args, **kwargs):
+        return rotation_matrix
+
     def rigid_transform(
         self,
         arr: TorchTensor,
@@ -317,6 +320,7 @@ class PytorchBackend(NumpyFFTWBackend):
         out_mask: TorchTensor = None,
         order: int = 1,
         cache: bool = False,
+        **kwargs,
     ):
         _mode_mapping = {0: "nearest", 1: "bilinear", 3: "bicubic"}
         mode = _mode_mapping.get(order, None)

tme/density.py

@@ -1763,12 +1763,13 @@ class Density:
                 axis=axis,
             )
 
-            arr_ft = np.fft.fftn(self.data)
+            mask, mask_ret = np.where(mask), np.where(mask_ret)
+
+            arr_ft = np.fft.fftn(self.data)[mask]
             arr_ft *= np.prod(ret_shape) / np.prod(self.shape)
             ret_ft = np.zeros(ret_shape, dtype=arr_ft.dtype)
-            ret_ft[mask_ret] = arr_ft[mask]
-            ret.data = np.real(np.fft.ifftn(ret_ft))
-
+            np.add.at(ret_ft, mask_ret, arr_ft)
+            ret.data = np.real(np.fft.ifftn(ret_ft)).astype(self.data.dtype)
         ret.sampling_rate = new_sampling_rate
         return ret
 

tme/filters/__init__.py

@@ -1,6 +1,6 @@
 from .ctf import CTF, CTFReconstructed
 from .compose import Compose, ComposableFilter
-from .bandpass import BandPassFilter
+from .bandpass import BandPass, BandPassReconstructed
 from .whitening import LinearWhiteningFilter
 from .wedge import Wedge, WedgeReconstructed
-from .reconstruction import ReconstructFromTilt
+from .reconstruction import ReconstructFromTilt, ShiftFourier