pytme 0.2.2__cp311-cp311-macosx_14_0_arm64.whl → 0.2.4__cp311-cp311-macosx_14_0_arm64.whl

tests/test_matching_cli.py

@@ -0,0 +1,283 @@
+import tempfile
+import subprocess
+from shutil import rmtree
+from os.path import exists
+from os import remove, makedirs
+
+import pytest
+import numpy as np
+from tme import Density
+from tme.backends import backend as be
+
+BACKEND_CLASSES = ["NumpyFFTWBackend", "PytorchBackend", "CupyBackend", "MLXBackend"]
+BACKENDS_TO_TEST = []
+
+test_gpu = (False,)
+for backend_class in BACKEND_CLASSES:
+    try:
+        BackendClass = getattr(
+            __import__("tme.backends", fromlist=[backend_class]), backend_class
+        )
+        BACKENDS_TO_TEST.append(BackendClass(device="cpu"))
+        if backend_class == "CupyBackend":
+            if BACKENDS_TO_TEST[-1].device_count() >= 1:
+                test_gpu = (False, True)
+    except ImportError:
+        print(f"Couldn't import {backend_class}. Skipping...")
+
+
+available_backends = (x for x in be.available_backends() if x != "mlx")
+
+
+def argdict_to_command(input_args, executable: str):
+    ret = []
+    for key, value in input_args.items():
+        if value is None:
+            continue
+        elif isinstance(value, bool):
+            if value:
+                ret.append(key)
+        else:
+            ret.extend([key, value])
+
+    ret = [str(x) for x in ret]
+    ret.insert(0, executable)
+    return " ".join(ret)
+
+
+class TestMatchTemplate:
+    @classmethod
+    def setup_class(cls):
+        np.random.seed(42)
+        target = np.random.rand(20, 20, 20)
+        template = np.random.rand(5, 5, 5)
+
+        target_mask = 1.0 * (target > 0.5)
+        template_mask = 1.0 * (template > 0.5)
+
+        cls.target_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mrc").name
+        cls.template_path = tempfile.NamedTemporaryFile(
+            delete=False, suffix=".mrc"
+        ).name
+        cls.target_mask_path = tempfile.NamedTemporaryFile(
+            delete=False, suffix=".mrc"
+        ).name
+        cls.template_mask_path = tempfile.NamedTemporaryFile(
+            delete=False, suffix=".mrc"
+        ).name
+
+        Density(target, sampling_rate=5).to_file(cls.target_path)
+        Density(template, sampling_rate=5).to_file(cls.template_path)
+        Density(target_mask, sampling_rate=5).to_file(cls.target_mask_path)
+        Density(template_mask, sampling_rate=5).to_file(cls.template_mask_path)
+
+    def teardown_class(cls):
+        cls.try_delete(cls.target_path)
+        cls.try_delete(cls.template_path)
+        cls.try_delete(cls.target_mask_path)
+        cls.try_delete(cls.template_mask_path)
+
+    @staticmethod
+    def try_delete(file_path: str):
+        try:
+            remove(file_path)
+        except Exception:
+            pass
+        try:
+            rmtree(file_path, ignore_errors=True)
+        except Exception:
+            pass
+
+    @staticmethod
+    def run_matching(
+        use_template_mask: bool,
+        test_filter: bool,
+        call_peaks: bool,
+        target_path: str,
+        template_path: str,
+        template_mask_path: str,
+        target_mask_path: str,
+        use_target_mask: bool = False,
+        backend: str = "numpyfftw",
+    ):
+        output_path = tempfile.NamedTemporaryFile(delete=False, suffix="pickle").name
+
+        argdict = {
+            "-m": target_path,
+            "-i": template_path,
+            "-n": 1,
+            "-a": 60,
+            "-o": output_path,
+            "--pad_edges": False,
+            "--pad_fourier": False,
+            "--backend": backend,
+        }
+
+        if use_template_mask:
+            argdict["--template_mask"] = template_mask_path
+
+        if use_target_mask:
+            argdict["--target_mask"] = target_mask_path
+
+        if backend in ("cupy", "pytorch") and True in test_gpu:
+            argdict["--use_gpu"] = True
+
+        if test_filter:
+            argdict["--lowpass"] = 30
+            argdict["--defocus"] = 3000
+            argdict["--tilt_angles"] = "40,40:10"
+            argdict["--wedge_axes"] = "0,2"
+            argdict["--whiten"] = True
+
+        if call_peaks:
+            argdict["-p"] = True
+
+        cmd = argdict_to_command(argdict, executable="match_template.py")
+        ret = subprocess.run(cmd, capture_output=True, shell=True)
+        print(ret)
+        assert ret.returncode == 0
+        return output_path
+
+    @pytest.mark.parametrize("backend", available_backends)
+    @pytest.mark.parametrize("call_peaks", (False, True))
+    @pytest.mark.parametrize("use_template_mask", (False, True))
+    @pytest.mark.parametrize("test_filter", (False, True))
+    def test_match_template(
+        self,
+        backend: bool,
+        call_peaks: bool,
+        use_template_mask: bool,
+        test_filter: bool,
+    ):
+        self.run_matching(
+            use_template_mask=use_template_mask,
+            use_target_mask=True,
+            backend=backend,
+            test_filter=test_filter,
+            call_peaks=call_peaks,
+            template_path=self.template_path,
+            target_path=self.target_path,
+            template_mask_path=self.template_mask_path,
+            target_mask_path=self.target_mask_path,
+        )
+
+
+class TestPostprocessing(TestMatchTemplate):
+    @classmethod
+    def setup_class(cls):
+        super().setup_class()
+
+        matching_kwargs = {
+            "use_template_mask": False,
+            "use_target_mask": False,
+            "test_filter": False,
+            "template_path": cls.template_path,
+            "target_path": cls.target_path,
+            "template_mask_path": cls.template_mask_path,
+            "target_mask_path": cls.target_mask_path,
+        }
+
+        cls.score_pickle = cls.run_matching(
+            call_peaks=False,
+            **matching_kwargs,
+        )
+        cls.peak_pickle = cls.run_matching(call_peaks=True, **matching_kwargs)
+        cls.tempdir = tempfile.TemporaryDirectory().name
+
+    @classmethod
+    def teardown_class(cls):
+        cls.try_delete(cls.score_pickle)
+        cls.try_delete(cls.peak_pickle)
+        cls.try_delete(cls.tempdir)
+
+    @pytest.mark.parametrize("distance_cutoff_strategy", (0, 1, 2, 3))
+    @pytest.mark.parametrize("score_cutoff", (None, (1,), (0, 1), (None, 1), (0, None)))
+    @pytest.mark.parametrize("peak_oversampling", (False, 4))
+    def test_postprocess_score_orientations(
+        self,
+        peak_oversampling,
+        score_cutoff,
+        distance_cutoff_strategy,
+    ):
+        self.try_delete(self.tempdir)
+        makedirs(self.tempdir, exist_ok=True)
+
+        argdict = {
+            "--input_file": self.score_pickle,
+            "--output_format": "orientations",
+            "--output_prefix": f"{self.tempdir}/temp",
+            "--peak_oversampling": peak_oversampling,
+            "--number_of_peaks": 3,
+        }
+
+        if score_cutoff is not None:
+            if len(score_cutoff) == 1:
+                argdict["--n_false_positives"] = 1
+            else:
+                min_score, max_score = score_cutoff
+                argdict["--minimum_score"] = min_score
+                argdict["--maximum_score"] = max_score
+
+        match distance_cutoff_strategy:
+            case 1:
+                argdict["--mask_edges"] = True
+            case 2:
+                argdict["--min_distance"] = 5
+            case 3:
+                argdict["--min_boundary_distance"] = 5
+
+        cmd = argdict_to_command(argdict, executable="postprocess.py")
+        ret = subprocess.run(cmd, capture_output=True, shell=True)
+        assert ret.returncode == 0
+
+    @pytest.mark.parametrize("input_format", ("score", "peaks"))
+    @pytest.mark.parametrize(
+        "output_format",
+        ("orientations", "alignment", "backmapping", "average", "relion"),
+    )
+    def test_postproces_score_formats(self, input_format, output_format):
+        self.try_delete(self.tempdir)
+        makedirs(self.tempdir, exist_ok=True)
+
+        input_file = self.score_pickle
+        if input_format == "peaks":
+            input_file = self.peak_pickle
+
+        argdict = {
+            "--input_file": input_file,
+            "--output_format": output_format,
+            "--output_prefix": f"{self.tempdir}/temp",
+            "--number_of_peaks": 3,
+            "--peak_caller": "PeakCallerMaximumFilter",
+        }
+        cmd = argdict_to_command(argdict, executable="postprocess.py")
+        ret = subprocess.run(cmd, capture_output=True, shell=True)
+
+        match output_format:
+            case "orientations":
+                assert exists(f"{self.tempdir}/temp.tsv")
+            case "alignment":
+                assert exists(f"{self.tempdir}/temp_0.mrc")
+            case "backmapping":
+                assert exists(f"{self.tempdir}/temp_backmapped.mrc")
+            case "average":
+                assert exists(f"{self.tempdir}/temp_average.mrc")
+            case "relion":
+                assert exists(f"{self.tempdir}/temp.star")
+
+        assert ret.returncode == 0
+
+    def test_postprocess_score_local_optimization(self):
+        self.try_delete(self.tempdir)
+        makedirs(self.tempdir, exist_ok=True)
+
+        argdict = {
+            "--input_file": self.score_pickle,
+            "--output_format": "orientations",
+            "--output_prefix": f"{self.tempdir}/temp",
+            "--number_of_peaks": 1,
+            "--local_optimization": True,
+        }
+        cmd = argdict_to_command(argdict, executable="postprocess.py")
+        ret = subprocess.run(cmd, capture_output=True, shell=True)
+        assert ret.returncode == 0

tests/test_matching_data.py

@@ -0,0 +1,162 @@
+from tempfile import mkstemp
+
+import pytest
+import numpy as np
+
+from tme.backends import backend as be
+from tme.matching_data import MatchingData
+
+
+class TestDensity:
+    def setup_method(self):
+        target = np.zeros((50, 50, 50))
+        target[20:30, 30:40, 12:17] = 1
+
+        self.target = target
+        template = np.zeros((50, 50, 50))
+        template[15:25, 20:30, 2:7] = 1
+        self.template = template
+        self.rotations = np.random.rand(100, target.ndim, target.ndim).astype(
+            np.float32
+        )
+
+    def teardown_method(self):
+        self.target = None
+        self.template = None
+        self.coordinates = None
+        self.coordinates_weights = None
+        self.rotations = None
+
+    def test_initialization(self):
+        _ = MatchingData(target=self.target, template=self.template)
+
+    @pytest.mark.parametrize("shape", [(10,), (10, 15), (10, 20, 30)])
+    def test__shape_to_slice(self, shape):
+        slices = MatchingData._shape_to_slice(shape=shape)
+        assert len(slices) == len(shape)
+        for i, k in enumerate(shape):
+            assert slices[i].start == 0
+            assert slices[i].stop == k
+
+    @pytest.mark.parametrize("shape", [(10,), (10, 15), (10, 20, 30)])
+    def test_slice_to_mesh(self, shape):
+        if shape is not None:
+            slices = MatchingData._shape_to_slice(shape=shape)
+
+        indices = MatchingData._slice_to_mesh(slice_variable=slices, shape=shape)
+        assert len(indices) == len(shape)
+        for i, k in enumerate(shape):
+            assert indices[i].min() == 0
+            assert indices[i].max() == k - 1
+
+        indices = MatchingData._slice_to_mesh(slice_variable=None, shape=shape)
+        assert len(indices) == len(shape)
+        for i, k in enumerate(shape):
+            assert indices[i].min() == 0
+            assert indices[i].max() == k - 1
+
+    def test__load_array(self):
+        arr = MatchingData._load_array(self.target)
+        assert np.allclose(arr, self.target)
+
+    def test__load_array_memmap(self):
+        _, filename = mkstemp()
+        shape, dtype = self.target.shape, self.target.dtype
+        arr_memmap = np.memmap(filename, mode="w+", dtype=dtype, shape=shape)
+        arr_memmap[:] = self.target[:]
+        arr_memmap.flush()
+
+        arr = MatchingData._load_array(arr_memmap)
+        assert np.allclose(arr, self.target)
+
+    def test_subset_array(self):
+        matching_data = MatchingData(target=self.target, template=self.template)
+        slices = MatchingData._shape_to_slice(
+            shape=np.divide(self.target.shape, 2).astype(int)
+        )
+        ret = matching_data.subset_array(
+            arr=self.target, arr_slice=slices, padding=(2, 2, 2)
+        )
+        assert np.allclose(
+            ret.shape, np.add(np.divide(self.target.shape, 2).astype(int), 2)
+        )
+
+    def test_subset_by_slice_none(self):
+        matching_data = MatchingData(target=self.target, template=self.template)
+        matching_data.rotations = self.rotations
+        matching_data.target_mask = self.target
+        matching_data.template_mask = self.template
+
+        ret = matching_data.subset_by_slice()
+
+        assert type(ret) == type(matching_data)
+        assert np.allclose(ret.target, matching_data.target)
+        assert np.allclose(ret.template, matching_data.template)
+        assert np.allclose(ret.target_mask, matching_data.target_mask)
+        assert np.allclose(ret.template_mask, matching_data.template_mask)
+
+    def test_subset_by_slice(self):
+        matching_data = MatchingData(target=self.target, template=self.template)
+        matching_data.rotations = self.rotations
+        matching_data.target_mask = self.target
+        matching_data.template_mask = self.template
+
+        target_slice = MatchingData._shape_to_slice(
+            shape=np.divide(self.target.shape, 2).astype(int)
+        )
+        template_slice = MatchingData._shape_to_slice(
+            shape=np.divide(self.template.shape, 2).astype(int)
+        )
+        ret = matching_data.subset_by_slice(
+            target_slice=target_slice, template_slice=template_slice
+        )
+        assert type(ret) == type(matching_data)
+
+        assert np.allclose(
+            ret.target.shape, np.divide(self.target.shape, 2).astype(int)
+        )
+        assert np.allclose(
+            ret.template.shape, np.divide(self.target.shape, 2).astype(int)[::-1]
+        )
+
+    def test_rotations(self):
+        matching_data = MatchingData(target=self.target, template=self.template)
+        matching_data.rotations = self.rotations
+        matching_data.target_mask = self.target
+
+        assert np.allclose(matching_data.rotations, self.rotations)
+
+        matching_data.rotations = np.random.rand(self.target.ndim, self.target.ndim)
+        assert np.allclose(
+            matching_data.rotations.shape, (1, self.target.ndim, self.target.ndim)
+        )
+
+    def test_target(self):
+        matching_data = MatchingData(target=self.target, template=self.template)
+
+        assert np.allclose(matching_data.target, self.target)
+
+    def test_template(self):
+        matching_data = MatchingData(target=self.target, template=self.template)
+
+        assert np.allclose(matching_data.template, be.reverse(self.template))
+
+    def test_target_mask(self):
+        matching_data = MatchingData(target=self.target, template=self.template)
+        matching_data.target_mask = self.target
+
+        assert np.allclose(matching_data.target_mask, self.target)
+
+    def test_template_mask(self):
+        matching_data = MatchingData(target=self.target, template=self.template)
+        matching_data.template_mask = self.template
+
+        assert np.allclose(matching_data.template_mask, be.reverse(self.template))
+
+    @pytest.mark.parametrize("jobs", range(1, 50, 5))
+    def test__split_rotations_on_jobs(self, jobs):
+        matching_data = MatchingData(target=self.target, template=self.template)
+        matching_data.rotations = self.rotations
+
+        ret = matching_data._split_rotations_on_jobs(n_jobs=jobs)
+        assert len(ret) == jobs

tests/test_matching_exhaustive.py

@@ -0,0 +1,162 @@
+import numpy as np
+import pytest
+
+from scipy.ndimage import laplace
+
+from tme.matching_data import MatchingData
+from tme.memory import MATCHING_MEMORY_REGISTRY
+from tme.matching_utils import get_rotation_matrices
+from tme.analyzer import MaxScoreOverRotations, PeakCallerSort
+from tme.matching_exhaustive import (
+    scan,
+    scan_subsets,
+    MATCHING_EXHAUSTIVE_REGISTER,
+    register_matching_exhaustive,
+)
+
+
+class TestMatchExhaustive:
+    def setup_method(self):
+        # To be valid for splitting, the template needs to be fully inside the object
+        target = np.zeros((80, 80, 80))
+        target[25:31, 22:28, 12:16] = 1
+
+        self.target = target
+        self.template = np.zeros((41, 41, 35))
+        self.template[20:26, 25:31, 17:21] = 1
+        self.template_mask = np.ones_like(self.template)
+        self.target_mask = np.ones_like(target)
+
+        self.rotations = get_rotation_matrices(60)[0,]
+        self.peak_position = np.array([25, 17, 12])
+
+    def teardown_method(self):
+        self.target = None
+        self.template = None
+        self.coordinates = None
+        self.coordinates_weights = None
+        self.rotations = None
+
+    @pytest.mark.parametrize("score", list(MATCHING_EXHAUSTIVE_REGISTER.keys()))
+    @pytest.mark.parametrize("n_jobs", (1, 2))
+    def test_scan(self, score: str, n_jobs: int):
+        matching_data = MatchingData(
+            target=self.target,
+            template=self.template,
+            target_mask=self.target_mask,
+            template_mask=self.template_mask,
+            rotations=self.rotations,
+        )
+        setup, process = MATCHING_EXHAUSTIVE_REGISTER[score]
+        ret = scan(
+            matching_data=matching_data,
+            matching_setup=setup,
+            matching_score=process,
+            n_jobs=n_jobs,
+            callback_class=MaxScoreOverRotations,
+        )
+        scores = ret[0]
+        peak = np.unravel_index(np.argmax(scores), scores.shape)
+
+        theoretical_score = 1
+        if score == "CC":
+            theoretical_score = self.template.sum()
+        elif score == "LCC":
+            theoretical_score = (laplace(self.template) * laplace(self.template)).sum()
+
+        assert np.allclose(peak, self.peak_position)
+        assert np.allclose(scores[peak], theoretical_score, rtol=0.05)
+
+    @pytest.mark.parametrize("evaluate_peak", (False, True))
+    @pytest.mark.parametrize("score", tuple(MATCHING_EXHAUSTIVE_REGISTER.keys()))
+    @pytest.mark.parametrize("job_schedule", ((2, 1), (1, 1)))
+    @pytest.mark.parametrize("pad_fourier", (True, False))
+    @pytest.mark.parametrize("pad_edge", (True, False))
+    def test_scan_subset(
+        self,
+        score: str,
+        job_schedule: int,
+        evaluate_peak: bool,
+        pad_fourier: bool,
+        pad_edge: bool,
+    ):
+        matching_data = MatchingData(
+            target=self.target,
+            template=self.template,
+            target_mask=self.target_mask,
+            template_mask=self.template_mask,
+            rotations=self.rotations,
+        )
+
+        setup, process = MATCHING_EXHAUSTIVE_REGISTER[score]
+
+        target_splits = {}
+        if job_schedule[0] == 2:
+            target_splits = {0: 2 if i != 0 else 2 for i in range(self.target.ndim)}
+
+        callback_class = PeakCallerSort
+        if evaluate_peak:
+            callback_class = MaxScoreOverRotations
+
+        ret = scan_subsets(
+            matching_data=matching_data,
+            matching_setup=setup,
+            matching_score=process,
+            target_splits=target_splits,
+            job_schedule=job_schedule,
+            callback_class=callback_class,
+            pad_target_edges=pad_edge,
+            pad_fourier=pad_fourier,
+        )
+
+        if evaluate_peak:
+            scores = ret[0]
+            peak = np.unravel_index(np.argmax(scores), scores.shape)
+            achieved_score = scores[tuple(peak)]
+        else:
+            try:
+                peak, achieved_score = ret[0][0], ret[2][0]
+            except Exception as e:
+                return None
+
+        if not pad_edge:
+            # To be valid, the match needs to be fully within the target subset
+            return None
+
+        theoretical_score = 1
+        if score == "CC":
+            theoretical_score = self.template.sum()
+        elif score == "LCC":
+            theoretical_score = (laplace(self.template) * laplace(self.template)).sum()
+
+        if not np.allclose(peak, self.peak_position):
+            print(peak)
+            assert False
+
+        assert np.allclose(achieved_score, theoretical_score, rtol=0.3)
+
+    def test_register_matching_exhaustive(self):
+        setup, matching = MATCHING_EXHAUSTIVE_REGISTER[
+            list(MATCHING_EXHAUSTIVE_REGISTER.keys())[0]
+        ]
+        memory_class = MATCHING_MEMORY_REGISTRY[
+            list(MATCHING_EXHAUSTIVE_REGISTER.keys())[0]
+        ]
+        register_matching_exhaustive(
+            matching="TEST",
+            matching_setup=setup,
+            matching_scoring=matching,
+            memory_class=memory_class,
+        )
+
+    def test_register_matching_exhaustive_error(self):
+        key = list(MATCHING_EXHAUSTIVE_REGISTER.keys())[0]
+        setup, matching = MATCHING_EXHAUSTIVE_REGISTER[key]
+        memory_class = MATCHING_MEMORY_REGISTRY[key]
+        with pytest.raises(ValueError):
+            register_matching_exhaustive(
+                matching=key,
+                matching_setup=setup,
+                matching_scoring=matching,
+                memory_class=memory_class,
+            )

tests/test_matching_memory.py

@@ -0,0 +1,30 @@
+import pytest
+from importlib_resources import files
+
+from tme import Density
+from tme.memory import MATCHING_MEMORY_REGISTRY, estimate_ram_usage
+
+BASEPATH = files("tests.data")
+
+
+class TestMatchingMemory:
+    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("analyzer_method", ["MaxScoreOverRotations", None])
+    @pytest.mark.parametrize("matching_method", list(MATCHING_MEMORY_REGISTRY.keys()))
+    @pytest.mark.parametrize("ncores", range(1, 10, 3))
+    def test_estimate_ram_usage(self, matching_method, ncores, analyzer_method):
+        estimate_ram_usage(
+            shape1=self.density.shape,
+            shape2=self.structure_density.shape,
+            matching_method=matching_method,
+            ncores=ncores,
+            analyzer_method=analyzer_method,
+        )