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

tme/mask.py

@@ -0,0 +1,341 @@
+"""
+Utility functions for generating template matching masks.
+
+Copyright (c) 2023 European Molecular Biology Laboratory
+
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+import numpy as np
+from typing import Tuple, Optional
+
+from .types import NDArray
+from scipy.ndimage import gaussian_filter
+from .matching_utils import rigid_transform
+
+__all__ = ["elliptical_mask", "tube_mask", "box_mask", "membrane_mask"]
+
+
+def elliptical_mask(
+    shape: Tuple[int],
+    radius: Tuple[float],
+    center: Optional[Tuple[float]] = None,
+    orientation: Optional[NDArray] = None,
+    sigma_decay: float = 0.0,
+    cutoff_sigma: float = 3,
+    **kwargs,
+) -> NDArray:
+    """
+    Creates an ellipsoidal mask.
+
+    Parameters
+    ----------
+    shape : tuple of ints
+        Shape of the mask to be created.
+    radius : tuple of floats
+        Radius of the mask.
+    center : tuple of floats, optional
+        Center of the mask, default to shape // 2.
+    orientation : NDArray, optional.
+        Orientation of the mask as rotation matrix with shape (d,d).
+
+    Returns
+    -------
+    NDArray
+        The created ellipsoidal mask.
+
+    Raises
+    ------
+    ValueError
+        If the length of center and radius is not one or the same as shape.
+
+    Examples
+    --------
+    >>> from tme.matching_utils import elliptical_mask
+    >>> mask = elliptical_mask(shape=(20,20), radius=(5,5), center=(10,10))
+    """
+    shape, radius = np.asarray(shape), np.asarray(radius)
+
+    shape = shape.astype(int)
+    if center is None:
+        center = np.divide(shape, 2).astype(int)
+
+    center = np.asarray(center, dtype=np.float32)
+    radius = np.repeat(radius, shape.size // radius.size)
+    center = np.repeat(center, shape.size // center.size)
+    if radius.size != shape.size:
+        raise ValueError("Length of radius has to be either one or match shape.")
+    if center.size != shape.size:
+        raise ValueError("Length of center has to be either one or match shape.")
+
+    n = shape.size
+    center = center.reshape((-1,) + (1,) * n)
+    radius = radius.reshape((-1,) + (1,) * n)
+
+    indices = np.indices(shape, dtype=np.float32) - center
+    if orientation is not None:
+        return_shape = indices.shape
+        indices = indices.reshape(n, -1)
+        rigid_transform(
+            coordinates=indices,
+            rotation_matrix=np.asarray(orientation),
+            out=indices,
+            translation=np.zeros(n),
+            use_geometric_center=False,
+        )
+        indices = indices.reshape(*return_shape)
+
+    dist = np.linalg.norm(indices / radius, axis=0)
+    if sigma_decay > 0:
+        sigma_decay = 2 * (sigma_decay / np.mean(radius)) ** 2
+        mask = np.maximum(0, dist - 1)
+        mask = np.exp(-(mask**2) / sigma_decay)
+        mask *= mask > np.exp(-(cutoff_sigma**2) / 2)
+    else:
+        mask = (dist <= 1).astype(int)
+    return mask
+
+
+def box_mask(
+    shape: Tuple[int],
+    center: Tuple[int],
+    size: Tuple[int],
+    sigma_decay: float = 0.0,
+    cutoff_sigma: float = 3.0,
+    **kwargs,
+) -> np.ndarray:
+    """
+    Creates a box mask centered around the provided center point.
+
+    Parameters
+    ----------
+    shape : tuple of ints
+        Shape of the output array.
+    center : tuple of ints
+        Center point coordinates of the box.
+    size : tuple of ints
+        Side length of the box along each axis.
+
+    Returns
+    -------
+    NDArray
+        The created box mask.
+
+    Raises
+    ------
+    ValueError
+        If ``shape`` and ``center`` do not have the same length.
+        If ``center`` and ``height`` do not have the same length.
+    """
+    if len(shape) != len(center) or len(center) != len(size):
+        raise ValueError("The length of shape, center, and height must be consistent.")
+
+    shape = tuple(int(x) for x in shape)
+    center, size = np.array(center, dtype=int), np.array(size, dtype=int)
+
+    half_heights = size // 2
+    starts = np.maximum(center - half_heights, 0)
+    stops = np.minimum(center + half_heights + np.mod(size, 2) + 1, shape)
+    slice_indices = tuple(slice(*coord) for coord in zip(starts, stops))
+
+    out = np.zeros(shape)
+    out[slice_indices] = 1
+
+    if sigma_decay > 0:
+        mask_filter = gaussian_filter(
+            out.astype(np.float32), sigma=sigma_decay, truncate=cutoff_sigma
+        )
+        out = np.add(out, (1 - out) * mask_filter)
+        out *= out > np.exp(-(cutoff_sigma**2) / 2)
+    return out
+
+
+def tube_mask(
+    shape: Tuple[int],
+    symmetry_axis: int,
+    center: Tuple[int],
+    inner_radius: float,
+    outer_radius: float,
+    height: int,
+    sigma_decay: float = 0.0,
+    cutoff_sigma: float = 3.0,
+    **kwargs,
+) -> NDArray:
+    """
+    Creates a tube mask.
+
+    Parameters
+    ----------
+    shape : tuple
+        Shape of the mask to be created.
+    symmetry_axis : int
+        The axis of symmetry for the tube.
+    base_center : tuple
+        Center of the tube.
+    inner_radius : float
+        Inner radius of the tube.
+    outer_radius : float
+        Outer radius of the tube.
+    height : int
+        Height of the tube.
+
+    Returns
+    -------
+    NDArray
+        The created tube mask.
+
+    Raises
+    ------
+    ValueError
+        If ``inner_radius`` is larger than ``outer_radius``.
+        If ``height`` is larger than the symmetry axis.
+        If ``base_center`` and ``shape`` do not have the same length.
+    """
+    if inner_radius > outer_radius:
+        raise ValueError("inner_radius should be smaller than outer_radius.")
+
+    if height > shape[symmetry_axis]:
+        raise ValueError(f"Height can be no larger than {shape[symmetry_axis]}.")
+
+    if symmetry_axis > len(shape):
+        raise ValueError(f"symmetry_axis can be not larger than {len(shape)}.")
+
+    if len(center) != len(shape):
+        raise ValueError("shape and base_center need to have the same length.")
+
+    shape = tuple(int(x) for x in shape)
+    circle_shape = tuple(b for ix, b in enumerate(shape) if ix != symmetry_axis)
+    circle_center = tuple(b for ix, b in enumerate(center) if ix != symmetry_axis)
+
+    inner_circle = np.zeros(circle_shape)
+    outer_circle = np.zeros_like(inner_circle)
+    if inner_radius > 0:
+        inner_circle = elliptical_mask(
+            shape=circle_shape,
+            radius=inner_radius,
+            center=circle_center,
+            sigma_decay=sigma_decay,
+            cutoff_sigma=cutoff_sigma,
+        )
+    if outer_radius > 0:
+        outer_circle = elliptical_mask(
+            shape=circle_shape,
+            radius=outer_radius,
+            center=circle_center,
+            sigma_decay=sigma_decay,
+            cutoff_sigma=cutoff_sigma,
+        )
+    circle = outer_circle - inner_circle
+    circle = np.expand_dims(circle, axis=symmetry_axis)
+
+    center = center[symmetry_axis]
+    start_idx = int(center - height // 2)
+    stop_idx = int(center + height // 2 + height % 2)
+
+    start_idx, stop_idx = max(start_idx, 0), min(stop_idx, shape[symmetry_axis])
+
+    slice_indices = tuple(
+        slice(None) if i != symmetry_axis else slice(start_idx, stop_idx)
+        for i in range(len(shape))
+    )
+    tube = np.zeros(shape)
+    tube[slice_indices] = circle
+
+    return tube
+
+
+def membrane_mask(
+    shape: Tuple[int],
+    radius: float,
+    thickness: float,
+    separation: float,
+    symmetry_axis: int = 2,
+    center: Optional[Tuple[float]] = None,
+    sigma_decay: float = 0.5,
+    cutoff_sigma: float = 3,
+    **kwargs,
+) -> NDArray:
+    """
+    Creates a membrane mask consisting of two parallel disks with Gaussian intensity profile.
+    Uses efficient broadcasting approach: flat disk mask × height profile.
+
+    Parameters
+    ----------
+    shape : tuple of ints
+        Shape of the mask to be created.
+    radius : float
+        Radius of the membrane disks.
+    thickness : float
+        Thickness of each disk in the membrane.
+    separation : float
+        Distance between the centers of the two disks.
+    symmetry_axis : int, optional
+        The axis perpendicular to the membrane disks, defaults to 2.
+    center : tuple of floats, optional
+        Center of the membrane (midpoint between the two disks), defaults to shape // 2.
+    sigma_decay : float, optional
+        Controls edge sharpness relative to radius, defaults to 0.5.
+    cutoff_sigma : float, optional
+        Cutoff for height profile in standard deviations, defaults to 3.
+
+    Returns
+    -------
+    NDArray
+        The created membrane mask with Gaussian intensity profile.
+
+    Raises
+    ------
+    ValueError
+        If ``thickness`` is negative.
+        If ``separation`` is negative.
+        If ``center`` and ``shape`` do not have the same length.
+        If ``symmetry_axis`` is out of bounds.
+
+    Examples
+    --------
+    >>> from tme.matching_utils import membrane_mask
+    >>> mask = membrane_mask(shape=(50,50,50), radius=10, thickness=2, separation=15)
+    """
+    shape = np.asarray(shape, dtype=int)
+
+    if center is None:
+        center = np.divide(shape, 2).astype(float)
+
+    center = np.asarray(center, dtype=np.float32)
+    center = np.repeat(center, shape.size // center.size)
+
+    if thickness < 0:
+        raise ValueError("thickness must be non-negative.")
+    if separation < 0:
+        raise ValueError("separation must be non-negative.")
+    if symmetry_axis >= len(shape):
+        raise ValueError(f"symmetry_axis must be less than {len(shape)}.")
+    if center.size != shape.size:
+        raise ValueError("Length of center has to be either one or match shape.")
+
+    disk_mask = elliptical_mask(
+        shape=[x for i, x in enumerate(shape) if i != symmetry_axis],
+        radius=radius,
+        sigma_decay=sigma_decay,
+        cutoff_sigma=cutoff_sigma,
+    )
+
+    axial_coord = np.arange(shape[symmetry_axis]) - center[symmetry_axis]
+    height_profile = np.zeros((shape[symmetry_axis],), dtype=np.float32)
+    for leaflet_pos in [-separation / 2, separation / 2]:
+        leaflet_profile = np.exp(
+            -((axial_coord - leaflet_pos) ** 2) / (2 * (thickness / 3) ** 2)
+        )
+        cutoff_threshold = np.exp(-(cutoff_sigma**2) / 2)
+        leaflet_profile *= leaflet_profile > cutoff_threshold
+
+        height_profile = np.maximum(height_profile, leaflet_profile)
+
+    disk_mask = disk_mask.reshape(
+        [x if i != symmetry_axis else 1 for i, x in enumerate(shape)]
+    )
+    height_profile = height_profile.reshape(
+        [1 if i != symmetry_axis else x for i, x in enumerate(shape)]
+    )
+
+    return disk_mask * height_profile

tme/matching_data.py

@@ -128,11 +128,8 @@ class MatchingData:
         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)
 
-        # We are deviating from our typical right_pad + mod here
-        # because cropping from full convolution mode to target shape
-        # is defined from the perspective of the origin
-        right_pad = np.divide(padding, 2).astype(int)
-        left_pad = np.add(right_pad, np.mod(padding, 2))
+        left_pad = np.divide(padding, 2).astype(int)
+        right_pad = np.add(left_pad, np.mod(padding, 2))
 
         data_voxels_left = np.minimum(slice_start, left_pad)
         data_voxels_right = np.minimum(
@@ -253,7 +250,7 @@ class MatchingData:
         target_offset[mask] = [x.start for x in target_slice]
         mask = np.subtract(1, self._target_batch).astype(bool)
         template_offset = np.zeros(len(self._output_template_shape), dtype=int)
-        template_offset[mask] = [x.start for x, b in zip(template_slice, mask) if b]
+        template_offset[mask] = [x.start for x in template_slice]
 
         translation_offset = tuple(x for x in target_offset)
 
@@ -299,7 +296,7 @@ class MatchingData:
 
     def set_matching_dimension(self, target_dim: int = None, template_dim: int = None):
         """
-        Sets matching dimensions for target and template.
+        Sets matching batch dimensions for target and template.
 
         Parameters
         ----------
@@ -490,8 +487,8 @@ class MatchingData:
     def _fourier_padding(
         target_shape: Tuple[int],
         template_shape: Tuple[int],
-        pad_target: bool = False,
         batch_mask: Tuple[int] = None,
+        **kwargs,
     ) -> Tuple[Tuple, Tuple, Tuple, Tuple]:
         if batch_mask is None:
             batch_mask = np.zeros_like(template_shape)
@@ -503,7 +500,7 @@ class MatchingData:
 
         # Avoid padding batch dimensions
         pad_shape = np.maximum(target_shape, template_shape)
-        pad_shape = np.maximum(target_shape, np.multiply(1 - batch_mask, pad_shape))
+        pad_shape = np.maximum(pad_shape, np.multiply(1 - batch_mask, pad_shape))
         ret = be.compute_convolution_shapes(pad_shape, fourier_pad)
         conv_shape, fast_shape, fast_ft_shape = ret
 
@@ -525,21 +522,13 @@ class MatchingData:
             shape_shift = np.multiply(np.add(shape_shift, offset), shape_mask)
             fourier_shift = np.subtract(fourier_shift, shape_shift).astype(int)
 
-        if pad_target:
-            fourier_shift = np.subtract(fourier_shift, np.subtract(1, template_mod))
-
         fourier_shift = tuple(np.multiply(fourier_shift, 1 - batch_mask).astype(int))
         return tuple(conv_shape), tuple(fast_shape), tuple(fast_ft_shape), fourier_shift
 
-    def fourier_padding(self, pad_target: bool = False) -> Tuple:
+    def fourier_padding(self, **kwargs) -> Tuple:
         """
         Computes efficient shape four Fourier transforms and potential associated shifts.
 
-        Parameters
-        ----------
-        pad_target : bool, optional
-            Whether the target has been padded to the full convolution shape.
-
         Returns
         -------
         Tuple[tuple of int, tuple of int, tuple of int, tuple of int]
@@ -553,9 +542,32 @@ class MatchingData:
             target_shape=be.to_numpy_array(self._output_target_shape),
             template_shape=be.to_numpy_array(self._output_template_shape),
             batch_mask=be.to_numpy_array(self._batch_mask),
-            pad_target=pad_target,
         )
 
+    def _score_mask(self, fast_shape: Tuple[int], shift: Tuple[int]) -> BackendArray:
+        """
+        Create a boolean mask to exclude scores derived from padding in template matching.
+        """
+        padding = self.target_padding(True)
+        offset = tuple(x // 2 for x in padding)
+        shape = tuple(y - x for x, y in zip(padding, self.target.shape))
+
+        subset = []
+        for i in range(len(offset)):
+            if self._batch_mask[i]:
+                subset.append(slice(None))
+            else:
+                subset.append(slice(offset[i], offset[i] + shape[i]))
+
+        score_mask = np.zeros(fast_shape, dtype=bool)
+        score_mask[tuple(subset)] = 1
+        score_mask = np.roll(
+            score_mask,
+            shift=tuple(-x for x in shift),
+            axis=tuple(i for i in range(len(shift))),
+        )
+        return be.to_backend_array(score_mask)
+
     def computation_schedule(
         self,
         matching_method: str = "FLCSphericalMask",

tme/matching_exhaustive.py

@@ -40,7 +40,7 @@ def _setup_template_filter_apply_target_filter(
     matching_data: MatchingData,
     fast_shape: Tuple[int],
     fast_ft_shape: Tuple[int],
-    pad_template_filter: bool = True,
+    pad_template_filter: bool = False,
 ):
     target_filter = None
     backend_arr = type(be.zeros((1), dtype=be._float_dtype))
@@ -146,11 +146,10 @@ def scan(
     matching_data: MatchingData,
     matching_setup: Callable,
     matching_score: Callable,
-    n_jobs: int = 4,
-    callback_class: CallbackClass = None,
+    callback_class: CallbackClass,
     callback_class_args: Dict = {},
+    n_jobs: int = 4,
     pad_target: bool = True,
-    pad_template_filter: bool = True,
     interpolation_order: int = 3,
     jobs_per_callback_class: int = 8,
     shm_handler=None,
@@ -172,20 +171,22 @@ def scan(
         Function pointer to scoring function.
     n_jobs : int, optional
         Number of parallel jobs. Default is 4.
-    callback_class : type, optional
+    callback_class : type
         Analyzer class pointer to operate on computed scores.
     callback_class_args : dict, optional
         Arguments passed to the callback_class. Default is an empty dictionary.
     pad_target: bool, optional
         Whether to pad target to the full convolution shape.
-    pad_template_filter: bool, optional
-        Whether to pad potential template filters to the full convolution shape.
     interpolation_order : int, optional
         Order of spline interpolation for rotations.
     jobs_per_callback_class : int, optional
         Number of jobs a callback_class instance is shared between, 8 by default.
     shm_handler : type, optional
         Manager for shared memory objects, None by default.
+    target_slice : tuple of slice, optional
+        Target subset to process.
+    template_slice : tuple of slice, optional
+        Template subset to process.
 
     Returns
     -------
@@ -220,13 +221,17 @@ def scan(
     template_shape = matching_data._batch_shape(
         matching_data.template.shape, matching_data._template_batch
     )
-    conv, fwd, inv, shift = matching_data.fourier_padding(pad_target=pad_target)
+    conv, fwd, inv, shift = matching_data.fourier_padding()
+
+    score_mask = be.full(shape=(1,), fill_value=1, dtype=bool)
+    if pad_target:
+        score_mask = matching_data._score_mask(fwd, shift)
 
     template_filter = _setup_template_filter_apply_target_filter(
         matching_data=matching_data,
         fast_shape=fwd,
         fast_ft_shape=inv,
-        pad_template_filter=pad_template_filter,
+        pad_template_filter=False,
     )
 
     default_callback_args = {
@@ -240,11 +245,10 @@ def scan(
         "thread_safe": n_jobs > 1,
         "convolution_mode": "valid" if pad_target else "same",
         "shm_handler": shm_handler,
-        "only_unique_rotations": True,
         "aggregate_axis": matching_data._batch_axis(matching_data._batch_mask),
         "n_rotations": matching_data.rotations.shape[0],
+        "inversion_mapping": n_jobs == 1,
     }
-    callback_class_args["inversion_mapping"] = n_jobs == 1
     default_callback_args.update(callback_class_args)
 
     setup = matching_setup(
@@ -254,22 +258,14 @@ def scan(
         fast_ft_shape=inv,
         shm_handler=shm_handler,
     )
-    setup["interpolation_order"] = interpolation_order
-    setup["template_filter"] = be.to_sharedarr(template_filter, shm_handler)
 
     matching_data._free_data()
-    be.free_cache()
-
     n_callback_classes = max(n_jobs // jobs_per_callback_class, 1)
     callback_classes = [
-        (
-            SharedAnalyzerProxy(
-                callback_class,
-                default_callback_args,
-                shm_handler=shm_handler if n_jobs > 1 else None,
-            )
-            if callback_class
-            else None
+        SharedAnalyzerProxy(
+            callback_class,
+            default_callback_args,
+            shm_handler=shm_handler if n_jobs > 1 else None,
         )
         for _ in range(n_callback_classes)
     ]
@@ -277,35 +273,33 @@ def scan(
         delayed(_wrap_backend(matching_score))(
             backend_name=be._backend_name,
             backend_args=be._backend_args,
+            fast_shape=fwd,
+            fast_ft_shape=inv,
             rotations=rotation,
             callback=callback_classes[index % n_callback_classes],
+            interpolation_order=interpolation_order,
+            template_filter=be.to_sharedarr(template_filter, shm_handler),
+            score_mask=be.to_sharedarr(score_mask, shm_handler),
             **setup,
         )
         for index, rotation in enumerate(matching_data._split_rotations_on_jobs(n_jobs))
     )
-    callbacks = [
-        callback.result(**default_callback_args)
-        for callback in ret[:n_callback_classes]
-        if callback
-    ]
     be.free_cache()
 
-    if callback_class:
-        ret = callback_class.merge(callbacks, **default_callback_args)
-    return ret
+    callbacks = [x.result(**default_callback_args) for x in ret[:n_callback_classes]]
+    return callback_class.merge(callbacks, **default_callback_args)
 
 
 def scan_subsets(
     matching_data: MatchingData,
     matching_score: Callable,
     matching_setup: Callable,
-    callback_class: CallbackClass = None,
+    callback_class: CallbackClass,
     callback_class_args: Dict = {},
     job_schedule: Tuple[int] = (1, 1),
     target_splits: Dict = {},
     template_splits: Dict = {},
     pad_target_edges: bool = False,
-    pad_template_filter: bool = True,
     interpolation_order: int = 3,
     jobs_per_callback_class: int = 8,
     backend_name: str = None,
@@ -325,7 +319,7 @@ def scan_subsets(
         Function pointer to setup function.
     matching_score : type
         Function pointer to scoring function.
-    callback_class : type, optional
+    callback_class : type
         Analyzer class pointer to operate on computed scores.
     callback_class_args : dict, optional
         Arguments passed to the callback_class. Default is an empty dictionary.
@@ -341,8 +335,6 @@ def scan_subsets(
         See :py:meth:`tme.matching_utils.compute_parallelization_schedule`.
     pad_target_edges : bool, optional
         Pad the target boundaries to avoid edge effects.
-    pad_template_filter: bool, optional
-        Whether to pad potential template filters to the full convolution shape.
     interpolation_order : int, optional
         Order of spline interpolation for rotations.
     jobs_per_callback_class : int, optional
@@ -424,18 +416,22 @@ def scan_subsets(
         )
     splits = tuple(product(target_splits, template_splits))
 
+    kwargs = {
+        "matching_data": matching_data,
+        "callback_class": callback_class,
+        "callback_class_args": callback_class_args,
+    }
+
     outer_jobs, inner_jobs = job_schedule
     if be._backend_name == "jax":
         func = be.scan
 
         corr_scoring = MATCHING_EXHAUSTIVE_REGISTER.get("CORR", (None, None))[1]
         results = func(
-            matching_data=matching_data,
             splits=splits,
             n_jobs=outer_jobs,
             rotate_mask=matching_score != corr_scoring,
-            callback_class=callback_class,
-            callback_class_args=callback_class_args,
+            **kwargs,
         )
     else:
         results = Parallel(n_jobs=outer_jobs, verbose=verbose)(
@@ -443,26 +439,20 @@ def scan_subsets(
                 delayed(_wrap_backend(scan))(
                     backend_name=be._backend_name,
                     backend_args=be._backend_args,
-                    matching_data=matching_data,
                     matching_score=matching_score,
                     matching_setup=matching_setup,
                     n_jobs=inner_jobs,
-                    callback_class=callback_class,
-                    callback_class_args=callback_class_args,
                     interpolation_order=interpolation_order,
                     pad_target=pad_target_edges,
                     gpu_index=index % outer_jobs,
-                    pad_template_filter=pad_template_filter,
                     target_slice=target_split,
                     template_slice=template_split,
+                    **kwargs,
                 )
                 for index, (target_split, template_split) in enumerate(splits)
             ]
         )
-    matching_data._free_data()
-    if callback_class is not None:
-        return callback_class.merge(results, **callback_class_args)
-    return None
+    return callback_class.merge(results, **callback_class_args)
 
 
 def register_matching_exhaustive(

tme/matching_optimization.py

@@ -1104,7 +1104,8 @@ def create_score_object(score: str, **kwargs) -> object:
     Examples
     --------
     >>> from tme import Density
-    >>> from tme.matching_utils import create_mask, euler_to_rotationmatrix
+    >>> from tme.mask import create_mask
+    >>> from tme.matching_utils import euler_to_rotationmatrix
     >>> from tme.matching_optimization import CrossCorrelation, optimize_match
     >>> translation, rotation = (5, -2, 7), (5, -10, 2)
     >>> target = create_mask(