pytme 0.2.0b0__cp311-cp311-macosx_14_0_arm64.whl → 0.2.2__cp311-cp311-macosx_14_0_arm64.whl

tme/orientations.py

@@ -1,5 +1,5 @@
 #!python3
-""" Handle template matching peaks and convert between formats.
+""" Handle template matching orientations and conversion between formats.
 
     Copyright (c) 2024 European Molecular Biology Laboratory
 
@@ -8,6 +8,7 @@
 import re
 from collections import deque
 from dataclasses import dataclass
+from string import ascii_lowercase
 from typing import List, Tuple, Dict
 
 import numpy as np
@@ -17,21 +18,85 @@ from scipy.spatial.transform import Rotation
 @dataclass
 class Orientations:
     """
-    Handle template matching peaks and convert between formats.
+    Handle template matching orientations and conversion between formats.
+
+    Examples
+    --------
+    The following achieves the minimal definition of an :py:class:`Orientations` instance
+
+    >>> import numpy as np
+    >>> from tme import Orientations
+    >>> translations = np.random.randint(low = 0, high = 100, size = (100,3))
+    >>> rotations = np.random.rand(100, 3)
+    >>> scores = np.random.rand(100)
+    >>> details = np.full((100,), fill_value = -1)
+    >>> orientations = Orientations(
+    >>>     translations = translations,
+    >>>     rotations = rotations,
+    >>>     scores = scores,
+    >>>     details = details,
+    >>> )
+
+    The created ``orientations`` object can be written to disk in a range of formats.
+    See :py:meth:`Orientations.to_file` for available formats. The following creates
+    a STAR file
+
+    >>> orientations.to_file("test.star")
+
+    :py:meth:`Orientations.from_file` can create :py:class:`Orientations` instances
+    from a range of formats, to enable conversion between formats
+
+    >>> orientations_star = Orientations.from_file("test.star")
+    >>> np.all(orientations.translations == orientations_star.translations)
+    True
+
+    Parameters
+    ----------
+    translations: np.ndarray
+        Array with translations of each orientations (n, d).
+    rotations: np.ndarray
+        Array with euler angles of each orientation in zxy convention (n, d).
+    scores: np.ndarray
+        Array with the score of each orientation (n, ).
+    details: np.ndarray
+        Array with additional orientation details (n, ).
     """
 
-    #: Return a numpy array with translations of each orientation (n x d).
+    #: Array with translations of each orientation (n, d).
     translations: np.ndarray
 
-    #: Return a numpy array with euler angles of each orientation in zxy format (n x d).
+    #: Array with zyx euler angles of each orientation (n, d).
     rotations: np.ndarray
 
-    #: Return a numpy array with the score of each orientation (n, ).
+    #: Array with scores of each orientation (n, ).
     scores: np.ndarray
 
-    #: Return a numpy array with additional orientation details (n, ).
+    #: Array with additional details of each orientation(n, ).
     details: np.ndarray
 
+    def __post_init__(self):
+        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)
+        n_orientations = set(
+            [
+                self.translations.shape[0],
+                self.rotations.shape[0],
+                self.scores.shape[0],
+                self.details.shape[0],
+            ]
+        )
+        if len(n_orientations) != 1:
+            raise ValueError(
+                "The first dimension of all parameters needs to be of equal length."
+            )
+        if self.translations.ndim != 2:
+            raise ValueError("Expected two dimensional translations parameter.")
+
+        if self.rotations.ndim != 2:
+            raise ValueError("Expected two dimensional rotations parameter.")
+
     def __iter__(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
         """
         Iterate over the current class instance. Each iteration returns a orientation
@@ -65,9 +130,21 @@ class Orientations:
             "scores",
             "details",
         )
-        kwargs = {attr: getattr(self, attr)[indices] for attr in attributes}
+        kwargs = {attr: getattr(self, attr)[indices].copy() for attr in attributes}
         return self.__class__(**kwargs)
 
+    def copy(self) -> "Orientations":
+        """
+        Create a copy of the current class instance.
+
+        Returns
+        -------
+        :py:class:`Orientations`
+            Copy of the class instance.
+        """
+        indices = np.arange(self.scores.size)
+        return self[indices]
+
     def to_file(self, filename: str, file_format: type = None, **kwargs) -> None:
         """
         Save the current class instance to a file in the specified format.
@@ -77,7 +154,17 @@ class Orientations:
         filename : str
             The name of the file where the orientations will be saved.
         file_format : type, optional
-            The format in which to save the orientations. Supported formats are 'text' and 'relion'.
+            The format in which to save the orientations. Defaults to None and infers
+            the file_format from the typical extension. Supported formats are
+
+            +---------------+----------------------------------------------------+
+            | text          | pytme's standard tab-separated orientations file   |
+            +---------------+----------------------------------------------------+
+            | relion        | Creates a STAR file of orientations                |
+            +---------------+----------------------------------------------------+
+            | dynamo        | Creates a dynamo table                             |
+            +---------------+----------------------------------------------------+
+
         **kwargs : dict
             Additional keyword arguments specific to the file format.
 
@@ -120,17 +207,23 @@ class Orientations:
         The file is saved with a header specifying each column: z, y, x, euler_z,
         euler_y, euler_x, score, detail. Each row in the file corresponds to an orientation.
         """
+        naming = ascii_lowercase[::-1]
         header = "\t".join(
-            ["z", "y", "x", "euler_z", "euler_y", "euler_x", "score", "detail"]
+            [
+                *list(naming[: self.translations.shape[1]]),
+                *[f"euler_{x}" for x in naming[: self.rotations.shape[1]]],
+                "score",
+                "detail",
+            ]
         )
         with open(filename, mode="w", encoding="utf-8") as ofile:
             _ = ofile.write(f"{header}\n")
             for translation, angles, score, detail in self:
-                translation_string = "\t".join([str(x) for x in translation])
-                angle_string = "\t".join([str(x) for x in angles])
-                _ = ofile.write(
-                    f"{translation_string}\t{angle_string}\t{score}\t{detail}\n"
+                out_string = (
+                    "\t".join([str(x) for x in (*translation, *angles, score, detail)])
+                    + "\n"
                 )
+                _ = ofile.write(out_string)
         return None
 
     def _to_dynamo_tbl(
@@ -158,9 +251,6 @@ class Orientations:
         References
         ----------
         .. [1]  https://wiki.dynamo.biozentrum.unibas.ch/w/index.php/Table
-
-        The file is saved with a standard header used in Dynamo STAR files.
-        Each row in the file corresponds to an orientation.
         """
         with open(filename, mode="w", encoding="utf-8") as ofile:
             for index, (translation, rotation, score, detail) in enumerate(self):
@@ -316,8 +406,18 @@ class Orientations:
         filename : str
             The name of the file from which to read the orientations.
         file_format : type, optional
-            The format of the file. Currently, only 'text' format is supported.
-        **kwargs : dict
+            The format of the file. Defaults to None and infers
+            the file_format from the typical extension. Supported formats are
+
+            +---------------+----------------------------------------------------+
+            | text          | pyTME's standard tab-separated orientations file   |
+            +---------------+----------------------------------------------------+
+            | relion        | Creates a STAR file of orientations                |
+            +---------------+----------------------------------------------------+
+            | dynamo        | Creates a dynamo table                             |
+            +---------------+----------------------------------------------------+
+
+        **kwargs
             Additional keyword arguments specific to the file format.
 
         Returns
@@ -330,11 +430,18 @@ class Orientations:
         ValueError
             If an unsupported file format is specified.
         """
-        mapping = {"text": cls._from_text, "relion": cls._from_relion_star}
+        mapping = {
+            "text": cls._from_text,
+            "relion": cls._from_relion_star,
+            "tbl": cls._from_tbl,
+        }
         if file_format is None:
             file_format = "text"
+
             if filename.lower().endswith(".star"):
                 file_format = "relion"
+            elif filename.lower().endswith(".tbl"):
+                file_format = "tbl"
 
         func = mapping.get(file_format, None)
         if func is None:
@@ -370,30 +477,61 @@ class Orientations:
 
         Notes
         -----
-        The text file is expected to have a header and data in columns corresponding to
-        z, y, x, euler_z, euler_y, euler_x, score, detail.
+        The text file is expected to have a header and data in columns. Colums containing
+        the name euler are considered to specify rotations. The second last and last
+        column correspond to score and detail. Its possible to only specify translations,
+        in this case the remaining columns will be filled with trivial values.
         """
         with open(filename, mode="r", encoding="utf-8") as infile:
             data = [x.strip().split("\t") for x in infile.read().split("\n")]
-            _ = data.pop(0)
 
+        header = data.pop(0)
         translation, rotation, score, detail = [], [], [], []
         for candidate in data:
             if len(candidate) <= 1:
                 continue
-            if len(candidate) != 8:
-                candidate.append(-1)
 
-            candidate = [float(x) for x in candidate]
-            translation.append((candidate[0], candidate[1], candidate[2]))
-            rotation.append((candidate[3], candidate[4], candidate[5]))
-            score.append(candidate[6])
-            detail.append(candidate[7])
+            translation.append(
+                tuple(
+                    candidate[i] for i, x in enumerate(header) if x in ascii_lowercase
+                )
+            )
+            rotation.append(
+                tuple(candidate[i] for i, x in enumerate(header) if "euler" in x)
+            )
+            score.append(candidate[-2])
+            detail.append(candidate[-1])
+
+        translation = np.vstack(translation)
+        rotation = np.vstack(rotation)
+        score = np.array(score)
+        detail = np.array(detail)
+
+        if translation.shape[1] == len(header):
+            rotation = np.zeros(translation.shape, dtype=np.float32)
+            score = np.zeros(translation.shape[0], dtype=np.float32)
+            detail = np.zeros(translation.shape[0], dtype=np.float32) - 1
+
+        if rotation.size == 0 and translation.shape[0] != 0:
+            rotation = np.zeros(translation.shape, dtype=np.float32)
+
+        header_order = tuple(x for x in header if x in ascii_lowercase)
+        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=True)
+        )
+        translation = translation[..., sort_order]
 
-        translation = np.vstack(translation).astype(int)
-        rotation = np.vstack(rotation).astype(float)
-        score = np.array(score).astype(float)
-        detail = np.array(detail).astype(float)
+        header_order = tuple(
+            x
+            for x in header
+            if "euler" in x and x.replace("euler_", "") in ascii_lowercase
+        )
+        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=True)
+        )
+        rotation = rotation[..., sort_order]
 
         return translation, rotation, score, detail
 
@@ -448,20 +586,15 @@ class Orientations:
         ret = cls._parse_star(filename=filename, delimiter=delimiter)
         ret = ret["data_particles"]
 
-        translation = (
-            np.vstack(
-                (ret["_rlnCoordinateZ"], ret["_rlnCoordinateY"], ret["_rlnCoordinateX"])
-            )
-            .astype(np.float32)
-            .astype(int)
-            .T
+        translation = np.vstack(
+            (ret["_rlnCoordinateZ"], ret["_rlnCoordinateY"], ret["_rlnCoordinateX"])
         )
+        translation = translation.astype(np.float32).T
 
-        rotation = (
-            np.vstack((ret["_rlnAngleRot"], ret["_rlnAngleTilt"], ret["_rlnAnglePsi"]))
-            .astype(np.float32)
-            .T
+        rotation = np.vstack(
+            (ret["_rlnAngleRot"], ret["_rlnAngleTilt"], ret["_rlnAnglePsi"])
         )
+        rotation = rotation.astype(np.float32).T
 
         rotation = Rotation.from_euler("xyx", rotation, degrees=True)
         rotation = rotation.as_euler(seq="zyx", degrees=True)
@@ -470,6 +603,33 @@ class Orientations:
 
         return translation, rotation, score, detail
 
+    @staticmethod
+    def _from_tbl(
+        filename: str, **kwargs
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        with open(filename, mode="r", encoding="utf-8") as infile:
+            data = infile.read().split("\n")
+        data = [x.strip().split(" ") for x in data if len(x.strip())]
+
+        if len(data[0]) != 38:
+            raise ValueError(
+                "Expected tbl file to have 38 columns generated by _to_tbl."
+            )
+
+        translations, rotations, scores, details = [], [], [], []
+        for peak in data:
+            rotation = Rotation.from_euler(
+                "xyx", (peak[6], peak[7], peak[8]), degrees=True
+            )
+            rotations.append(rotation.as_euler(seq="zyx", degrees=True))
+            scores.append(peak[9])
+            details.append(-1)
+            translations.append((peak[25], peak[24], peak[23]))
+
+        translations, rotations = np.array(translations), np.array(rotations)
+        scores, details = np.array(scores), np.array(details)
+        return translations, rotations, scores, details
+
     def get_extraction_slices(
         self,
         target_shape: Tuple[int],
@@ -504,55 +664,55 @@ class Orientations:
         SystemExit
             If no peak remains after filtering, indicating an error.
         """
-        left_pad = np.divide(extraction_shape, 2).astype(int)
-        right_pad = np.add(left_pad, np.mod(extraction_shape, 2)).astype(int)
+        right_pad = np.divide(extraction_shape, 2).astype(int)
+        left_pad = np.add(right_pad, np.mod(extraction_shape, 2)).astype(int)
+
+        peaks = self.translations.astype(int)
+        obs_beg = np.subtract(peaks, left_pad)
+        obs_end = np.add(peaks, right_pad)
 
-        obs_start = np.subtract(self.translations, left_pad)
-        obs_stop = np.add(self.translations, right_pad)
+        obs_beg = np.maximum(obs_beg, 0)
+        obs_end = np.minimum(obs_end, target_shape)
 
-        cand_start = np.subtract(np.maximum(obs_start, 0), obs_start)
-        cand_stop = np.subtract(obs_stop, np.minimum(obs_stop, target_shape))
-        cand_stop = np.subtract(extraction_shape, cand_stop)
-        obs_start = np.maximum(obs_start, 0)
-        obs_stop = np.minimum(obs_stop, target_shape)
+        cand_beg = left_pad - np.subtract(peaks, obs_beg)
+        cand_end = left_pad + np.subtract(obs_end, peaks)
 
         subset = self
         if drop_out_of_box:
-            stops = np.subtract(cand_stop, extraction_shape)
+            stops = np.subtract(cand_end, extraction_shape)
             keep_peaks = (
                 np.sum(
-                    np.multiply(cand_start == 0, stops == 0),
+                    np.multiply(cand_beg == 0, stops == 0),
                     axis=1,
                 )
-                == self.translations.shape[1]
+                == peaks.shape[1]
             )
             n_remaining = keep_peaks.sum()
             if n_remaining == 0:
                 print(
                     "No peak remaining after filtering. Started with"
-                    f" {self.translations.shape[0]} filtered to {n_remaining}."
+                    f" {peaks.shape[0]} filtered to {n_remaining}."
                     " Consider reducing min_distance, increase num_peaks or use"
                     " a different peak caller."
                 )
-                exit(-1)
 
-            cand_start = cand_start[keep_peaks,]
-            cand_stop = cand_stop[keep_peaks,]
-            obs_start = obs_start[keep_peaks,]
-            obs_stop = obs_stop[keep_peaks,]
+            cand_beg = cand_beg[keep_peaks,]
+            cand_end = cand_end[keep_peaks,]
+            obs_beg = obs_beg[keep_peaks,]
+            obs_end = obs_end[keep_peaks,]
             subset = self[keep_peaks]
 
-        cand_start, cand_stop = cand_start.astype(int), cand_stop.astype(int)
-        obs_start, obs_stop = obs_start.astype(int), obs_stop.astype(int)
+        cand_beg, cand_end = cand_beg.astype(int), cand_end.astype(int)
+        obs_beg, obs_end = obs_beg.astype(int), obs_end.astype(int)
 
         candidate_slices = [
             tuple(slice(s, e) for s, e in zip(start_row, stop_row))
-            for start_row, stop_row in zip(cand_start, cand_stop)
+            for start_row, stop_row in zip(cand_beg, cand_end)
         ]
 
         observation_slices = [
             tuple(slice(s, e) for s, e in zip(start_row, stop_row))
-            for start_row, stop_row in zip(obs_start, obs_stop)
+            for start_row, stop_row in zip(obs_beg, obs_end)
         ]
 
         if return_orientations:

tme/parser.py

@@ -137,8 +137,7 @@ class Parser(ABC):
 
 class PDBParser(Parser):
     """
-    A Parser subclass for converting PDB file data into a dictionary representation.
-    This class is specifically designed to work with PDB file format.
+    Convert PDB file data into a dictionary representation [1]_.
 
     References
     ----------
@@ -228,8 +227,8 @@ class PDBParser(Parser):
 
 class MMCIFParser(Parser):
     """
-    A Parser subclass for converting MMCIF file data into a dictionary representation.
-    This implementation heavily relies on the atomium library [1]_.
+    Convert MMCIF file data into a dictionary representation. This implementation
+    heavily relies on the atomium library [1]_.
 
     References
     ----------

tme/preprocessing/__init__.py

@@ -0,0 +1,2 @@
+from .compose import Compose
+from .frequency_filters import BandPassFilter, LinearWhiteningFilter

tme/preprocessing/_utils.py

@@ -0,0 +1,217 @@
+""" Utilities for the generation of frequency grids.
+
+    Copyright (c) 2024 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Tuple, List
+
+import numpy as np
+
+from ..backends import backend as be
+from ..backends import NumpyFFTWBackend
+from ..types import BackendArray, NDArray
+from ..matching_utils import euler_to_rotationmatrix
+
+
+def compute_tilt_shape(shape: Tuple[int], opening_axis: int, reduce_dim: bool = False):
+    """
+    Given an opening_axis, computes the shape of the remaining dimensions.
+
+    Parameters:
+    -----------
+    shape : Tuple[int]
+        The shape of the input array.
+    opening_axis : int
+        The axis along which the array will be tilted.
+    reduce_dim : bool, optional (default=False)
+        Whether to reduce the dimensionality after tilting.
+
+    Returns:
+    --------
+    Tuple[int]
+        The shape of the array after tilting.
+    """
+    tilt_shape = tuple(x if i != opening_axis else 1 for i, x in enumerate(shape))
+    if reduce_dim:
+        tilt_shape = tuple(x for i, x in enumerate(shape) if i != opening_axis)
+
+    return tilt_shape
+
+
+def centered_grid(shape: Tuple[int]) -> NDArray:
+    """
+    Generate an integer valued grid centered around size // 2
+
+    Parameters:
+    -----------
+    shape : Tuple[int]
+        The shape of the grid.
+
+    Returns:
+    --------
+    NDArray
+        The centered grid.
+    """
+    index_grid = np.array(
+        np.meshgrid(*[np.arange(size) - size // 2 for size in shape], indexing="ij")
+    )
+    return index_grid
+
+
+def frequency_grid_at_angle(
+    shape: Tuple[int],
+    angle: float,
+    sampling_rate: Tuple[float],
+    opening_axis: int = None,
+    tilt_axis: int = None,
+) -> NDArray:
+    """
+    Generate a frequency grid from 0 to 1/(2 * sampling_rate) in each axis.
+
+    Parameters:
+    -----------
+    shape : Tuple[int]
+        The shape of the grid.
+    angle : float
+        The angle at which to generate the grid.
+    sampling_rate : Tuple[float]
+        The sampling rate for each dimension.
+    opening_axis : int, optional
+        The axis to be opened, defaults to None.
+    tilt_axis : int, optional
+        The axis along which the grid is tilted, defaults to None.
+
+    Returns:
+    --------
+    NDArray
+        The frequency grid.
+    """
+    sampling_rate = np.array(sampling_rate)
+    sampling_rate = np.repeat(sampling_rate, len(shape) // sampling_rate.size)
+
+    tilt_shape = compute_tilt_shape(
+        shape=shape, opening_axis=opening_axis, reduce_dim=False
+    )
+
+    if angle == 0:
+        index_grid = fftfreqn(
+            tuple(x for x in tilt_shape if x != 1),
+            sampling_rate=1,
+            compute_euclidean_norm=True,
+        )
+
+    if angle != 0:
+        angles = np.zeros(len(shape))
+        angles[tilt_axis] = angle
+        rotation_matrix = euler_to_rotationmatrix(np.roll(angles, opening_axis - 1))
+
+        index_grid = fftfreqn(tilt_shape, sampling_rate=None)
+        index_grid = np.einsum("ij,j...->i...", rotation_matrix, index_grid)
+        norm = np.multiply(sampling_rate, shape).astype(int)
+
+        index_grid = np.divide(index_grid.T, norm).T
+        index_grid = np.squeeze(index_grid)
+        index_grid = np.linalg.norm(index_grid, axis=(0))
+
+    return index_grid
+
+
+def fftfreqn(
+    shape: Tuple[int],
+    sampling_rate: Tuple[float],
+    compute_euclidean_norm: bool = False,
+    shape_is_real_fourier: bool = False,
+    return_sparse_grid: bool = False,
+) -> NDArray:
+    """
+    Generate the n-dimensional discrete Fourier transform sample frequencies.
+
+    Parameters:
+    -----------
+    shape : Tuple[int]
+        The shape of the data.
+    sampling_rate : float or Tuple[float]
+        The sampling rate.
+    compute_euclidean_norm : bool, optional
+        Whether to compute the Euclidean norm, defaults to False.
+    shape_is_real_fourier : bool, optional
+        Whether the shape corresponds to a real Fourier transform, defaults to False.
+
+    Returns:
+    --------
+    NDArray
+        The sample frequencies.
+    """
+    # There is no real need to have these operations on GPU right now
+    temp_backend = NumpyFFTWBackend()
+    norm = temp_backend.full(len(shape), fill_value=1)
+    center = temp_backend.astype(temp_backend.divide(shape, 2), temp_backend._int_dtype)
+    if sampling_rate is not None:
+        norm = temp_backend.astype(temp_backend.multiply(shape, sampling_rate), int)
+
+    if shape_is_real_fourier:
+        center[-1], norm[-1] = 0, 1
+        if sampling_rate is not None:
+            norm[-1] = (shape[-1] - 1) * 2 * sampling_rate
+
+    grids = []
+    for i, x in enumerate(shape):
+        baseline_dims = tuple(1 if i != t else x for t in range(len(shape)))
+        grid = (temp_backend.arange(x) - center[i]) / norm[i]
+        grids.append(temp_backend.reshape(grid, baseline_dims))
+
+    if compute_euclidean_norm:
+        grids = sum(temp_backend.square(x) for x in grids)
+        grids = temp_backend.sqrt(grids, out=grids)
+        return grids
+
+    if return_sparse_grid:
+        return grids
+
+    grid_flesh = temp_backend.full(shape, fill_value=1)
+    grids = temp_backend.stack(tuple(grid * grid_flesh for grid in grids))
+
+    return grids
+
+
+def crop_real_fourier(data: BackendArray) -> BackendArray:
+    """
+    Crop the real part of a Fourier transform.
+
+    Parameters:
+    -----------
+    data : BackendArray
+        The Fourier transformed data.
+
+    Returns:
+    --------
+    BackendArray
+        The cropped data.
+    """
+    stop = 1 + (data.shape[-1] // 2)
+    return data[..., :stop]
+
+
+def compute_fourier_shape(
+    shape: Tuple[int], shape_is_real_fourier: bool = False
+) -> List[int]:
+    if shape_is_real_fourier:
+        return shape
+    shape = [int(x) for x in shape]
+    shape[-1] = 1 + shape[-1] // 2
+    return shape
+
+
+def shift_fourier(
+    data: BackendArray, shape_is_real_fourier: bool = False
+) -> BackendArray:
+    shape = be.to_backend_array(data.shape)
+    shift = be.add(be.divide(shape, 2), be.mod(shape, 2))
+    shift = [int(x) for x in shift]
+    if shape_is_real_fourier:
+        shift[-1] = 0
+
+    data = be.roll(data, shift, tuple(i for i in range(len(shift))))
+    return data