pytme 0.1.5__cp311-cp311-macosx_14_0_arm64.whl

tme/matching_data.py

@@ -0,0 +1,377 @@
+""" Data class for holding template matching data.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Tuple, List
+
+import numpy as np
+from numpy.typing import NDArray
+
+from . import Density
+from .backends import backend
+from .matching_utils import compute_full_convolution_index
+
+
+class MatchingData:
+    """
+    Contains data required for template matching.
+
+    Parameters
+    ----------
+    target : np.ndarray or Density
+        Target data array for template matching.
+    template : np.ndarray or Density
+        Template data array for template matching.
+
+    """
+
+    def __init__(self, target: NDArray, template: NDArray):
+        self._default_dtype = np.float32
+        self._complex_dtype = np.complex64
+
+        self._target = target
+        self._target_mask = None
+        self._template_mask = None
+        self._translation_offset = np.zeros(len(target.shape), dtype=int)
+
+        self.template = template
+
+        self._target_pad = np.zeros(len(target.shape), dtype=int)
+        self._template_pad = np.zeros(len(template.shape), dtype=int)
+
+        self.template_filter = {}
+        self.target_filter = {}
+
+        self._invert_target = False
+
+    @staticmethod
+    def _shape_to_slice(shape: Tuple[int]):
+        return tuple(slice(0, dim) for dim in shape)
+
+    @classmethod
+    def _slice_to_mesh(cls, slice_variable: (slice,), shape: (int,)):
+        if slice_variable is None:
+            slice_variable = cls._shape_to_slice(shape)
+        ranges = [range(slc.start, slc.stop) for slc in slice_variable]
+        indices = np.meshgrid(*ranges, sparse=True, indexing="ij")
+        return indices
+
+    @staticmethod
+    def _load_array(arr: NDArray):
+        """
+        Load ``arr``,  If ``arr`` type is memmap, reload from disk.
+
+        Parameters
+        ----------
+        arr : NDArray
+            Array to load.
+
+        Returns
+        -------
+        NDArray
+            Loaded array.
+        """
+
+        if type(arr) == np.memmap:
+            return np.memmap(arr.filename, mode="r", shape=arr.shape, dtype=arr.dtype)
+        return arr
+
+    def subset_array(
+        self, arr: NDArray, arr_slice: Tuple[slice], padding: NDArray
+    ) -> NDArray:
+        """
+        Extract a subset of the input array according to the given slice and
+        apply padding.
+
+        Parameters
+        ----------
+        arr : NDArray
+            The input array from which a subset is extracted.
+        arr_slice : tuple of slice
+            Defines the region of the input array to be extracted.
+        padding : NDArray
+            Padding values for each dimension. If the padding exceeds the array
+            dimensions, the extra regions are filled with the mean of the array
+            values, otherwise, the
+            values in ``arr`` are used.
+
+        Returns
+        -------
+        NDArray
+            Subset of the input array with padding applied.
+        """
+        padding = np.maximum(padding, 0)
+
+        slice_start = np.array([x.start for x in arr_slice], dtype=int)
+        slice_stop = np.array([x.stop for x in arr_slice], dtype=int)
+        slice_shape = np.subtract(slice_stop, slice_start)
+
+        padding = np.add(padding, np.mod(padding, 2))
+        left_pad = right_pad = np.divide(padding, 2).astype(int)
+
+        data_voxels_left = np.minimum(slice_start, left_pad)
+        data_voxels_right = np.minimum(
+            np.subtract(arr.shape, slice_stop), right_pad
+        ).astype(int)
+
+        ret_shape = np.add(slice_shape, padding)
+        arr_start = np.subtract(slice_start, data_voxels_left)
+        arr_stop = np.add(slice_stop, data_voxels_right)
+        arr_slice = tuple(slice(*pos) for pos in zip(arr_start, arr_stop))
+        arr_mesh = self._slice_to_mesh(arr_slice, arr.shape)
+
+        subset_start = np.subtract(left_pad, data_voxels_left)
+        subset_stop = np.add(subset_start, np.subtract(arr_stop, arr_start))
+        subset_slice = tuple(slice(*prod) for prod in zip(subset_start, subset_stop))
+        subset_mesh = self._slice_to_mesh(subset_slice, ret_shape)
+
+        if type(arr) == Density:
+            if type(arr.data) == np.memmap:
+                arr = Density.from_file(arr.data.filename, subset=arr_slice).data
+            else:
+                arr = np.asarray(arr.data[*arr_mesh])
+        else:
+            if type(arr) == np.memmap:
+                arr = np.memmap(
+                    arr.filename, mode="r", shape=arr.shape, dtype=arr.dtype
+                )
+            arr = np.asarray(arr[*arr_mesh])
+        ret = np.full(
+            shape=np.add(slice_shape, padding), fill_value=arr.mean(), dtype=arr.dtype
+        )
+        ret[*subset_mesh] = arr
+
+        return ret
+
+    def subset_by_slice(
+        self,
+        target_slice: Tuple[slice] = None,
+        template_slice: Tuple[slice] = None,
+        target_pad: NDArray = None,
+        template_pad: NDArray = None,
+        invert_target: bool = False,
+    ) -> "MatchingData":
+        """
+        Slice the instance arrays based on the provided slices.
+
+        Parameters
+        ----------
+        target_slice : tuple of slice, optional
+            Slices for the target. If not provided, the full shape is used.
+        template_slice : tuple of slice, optional
+            Slices for the template. If not provided, the full shape is used.
+        target_pad : NDArray, optional
+            Padding for target. Defaults to zeros. If padding exceeds target,
+            pad with mean.
+        template_pad : NDArray, optional
+            Padding for template. Defaults to zeros. If padding exceeds template,
+            pad with mean.
+
+        Returns
+        -------
+        MatchingData
+            Newly allocated sliced class instance.
+        """
+        target_shape = self.target.shape
+        template_shape = self._template.shape
+
+        if target_slice is None:
+            target_slice = self._shape_to_slice(target_shape)
+        if template_slice is None:
+            template_slice = self._shape_to_slice(template_shape)
+
+        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,
+        )
+
+        target_subset = self.subset_array(
+            arr=self._target, arr_slice=target_slice, padding=target_pad
+        )
+        if self._invert_target:
+            target_subset *= -1
+            target_min, target_max = target_subset.min(), target_subset.max()
+            target_subset = (target_subset - target_min) / (target_max - target_min)
+        template_subset = self.subset_array(
+            arr=self._template,
+            arr_slice=template_slice,
+            padding=template_pad,
+        )
+        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
+
+        ret.rotations, ret.indices = self.rotations, indices
+        ret._target_pad, ret._template_pad = target_pad, template_pad
+        ret._invert_target = self._invert_target
+
+        if self._target_mask is not None:
+            ret.target_mask = self.subset_array(
+                arr=self._target_mask, arr_slice=target_slice, padding=target_pad
+            )
+        if self._template_mask is not None:
+            ret.template_mask = self.subset_array(
+                arr=self._template_mask,
+                arr_slice=template_slice,
+                padding=template_pad,
+            )
+
+        return ret
+
+    def to_backend(self) -> None:
+        """
+        Transfer the class instance's numpy arrays to the current backend.
+        """
+        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)
+
+        self._default_dtype = backend._default_dtype
+        self._complex_dtype = backend._complex_dtype
+
+    @property
+    def rotations(self):
+        """Return rotation matrices used for fitting."""
+        return self._rotations
+
+    @rotations.setter
+    def rotations(self, rotations: NDArray):
+        """
+        Set and reshape the rotation matrices for fitting.
+
+        Parameters
+        ----------
+        rotations : NDArray
+            Rotations in shape (3 x 3), (1 x 3 x 3), or (n x k x k).
+        """
+        if rotations.__class__ != np.ndarray:
+            raise ValueError("Rotation set has to be of type numpy ndarray.")
+        if rotations.ndim == 2:
+            print("Reshaping rotations array to rank 3.")
+            rotations = rotations.reshape(1, *rotations.shape)
+        elif rotations.ndim == 3:
+            pass
+        else:
+            raise ValueError("Rotations have to be a rank 2 or 3 array.")
+        self._rotations = rotations.astype(self._default_dtype)
+
+    @property
+    def target(self):
+        """Returns the target NDArray."""
+        if type(self._target) == Density:
+            return self._target.data
+        return self._target
+
+    @property
+    def template(self):
+        """Returns the reversed template NDArray."""
+        if type(self._template) == Density:
+            return backend.reverse(self._template.data)
+        return backend.reverse(self._template)
+
+    @template.setter
+    def template(self, template: NDArray):
+        """
+        Set the template array.
+
+        Parameters
+        ----------
+        template : NDArray
+            Array to set as the template.
+        """
+        if type(template) == Density:
+            template.data = template.data.astype(self._default_dtype, copy=False)
+            self._template = template
+            self._templateshape = self._template.shape[::-1]
+            return None
+        self._template = template.astype(self._default_dtype, copy=False)
+        self._templateshape = self._template.shape[::-1]
+
+    @property
+    def target_mask(self):
+        """Returns the target mask NDArray."""
+        if type(self._target_mask) == Density:
+            return self._target_mask.data
+        return self._target_mask
+
+    @target_mask.setter
+    def target_mask(self, mask: NDArray):
+        """Sets the target mask."""
+        if not np.all(self.target.shape == mask.shape):
+            raise ValueError("Target and its mask have to have the same shape.")
+
+        if type(mask) == Density:
+            mask.data = mask.data.astype(self._default_dtype, copy=False)
+            self._target_mask = mask
+            self._targetmaskshape = self._target_mask.shape[::-1]
+            return None
+        self._target_mask = mask.astype(self._default_dtype, copy=False)
+        self._targetmaskshape = self._target_mask.shape
+
+    @property
+    def template_mask(self):
+        """
+        Set the template mask array after reversing it.
+
+        Parameters
+        ----------
+        template : NDArray
+            Array to set as the template.
+        """
+        if type(self._template_mask) == Density:
+            return backend.reverse(self._template_mask.data)
+        return backend.reverse(self._template_mask)
+
+    @template_mask.setter
+    def template_mask(self, mask: NDArray):
+        """Returns the reversed template mask NDArray."""
+        if not np.all(self._template.shape == mask.shape):
+            raise ValueError("Target and its mask have to have the same shape.")
+
+        if type(mask) == Density:
+            mask.data = mask.data.astype(self._default_dtype, copy=False)
+            self._template_mask = mask
+            self._templatemaskshape = self._template_mask.shape[::-1]
+            return None
+
+        self._template_mask = mask.astype(self._default_dtype, copy=False)
+        self._templatemaskshape = self._template_mask.shape[::-1]
+
+    def _split_rotations_on_jobs(self, n_jobs: int) -> List[NDArray]:
+        """
+        Split the rotation matrices into parts based on the number of jobs.
+
+        Parameters
+        ----------
+        n_jobs : int
+            Number of jobs for splitting.
+
+        Returns
+        -------
+        list of NDArray
+            List of split rotation matrices.
+        """
+        nrot_per_job = self.rotations.shape[0] // n_jobs
+        rot_list = []
+        for n in range(n_jobs):
+            init_rot = n * nrot_per_job
+            end_rot = init_rot + nrot_per_job
+            if n == n_jobs - 1:
+                end_rot = None
+            rot_list.append(self.rotations[init_rot:end_rot])
+        return rot_list