pytme 0.2.9.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.0__cp311-cp311-macosx_15_0_arm64.whl

tme/matching_utils.py

@@ -1,8 +1,9 @@
-""" Utility functions for template matching.
+"""
+Utility functions for template matching.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import os
@@ -519,7 +520,7 @@ def apply_convolution_mode(
     elif convolution_mode == "same":
         return func(arr, s1)
     elif convolution_mode == "valid":
-        valid_shape = [s1[i] - s2[i] + s2[i] % 2 for i in range(arr.ndim)]
+        valid_shape = [s1[i] - s2[i] + 1 for i in range(arr.ndim)]
         return func(arr, valid_shape)
 
 
@@ -724,13 +725,15 @@ def create_mask(mask_type: str, sigma_decay: float = 0, **kwargs) -> NDArray:
     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`)        |
-            +---------+----------------------------------------------------------+
+            +----------+---------------------------------------------------------+
+            | 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
@@ -746,16 +749,16 @@ def create_mask(mask_type: str, sigma_decay: float = 0, **kwargs) -> NDArray:
     ValueError
         If the mask_type is invalid.
     """
-    mapping = {"ellipse": elliptical_mask, "box": box_mask, "tube": tube_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)
-    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
-
+    mask = mapping[mask_type](**kwargs, sigma_decay=sigma_decay)
     return mask
 
 
@@ -764,6 +767,8 @@ def elliptical_mask(
     radius: Tuple[float],
     center: Optional[Tuple[float]] = None,
     orientation: Optional[NDArray] = None,
+    sigma_decay: float = 0.0,
+    cutoff_sigma: float = 3,
 ) -> NDArray:
     """
     Creates an ellipsoidal mask.
@@ -825,9 +830,14 @@ def elliptical_mask(
         )
         indices = indices.reshape(*return_shape)
 
-    mask = np.linalg.norm(indices / radius, axis=0)
-    mask = (mask <= 1).astype(int)
-
+    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
 
 
@@ -925,7 +935,13 @@ def tube_mask2(
     return mask
 
 
-def box_mask(shape: Tuple[int], center: Tuple[int], height: Tuple[int]) -> np.ndarray:
+def box_mask(
+    shape: Tuple[int],
+    center: Tuple[int],
+    height: Tuple[int],
+    sigma_decay: float = 0.0,
+    cutoff_sigma: float = 0.0,
+) -> np.ndarray:
     """
     Creates a box mask centered around the provided center point.
 
@@ -962,6 +978,11 @@ def box_mask(shape: Tuple[int], center: Tuple[int], height: Tuple[int]) -> np.nd
 
     out = np.zeros(shape)
     out[slice_indices] = 1
+
+    if sigma_decay > 0:
+        mask_filter = gaussian_filter(out.astype(np.float32), sigma=sigma_decay)
+        out = np.add(out, (1 - out) * mask_filter)
+        out *= out > np.exp(-(cutoff_sigma**2) / 2)
     return out
 
 
@@ -972,6 +993,8 @@ def tube_mask(
     inner_radius: float,
     outer_radius: float,
     height: int,
+    sigma_decay: float = 0.0,
+    **kwargs,
 ) -> NDArray:
     """
     Creates a tube mask.
@@ -1027,6 +1050,7 @@ def tube_mask(
             shape=circle_shape,
             radius=inner_radius,
             center=circle_center,
+            sigma_decay=sigma_decay,
         )
     if outer_radius > 0:
         outer_circle = create_mask(
@@ -1034,6 +1058,7 @@ def tube_mask(
             shape=circle_shape,
             radius=outer_radius,
             center=circle_center,
+            sigma_decay=sigma_decay,
         )
     circle = outer_circle - inner_circle
     circle = np.expand_dims(circle, axis=symmetry_axis)
@@ -1054,9 +1079,106 @@ def tube_mask(
     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
+
+
 def scramble_phases(
     arr: NDArray,
-    noise_proportion: float = 0.5,
+    noise_proportion: float = 1.0,
     seed: int = 42,
     normalize_power: bool = False,
 ) -> NDArray:
@@ -1068,7 +1190,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
@@ -1079,15 +1201,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)))
+    # 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
+    )
+
+    ph_noise = np.random.permutation(ph[mask])
+    ph[mask] = ph[mask] * (1 - noise_proportion) + ph_noise * noise_proportion
+    ret = np.real(np.fft.ifftn(amp * np.exp(1j * ph)))
 
     if normalize_power:
         np.divide(ret - ret.min(), ret.max() - ret.min(), out=ret)
@@ -1150,3 +1279,40 @@ def compute_extraction_box(
     keep = be.multiply(keep, clamp_change == 0)
 
     return obs_beg_clamp, obs_end_clamp, cand_beg, cand_end, keep
+
+
+class TqdmParallel(Parallel):
+    """
+    A minimal Parallel implementation using tqdm for progress reporting.
+
+    Parameters:
+    -----------
+    tqdm_args : dict, optional
+        Dictionary of arguments passed to tqdm.tqdm
+    *args, **kwargs:
+        Arguments to pass to joblib.Parallel
+    """
+
+    def __init__(self, tqdm_args: Dict = {}, *args, **kwargs):
+        from tqdm import tqdm
+
+        super().__init__(*args, **kwargs)
+        self.pbar = tqdm(**tqdm_args)
+
+    def __call__(self, iterable, *args, **kwargs):
+        self.n_tasks = len(iterable) if hasattr(iterable, "__len__") else None
+        return super().__call__(iterable, *args, **kwargs)
+
+    def print_progress(self):
+        if self.n_tasks is None:
+            return super().print_progress()
+
+        if self.n_tasks != self.pbar.total:
+            self.pbar.total = self.n_tasks
+            self.pbar.refresh()
+
+        self.pbar.n = self.n_completed_tasks
+        self.pbar.refresh()
+
+        if self.n_completed_tasks >= self.n_tasks:
+            self.pbar.close()

tme/memory.py

@@ -1,8 +1,9 @@
-""" Compute memory consumption of template matching components.
+"""
+Compute memory consumption of template matching components.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from abc import ABC, abstractmethod

tme/orientations.py

@@ -1,10 +1,11 @@
-#!python3
-""" Handle template matching orientations and conversion between formats.
+"""
+Handle template matching orientations and conversion between formats.
 
-    Copyright (c) 2024 European Molecular Biology Laboratory
+Copyright (c) 2024 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
+
 from typing import List, Tuple
 from dataclasses import dataclass
 from string import ascii_lowercase, ascii_uppercase
@@ -14,7 +15,7 @@ import numpy as np
 from .parser import StarParser
 from .matching_utils import compute_extraction_box
 
-# Exceeds available numpy dimensions for default installations.
+# Exceeds available numpy dimensions for default installations
 NAMES = ["x", "y", "z", *ascii_lowercase[:-3], *ascii_uppercase]
 
 
@@ -81,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).astype(np.float32)
+        self.details = np.array(self.details)
         n_orientations = set(
             [
                 self.translations.shape[0],
@@ -324,6 +325,7 @@ class Orientations:
             "_rlnAngleRot",
             "_rlnAngleTilt",
             "_rlnAnglePsi",
+            "_rlnClassNumber",
         ]
         if source_path is not None:
             header.append("_rlnMicrographName")
@@ -339,6 +341,7 @@ class Orientations:
             for index, (translation, rotation, score, detail) in enumerate(self):
                 line = [str(x) for x in translation]
                 line.extend([str(x) for x in rotation])
+                line.extend([str(detail)])
 
                 if source_path is not None:
                     line.append(source_path)
@@ -489,9 +492,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}.")
 
@@ -500,13 +506,20 @@ 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
+
+        scores = ret.get("_pytmeScore", default)
+        return translation, rotation, scores, default
 
     @staticmethod
     def _from_tbl(

tme/parser.py

@@ -1,8 +1,9 @@
-""" Implements parsers for atomic structure file formats.
+"""
+Implements parsers for atomic structure file formats.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import re
@@ -15,7 +16,14 @@ from typing import List, Dict, Union
 
 import numpy as np
 
-__all__ = ["PDBParser", "MMCIFParser", "GROParser", "StarParser", "XMLParser"]
+__all__ = [
+    "PDBParser",
+    "MMCIFParser",
+    "GROParser",
+    "StarParser",
+    "XMLParser",
+    "MDOCParser",
+]
 
 
 class Parser(ABC):
@@ -84,6 +92,34 @@ class Parser(ABC):
         """
         return key in self._data
 
+    def __repr__(self) -> str:
+        """
+        String representation of the Parser showing available keys and their lengths.
+
+        Returns
+        -------
+        str
+            A formatted string showing each key and the length of its value.
+        """
+        if not self._data:
+            return f"{self.__class__.__name__}(empty)"
+
+        lines = [f"{self.__class__.__name__}:"]
+        try:
+            for key, value in sorted(self._data.items()):
+                if isinstance(value, (list, tuple)):
+                    lines.append(f"  {key}: length {len(value)}")
+                elif isinstance(value, dict):
+                    lines.append(f"  {key}: dict with {len(value)} keys")
+                elif isinstance(value, str):
+                    lines.append(f"  {key}: str")
+                else:
+                    lines.append(f"  {key}: {type(value).__name__}")
+        except Exception:
+            pass
+
+        return "\n".join(lines)
+
     def get(self, key, default=None):
         """
         Retrieve a value from the internal data using a given key. If the
@@ -247,19 +283,6 @@ class MMCIFParser(Parser):
     """
 
     def parse_input(self, lines: deque) -> Dict:
-        """
-        Parse a list of lines from an MMCIF file and convert the data into a dictionary.
-
-        Parameters
-        ----------
-        lines : deque of str
-            The lines of an MMCIF file to parse.
-
-        Returns
-        -------
-        dict
-            A dictionary containing the parsed data from the MMCIF file.
-        """
         lines = self._consolidate_strings(lines)
         blocks = self._split_in_blocks(lines)
         mmcif_dict = {}
@@ -448,21 +471,6 @@ class GROParser(Parser):
     """
 
     def parse_input(self, lines, **kwargs) -> Dict:
-        """
-        Parse a list of lines from a GRO file and convert the data into a dictionary.
-
-        Parameters
-        ----------
-        lines : deque of str
-            The lines of a GRO file to parse.
-        kwargs : Dict, optional
-            Optional keyword arguments.
-
-        Returns
-        -------
-        dict
-            A dictionary containing the parsed data from the GRO file.
-        """
         data = {
             "title": [],
             "num_atoms": [],
@@ -560,7 +568,7 @@ class StarParser(MMCIFParser):
     .. [1]  https://www.iucr.org/__data/assets/file/0013/11416/star.5.html
     """
 
-    def parse_input(self, lines: List[str], delimiter: str = "\t") -> Dict:
+    def parse_input(self, lines: List[str], delimiter: str = None) -> Dict:
         pattern = re.compile(r"\s*#.*")
 
         ret, category, block = {}, None, []
@@ -683,3 +691,76 @@ class XMLParser(Parser):
             pass
 
         return value_str
+
+
+class MDOCParser(Parser):
+    """
+    Convert MDOC file (SerialEM metadata) into a dictionary representation.
+
+    MDOC files contain global parameters and per-tilt metadata for cryo-ET
+    tilt series, with sections marked by [ZValue = N] for individual tilts.
+    """
+
+    def parse_input(self, lines: deque, **kwargs) -> Dict:
+        data = {}
+        global_params = {}
+        in_zvalue_section = False
+        zvalue_pattern = re.compile(r"\[ZValue\s*=\s*(\d+)\]")
+        section_pattern = re.compile(r"\[T\s*=\s*(.*?)\]")
+
+        if not lines:
+            return data
+
+        while lines:
+            line = lines.popleft().strip()
+
+            if not line:
+                continue
+
+            # Check for ZValue section header
+            zvalue_match = zvalue_pattern.match(line)
+            if zvalue_match:
+                in_zvalue_section = True
+
+                zvalue = int(zvalue_match.group(1))
+                if "ZValue" not in data:
+                    data["ZValue"] = []
+                data["ZValue"].append(zvalue)
+                continue
+
+            # Check for T section header (comments/metadata)
+            section_match = section_pattern.match(line)
+            if section_match:
+                section_content = section_match.group(1)
+                if "sections" not in global_params:
+                    global_params["sections"] = []
+                global_params["sections"].append(section_content)
+                continue
+
+            # Parse key-value pairs
+            if "=" in line:
+                try:
+                    key, value = line.split("=", 1)
+                    key = key.strip()
+                    value = value.strip()
+
+                    try:
+                        if "." not in value and "e" not in value.lower():
+                            parsed_value = int(value)
+                        else:
+                            parsed_value = float(value)
+                    except ValueError:
+                        parsed_value = value
+
+                    if not in_zvalue_section:
+                        global_params[key] = parsed_value
+                    else:
+                        if key not in data:
+                            data[key] = []
+                        data[key].append(parsed_value)
+
+                except ValueError:
+                    continue
+
+        data.update(global_params)
+        return data

tme/preprocessor.py

@@ -1,8 +1,9 @@
-""" Implements Preprocessor class for filtering operations.
+"""
+Implements Preprocessor class for filtering operations.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import os
@@ -669,9 +670,9 @@ class Preprocessor:
         NDArray
             Bandpass filtered.
         """
-        from .filters import BandPassFilter
+        from .filters import BandPassReconstructed
 
-        return BandPassFilter(
+        return BandPassReconstructed(
             sampling_rate=sampling_rate,
             lowpass=lowpass,
             highpass=highpass,

tme/rotations.py

@@ -1,8 +1,9 @@
-""" Implements various means of generating rotation matrices.
+"""
+Implements various means of generating rotation matrices.
 
-    Copyright (c) 2023-2025 European Molecular Biology Laboratory
+Copyright (c) 2023-2025 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import yaml
@@ -183,12 +184,14 @@ def euler_to_rotationmatrix(angles: Tuple[float], seq: str = "zyz") -> NDArray:
     NDArray
         The generated rotation matrix.
     """
+    angles = np.asarray(angles)
+
     n_angles = len(angles)
-    angle_convention = seq[:n_angles]
-    if n_angles == 1:
-        angles = (angles, 0, 0)
+    if angles.ndim == 2:
+        n_angles = angles.shape[1]
+
     rotation_matrix = Rotation.from_euler(
-        seq=angle_convention, angles=angles, degrees=True
+        seq=seq[:n_angles], angles=angles, degrees=True
     )
     return rotation_matrix.as_matrix().astype(np.float32)
 

tme/structure.py

@@ -1,8 +1,9 @@
-""" Implements class Structure to represent atomic structures.
+"""
+Implements class Structure to represent atomic structures.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import warnings