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

tme/filters/ctf.py

@@ -36,7 +36,7 @@ class CTF(ComposableFilter):
 
     #: The shape of the to-be created mask.
     shape: Tuple[int] = None
-    #: The defocus value in x direction (in units of sampling rate).
+    #: The defocus in x direction (in units of sampling rate).
     defocus_x: Tuple[float] = None
     #: The tilt angles.
     angles: Tuple[float] = None
@@ -164,7 +164,7 @@ class CTF(ComposableFilter):
         shape : tuple of int
             The shape of the CTF.
         defocus_x : tuple of float
-            The defocus value in x direction.
+            The defocus in x direction (in units of sampling rate).
         angles : tuple of float
             The tilt angles.
         opening_axis : int, optional
@@ -178,7 +178,7 @@ class CTF(ComposableFilter):
         defocus_angle : tuple of float, optional
             The defocus angle in radians, defaults to 0.
         defocus_y : tuple of float, optional
-            The defocus value in y direction, defaults to None.
+            The defocus in x direction (in units of sampling rate).
         correct_defocus_gradient : bool, optional
             Whether to correct defocus gradient, defaults to False.
         sampling_rate : tuple of float, optional
@@ -219,14 +219,12 @@ class CTF(ComposableFilter):
                 corrected_tilt_axis -= 1
 
         for index, angle in enumerate(angles):
-            defocus_x, defocus_y = defoci_x[index], defoci_y[index]
-
             correction = correct_defocus_gradient and angle is not None
             chi = create_ctf(
                 angle=angle,
                 shape=ctf_shape,
-                defocus_x=defocus_x,
-                defocus_y=defocus_y,
+                defocus_x=defoci_x[index],
+                defocus_y=defoci_y[index],
                 sampling_rate=sampling_rate,
                 acceleration_voltage=acceleration_voltage[index],
                 correct_defocus_gradient=correction,
@@ -243,12 +241,10 @@ class CTF(ComposableFilter):
             stack[index] = chi
 
         # Avoid contrast inversion
-        np.negative(stack, out=stack)
+        stack = np.negative(stack, out=stack)
         if flip_phase:
-            np.abs(stack, out=stack)
-
-        stack = be.to_backend_array(np.squeeze(stack))
-        return stack
+            stack = np.abs(stack, out=stack)
+        return be.to_backend_array(np.squeeze(stack))
 
 
 class CTFReconstructed(CTF):
@@ -281,7 +277,7 @@ class CTFReconstructed(CTF):
         shape : tuple of int
             The shape of the CTF.
         defocus_x : tuple of float
-            The defocus value in x direction.
+            The defocus in x direction in units of sampling rate.
         opening_axis : int, optional
             The axis around which the wedge is opened, defaults to 2.
         amplitude_contrast : float, optional
@@ -291,7 +287,7 @@ class CTFReconstructed(CTF):
         defocus_angle : tuple of float, optional
             The defocus angle in radians, defaults to 0.
         defocus_y : tuple of float, optional
-            The defocus value in y direction, defaults to None.
+            The defocus in y direction in units of sampling rate.
         sampling_rate : tuple of float, optional
             The sampling rate, defaults to 1.
         acceleration_voltage : float, optional
@@ -321,18 +317,15 @@ class CTFReconstructed(CTF):
             defocus_angle=defocus_angle,
             amplitude_contrast=amplitude_contrast,
         )
-        stack = shift_fourier(data=stack, shape_is_real_fourier=False)
-
         # Avoid contrast inversion
         np.negative(stack, out=stack)
         if flip_phase:
             np.abs(stack, out=stack)
 
-        stack = be.to_backend_array(np.squeeze(stack))
+        stack = shift_fourier(data=stack, shape_is_real_fourier=False)
         if return_real_fourier:
             stack = crop_real_fourier(stack)
-
-        return stack
+        return be.to_backend_array(np.squeeze(stack))
 
 
 def _from_xml(filename: str) -> Dict:
@@ -436,6 +429,9 @@ def _from_ctffind(filename: str) -> Dict:
         output[key] = np.array(output[key])
 
     output["additional_phase_shift"] = np.degrees(output["additional_phase_shift"])
+    cs = output.get("spherical_aberration")
+    if cs is not None:
+        output["spherical_aberration"] = float(cs) * 1e7
     return output
 
 
@@ -566,7 +562,7 @@ def create_ctf(
     amplitude_contrast : float, optional
         Amplitude contrast of microscope, defaults to 0.07.
     spherical_aberration : float, optional
-        Spherical aberration of microscope in Angstrom.
+        Spherical aberration of microscope in units of sampling rate.
     angle : float, optional
         Assume the created CTF is a projection over opening_axis observed at angle.
     opening_axis : int, optional
@@ -590,10 +586,14 @@ def create_ctf(
     electron_wavelength = _compute_electron_wavelength(acceleration_voltage)
     electron_wavelength /= sampling_rate
     aberration = (spherical_aberration / sampling_rate) * electron_wavelength**2
+
+    defocus_x = defocus_x / sampling_rate if defocus_x is not None else None
+    defocus_y = defocus_y / sampling_rate if defocus_y is not None else None
     if correct_defocus_gradient or defocus_y is not None:
         if len(shape) < 2:
             raise ValueError(f"Length of shape needs to be at least 2, got {shape}")
 
+        # Axial distance from grid center in multiples of sampling rate
         sampling = tuple(float(x) for x in np.divide(sampling_rate, shape))
         grid = fftfreqn(
             shape=shape,
@@ -619,6 +619,7 @@ def create_ctf(
         defocus_sum = np.add(defocus_x, defocus_y)
         defocus_difference = np.subtract(defocus_x, defocus_y)
 
+        # Reusing grid, but in principle pure frequencies would suffice
         angular_grid = np.arctan2(grid[1], grid[0])
         defocus_difference = np.multiply(
             defocus_difference,
@@ -627,7 +628,7 @@ def create_ctf(
         defocus_x = np.add(defocus_sum, defocus_difference)
         defocus_x *= 0.5
 
-    frequency_grid = fftfreqn(shape, sampling_rate=True, compute_euclidean_norm=True)
+    frequency_grid = fftfreqn(shape, sampling_rate=1, compute_euclidean_norm=True)
     if angle is not None and opening_axis is not None and full_shape is not None:
         frequency_grid = frequency_grid_at_angle(
             shape=full_shape,

tme/filters/wedge.py

@@ -15,7 +15,7 @@ import numpy as np
 from ..types import NDArray
 from ..backends import backend as be
 from .compose import ComposableFilter
-from ..matching_utils import centered
+from ..matching_utils import center_slice
 from ..rotations import euler_to_rotationmatrix
 from ..parser import XMLParser, StarParser, MDOCParser
 from ._utils import (
@@ -207,11 +207,10 @@ class Wedge(ComposableFilter):
             )
             sigma = np.sqrt(self.weights[index] * 4 / (8 * np.pi**2))
             sigma = -2 * np.pi**2 * sigma**2
-            np.square(frequency_grid, out=frequency_grid)
-            np.multiply(sigma, frequency_grid, out=frequency_grid)
-            np.exp(frequency_grid, out=frequency_grid)
-            np.multiply(frequency_grid, np.cos(np.radians(angle)), out=frequency_grid)
-            wedges[index] = frequency_grid
+            frequency_grid = np.square(frequency_grid, out=frequency_grid)
+            frequency_grid = np.multiply(sigma, frequency_grid, out=frequency_grid)
+            frequency_grid = np.exp(frequency_grid, out=frequency_grid)
+            wedges[index] = np.multiply(frequency_grid, np.cos(np.radians(angle)))
 
         return wedges
 
@@ -490,7 +489,11 @@ class WedgeReconstructed:
             )
             wedge_volume += plane_rotated * weights[index]
 
-        wedge_volume = centered(wedge_volume, (shape[opening_axis], shape[tilt_axis]))
+        subset = center_slice(
+            wedge_volume.shape, (shape[opening_axis], shape[tilt_axis])
+        )
+        wedge_volume = wedge_volume[subset]
+
         np.fmin(wedge_volume, np.max(weights), wedge_volume)
 
         if opening_axis > tilt_axis:

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,7 +542,6 @@ 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 computation_schedule(