pytme 0.2.0b0__cp311-cp311-macosx_14_0_arm64.whl → 0.2.2__cp311-cp311-macosx_14_0_arm64.whl

tme/matching_optimization.py

@@ -10,21 +10,23 @@ from typing import Tuple, Dict
 from abc import ABC, abstractmethod
 
 import numpy as np
-from numpy.typing import NDArray
+from scipy.spatial import KDTree
+from scipy.ndimage import laplace, map_coordinates
 from scipy.optimize import (
-    differential_evolution,
-    LinearConstraint,
-    basinhopping,
     minimize,
+    basinhopping,
+    LinearConstraint,
+    differential_evolution,
 )
-from scipy.ndimage import laplace, map_coordinates
-from scipy.spatial import KDTree
 
-from .types import ArrayLike
-from .backends import backend
+from .backends import backend as be
+from .types import ArrayLike, NDArray
 from .matching_data import MatchingData
-from .matching_utils import rigid_transform, euler_to_rotationmatrix
-from .matching_exhaustive import normalize_under_mask
+from .matching_utils import (
+    rigid_transform,
+    euler_to_rotationmatrix,
+    normalize_template,
+)
 
 
 def _format_rigid_transform(x: Tuple[float]) -> Tuple[ArrayLike, ArrayLike]:
@@ -45,9 +47,9 @@ def _format_rigid_transform(x: Tuple[float]) -> Tuple[ArrayLike, ArrayLike]:
     split = len(x) // 2
     translation, angles = x[:split], x[split:]
 
-    translation = backend.to_backend_array(translation)
-    rotation_matrix = euler_to_rotationmatrix(backend.to_numpy_array(angles))
-    rotation_matrix = backend.to_backend_array(rotation_matrix)
+    translation = be.to_backend_array(translation)
+    rotation_matrix = euler_to_rotationmatrix(be.to_numpy_array(angles))
+    rotation_matrix = be.to_backend_array(rotation_matrix)
 
     return translation, rotation_matrix
 
@@ -91,6 +93,7 @@ class _MatchDensityToDensity(ABC):
         negate_score: bool = True,
         **kwargs: Dict,
     ):
+        self.eps = be.eps(target.dtype)
         self.rotate_mask = rotate_mask
         self.interpolation_order = interpolation_order
 
@@ -100,36 +103,29 @@ class _MatchDensityToDensity(ABC):
         if target_mask is not None:
             matching_data.target_mask = target_mask
 
-        target_pad = matching_data.target_padding(pad_target=pad_target_edges)
-        matching_data = matching_data.subset_by_slice(target_pad=target_pad)
+        self.target, self.target_mask = matching_data.target, matching_data.target_mask
 
-        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 = matching_data._template
+        self.template_rot = be.zeros(template.shape, be._float_dtype)
 
         self.template_mask, self.template_mask_rot = 1, 1
-        rotate_mask = False if matching_data.template_mask is None else rotate_mask
+        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.template_mask = matching_data._template_mask
+            self.template_mask_rot = be.topleft_pad(
+                matching_data._template_mask, self.template_mask.shape
             )
 
+        self.template_slices = tuple(slice(None) for _ in self.template.shape)
+        self.target_slices = tuple(slice(0, x) for x in self.template.shape)
+
         self.score_sign = -1 if negate_score else 1
 
         if hasattr(self, "_post_init"):
             self._post_init(**kwargs)
 
-    @staticmethod
-    def rigid_transform(
+    def rotate_array(
+        self,
         arr,
         rotation_matrix,
         translation,
@@ -137,28 +133,39 @@ class _MatchDensityToDensity(ABC):
         out=None,
         out_mask=None,
         order: int = 1,
-        use_geometric_center: bool = False,
+        **kwargs,
     ):
         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 not hasattr(self, "_previous_center"):
+            self._previous_center = arr.shape
+
+        if not hasattr(self, "grid") or not np.allclose(self._previous_center, center):
+            self.grid = np.indices(arr.shape, dtype=np.float32).reshape(arr.ndim, -1)
+            np.subtract(self.grid, center, out=self.grid)
+            self.grid_out = np.zeros_like(self.grid)
+            self._previous_center = center
+
+        np.matmul(rotation_matrix.T, self.grid, out=self.grid_out)
+        translation = np.add(translation[:, None], center)
+        np.add(self.grid_out, translation, out=self.grid_out)
 
         if out is None:
             out = np.zeros_like(arr)
 
-        map_coordinates(arr, grid, order=order, output=out.ravel())
+        map_coordinates(arr, self.grid_out, 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())
+            map_coordinates(
+                arr_mask, self.grid_out, order=order, output=out_mask.ravel()
+            )
 
         match return_type:
             case 0:
@@ -218,19 +225,48 @@ class _MatchDensityToDensity(ABC):
             The matching score obtained for the transformation.
         """
         translation, rotation_matrix = _format_rigid_transform(x)
+        self.template_rot.fill(0)
+
+        voxel_translation = be.astype(translation, be._int_dtype)
+        subvoxel_translation = be.subtract(translation, voxel_translation)
+
+        center = be.astype(be.divide(self.template.shape, 2), be._int_dtype)
+        right_pad = be.subtract(self.template.shape, center)
+
+        translated_center = be.add(voxel_translation, center)
+
+        target_starts = be.subtract(translated_center, center)
+        target_stops = be.add(translated_center, right_pad)
+
+        template_starts = be.subtract(be.maximum(target_starts, 0), target_starts)
+        template_stops = be.subtract(
+            target_stops, be.minimum(target_stops, self.target.shape)
+        )
+        template_stops = be.subtract(self.template.shape, template_stops)
+
+        target_starts = be.maximum(target_starts, 0)
+        target_stops = be.minimum(target_stops, self.target.shape)
+
+        cand_start, cand_stop = template_starts.astype(int), template_stops.astype(int)
+        obs_start, obs_stop = target_starts.astype(int), target_stops.astype(int)
+
+        self.template_slices = tuple(slice(s, e) for s, e in zip(cand_start, cand_stop))
+        self.target_slices = tuple(slice(s, e) for s, e in zip(obs_start, obs_stop))
+
         kw_dict = {
             "arr": self.template,
             "rotation_matrix": rotation_matrix,
-            "translation": translation,
+            "translation": subvoxel_translation,
             "out": self.template_rot,
-            "use_geometric_center": False,
             "order": self.interpolation_order,
+            "use_geometric_center": True,
         }
         if self.rotate_mask:
+            self.template_mask_rot.fill(0)
             kw_dict["arr_mask"] = self.template_mask
             kw_dict["out_mask"] = self.template_mask_rot
 
-        self.rigid_transform(**kw_dict)
+        self.rotate_array(**kw_dict)
 
         return self()
 
@@ -246,11 +282,11 @@ class _MatchCoordinatesToDensity(_MatchDensityToDensity):
     target : NDArray
         A d-dimensional target to match the template coordinate set to.
     template_coordinates : NDArray
-        Template coordinate array with shape [d x N].
+        Template coordinate array with shape (d,n).
     template_weights : NDArray
-        Template weight array with shape [N].
+        Template weight array with shape (n,).
     template_mask_coordinates : NDArray, optional
-        Template mask coordinates with shape [d x N].
+        Template mask coordinates with shape (d,n).
     target_mask : NDArray, optional
         A d-dimensional mask to be applied to the target.
     negate_score : bool, optional
@@ -269,6 +305,7 @@ class _MatchCoordinatesToDensity(_MatchDensityToDensity):
         negate_score: bool = True,
         **kwargs: Dict,
     ):
+        self.eps = be.eps(target.dtype)
         self.target_density = target
         self.target_mask_density = target_mask
 
@@ -277,6 +314,8 @@ class _MatchCoordinatesToDensity(_MatchDensityToDensity):
         self.template_coordinates_rotated = np.copy(self.template_coordinates).astype(
             np.float32
         )
+        if template_mask_coordinates is None:
+            template_mask_coordinates = template_coordinates.copy()
 
         self.template_mask_coordinates = template_mask_coordinates
         self.template_mask_coordinates_rotated = template_mask_coordinates
@@ -291,9 +330,9 @@ class _MatchCoordinatesToDensity(_MatchDensityToDensity):
         self.in_volume, self.in_volume_mask = self.map_coordinates_to_array(
             coordinates=self.template_coordinates_rotated,
             coordinates_mask=self.template_mask_coordinates_rotated,
-            array_origin=backend.zeros(target.ndim),
+            array_origin=be.zeros(target.ndim),
             array_shape=self.target_density.shape,
-            sampling_rate=backend.full(target.ndim, fill_value=1),
+            sampling_rate=be.full(target.ndim, fill_value=1),
         )
 
         if hasattr(self, "_post_init"):
@@ -329,9 +368,9 @@ class _MatchCoordinatesToDensity(_MatchDensityToDensity):
         self.in_volume, self.in_volume_mask = self.map_coordinates_to_array(
             coordinates=self.template_coordinates_rotated,
             coordinates_mask=self.template_mask_coordinates_rotated,
-            array_origin=backend.zeros(rotation_matrix.shape[0]),
+            array_origin=be.zeros(rotation_matrix.shape[0]),
             array_shape=self.target_density.shape,
-            sampling_rate=backend.full(rotation_matrix.shape[0], fill_value=1),
+            sampling_rate=be.full(rotation_matrix.shape[0], fill_value=1),
         )
 
         return self()
@@ -521,55 +560,65 @@ class _MatchCoordinatesToCoordinates(_MatchDensityToDensity):
 
 
 class FLC(_MatchDensityToDensity):
+    """
+    Computes a normalized cross-correlation score of a target f a template g
+    and a mask m:
+
+    .. math::
+
+        \\frac{CC(f, \\frac{g*m - \\overline{g*m}}{\\sigma_{g*m}})}
+        {N_m * \\sqrt{
+            \\frac{CC(f^2, m)}{N_m} - (\\frac{CC(f, m)}{N_m})^2}
+        }
+
+    Where:
+
+    .. math::
+
+        CC(f,g) = \\mathcal{F}^{-1}(\\mathcal{F}(f) \\cdot \\mathcal{F}(g)^*)
+
+    and Nm is the number of voxels within the template mask m.
+    """
+
     __doc__ += _MatchDensityToDensity.__doc__
 
     def _post_init(self, **kwargs: Dict):
         if self.target_mask is not None:
-            backend.multiply(self.target, self.target_mask, out=self.target)
+            be.multiply(self.target, self.target_mask, out=self.target)
 
-        self.target_square = backend.square(self.target)
+        self.target_square = be.square(self.target)
 
-        normalize_under_mask(
+        normalize_template(
             template=self.template,
             mask=self.template_mask,
-            mask_intensity=backend.sum(self.template_mask),
+            n_observations=be.sum(self.template_mask),
         )
 
-        self.template = backend.reverse(self.template)
-        self.template_mask = backend.reverse(self.template_mask)
-
     def __call__(self) -> float:
         """Returns the score of the current configuration."""
-        n_observations = backend.sum(self.template_mask_rot)
+        n_obs = be.sum(self.template_mask_rot)
 
-        normalize_under_mask(
+        normalize_template(
             template=self.template_rot,
             mask=self.template_mask_rot,
-            mask_intensity=n_observations,
+            n_observations=n_obs,
         )
-
-        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,
+        overlap = be.sum(
+            be.multiply(
+                self.template_rot[self.template_slices], self.target[self.target_slices]
             )
         )
 
-        denominator = backend.maximum(backend.subtract(ex2, e2x), 0.0)
-        denominator = backend.sqrt(denominator)
-        denominator = backend.multiply(denominator, n_observations)
+        mask_rot = self.template_mask_rot[self.template_slices]
+        exp_sq = be.sum(self.target_square[self.target_slices] * mask_rot) / n_obs
+        sq_exp = be.square(be.sum(self.target[self.target_slices] * mask_rot) / n_obs)
 
-        overlap = backend.sum(backend.multiply(self.template_rot, self.target))
+        denominator = be.maximum(be.subtract(exp_sq, sq_exp), 0.0)
+        denominator = be.sqrt(denominator)
+        if denominator < self.eps:
+            return 0
 
-        score = backend.divide(overlap, denominator) * self.score_sign
+        score = be.divide(overlap, denominator * n_obs) * self.score_sign
         return score
 
 
@@ -586,29 +635,12 @@ class CrossCorrelation(_MatchCoordinatesToDensity):
 
     def __call__(self) -> float:
         """Returns the score of the current configuration."""
-        try:
-            score = np.dot(
-                self.target_density[
-                    tuple(
-                        self.template_coordinates_rotated[:, self.in_volume].astype(int)
-                    )
-                ],
-                self.template_weights[self.in_volume],
-            )
-        except:
-            print(self.template_coordinates_rotated[:, self.in_volume].astype(int))
-            print(self.target_density.shape)
-            print(self.in_volume)
-            coordinates = self.template_coordinates_rotated[:, self.in_volume].astype(
-                int
-            )
-            in_volume = np.logical_and(
-                coordinates < np.array(self.target_density.shape)[:, None],
-                coordinates >= 0,
-            ).min(axis=0)
-            print(in_volume)
-
-            raise ValueError()
+        score = np.dot(
+            self.target_density[
+                tuple(self.template_coordinates_rotated[:, self.in_volume].astype(int))
+            ],
+            self.template_weights[self.in_volume],
+        )
         score /= self.denominator
         return score * self.score_sign
 
@@ -663,10 +695,24 @@ class NormalizedCrossCorrelation(CrossCorrelation):
 
     __doc__ += _MatchCoordinatesToDensity.__doc__
 
-    def _post_init(self, **kwargs):
-        target_norm = np.linalg.norm(self.target_density[self.target_density != 0])
-        template_norm = np.linalg.norm(self.template_weights)
-        self.denominator = np.fmax(target_norm * template_norm, np.finfo(float).eps)
+    def __call__(self) -> float:
+        n_observations = be.sum(self.in_volume_mask)
+        target_coordinates = be.astype(
+            self.template_mask_coordinates_rotated[:, self.in_volume_mask], int
+        )
+        target_weight = self.target_density[tuple(target_coordinates)]
+        ex2 = be.divide(be.sum(be.square(target_weight)), n_observations)
+        e2x = be.square(be.divide(be.sum(target_weight), n_observations))
+
+        denominator = be.maximum(be.subtract(ex2, e2x), 0.0)
+        denominator = be.sqrt(denominator)
+        denominator = be.multiply(denominator, n_observations)
+
+        if denominator <= self.eps:
+            return 0.0
+
+        self.denominator = denominator
+        return super().__call__()
 
 
 class NormalizedCrossCorrelationMean(NormalizedCrossCorrelation):
@@ -1037,7 +1083,7 @@ def register_matching_optimization(match_name: str, match_class: type):
 
 def create_score_object(score: str, **kwargs) -> object:
     """
-    Initialize score object with name ``score`` using `**kwargs``.
+    Initialize score object with name ``score`` using ``**kwargs``.
 
     Parameters
     ----------
@@ -1059,6 +1105,32 @@ def create_score_object(score: str, **kwargs) -> object:
     See Also
     --------
     :py:meth:`register_matching_optimization`
+
+    Examples
+    --------
+    >>> from tme import Density
+    >>> from tme.matching_utils import create_mask, euler_to_rotationmatrix
+    >>> from tme.matching_optimization import CrossCorrelation, optimize_match
+    >>> translation, rotation = (5, -2, 7), (5, -10, 2)
+    >>> target = create_mask(
+    >>>     mask_type="ellipse",
+    >>>     radius=(5,5,5),
+    >>>     shape=(51,51,51),
+    >>>     center=(25,25,25),
+    >>> ).astype(float)
+    >>> template = Density(data=target)
+    >>> template = template.rigid_transform(
+    >>>     translation=translation,
+    >>>     rotation_matrix=euler_to_rotationmatrix(rotation),
+    >>> )
+    >>> template_coordinates = template.to_pointcloud(0)
+    >>> template_weights = template.data[tuple(template_coordinates)]
+    >>> score_object = CrossCorrelation(
+    >>>     target=target,
+    >>>     template_coordinates=template_coordinates,
+    >>>     template_weights=template_weights,
+    >>>     negate_score=True # Multiply returned score with -1 for minimization
+    >>> )
     """
 
     score_object = MATCHING_OPTIMIZATION_REGISTER.get(score, None)
@@ -1078,10 +1150,11 @@ def optimize_match(
     bounds_translation: Tuple[Tuple[float]] = None,
     bounds_rotation: Tuple[Tuple[float]] = None,
     optimization_method: str = "basinhopping",
-    maxiter: int = 500,
+    maxiter: int = 50,
+    x0: Tuple[float] = None,
 ) -> Tuple[ArrayLike, ArrayLike, float]:
     """
-    Find the translation and rotation optimizing the score returned by `score_object`
+    Find the translation and rotation optimizing the score returned by ``score_object``
     with respect to provided bounds.
 
     Parameters
@@ -1098,36 +1171,68 @@ def optimize_match(
         Bounds on the evaluated zyx Euler angles. Has to be specified per dimension
         as tuple of (min, max). Default is None.
     optimization_method : str, optional
-        Optimizer that will be used, by default basinhopping. For further
+        Optimizer that will be used, basinhopping by default. For further
         information refer to :doc:`scipy:reference/optimize`.
 
-        +--------------------------+-----------------------------------------+
-        | 'differential_evolution' | Highest accuracy but long runtime.      |
-        |                          | Requires bounds on translation.         |
-        +--------------------------+-----------------------------------------+
-        | 'basinhopping'           | Decent accuracy, medium runtime.        |
-        +--------------------------+-----------------------------------------+
-        | 'minimize'               | If initial values are closed to optimum |
-        |                          | decent performance, short runtime.      |
-        +--------------------------+-----------------------------------------+
+        +------------------------+-------------------------------------------+
+        | differential_evolution | Highest accuracy but long runtime.        |
+        |                        | Requires bounds on translation.           |
+        +------------------------+-------------------------------------------+
+        | basinhopping           | Decent accuracy, medium runtime.          |
+        +------------------------+-------------------------------------------+
+        | minimize               | If initial values are closed to optimum   |
+        |                        | acceptable accuracy and short runtime     |
+        +------------------------+-------------------------------------------+
+
     maxiter : int, optional
-        The maximum number of iterations. Default is 500. Not considered for
-        `optimization_method` 'minimize'.
+        The maximum number of iterations, 50 by default.
+    x0 : tuple of floats, optional
+        Initial values for the optimizer, zero by default.
 
     Returns
     -------
     Tuple[ArrayLike, ArrayLike, float]
-        Translation and rotation matrix yielding final score.
+        Optimal translation, rotation matrix and corresponding score.
 
     Raises
     ------
     ValueError
-        If `optimization_method` is not supported.
+        If ``optimization_method`` is not supported.
 
     Notes
     -----
     This function currently only supports three-dimensional optimization and
-    `score_object` will be modified during this operation.
+    ``score_object`` will be modified during this operation.
+
+    Examples
+    --------
+    Having defined ``score_object``, for instance via :py:meth:`create_score_object`,
+    non-exhaustive template matching can be performed as follows
+
+    >>> translation_fit, rotation_fit, score = optimize_match(score_object)
+
+    `translation_fit` and `rotation_fit` correspond to the inverse of the applied
+    translation and rotation, so the following statements should hold within tolerance
+
+    >>> np.allclose(translation, -translation_fit, atol = 1) # True
+    >>> np.allclose(rotation, np.linalg.inv(rotation_fit), rtol = .1) # True
+
+    Bounds on translation and rotation can be defined as follows
+
+    >>> translation_fit, rotation_fit, score = optimize_match(
+    >>>     score_object=score_object,
+    >>>     bounds_translation=((-5,5),(-2,2),(0,0)),
+    >>>     bounds_rotation=((-10,10), (-5,5), (0,0)),
+    >>> )
+
+    The optimization scheme and the initial parameter estimates can also be adapted
+
+    >>> translation_fit, rotation_fit, score = optimize_match(
+    >>>     score_object=score_object,
+    >>>     optimization_method="minimize",
+    >>>     x0=(0,0,0,5,3,-5),
+    >>> )
+
     """
     ndim = 3
     _optimization_method = {
@@ -1164,10 +1269,12 @@ def optimize_match(
             np.eye(len(bounds)), np.min(bounds, axis=1), np.max(bounds, axis=1)
         )
 
-    initial_score = score_object()
+    x0 = np.zeros(2 * ndim) if x0 is None else x0
+
+    initial_score = score_object.score(x=x0)
     if optimization_method == "basinhopping":
         result = basinhopping(
-            x0=np.zeros(2 * ndim),
+            x0=x0,
             func=score_object.score,
             niter=maxiter,
             minimizer_kwargs={"method": "COBYLA", "constraints": linear_constraint},
@@ -1180,11 +1287,13 @@ def optimize_match(
             maxiter=maxiter,
         )
     elif optimization_method == "minimize":
+        print(maxiter)
         result = minimize(
-            x0=np.zeros(2 * ndim),
+            x0=x0,
             fun=score_object.score,
             bounds=bounds,
             constraints=linear_constraint,
+            options={"maxiter": maxiter},
         )
     print(f"Niter: {result.nit}, success : {result.success} ({result.message}).")
     print(f"Initial score: {initial_score} - Refined score: {result.fun}")
@@ -1194,7 +1303,3 @@ def optimize_match(
     translation, rotation = result.x[:ndim], result.x[ndim:]
     rotation_matrix = euler_to_rotationmatrix(rotation)
     return translation, rotation_matrix, result.fun
-
-
-class FitRefinement:
-    pass