pytme 0.1.8__cp311-cp311-macosx_14_0_arm64.whl → 0.2.0__cp311-cp311-macosx_14_0_arm64.whl

tme/matching_constrained.py

@@ -0,0 +1,195 @@
+import numpy as np
+from typing import Tuple, Dict
+
+from scipy.ndimage import map_coordinates
+
+from tme.types import ArrayLike
+from tme.backends import backend
+from tme.matching_data import MatchingData
+from tme.matching_exhaustive import _normalize_under_mask
+
+
+class MatchDensityToDensity:
+    def __init__(
+        self,
+        matching_data: "MatchingData",
+        pad_target_edges: bool = False,
+        pad_fourier: bool = False,
+        rotate_mask: bool = True,
+        interpolation_order: int = 1,
+        negate_score: bool = False,
+    ):
+        self.rotate_mask = rotate_mask
+        self.interpolation_order = interpolation_order
+
+        target_pad = matching_data.target_padding(pad_target=pad_target_edges)
+        matching_data = matching_data.subset_by_slice(target_pad=target_pad)
+
+        fast_shape, fast_ft_shape, fourier_shift = matching_data.fourier_padding(
+            pad_fourier=pad_fourier
+        )
+
+        self.target = backend.topleft_pad(matching_data.target, fast_shape)
+        self.target_mask = matching_data.target_mask
+
+        self.template = matching_data.template
+        self.template_rot = backend.preallocate_array(
+            fast_shape, backend._default_dtype
+        )
+
+        self.template_mask, self.template_mask_rot = 1, 1
+        rotate_mask = False if matching_data.template_mask is None else rotate_mask
+        if matching_data.template_mask is not None:
+            self.template_mask = matching_data.template_mask
+            self.template_mask_rot = backend.topleft_pad(
+                matching_data.template_mask, fast_shape
+            )
+
+        self.score_sign = -1 if negate_score else 1
+
+    @staticmethod
+    def rigid_transform(
+        arr,
+        rotation_matrix,
+        translation,
+        arr_mask=None,
+        out=None,
+        out_mask=None,
+        order: int = 1,
+        use_geometric_center: bool = False,
+    ):
+        rotate_mask = arr_mask is not None
+        return_type = (out is None) + 2 * rotate_mask * (out_mask is None)
+        translation = np.zeros(arr.ndim) if translation is None else translation
+
+        center = np.floor(np.array(arr.shape) / 2)[:, None]
+        grid = np.indices(arr.shape, dtype=np.float32).reshape(arr.ndim, -1)
+        np.subtract(grid, center, out=grid)
+        np.matmul(rotation_matrix.T, grid, out=grid)
+        np.add(grid, center, out=grid)
+
+        if out is None:
+            out = np.zeros_like(arr)
+
+        map_coordinates(arr, grid, order=order, output=out.ravel())
+
+        if out_mask is None and arr_mask is not None:
+            out_mask = np.zeros_like(arr_mask)
+
+        if arr_mask is not None:
+            map_coordinates(arr_mask, grid, order=order, output=out_mask.ravel())
+
+        match return_type:
+            case 0:
+                return None
+            case 1:
+                return out
+            case 2:
+                return out_mask
+            case 3:
+                return out, out_mask
+
+    @staticmethod
+    def angles_to_rotationmatrix(angles: Tuple[float]) -> ArrayLike:
+        angles = backend.to_numpy_array(angles)
+        rotation_matrix = euler_to_rotationmatrix(angles)
+        return backend.to_backend_array(rotation_matrix)
+
+    def format_translation(self, translation: Tuple[float] = None) -> ArrayLike:
+        if translation is None:
+            return backend.zeros(self.template.ndim, backend._default_dtype)
+
+        return backend.to_backend_array(translation)
+
+    def score_translation(self, x: Tuple[float]) -> float:
+        translation = self.format_translation(x)
+        rotation_matrix = self.angles_to_rotationmatrix((0, 0, 0))
+
+        return self(translation=translation, rotation_matrix=rotation_matrix)
+
+    def score_angles(self, x: Tuple[float]) -> float:
+        translation = self.format_translation(None)
+        rotation_matrix = self.angles_to_rotationmatrix(x)
+
+        return self(translation=translation, rotation_matrix=rotation_matrix)
+
+    def score(self, x: Tuple[float]) -> float:
+        split = len(x) // 2
+        translation, angles = x[:split], x[split:]
+
+        translation = self.format_translation(translation)
+        rotation_matrix = self.angles_to_rotationmatrix(angles)
+
+        return self(translation=translation, rotation_matrix=rotation_matrix)
+
+
+class FLC(MatchDensityToDensity):
+    def __init__(self, **kwargs: Dict):
+        super().__init__(**kwargs)
+
+        if self.target_mask is not None:
+            backend.multiply(self.target, self.target_mask, out=self.target)
+
+        self.target_square = backend.square(self.target)
+
+        _normalize_under_mask(
+            template=self.template,
+            mask=self.template_mask,
+            mask_intensity=backend.sum(self.template_mask),
+        )
+
+        self.template = backend.reverse(self.template)
+        self.template_mask = backend.reverse(self.template_mask)
+
+    def __call__(self, translation: ArrayLike, rotation_matrix: ArrayLike) -> float:
+        if self.rotate_mask:
+            self.rigid_transform(
+                arr=self.template,
+                arr_mask=self.template_mask,
+                rotation_matrix=rotation_matrix,
+                translation=translation,
+                out=self.template_rot,
+                out_mask=self.template_mask_rot,
+                use_geometric_center=False,
+                order=self.interpolation_order,
+            )
+        else:
+            self.rigid_transform(
+                arr=self.template,
+                rotation_matrix=rotation_matrix,
+                translation=translation,
+                out=self.template_rot,
+                use_geometric_center=False,
+                order=self.interpolation_order,
+            )
+        n_observations = backend.sum(self.template_mask_rot)
+
+        _normalize_under_mask(
+            template=self.template_rot,
+            mask=self.template_mask_rot,
+            mask_intensity=n_observations,
+        )
+
+        ex2 = backend.sum(
+            backend.divide(
+                backend.sum(
+                    backend.multiply(self.target_square, self.template_mask_rot),
+                ),
+                n_observations,
+            )
+        )
+        e2x = backend.square(
+            backend.divide(
+                backend.sum(backend.multiply(self.target, self.template_mask_rot)),
+                n_observations,
+            )
+        )
+
+        denominator = backend.maximum(backend.subtract(ex2, e2x), 0.0)
+        denominator = backend.sqrt(denominator)
+        denominator = backend.multiply(denominator, n_observations)
+
+        overlap = backend.sum(backend.multiply(self.template_rot, self.target))
+
+        score = backend.divide(overlap, denominator) * self.score_sign
+        return score

tme/matching_data.py

@@ -47,6 +47,9 @@ class MatchingData:
         self.target_filter = {}
 
         self._invert_target = False
+        self._rotations = None
+
+        self._set_batch_dimension()
 
     @staticmethod
     def _shape_to_slice(shape: Tuple[int]):
@@ -149,13 +152,13 @@ class MatchingData:
         def _warn_on_mismatch(
             expectation: float, computation: float, name: str
         ) -> float:
-            expectation, computation = float(expectation), float(computation)
             if expectation is None:
                 expectation = computation
+            expectation, computation = float(expectation), float(computation)
 
             if abs(computation) > abs(expectation):
                 warnings.warn(
-                    f"Computed {name} value is more extreme than value specified in file"
+                    f"Computed {name} value is more extreme than value in file header"
                     f" (|{computation}| > |{expectation}|). This may lead to issues"
                     " with padding and contrast inversion."
                 )
@@ -175,9 +178,10 @@ class MatchingData:
             arr_min = _warn_on_mismatch(arr_min, arr.min(), "min")
             arr_max = _warn_on_mismatch(arr_max, arr.max(), "max")
 
-            np.subtract(-ret, arr_min, out=ret)
-            np.divide(ret, arr_max - arr_min, out=ret)
-
+            # Avoid in-place operation in case ret is not floating point
+            ret = (
+                -np.divide(np.subtract(ret, arr_min), np.subtract(arr_max, arr_min)) + 1
+            )
         return ret
 
     def subset_by_slice(
@@ -220,14 +224,16 @@ class MatchingData:
         if target_pad is None:
             target_pad = np.zeros(len(self._target.shape), dtype=int)
         if template_pad is None:
-            template_pad = np.zeros(len(self._target.shape), dtype=int)
-
-        indices = compute_full_convolution_index(
-            outer_shape=self._target.shape,
-            inner_shape=self._template.shape,
-            outer_split=target_slice,
-            inner_split=template_slice,
-        )
+            template_pad = np.zeros(len(self._template.shape), dtype=int)
+
+        indices = None
+        if len(self._target.shape) == len(self._template.shape):
+            indices = compute_full_convolution_index(
+                outer_shape=self._target.shape,
+                inner_shape=self._template.shape,
+                outer_split=target_slice,
+                inner_split=template_slice,
+            )
 
         target_subset = self.subset_array(
             arr=self._target,
@@ -243,13 +249,19 @@ class MatchingData:
         )
         ret = self.__class__(target=target_subset, template=template_subset)
 
-        ret._translation_offset = np.add(
-            [x.start for x in target_slice],
-            [x.start for x in template_slice],
-        )
-        ret.template_filter = self.template_filter
+        target_offset = np.zeros(len(self._output_target_shape), dtype=int)
+        target_offset[(target_offset.size - len(target_slice)) :] = [
+            x.start for x in target_slice
+        ]
+        template_offset = np.zeros(len(self._output_target_shape), dtype=int)
+        template_offset[(template_offset.size - len(template_slice)) :] = [
+            x.start for x in template_slice
+        ]
+        ret._translation_offset = np.add(target_offset, template_offset)
 
-        ret.rotations, ret.indices = self.rotations, indices
+        ret.template_filter = self.template_filter
+        ret.target_filter = self.target_filter
+        ret._rotations, ret.indices = self.rotations, indices
         ret._target_pad, ret._template_pad = target_pad, template_pad
         ret._invert_target = self._invert_target
 
@@ -279,10 +291,14 @@ class MatchingData:
         """
         Transfer the class instance's numpy arrays to the current backend.
         """
+        backend_arr = type(backend.zeros((1), dtype=backend._default_dtype))
         for attr_name, attr_value in vars(self).items():
             if isinstance(attr_value, np.ndarray):
                 converted_array = backend.to_backend_array(attr_value.copy())
                 setattr(self, attr_name, converted_array)
+            elif isinstance(attr_value, backend_arr):
+                converted_array = backend.to_backend_array(attr_value)
+                setattr(self, attr_name, converted_array)
 
         self._default_dtype = backend._default_dtype
         self._complex_dtype = backend._complex_dtype
@@ -424,13 +440,13 @@ class MatchingData:
             An array indicating the padding for each dimension of the target.
         """
         target_padding = backend.zeros(
-            len(self.target.shape), dtype=backend._default_dtype_int
+            len(self._output_target_shape), dtype=backend._default_dtype_int
         )
 
         if pad_target:
             backend.subtract(
-                self._template.shape,
-                backend.mod(self._template.shape, 2),
+                self._output_template_shape,
+                backend.mod(self._output_template_shape, 2),
                 out=target_padding,
             )
             if hasattr(self, "_is_target_batch"):
@@ -472,22 +488,34 @@ class MatchingData:
         fourier_shift = backend.zeros(len(fourier_pad))
 
         if not pad_fourier:
-            fourier_pad = backend.full(shape=len(fourier_pad), fill_value=1, dtype=int)
+            fourier_pad = backend.full(
+                shape=(len(fourier_pad),),
+                fill_value=1,
+                dtype=backend._default_dtype_int,
+            )
 
         fourier_pad = backend.to_backend_array(fourier_pad)
         if hasattr(self, "_batch_mask"):
             batch_mask = backend.to_backend_array(self._batch_mask)
             fourier_pad[batch_mask] = 1
 
-        ret = backend.compute_convolution_shapes(target_shape, fourier_pad)
+        pad_shape = backend.maximum(target_shape, template_shape)
+        ret = backend.compute_convolution_shapes(pad_shape, fourier_pad)
         convolution_shape, fast_shape, fast_ft_shape = ret
         if not pad_fourier:
             fourier_shift = 1 - backend.astype(backend.divide(template_shape, 2), int)
             fourier_shift -= backend.mod(template_shape, 2)
             shape_diff = backend.subtract(fast_shape, convolution_shape)
             shape_diff = backend.astype(backend.divide(shape_diff, 2), int)
+
+            if hasattr(self, "_batch_mask"):
+                batch_mask = backend.to_backend_array(self._batch_mask)
+                shape_diff[batch_mask] = 0
+
             backend.add(fourier_shift, shape_diff, out=fourier_shift)
 
+        fourier_shift = backend.astype(fourier_shift, backend._default_dtype_int)
+
         return fast_shape, fast_ft_shape, fourier_shift
 
     @property
@@ -520,15 +548,21 @@ class MatchingData:
     def target(self):
         """Returns the target NDArray."""
         if isinstance(self._target, Density):
-            return self._target.data
-        return self._target
+            target = self._target.data
+        else:
+            target = self._target
+        out_shape = backend.to_numpy_array(self._output_target_shape)
+        return target.reshape(tuple(int(x) for x in out_shape))
 
     @property
     def template(self):
         """Returns the reversed template NDArray."""
+        template = self._template
         if isinstance(self._template, Density):
-            return backend.reverse(self._template.data)
-        return backend.reverse(self._template)
+            template = self._template.data
+        template = backend.reverse(template)
+        out_shape = backend.to_numpy_array(self._output_template_shape)
+        return template.reshape(tuple(int(x) for x in out_shape))
 
     @template.setter
     def template(self, template: NDArray):
@@ -555,9 +589,15 @@ class MatchingData:
     @property
     def target_mask(self):
         """Returns the target mask NDArray."""
+        target_mask = self._target_mask
         if isinstance(self._target_mask, Density):
-            return self._target_mask.data
-        return self._target_mask
+            target_mask = self._target_mask.data
+
+        if target_mask is not None:
+            out_shape = backend.to_numpy_array(self._output_target_shape)
+            target_mask = target_mask.reshape(tuple(int(x) for x in out_shape))
+
+        return target_mask
 
     @target_mask.setter
     def target_mask(self, mask: NDArray):
@@ -584,9 +624,15 @@ class MatchingData:
         template : NDArray
             Array to set as the template.
         """
+        mask = self._template_mask
         if isinstance(self._template_mask, Density):
-            return backend.reverse(self._template_mask.data)
-        return backend.reverse(self._template_mask)
+            mask = self._template_mask.data
+
+        if mask is not None:
+            mask = backend.reverse(mask)
+            out_shape = backend.to_numpy_array(self._output_template_shape)
+            mask = mask.reshape(tuple(int(x) for x in out_shape))
+        return mask
 
     @template_mask.setter
     def template_mask(self, mask: NDArray):