pytme 0.2.9__cp311-cp311-macosx_15_0_arm64.whl

tests/test_rotations.py

@@ -0,0 +1,153 @@
+from importlib_resources import files
+from itertools import combinations, chain, product
+
+import pytest
+import numpy as np
+
+from tme import Density
+from scipy.spatial.transform import Rotation
+from scipy.signal import correlate
+
+from tme.rotations import (
+    euler_from_rotationmatrix,
+    euler_to_rotationmatrix,
+    get_cone_rotations,
+    align_vectors,
+    get_rotation_matrices,
+)
+from tme.matching_utils import (
+    elliptical_mask,
+    split_shape,
+    compute_full_convolution_index,
+)
+
+BASEPATH = files("tests.data")
+
+
+class TestRotations:
+    def setup_method(self):
+        self.density = Density.from_file(str(BASEPATH.joinpath("Raw/em_map.map")))
+        self.structure_density = Density.from_structure(
+            filename_or_structure=str(BASEPATH.joinpath("Structures/5khe.cif")),
+            origin=self.density.origin,
+            shape=self.density.shape,
+            sampling_rate=self.density.sampling_rate,
+        )
+
+    @pytest.mark.parametrize(
+        "initial_vector, target_vector, convention",
+        [
+            ([1, 0, 0], [0, 1, 0], None),
+            ([0, 1, 0], [0, 0, 1], "zyx"),
+            ([1, 1, 1], [1, 0, 0], "xyz"),
+        ],
+    )
+    def test_align_vectors(self, initial_vector, target_vector, convention):
+        result = align_vectors(initial_vector, target_vector, convention)
+
+        assert isinstance(result, np.ndarray)
+        if convention is None:
+            assert result.shape == (3, 3)
+            assert np.allclose(np.dot(result, result.T), np.eye(3), atol=1e-6)
+        else:
+            assert len(result) == 3
+            result = Rotation.from_euler(convention, result, degrees=True).as_matrix()
+            assert np.allclose(np.dot(result, result.T), np.eye(3), atol=1e-6)
+
+        rotated = np.dot(Rotation.from_matrix(result).as_matrix(), initial_vector)
+        assert np.allclose(
+            rotated / np.linalg.norm(rotated),
+            target_vector / np.linalg.norm(target_vector),
+            atol=1e-6,
+        )
+
+    @pytest.mark.parametrize(
+        "cone_angle, cone_sampling, axis_angle, axis_sampling, vector, n_symmetry, convention",
+        [
+            (30, 5, 360, None, (1, 0, 0), 1, None),
+            (45, 10, 180, 15, (0, 1, 0), 2, "zyx"),
+            (60, 15, 90, 30, (0, 0, 1), 4, "xyz"),
+        ],
+    )
+    def test_get_cone_rotations(
+        self,
+        cone_angle,
+        cone_sampling,
+        axis_angle,
+        axis_sampling,
+        vector,
+        n_symmetry,
+        convention,
+    ):
+        result = get_cone_rotations(
+            cone_angle=cone_angle,
+            cone_sampling=cone_sampling,
+            axis_angle=axis_angle,
+            axis_sampling=axis_sampling,
+            reference=vector,
+            n_symmetry=n_symmetry,
+            seq=convention,
+        )
+
+        assert isinstance(result, np.ndarray)
+        if convention is None:
+            assert result.shape[1:] == (3, 3)
+        else:
+            assert result.shape[1] == 3
+
+    def test_euler_conversion(self):
+        rotation_matrix_initial = np.array(
+            [
+                [0.35355339, 0.61237244, -0.70710678],
+                [-0.8660254, 0.5, -0.0],
+                [0.35355339, 0.61237244, 0.70710678],
+            ]
+        )
+        euler_angles = euler_from_rotationmatrix(rotation_matrix_initial)
+        rotation_matrix_converted = euler_to_rotationmatrix(euler_angles)
+        assert np.allclose(
+            rotation_matrix_initial, rotation_matrix_converted, atol=1e-6
+        )
+
+    @pytest.mark.parametrize("dim", range(1, 3, 5))
+    @pytest.mark.parametrize("angular_sampling", [10, 15, 20])
+    def test_get_rotation_matrices(self, dim, angular_sampling):
+        rotation_matrices = get_rotation_matrices(
+            angular_sampling=angular_sampling, dim=dim
+        )
+        assert np.allclose(rotation_matrices[0] @ rotation_matrices[0].T, np.eye(dim))
+
+    def test_split_correlation(self):
+        arr1 = elliptical_mask(shape=(50, 51), center=(20, 30), radius=5)
+
+        arr2 = elliptical_mask(shape=(41, 36), center=(25, 20), radius=5)
+        s = range(arr1.ndim)
+        outer_split = chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))
+
+        s = range(arr2.ndim)
+        inner_split = chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))
+
+        outer_splits = [dict(zip(i, [2] * len(i))) for i in list(outer_split)]
+        inner_splits = [dict(zip(i, [2] * len(i))) for i in list(inner_split)]
+
+        for outer_split, inner_split in product(outer_splits, inner_splits):
+            splits1 = split_shape(
+                shape=arr1.shape, splits=outer_split, equal_shape=False
+            )
+            splits2 = split_shape(
+                shape=arr2.shape, splits=inner_split, equal_shape=False
+            )
+
+            full = correlate(arr1, arr2, method="direct", mode="full")
+            temp = np.zeros_like(full)
+
+            for arr1_split, arr2_split in product(splits1, splits2):
+                correlation = correlate(
+                    arr1[arr1_split], arr2[arr2_split], method="direct", mode="full"
+                )
+                score_slice = compute_full_convolution_index(
+                    arr1.shape, arr2.shape, arr1_split, arr2_split
+                )
+                temp[score_slice] += correlation
+
+            assert np.allclose(temp, full)

tests/test_structure.py

@@ -0,0 +1,247 @@
+from os import remove
+from tempfile import mkstemp
+from importlib_resources import files
+
+import pytest
+import numpy as np
+
+from tme import Structure
+from tme.matching_utils import minimum_enclosing_box
+from tme.rotations import euler_to_rotationmatrix
+
+
+STRUCTURE_ATTRIBUTES = [
+    "record_type",
+    "atom_serial_number",
+    "atom_name",
+    "atom_coordinate",
+    "alternate_location_indicator",
+    "residue_name",
+    "chain_identifier",
+    "residue_sequence_number",
+    "code_for_residue_insertion",
+    "occupancy",
+    "temperature_factor",
+    "segment_identifier",
+    "element_symbol",
+    "charge",
+    "metadata",
+]
+
+
+class TestStructure:
+    def setup_method(self):
+        self.structure = Structure.from_file(
+            str(files("tests.data").joinpath("Structures/5khe.cif"))
+        )
+        _, self.path = mkstemp()
+
+    def teardown_method(self):
+        del self.structure
+        remove(self.path)
+
+    def compare_structures(self, structure1, structure2, exclude_attributes=[]):
+        for attribute in STRUCTURE_ATTRIBUTES:
+            if attribute in exclude_attributes:
+                continue
+            value = getattr(structure1, attribute)
+            value_comparison = getattr(structure2, attribute)
+            if isinstance(value, np.ndarray):
+                assert np.all(value_comparison == value)
+            else:
+                assert value == value_comparison
+
+    def test_initialization(self):
+        structure = Structure(
+            record_type=self.structure.record_type,
+            atom_serial_number=self.structure.atom_serial_number,
+            atom_name=self.structure.atom_name,
+            atom_coordinate=self.structure.atom_coordinate,
+            alternate_location_indicator=self.structure.alternate_location_indicator,
+            residue_name=self.structure.residue_name,
+            chain_identifier=self.structure.chain_identifier,
+            residue_sequence_number=self.structure.residue_sequence_number,
+            code_for_residue_insertion=self.structure.code_for_residue_insertion,
+            occupancy=self.structure.occupancy,
+            temperature_factor=self.structure.temperature_factor,
+            segment_identifier=self.structure.segment_identifier,
+            element_symbol=self.structure.element_symbol,
+            charge=self.structure.charge,
+            metadata=self.structure.metadata,
+        )
+
+        for attribute in STRUCTURE_ATTRIBUTES:
+            value = getattr(self.structure, attribute)
+            value_comparison = getattr(structure, attribute)
+            if isinstance(value, np.ndarray):
+                assert np.all(value_comparison == value)
+            else:
+                assert value == value_comparison
+
+    @pytest.mark.parametrize(
+        "modified_attribute",
+        [
+            ("record_type"),
+            ("atom_serial_number"),
+            ("atom_name"),
+            ("atom_coordinate"),
+            ("alternate_location_indicator"),
+            ("residue_name"),
+            ("chain_identifier"),
+            ("residue_sequence_number"),
+            ("code_for_residue_insertion"),
+            ("occupancy"),
+            ("temperature_factor"),
+            ("segment_identifier"),
+            ("element_symbol"),
+        ],
+    )
+    def test_initialization_errors(self, modified_attribute):
+        kwargs = {
+            attribute: getattr(self.structure, attribute)
+            for attribute in STRUCTURE_ATTRIBUTES
+            if attribute != modified_attribute
+        }
+        kwargs[modified_attribute] = getattr(self.structure, modified_attribute)[:1]
+
+        with pytest.raises(ValueError):
+            Structure(**kwargs)
+
+    def test__getitem__(self):
+        ret_single_index = self.structure[1]
+        ret = self.structure[[1]]
+        self.compare_structures(ret_single_index, ret)
+
+        ret = self.structure[self.structure.record_type == "ATOM"]
+        assert np.all(ret.record_type == "ATOM")
+
+        ret = self.structure[self.structure.element_symbol == "C"]
+        assert np.all(ret.element_symbol == "C")
+
+    def test__repr__(self):
+        unique_chains = "-".join(
+            [
+                ",".join([str(x) for x in entity])
+                for entity in self.structure.metadata["unique_chains"]
+            ]
+        )
+
+        min_atom = np.min(self.structure.atom_serial_number)
+        max_atom = np.max(self.structure.atom_serial_number)
+        n_atom = self.structure.atom_serial_number.size
+
+        min_residue = np.min(self.structure.residue_sequence_number)
+        max_residue = np.max(self.structure.residue_sequence_number)
+        n_residue = np.unique(self.structure.residue_sequence_number).size
+
+        repr_str = (
+            f"Structure object at {id(self.structure)}\n"
+            f"Unique Chains: {unique_chains}, "
+            f"Atom Range: {min_atom}-{max_atom} [N = {n_atom}], "
+            f"Residue Range: {min_residue}-{max_residue} [N = {n_residue}]"
+        )
+        assert repr_str == self.structure.__repr__()
+
+    @pytest.mark.parametrize(
+        "path",
+        [
+            str(files("tests.data").joinpath("Structures/5khe.cif")),
+            str(files("tests.data").joinpath("Structures/5khe.pdb")),
+        ],
+    )
+    def test_fromfile(self, path):
+        _ = Structure.from_file(path)
+
+    def test_fromfile_error(self):
+        with pytest.raises(NotImplementedError):
+            _ = Structure.from_file("madeup.extension")
+
+    @pytest.mark.parametrize("file_format", ["cif", "pdb", "gro"])
+    def test_to_file(self, file_format):
+        _, path = mkstemp()
+        path = f"{path}.{file_format}"
+        self.structure.to_file(path)
+        read = self.structure.from_file(path)
+        comparison = self.structure.copy()
+
+        if file_format != "gro":
+            self.compare_structures(comparison, read, exclude_attributes=["metadata"])
+        else:
+            assert np.allclose(comparison.atom_coordinate, read.atom_coordinate)
+
+    def test_to_file_error(self):
+        _, path = mkstemp()
+        path = f"{path}.RAISERROR"
+        with pytest.raises(NotImplementedError):
+            self.structure.to_file(path)
+
+    def test_subset_by_chain(self):
+        chain = "A"
+        ret = self.structure.subset_by_chain(chain=chain)
+        assert np.all(ret.chain_identifier == chain)
+
+    def test_subset_by_chain_range(self):
+        chain, start, stop = "A", 0, 20
+        ret = self.structure.subset_by_range(chain=chain, start=start, stop=stop)
+        assert np.all(ret.chain_identifier == chain)
+        assert np.all(
+            np.logical_and(
+                ret.residue_sequence_number >= start,
+                ret.residue_sequence_number <= stop,
+            )
+        )
+
+    def test_center_of_mass(self):
+        center_of_mass = self.structure.center_of_mass()
+        assert center_of_mass.shape[0] == self.structure.atom_coordinate.shape[1]
+        assert np.allclose(center_of_mass, [-0.89391639, 29.94908928, -2.64736741])
+
+    def test_centered(self):
+        ret, translation = self.structure.centered()
+        box = minimum_enclosing_box(coordinates=self.structure.atom_coordinate.T)
+        assert np.allclose(ret.center_of_mass(), np.divide(box, 2), atol=1)
+
+    def test__get_atom_weights_error(self):
+        with pytest.raises(NotImplementedError):
+            self.structure._get_atom_weights(
+                self.structure.atom_name, weight_type="RAISEERROR"
+            )
+
+    def test_compare_structures(self):
+        rmsd = Structure.compare_structures(self.structure, self.structure)
+        assert rmsd == 0
+
+        rmsd = Structure.compare_structures(
+            self.structure, self.structure, weighted=True
+        )
+        assert rmsd == 0
+
+        translation = (3, 0, 0)
+        structure_transform = self.structure.rigid_transform(
+            translation=translation,
+            rotation_matrix=np.eye(self.structure.atom_coordinate.shape[1]),
+        )
+        rmsd = Structure.compare_structures(self.structure, structure_transform)
+        assert np.allclose(rmsd, np.linalg.norm(translation))
+
+    def test_comopare_structures_error(self):
+        ret = self.structure[[1, 2, 3, 4, 5]]
+        with pytest.raises(ValueError):
+            Structure.compare_structures(self.structure, ret)
+
+    def test_align_structures(self):
+        rotation_matrix = euler_to_rotationmatrix((20, -10, 45))
+        translation = (10, 0, -15)
+
+        structure_transform = self.structure.rigid_transform(
+            rotation_matrix=rotation_matrix, translation=translation
+        )
+        aligned, final_rmsd = Structure.align_structures(
+            self.structure, structure_transform
+        )
+        assert final_rmsd <= 0.1
+
+        aligned, final_rmsd = Structure.align_structures(
+            self.structure, structure_transform, sampling_rate=1
+        )
+        assert final_rmsd <= 1

tme/__init__.py

@@ -0,0 +1,6 @@
+from .__version__ import __version__
+from .density import Density
+from .preprocessor import Preprocessor
+from .structure import Structure
+from .orientations import Orientations
+from .matching_data import MatchingData

tme/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.2.9"

tme/analyzer/__init__.py

@@ -0,0 +1,2 @@
+from .peaks import *
+from .aggregation import *

tme/analyzer/_utils.py

@@ -0,0 +1,186 @@
+""" Analyzer utility functions.
+
+    Copyright (c) 2023-2025 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Tuple
+
+from ..types import BackendArray
+from ..backends import backend as be
+
+__all__ = ["cart_to_score", "score_to_cart"]
+
+
+def _convmode_to_shape(
+    convolution_mode: str,
+    targetshape: BackendArray,
+    templateshape: BackendArray,
+    convolution_shape: BackendArray,
+) -> BackendArray:
+    """
+    Calculate convolution shape based on convolution mode.
+
+    Parameters
+    ----------
+    convolution_mode : str
+        Mode of convolution. Supported values are:
+        - 'same': Output shape will match target shape
+        - 'valid': Output shape will be target shape minus template shape plus
+                   template shape modulo 2
+        - Other: Output shape will be equal to convolution_shape
+    targetshape : BackendArray
+        Shape of the target array.
+    templateshape : BackendArray
+        Shape of the template array.
+    convolution_shape : BackendArray
+        Shape of the convolution output.
+
+    Returns
+    -------
+    BackendArray
+        Convolution shape.
+    """
+    output_shape = convolution_shape
+    if convolution_mode == "same":
+        output_shape = targetshape
+    elif convolution_mode == "valid":
+        output_shape = be.add(
+            be.subtract(targetshape, templateshape),
+            be.mod(templateshape, 2),
+        )
+    return be.to_backend_array(output_shape)
+
+
+def cart_to_score(
+    positions: BackendArray,
+    fast_shape: Tuple[int],
+    targetshape: Tuple[int],
+    templateshape: Tuple[int],
+    convolution_shape: Tuple[int] = None,
+    fourier_shift: Tuple[int] = None,
+    convolution_mode: str = None,
+    **kwargs,
+) -> Tuple[BackendArray]:
+    """
+    Maps peak positions from cartesian to padded score space coordinates.
+
+    Parameters
+    ----------
+    positions : BackendArray
+        Positions in cartesian coordinates.
+    fast_shape : tuple of int
+        Shape of the score space padded to efficient Fourier shape.
+    targetshape : tuple of int
+        Shape of the target array.
+    templateshape : tuple of int
+        Shape of the template array.
+    convolution_shape : tuple of int, optional
+        Non-padded convolution_shape of template and target.
+    fourier_shift : tuple of int, optional
+        Translation offset of coordinates.
+    convolution_mode : str, optional
+        Mode of convolution ('same', 'valid', or 'full')
+
+    Returns
+    -------
+    Tuple of BackendArray
+        Adjusted positions. and boolean array indicating whether corresponding
+        positions are valid positions w.r.t. to supplied bounds.
+    """
+    positions = be.to_backend_array(positions)
+    fast_shape = be.to_backend_array(fast_shape)
+    targetshape = be.to_backend_array(targetshape)
+    templateshape = be.to_backend_array(templateshape)
+    convolution_shape = be.to_backend_array(convolution_shape)
+
+    # Compute removed padding
+    output_shape = _convmode_to_shape(
+        convolution_mode=convolution_mode,
+        targetshape=targetshape,
+        templateshape=templateshape,
+        convolution_shape=convolution_shape,
+    )
+    valid_positions = be.multiply(positions >= 0, positions < output_shape)
+    valid_positions = be.sum(valid_positions, axis=1) == positions.shape[1]
+
+    starts = be.astype(
+        be.divide(be.subtract(convolution_shape, output_shape), 2),
+        be._int_dtype,
+    )
+
+    positions = be.add(positions, starts)
+    if fourier_shift is not None:
+        fourier_shift = be.to_backend_array(fourier_shift)
+        positions = be.subtract(positions, fourier_shift)
+        positions = be.mod(positions, fast_shape)
+
+    return positions, valid_positions
+
+
+def score_to_cart(
+    positions,
+    fast_shape: Tuple[int] = None,
+    targetshape: Tuple[int] = None,
+    templateshape: Tuple[int] = None,
+    convolution_shape: Tuple[int] = None,
+    fourier_shift: Tuple[int] = None,
+    convolution_mode: str = None,
+    **kwargs,
+) -> Tuple[BackendArray]:
+    """
+    Maps peak positions from padded score to cartesian coordinates.
+
+    Parameters
+    ----------
+    positions : BackendArray
+        Positions in padded Fourier space system.
+    fast_shape : tuple of int
+        Shape of the score space padded to efficient Fourier shape.
+    targetshape : tuple of int
+        Shape of the target array.
+    templateshape : tuple of int
+        Shape of the template array.
+    convolution_shape : tuple of int, optional
+        Non-padded convolution_shape of template and target.
+    fourier_shift : tuple of int, optional
+        Translation offset of coordinates.
+    convolution_mode : str, optional
+        Mode of convolution ('same', 'valid', or 'full')
+
+    Returns
+    -------
+    Tuple of BackendArray
+        Adjusted positions. and boolean array indicating whether corresponding
+        positions are valid positions w.r.t. to supplied bounds.
+    """
+    positions = be.to_backend_array(positions)
+    convolution_shape = be.to_backend_array(convolution_shape)
+    fast_shape = be.to_backend_array(fast_shape)
+    targetshape = be.to_backend_array(targetshape)
+    templateshape = be.to_backend_array(templateshape)
+
+    # Wrap peaks around score space
+    if fourier_shift is not None:
+        fourier_shift = be.to_backend_array(fourier_shift)
+        positions = be.add(positions, fourier_shift)
+        positions = be.mod(positions, fast_shape)
+
+    output_shape = _convmode_to_shape(
+        convolution_mode=convolution_mode,
+        targetshape=targetshape,
+        templateshape=templateshape,
+        convolution_shape=convolution_shape,
+    )
+    starts = be.astype(
+        be.divide(be.subtract(convolution_shape, output_shape), 2),
+        be._int_dtype,
+    )
+    stops = be.add(starts, output_shape)
+
+    valid_positions = be.multiply(positions >= starts, positions < stops)
+    valid_positions = be.sum(valid_positions, axis=1) == positions.shape[1]
+    positions = be.subtract(positions, starts)
+
+    return positions, valid_positions