pytme 0.3b0__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1__cp311-cp311-macosx_15_0_arm64.whl

tme/matching_utils.py

@@ -12,40 +12,36 @@ from shutil import move
 from joblib import Parallel
 from tempfile import mkstemp
 from itertools import product
+from gzip import open as gzip_open
+from typing import Tuple, Dict, Callable
 from concurrent.futures import ThreadPoolExecutor
-from typing import Tuple, Dict, Callable, Optional
 
 import numpy as np
 from scipy.spatial import ConvexHull
-from scipy.ndimage import gaussian_filter
 
 from .backends import backend as be
 from .memory import estimate_memory_usage
 from .types import NDArray, BackendArray
 
 
-def noop(*args, **kwargs):
-    pass
-
-
-def identity(arr, *args):
+def identity(arr, *args, **kwargs):
     return arr
 
 
 def conditional_execute(
     func: Callable,
-    execute_operation: bool,
-    alt_func: Callable = noop,
+    execute_operation: bool = False,
+    alt_func: Callable = identity,
 ) -> Callable:
     """
-    Return the given function or a no-op function based on execute_operation.
+    Return the given function or alternative function based on execute_operation.
 
     Parameters
     ----------
     func : Callable
         Callable.
     alt_func : Callable
-        Callable to return if ``execute_operation`` is False, no-op by default.
+        Callable to return if ``execute_operation`` is False, identity by default.
     execute_operation : bool
         Whether to return ``func`` or a ``alt_func`` function.
 
@@ -175,6 +171,12 @@ def memmap_to_array(arr: NDArray) -> NDArray:
     return arr
 
 
+def is_gzipped(filename: str) -> bool:
+    """Check if a file is a gzip file by reading its magic number."""
+    with open(filename, "rb") as f:
+        return f.read(2) == b"\x1f\x8b"
+
+
 def write_pickle(data: object, filename: str) -> None:
     """
     Serialize and write data to a file invalidating the input data.
@@ -242,7 +244,11 @@ def load_pickle(filename: str) -> object:
         return ret
 
     items = []
-    with open(filename, "rb") as ifile:
+    func = open
+    if is_gzipped(filename):
+        func = gzip_open
+
+    with func(filename, "rb") as ifile:
         for data in _load_pickle(ifile):
             if isinstance(data, tuple):
                 if _is_pickle_memmap(data):
@@ -409,7 +415,7 @@ def compute_parallelization_schedule(
     return splits, core_assignment
 
 
-def _center_slice(current_shape: Tuple[int], new_shape: Tuple[int]) -> Tuple[slice]:
+def center_slice(current_shape: Tuple[int], new_shape: Tuple[int]) -> Tuple[slice]:
     """Extract the center slice of ``current_shape`` to retrieve ``new_shape``."""
     new_shape = tuple(int(x) for x in new_shape)
     current_shape = tuple(int(x) for x in current_shape)
@@ -419,60 +425,12 @@ def _center_slice(current_shape: Tuple[int], new_shape: Tuple[int]) -> Tuple[sli
     return box
 
 
-def centered(arr: BackendArray, new_shape: Tuple[int]) -> BackendArray:
-    """
-    Extract the centered portion of an array based on a new shape.
-
-    Parameters
-    ----------
-    arr : BackendArray
-        Input data.
-    new_shape : tuple of ints
-        Desired shape for the central portion.
-
-    Returns
-    -------
-    BackendArray
-        Central portion of the array with shape ``new_shape``.
-
-    References
-    ----------
-    .. [1] https://github.com/scipy/scipy/blob/v1.11.2/scipy/signal/_signaltools.py#L388
-    """
-    box = _center_slice(arr.shape, new_shape=new_shape)
-    return arr[box]
-
-
-def centered_mask(arr: BackendArray, new_shape: Tuple[int]) -> BackendArray:
-    """
-    Mask the centered portion of an array based on a new shape.
-
-    Parameters
-    ----------
-    arr : BackendArray
-        Input data.
-    new_shape : tuple of ints
-        Desired shape for the mask.
-
-    Returns
-    -------
-    BackendArray
-        Array with central portion unmasked and the rest set to 0.
-    """
-    box = _center_slice(arr.shape, new_shape=new_shape)
-    mask = np.zeros_like(arr)
-    mask[box] = 1
-    arr *= mask
-    return arr
-
-
 def apply_convolution_mode(
     arr: BackendArray,
     convolution_mode: str,
     s1: Tuple[int],
     s2: Tuple[int],
     convolution_shape: Tuple[int] = None,
-    mask_output: bool = False,
 ) -> BackendArray:
     """
     Applies convolution_mode to ``arr``.
@@ -497,9 +455,6 @@ def apply_convolution_mode(
         Tuple of integers corresponding to shape of convolution array 2.
     convolution_shape : tuple of ints, optional
         Size of the actually computed convolution. s1 + s2 - 1 by default.
-    mask_output : bool, optional
-        Whether to mask values outside of convolution_mode rather than
-        removing them. Defaults to False.
 
     Returns
     -------
@@ -514,14 +469,13 @@ def apply_convolution_mode(
     if convolution_mode not in ("full", "same", "valid"):
         raise ValueError("Supported convolution_mode are 'full', 'same' and 'valid'.")
 
-    func = centered_mask if mask_output else centered
     if convolution_mode == "full":
-        return arr
+        subset = ...
     elif convolution_mode == "same":
-        return func(arr, s1)
+        subset = center_slice(arr.shape, s1)
     elif convolution_mode == "valid":
-        valid_shape = [s1[i] - s2[i] + s2[i] % 2 for i in range(arr.ndim)]
-        return func(arr, valid_shape)
+        subset = center_slice(arr.shape, [x - y + 1 for x, y in zip(s1, s2)])
+    return arr[subset]
 
 
 def compute_full_convolution_index(
@@ -716,350 +670,8 @@ def minimum_enclosing_box(
     return shape
 
 
-def create_mask(mask_type: str, sigma_decay: float = 0, **kwargs) -> NDArray:
-    """
-    Creates a mask of the specified type.
-
-    Parameters
-    ----------
-    mask_type : str
-        Type of the mask to be created. Can be one of:
-
-            +---------+----------------------------------------------------------+
-            | box     | Box mask (see :py:meth:`box_mask`)                       |
-            +---------+----------------------------------------------------------+
-            | tube    | Cylindrical mask (see :py:meth:`tube_mask`)              |
-            +---------+----------------------------------------------------------+
-            | ellipse | Ellipsoidal mask (see :py:meth:`elliptical_mask`)        |
-            +---------+----------------------------------------------------------+
-    sigma_decay : float, optional
-        Smoothing along mask edges using a Gaussian filter, 0 by default.
-    kwargs : dict
-        Parameters passed to the indivdual mask creation funcitons.
-
-    Returns
-    -------
-    NDArray
-        The created mask.
-
-    Raises
-    ------
-    ValueError
-        If the mask_type is invalid.
-    """
-    mapping = {"ellipse": elliptical_mask, "box": box_mask, "tube": tube_mask}
-    if mask_type not in mapping:
-        raise ValueError(f"mask_type has to be one of {','.join(mapping.keys())}")
-
-    mask = mapping[mask_type](**kwargs)
-    if sigma_decay > 0:
-        mask_filter = gaussian_filter(mask.astype(np.float32), sigma=sigma_decay)
-        mask = np.add(mask, (1 - mask) * mask_filter)
-        mask[mask < np.exp(-np.square(sigma_decay))] = 0
-
-    return mask
-
-
-def elliptical_mask(
-    shape: Tuple[int],
-    radius: Tuple[float],
-    center: Optional[Tuple[float]] = None,
-    orientation: Optional[NDArray] = None,
-) -> 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)
-
-    mask = np.linalg.norm(indices / radius, axis=0)
-    mask = (mask <= 1).astype(int)
-
-    return mask
-
-
-def tube_mask2(
-    shape: Tuple[int],
-    inner_radius: float,
-    outer_radius: float,
-    height: int,
-    symmetry_axis: Optional[int] = 2,
-    center: Optional[Tuple[float]] = None,
-    orientation: Optional[NDArray] = None,
-    epsilon: float = 0.5,
-) -> NDArray:
-    """
-    Creates a tube mask.
-
-    Parameters
-    ----------
-    shape : tuple
-        Shape of the mask to be created.
-    inner_radius : float
-        Inner radius of the tube.
-    outer_radius : float
-        Outer radius of the tube.
-    height : int
-        Height of the tube.
-    symmetry_axis : int, optional
-        The axis of symmetry for the tube, defaults to 2.
-    center : tuple of float, optional.
-        Center of the mask, defaults to shape // 2.
-    orientation : NDArray, optional.
-        Orientation of the mask as rotation matrix with shape (d,d).
-    epsilon : float, optional
-        Tolerance to handle discretization errors, defaults to 0.5.
-
-    Returns
-    -------
-    NDArray
-        The created tube mask.
-
-    Raises
-    ------
-    ValueError
-        If ``inner_radius`` is larger than ``outer_radius``.
-        If ``center`` and ``shape`` do not have the same length.
-    """
-    shape = np.asarray(shape, dtype=int)
-
-    if center is None:
-        center = np.divide(shape, 2).astype(int)
-
-    center = np.asarray(center, dtype=np.float32)
-    center = np.repeat(center, shape.size // center.size)
-    if inner_radius > outer_radius:
-        raise ValueError("inner_radius should be smaller than outer_radius.")
-    if symmetry_axis > len(shape):
-        raise ValueError(f"symmetry_axis can be not larger than {len(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)
-    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)
-
-    mask = np.zeros(shape, dtype=bool)
-    sq_dist = np.zeros(shape)
-    for i in range(len(shape)):
-        if i == symmetry_axis:
-            continue
-        sq_dist += indices[i] ** 2
-
-    sym_coord = indices[symmetry_axis]
-    half_height = height / 2
-    height_mask = np.abs(sym_coord) <= half_height
-
-    inner_mask = 1
-    if inner_radius > epsilon:
-        inner_mask = sq_dist >= ((inner_radius) ** 2 - epsilon)
-
-    height_mask = np.abs(sym_coord) <= (half_height + epsilon)
-    outer_mask = sq_dist <= ((outer_radius) ** 2 + epsilon)
-
-    mask = height_mask & inner_mask & outer_mask
-    return mask
-
-
-def box_mask(shape: Tuple[int], center: Tuple[int], height: Tuple[int]) -> 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.
-    height : tuple of ints
-        Height (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(height):
-        raise ValueError("The length of shape, center, and height must be consistent.")
-
-    shape = tuple(int(x) for x in shape)
-    center, height = np.array(center, dtype=int), np.array(height, dtype=int)
-
-    half_heights = height // 2
-    starts = np.maximum(center - half_heights, 0)
-    stops = np.minimum(center + half_heights + np.mod(height, 2) + 1, shape)
-    slice_indices = tuple(slice(*coord) for coord in zip(starts, stops))
-
-    out = np.zeros(shape)
-    out[slice_indices] = 1
-    return out
-
-
-def tube_mask(
-    shape: Tuple[int],
-    symmetry_axis: int,
-    base_center: Tuple[int],
-    inner_radius: float,
-    outer_radius: float,
-    height: int,
-) -> 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(base_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(base_center) if ix != symmetry_axis)
-
-    inner_circle = np.zeros(circle_shape)
-    outer_circle = np.zeros_like(inner_circle)
-    if inner_radius > 0:
-        inner_circle = create_mask(
-            mask_type="ellipse",
-            shape=circle_shape,
-            radius=inner_radius,
-            center=circle_center,
-        )
-    if outer_radius > 0:
-        outer_circle = create_mask(
-            mask_type="ellipse",
-            shape=circle_shape,
-            radius=outer_radius,
-            center=circle_center,
-        )
-    circle = outer_circle - inner_circle
-    circle = np.expand_dims(circle, axis=symmetry_axis)
-
-    center = base_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 scramble_phases(
-    arr: NDArray,
-    noise_proportion: float = 0.5,
-    seed: int = 42,
-    normalize_power: bool = False,
+    arr: NDArray, noise_proportion: float = 1.0, seed: int = 42, **kwargs
 ) -> NDArray:
     """
     Perform random phase scrambling of ``arr``.
@@ -1069,7 +681,7 @@ def scramble_phases(
     arr : NDArray
         Input data.
     noise_proportion : float, optional
-        Proportion of scrambled phases, 0.5 by default.
+        Proportion of scrambled phases, 1.0 by default.
     seed : int, optional
         The seed for the random phase scrambling, 42 by default.
     normalize_power : bool, optional
@@ -1080,24 +692,22 @@ def scramble_phases(
     NDArray
         Phase scrambled version of ``arr``.
     """
+    from tme.filters._utils import fftfreqn
+
     np.random.seed(seed)
     noise_proportion = max(min(noise_proportion, 1), 0)
 
     arr_fft = np.fft.fftn(arr)
     amp, ph = np.abs(arr_fft), np.angle(arr_fft)
 
-    ph_noise = np.random.permutation(ph)
-    ph_new = ph * (1 - noise_proportion) + ph_noise * noise_proportion
-    ret = np.real(np.fft.ifftn(amp * np.exp(1j * ph_new)))
-
-    if normalize_power:
-        np.divide(ret - ret.min(), ret.max() - ret.min(), out=ret)
-        np.multiply(ret, np.subtract(arr.max(), arr.min()), out=ret)
-        np.add(ret, arr.min(), out=ret)
-        scaling = np.divide(np.abs(arr).sum(), np.abs(ret).sum())
-        np.multiply(ret, scaling, out=ret)
+    # Scrambling up to nyquist gives more uniform noise distribution
+    mask = np.fft.ifftshift(
+        fftfreqn(arr_fft.shape, sampling_rate=1, compute_euclidean_norm=True) <= 0.5
+    )
 
-    return ret
+    ph_noise = np.random.permutation(ph[mask])
+    ph[mask] = ph[mask] * (1 - noise_proportion) + ph_noise * noise_proportion
+    return np.real(np.fft.ifftn(amp * np.exp(1j * ph)))
 
 
 def compute_extraction_box(
@@ -1153,6 +763,54 @@ def compute_extraction_box(
     return obs_beg_clamp, obs_end_clamp, cand_beg, cand_end, keep
 
 
+def create_mask(mask_type: str, sigma_decay: float = 0, **kwargs) -> NDArray:
+    """
+    Creates a mask of the specified type.
+
+    Parameters
+    ----------
+    mask_type : str
+        Type of the mask to be created. Can be one of:
+
+            +----------+---------------------------------------------------------+
+            | box      | Box mask (see :py:meth:`box_mask`)                      |
+            +----------+---------------------------------------------------------+
+            | tube     | Cylindrical mask (see :py:meth:`tube_mask`)             |
+            +----------+---------------------------------------------------------+
+            | membrane | Cylindrical mask (see :py:meth:`membrane_mask`)         |
+            +----------+---------------------------------------------------------+
+            | ellipse  | Ellipsoidal mask (see :py:meth:`elliptical_mask`)       |
+            +----------+---------------------------------------------------------+
+    sigma_decay : float, optional
+        Smoothing along mask edges using a Gaussian filter, 0 by default.
+    kwargs : dict
+        Parameters passed to the indivdual mask creation funcitons.
+
+    Returns
+    -------
+    NDArray
+        The created mask.
+
+    Raises
+    ------
+    ValueError
+        If the mask_type is invalid.
+    """
+    from .mask import elliptical_mask, box_mask, tube_mask, membrane_mask
+
+    mapping = {
+        "ellipse": elliptical_mask,
+        "box": box_mask,
+        "tube": tube_mask,
+        "membrane": membrane_mask,
+    }
+    if mask_type not in mapping:
+        raise ValueError(f"mask_type has to be one of {','.join(mapping.keys())}")
+
+    mask = mapping[mask_type](**kwargs, sigma_decay=sigma_decay)
+    return mask
+
+
 class TqdmParallel(Parallel):
     """
     A minimal Parallel implementation using tqdm for progress reporting.

tme/memory.py

@@ -244,7 +244,7 @@ MATCHING_MEMORY_REGISTRY = {
     "PeakCallerMaximumFilter": PeakCallerMaximumFilterMemoryUsage,
     "cupy": CupyBackendMemoryUsage,
     "pytorch": CupyBackendMemoryUsage,
-    "batchFLCSpherical": FLCSphericalMaskMemoryUsage,
+    "batchFLCSphericalMask": FLCSphericalMaskMemoryUsage,
     "batchFLC": FLCMemoryUsage,
 }
 

tme/orientations.py

@@ -82,7 +82,7 @@ class Orientations:
         self.translations = np.array(self.translations).astype(np.float32)
         self.rotations = np.array(self.rotations).astype(np.float32)
         self.scores = np.array(self.scores).astype(np.float32)
-        self.details = np.array(self.details)
+        self.details = np.array(self.details).astype(np.float32)
         n_orientations = set(
             [
                 self.translations.shape[0],
@@ -471,10 +471,7 @@ class Orientations:
             rotation = np.zeros(translation.shape, dtype=np.float32)
 
         header_order = tuple(x for x in header if x in NAMES)
-        header_order = zip(header_order, range(len(header_order)))
-        sort_order = tuple(
-            x[1] for x in sorted(header_order, key=lambda x: x[0], reverse=False)
-        )
+        sort_order = tuple(NAMES.index(x) for x in header_order)
         translation = translation[..., sort_order]
 
         header_order = tuple(
@@ -492,9 +489,12 @@ class Orientations:
     def _from_star(
         cls, filename: str, delimiter: str = "\t"
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-        ret = StarParser(filename, delimiter=delimiter)
+        parser = StarParser(filename, delimiter=delimiter)
+
+        ret = parser.get("data_particles", None)
+        if ret is None:
+            ret = parser.get("data_", None)
 
-        ret = ret.get("data_particles", None)
         if ret is None:
             raise ValueError(f"No data_particles section found in {filename}.")
 
@@ -503,13 +503,21 @@ class Orientations:
         )
         translation = translation.astype(np.float32).T
 
+        default_angle = np.zeros(translation.shape[0], dtype=np.float32)
+        for x in ("_rlnAngleRot", "_rlnAngleTilt", "_rlnAnglePsi"):
+            if x not in ret:
+                ret[x] = default_angle
+
         rotation = np.vstack(
             (ret["_rlnAngleRot"], ret["_rlnAngleTilt"], ret["_rlnAnglePsi"])
         )
         rotation = rotation.astype(np.float32).T
 
         default = np.zeros(translation.shape[0])
-        return translation, rotation, default, default
+
+        details = ret.get("_rlnClassNumber", default)
+        scores = ret.get("_pytmeScore", default)
+        return translation, rotation, scores, details
 
     @staticmethod
     def _from_tbl(