pytme 0.2.9__cp311-cp311-macosx_15_0_arm64.whl

tests/test_matching_optimization.py

@@ -0,0 +1,226 @@
+import numpy as np
+import pytest
+
+from tme.rotations import euler_from_rotationmatrix
+from tme.matching_optimization import (
+    MATCHING_OPTIMIZATION_REGISTER,
+    register_matching_optimization,
+    _MatchCoordinatesToDensity,
+    _MatchCoordinatesToCoordinates,
+    optimize_match,
+    create_score_object,
+)
+
+density_to_density = ["FLC"]
+
+coordinate_to_density = [
+    k
+    for k, v in MATCHING_OPTIMIZATION_REGISTER.items()
+    if issubclass(v, _MatchCoordinatesToDensity)
+]
+
+coordinate_to_coordinate = [
+    k
+    for k, v in MATCHING_OPTIMIZATION_REGISTER.items()
+    if issubclass(v, _MatchCoordinatesToCoordinates)
+]
+
+
+class TestMatchDensityToDensity:
+    def setup_method(self):
+        target = np.zeros((50, 50, 50))
+        target[20:30, 30:40, 12:17] = 1
+        self.target = target
+        self.template = target.copy()
+        self.template_mask = np.ones_like(target)
+
+    def teardown_method(self):
+        self.target = None
+        self.template = None
+        self.template_mask = None
+
+    @pytest.mark.parametrize("method", density_to_density)
+    def test_initialization(self, method: str, notest: bool = False):
+        instance = create_score_object(
+            score=method,
+            target=self.target,
+            template=self.template,
+            template_mask=self.template_mask,
+        )
+        if notest:
+            return instance
+
+    @pytest.mark.parametrize("method", density_to_density)
+    def test_call(self, method):
+        instance = self.test_initialization(method=method, notest=True)
+        score = instance()
+        assert isinstance(score, float)
+
+
+class TestMatchDensityToCoordinates:
+    def setup_method(self):
+        data = np.zeros((50, 50, 50))
+        data[20:30, 30:40, 12:17] = 1
+        self.target = data
+        self.target_mask_density = data > 0
+        self.coordinates = np.array(np.where(self.target > 0))
+        self.coordinates_weights = self.target[tuple(self.coordinates)]
+
+        np.random.seed(42)
+        random_pixels = np.random.choice(
+            range(self.coordinates.shape[1]), self.coordinates.shape[1] // 2
+        )
+        self.coordinates_mask = self.coordinates[:, random_pixels]
+
+        self.origin = np.zeros(self.coordinates.shape[0])
+        self.sampling_rate = np.ones(self.coordinates.shape[0])
+
+    def teardown_method(self):
+        self.target = None
+        self.target_mask_density = None
+        self.coordinates = None
+        self.coordinates_weights = None
+        self.coordinates_mask = None
+
+    @pytest.mark.parametrize("method", coordinate_to_density)
+    def test_initialization(self, method: str, notest: bool = False):
+        instance = create_score_object(
+            score=method,
+            target=self.target,
+            target_mask=self.target_mask_density,
+            template_coordinates=self.coordinates,
+            template_weights=self.coordinates_weights,
+            template_mask_coordinates=self.coordinates_mask,
+        )
+        if notest:
+            return instance
+
+    @pytest.mark.parametrize("method", coordinate_to_density)
+    def test_call(self, method):
+        instance = self.test_initialization(method=method, notest=True)
+        score = instance()
+        assert isinstance(score, float)
+
+
+class TestMatchCoordinateToCoordinates:
+    def setup_method(self):
+        data = np.zeros((50, 50, 50))
+        data[20:30, 30:40, 12:17] = 1
+        self.target_coordinates = np.array(np.where(data > 0))
+        self.target_weights = data[tuple(self.target_coordinates)]
+
+        self.coordinates = np.array(np.where(data > 0))
+        self.coordinates_weights = data[tuple(self.coordinates)]
+
+        self.origin = np.zeros(self.coordinates.shape[0])
+        self.sampling_rate = np.ones(self.coordinates.shape[0])
+
+    def teardown_method(self):
+        self.target_coordinates = None
+        self.target_weights = None
+        self.coordinates = None
+        self.coordinates_weights = None
+
+    @pytest.mark.parametrize("method", coordinate_to_coordinate)
+    def test_initialization(self, method: str, notest: bool = False):
+        instance = create_score_object(
+            score=method,
+            target_coordinates=self.target_coordinates,
+            target_weights=self.target_weights,
+            template_coordinates=self.coordinates,
+            template_weights=self.coordinates_weights,
+        )
+        if notest:
+            return instance
+
+    @pytest.mark.parametrize("method", coordinate_to_coordinate)
+    def test_call(self, method):
+        instance = self.test_initialization(method=method, notest=True)
+        score = instance()
+        assert isinstance(score, float)
+
+
+class TestOptimizeMatch:
+    def setup_method(self):
+        data = np.zeros((50, 50, 50))
+        data[20:30, 30:40, 12:17] = 1
+        self.target = data
+        self.coordinates = np.array(np.where(self.target > 0))
+        self.coordinates_weights = self.target[tuple(self.coordinates)]
+
+        self.origin = np.zeros(self.coordinates.shape[0])
+        self.sampling_rate = np.ones(self.coordinates.shape[0])
+
+        self.score_object = MATCHING_OPTIMIZATION_REGISTER["CrossCorrelation"]
+        self.score_object = self.score_object(
+            target=self.target,
+            template_coordinates=self.coordinates,
+            template_weights=self.coordinates_weights,
+        )
+
+    def teardown_method(self):
+        self.target = None
+        self.coordinates = None
+        self.coordinates_weights = None
+
+    @pytest.mark.parametrize(
+        "method", ("differential_evolution", "basinhopping", "minimize")
+    )
+    @pytest.mark.parametrize("bound_translation", (True, False))
+    @pytest.mark.parametrize("bound_rotation", (True, False))
+    def test_call(self, method, bound_translation, bound_rotation):
+        if bound_rotation:
+            bound_rotation = tuple((-90, 90) for _ in range(self.target.ndim))
+        else:
+            bound_rotation = None
+
+        if bound_translation:
+            bound_translation = tuple((-5, 5) for _ in range(self.target.ndim))
+        else:
+            bound_translation = None
+
+        translation, rotation, score = optimize_match(
+            score_object=self.score_object,
+            optimization_method=method,
+            bounds_rotation=bound_rotation,
+            bounds_translation=bound_translation,
+            maxiter=10,
+        )
+        assert translation.size == self.target.ndim
+        assert rotation.shape[0] == self.target.ndim
+        assert rotation.shape[1] == self.target.ndim
+        assert isinstance(score, float)
+
+        if bound_translation is not None:
+            lower_bound = np.array([x[0] for x in bound_translation])
+            upper_bound = np.array([x[1] for x in bound_translation])
+            assert np.all(
+                np.logical_and(translation >= lower_bound, translation <= upper_bound)
+            )
+
+        if bound_rotation is not None:
+            angles = euler_from_rotationmatrix(rotation)
+            lower_bound = np.array([x[0] for x in bound_rotation])
+            upper_bound = np.array([x[1] for x in bound_rotation])
+            assert np.all(np.logical_and(angles >= lower_bound, angles <= upper_bound))
+
+    def test_call_error(self):
+        with pytest.raises(ValueError):
+            translation, rotation, score = optimize_match(
+                score_object=self.score_object,
+                optimization_method="RAISERROR",
+                maxiter=10,
+            )
+
+
+class TestUtils:
+    def test_register_matching_optimization(self):
+        new_class = list(MATCHING_OPTIMIZATION_REGISTER.keys())[0]
+        register_matching_optimization(
+            match_name="new_score",
+            match_class=MATCHING_OPTIMIZATION_REGISTER[new_class],
+        )
+
+    def test_register_matching_optimization_error(self):
+        with pytest.raises(ValueError):
+            register_matching_optimization(match_name="new_score", match_class=None)

tests/test_matching_utils.py

@@ -0,0 +1,189 @@
+from tempfile import mkstemp
+from importlib_resources import files
+
+import pytest
+import numpy as np
+from scipy.signal import correlate
+
+from tme import Density
+from tme.backends import backend as be
+from tme.memory import MATCHING_MEMORY_REGISTRY
+from tme.matching_utils import (
+    compute_parallelization_schedule,
+    elliptical_mask,
+    box_mask,
+    tube_mask,
+    create_mask,
+    scramble_phases,
+    apply_convolution_mode,
+    write_pickle,
+    load_pickle,
+    _normalize_template_overflow_safe,
+)
+
+BASEPATH = files("tests.data")
+
+
+class TestMatchingUtils:
+    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("matching_method", list(MATCHING_MEMORY_REGISTRY.keys()))
+    @pytest.mark.parametrize("max_cores", range(1, 10, 3))
+    @pytest.mark.parametrize("max_ram", [1e5, 1e7, 1e9])
+    def test_compute_parallelization_schedule(
+        self, matching_method, max_cores, max_ram
+    ):
+        max_cores, max_ram = int(max_cores), int(max_ram)
+        compute_parallelization_schedule(
+            self.density.shape,
+            self.structure_density.shape,
+            matching_method=matching_method,
+            max_cores=max_cores,
+            max_ram=max_ram,
+            max_splits=256,
+        )
+
+    def test_create_mask(self):
+        create_mask(
+            mask_type="ellipse",
+            shape=self.density.shape,
+            radius=5,
+            center=np.divide(self.density.shape, 2),
+        )
+
+    def test_create_mask_error(self):
+        with pytest.raises(ValueError):
+            create_mask(mask_type=None)
+
+    def test_elliptical_mask(self):
+        elliptical_mask(
+            shape=self.density.shape,
+            radius=5,
+            center=np.divide(self.density.shape, 2),
+        )
+
+    def test_box_mask(self):
+        box_mask(
+            shape=self.density.shape,
+            height=[5, 10, 20],
+            center=np.divide(self.density.shape, 2),
+        )
+
+    def test_tube_mask(self):
+        tube_mask(
+            shape=self.density.shape,
+            outer_radius=10,
+            inner_radius=5,
+            height=5,
+            base_center=np.divide(self.density.shape, 2),
+            symmetry_axis=1,
+        )
+
+    def test_tube_mask_error(self):
+        with pytest.raises(ValueError):
+            tube_mask(
+                shape=self.density.shape,
+                outer_radius=5,
+                inner_radius=10,
+                height=5,
+                base_center=np.divide(self.density.shape, 2),
+                symmetry_axis=1,
+            )
+
+        with pytest.raises(ValueError):
+            tube_mask(
+                shape=self.density.shape,
+                outer_radius=5,
+                inner_radius=10,
+                height=10 * np.max(self.density.shape),
+                base_center=np.divide(self.density.shape, 2),
+                symmetry_axis=1,
+            )
+
+        with pytest.raises(ValueError):
+            tube_mask(
+                shape=self.density.shape,
+                outer_radius=5,
+                inner_radius=10,
+                height=10 * np.max(self.density.shape),
+                base_center=np.divide(self.density.shape, 2),
+                symmetry_axis=len(self.density.shape) + 1,
+            )
+
+    def test_scramble_phases(self):
+        scramble_phases(arr=self.density.data, noise_proportion=0.5)
+
+    @pytest.mark.parametrize("convolution_mode", ["full", "valid", "same"])
+    def test_apply_convolution_mode(self, convolution_mode):
+        correlation = correlate(
+            self.density.data, self.structure_density.data, method="direct", mode="full"
+        )
+        ret = apply_convolution_mode(
+            arr=correlation,
+            convolution_mode=convolution_mode,
+            s1=self.density.shape,
+            s2=self.structure_density.shape,
+        )
+        if convolution_mode == "full":
+            expected_size = correlation.shape
+        elif convolution_mode == "same":
+            expected_size = self.density.shape
+        else:
+            expected_size = np.subtract(
+                self.density.shape, self.structure_density.shape
+            )
+            expected_size += np.mod(self.structure_density.shape, 2)
+        assert np.allclose(ret.shape, expected_size)
+
+    def test_apply_convolution_mode_error(self):
+        correlation = correlate(
+            self.density.data, self.structure_density.data, method="direct", mode="full"
+        )
+        with pytest.raises(ValueError):
+            _ = apply_convolution_mode(
+                arr=correlation,
+                convolution_mode=None,
+                s1=self.density.shape,
+                s2=self.structure_density.shape,
+            )
+
+    def test_pickle_io(self):
+        _, filename = mkstemp()
+
+        data = ["Hello", 123, np.array([1, 2, 3])]
+        write_pickle(data=data, filename=filename)
+        loaded_data = load_pickle(filename)
+        assert all([np.array_equal(a, b) for a, b in zip(data, loaded_data)])
+
+        data = 42
+        write_pickle(data=data, filename=filename)
+        loaded_data = load_pickle(filename)
+        assert loaded_data == data
+
+        _, filename = mkstemp()
+        data = np.memmap(filename, dtype="float32", mode="w+", shape=(3,))
+        data[:] = [1.1, 2.2, 3.3]
+        data.flush()
+        data = np.memmap(filename, dtype="float32", mode="r+", shape=(3,))
+        _, filename = mkstemp()
+        write_pickle(data=data, filename=filename)
+        loaded_data = load_pickle(filename)
+        assert np.array_equal(loaded_data, data)
+
+    def test_normalize_template_overflow_safe(self):
+        template = be.random.random((10, 10)).astype(be.float32)
+        mask = be.ones_like(template)
+        n_observations = 100.0
+
+        result = _normalize_template_overflow_safe(template, mask, n_observations)
+        assert result.shape == template.shape
+        assert result.dtype == template.dtype
+        assert np.allclose(result.mean(), 0, atol=0.1)
+        assert np.allclose(result.std(), 1, atol=0.1)

tests/test_orientations.py

@@ -0,0 +1,175 @@
+from tempfile import mkstemp
+
+import pytest
+import numpy as np
+
+from tme import Orientations
+
+
+class TestDensity:
+    def setup_method(self):
+        self.translations = np.random.rand(100, 3).astype(np.float32)
+        self.rotations = np.random.rand(100, 3).astype(np.float32)
+        self.scores = np.random.rand(100).astype(np.float32)
+        self.details = np.random.rand(100).astype(np.float32)
+
+        self.orientations = Orientations(
+            translations=self.translations,
+            rotations=self.rotations,
+            scores=self.scores,
+            details=self.details,
+        )
+
+    def teardown_method(self):
+        self.translations = None
+        self.rotations = None
+        self.scores = None
+        self.details = None
+        self.orientations = None
+
+    def test_initialization(self):
+        orientations = Orientations(
+            translations=self.translations,
+            rotations=self.rotations,
+            scores=self.scores,
+            details=self.details,
+        )
+
+        assert np.array_equal(self.translations, orientations.translations)
+        assert np.array_equal(self.rotations, orientations.rotations)
+        assert np.array_equal(self.scores, orientations.scores)
+        assert np.array_equal(self.details, orientations.details)
+
+    def test_initialization_type(self):
+        orientations = Orientations(
+            translations=self.translations.astype(int),
+            rotations=self.rotations.astype(int),
+            scores=self.scores.astype(int),
+            details=self.details.astype(int),
+        )
+        assert np.issubdtype(orientations.translations.dtype, np.floating)
+        assert np.issubdtype(orientations.rotations.dtype, np.floating)
+        assert np.issubdtype(orientations.scores.dtype, np.floating)
+        assert np.issubdtype(orientations.details.dtype, np.floating)
+
+    def test_initialization_error(self):
+        with pytest.raises(ValueError):
+            _ = Orientations(
+                translations=self.translations,
+                rotations=np.random.rand(self.translations.shape[0] + 1),
+                scores=self.scores,
+                details=self.details,
+            )
+
+        with pytest.raises(ValueError):
+            _ = Orientations(
+                translations=np.random.rand(self.translations.shape[0]),
+                rotations=np.random.rand(self.translations.shape[0] + 1),
+                scores=self.scores,
+                details=self.details,
+            )
+            _ = Orientations(
+                translations=self.translations,
+                rotations=np.random.rand(self.translations.shape[0]),
+                scores=self.scores,
+                details=self.details,
+            )
+
+        assert True
+
+    @pytest.mark.parametrize("file_format", ("text", "relion", "tbl"))
+    def test_to_file(self, file_format: str):
+        _, output_file = mkstemp(suffix=f".{file_format}")
+        self.orientations.to_file(output_file)
+        assert True
+
+    @pytest.mark.parametrize("file_format", ("text", "star", "tbl"))
+    def test_from_file(self, file_format: str):
+        _, output_file = mkstemp(suffix=f".{file_format}")
+        self.orientations.to_file(output_file)
+        orientations_new = Orientations.from_file(output_file)
+
+        assert np.allclose(
+            self.orientations.translations, orientations_new.translations
+        )
+        assert np.allclose(
+            self.orientations.rotations, orientations_new.rotations, atol=1e-3
+        )
+
+    @pytest.mark.parametrize("input_format", ("text", "star", "tbl"))
+    @pytest.mark.parametrize("output_format", ("text", "star", "tbl"))
+    def test_file_format_io(self, input_format: str, output_format: str):
+        _, output_file = mkstemp(suffix=f".{input_format}")
+        _, output_file2 = mkstemp(suffix=f".{output_format}")
+
+        self.orientations.to_file(output_file)
+        orientations_new = Orientations.from_file(output_file)
+        orientations_new.to_file(output_file2)
+
+        assert True
+
+    @pytest.mark.parametrize("drop_oob", (True, False))
+    @pytest.mark.parametrize("shape", (10, 40, 80))
+    @pytest.mark.parametrize("odd", (True, False))
+    def test_extraction(self, shape: int, drop_oob: bool, odd: bool):
+        if odd:
+            shape = shape + (1 - shape % 2)
+
+        data = np.random.rand(50, 50, 50)
+        translations = np.array([[25, 25, 25], [15, 25, 35], [35, 25, 15], [0, 15, 49]])
+        orientations = Orientations(
+            translations=translations,
+            rotations=np.random.rand(*translations.shape),
+            scores=np.random.rand(translations.shape[0]),
+            details=np.random.rand(translations.shape[0]),
+        )
+        extraction_shape = np.repeat(np.array(shape), data.ndim)
+        orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
+            target_shape=data.shape,
+            extraction_shape=extraction_shape,
+            drop_out_of_box=drop_oob,
+            return_orientations=True,
+        )
+        assert orientations.translations.shape[0] == len(cand_slices)
+        assert len(cand_slices) == len(obs_slices)
+
+        cand_slices2, obs_slices2 = orientations.get_extraction_slices(
+            target_shape=data.shape,
+            extraction_shape=extraction_shape,
+            drop_out_of_box=drop_oob,
+            return_orientations=False,
+        )
+        assert cand_slices == cand_slices2
+        assert obs_slices == obs_slices2
+
+        # Check whether extraction slices are pasted in center
+        out = np.zeros(extraction_shape, dtype=data.dtype)
+        center = np.divide(extraction_shape, 2).astype(int)
+        for index, (cand_slice, obs_slice) in enumerate(zip(cand_slices, obs_slices)):
+            out[cand_slice] = data[obs_slice]
+            assert np.allclose(
+                out[tuple(center)],
+                data[tuple(orientations.translations[index].astype(int))],
+            )
+
+    @pytest.mark.parametrize(
+        "order", (("x", "y", "z"), ("z", "y", "x"), ("y", "x", "z"))
+    )
+    def test_txt_sort(self, order: str):
+        _, output_file = mkstemp(suffix=".tsv")
+        translations = ((50, 30, 20), (10, 5, 30))
+
+        with open(output_file, mode="w", encoding="utf-8") as ofile:
+            ofile.write("\t".join([str(x) for x in order]) + "\n")
+            for translation in translations:
+                ofile.write("\t".join([str(x) for x in translation]) + "\n")
+
+        translations = np.array(translations).astype(np.float32)
+        orientations = Orientations.from_file(output_file)
+
+        out_order = zip(order, range(len(order)))
+        out_order = tuple(
+            x[1] for x in sorted(out_order, key=lambda x: x[0], reverse=False)
+        )
+
+        assert np.array_equal(translations[..., out_order], orientations.translations)

tests/test_parser.py

@@ -0,0 +1,33 @@
+import pytest
+from importlib_resources import files
+
+from tme.parser import Parser, PDBParser
+
+
+class TestParser:
+    def setup_method(self):
+        self.pdb_file = str(files("tests.data").joinpath("Structures/5khe.pdb"))
+
+    def teardown_method(self):
+        self.pdb_file = None
+
+    def test_initialize_parser_error(self):
+        with pytest.raises(TypeError):
+            _ = Parser(self.pdb_file)
+
+    def test_parser_get(self):
+        parser = PDBParser(self.pdb_file)
+        assert parser.get("NOTPRESENT", None) is None
+        assert parser.get("record_type", None) is not None
+
+    def test_parser_keys(self):
+        parser = PDBParser(self.pdb_file)
+        assert parser.keys() == parser._data.keys()
+
+    def test_parser_values(self):
+        parser = PDBParser(self.pdb_file)
+        assert str(parser.values()) == str(parser._data.values())
+
+    def test_parser_items(self):
+        parser = PDBParser(self.pdb_file)
+        assert parser.items() == parser._data.items()