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

tests/preprocessing/test_frequency_filters.py

@@ -4,13 +4,15 @@ from typing import Tuple
 
 from tme.backends import backend as be
 from tme.filters._utils import compute_fourier_shape
-from tme.filters import BandPassFilter, LinearWhiteningFilter
+from tme.filters import BandPassReconstructed, LinearWhiteningFilter
+from tme.filters.bandpass import gaussian_bandpass, discrete_bandpass
+from tme.filters._utils import fftfreqn
 
 
 class TestBandPassFilter:
     @pytest.fixture
     def band_pass_filter(self):
-        return BandPassFilter()
+        return BandPassReconstructed()
 
     @pytest.mark.parametrize(
         "shape, lowpass, highpass, sampling_rate",
@@ -24,9 +26,13 @@ class TestBandPassFilter:
     def test_discrete_bandpass(
         self, shape: Tuple[int], lowpass: float, highpass: float, sampling_rate: float
     ):
-        result = BandPassFilter.discrete_bandpass(
-            shape, lowpass, highpass, sampling_rate
+        grid = fftfreqn(
+            shape=shape,
+            sampling_rate=0.5,
+            shape_is_real_fourier=False,
+            compute_euclidean_norm=True,
         )
+        result = discrete_bandpass(grid, lowpass, highpass, sampling_rate)
         assert isinstance(result, type(be.ones((1,))))
         assert result.shape == shape
         assert np.all((result >= 0) & (result <= 1))
@@ -43,9 +49,13 @@ class TestBandPassFilter:
     def test_gaussian_bandpass(
         self, shape: Tuple[int], lowpass: float, highpass: float, sampling_rate: float
     ):
-        result = BandPassFilter.gaussian_bandpass(
-            shape, lowpass, highpass, sampling_rate
+        grid = fftfreqn(
+            shape=shape,
+            sampling_rate=0.5,
+            shape_is_real_fourier=False,
+            compute_euclidean_norm=True,
         )
+        result = gaussian_bandpass(grid, lowpass, highpass, sampling_rate)
         assert isinstance(result, type(be.ones((1,))))
         assert result.shape == shape
         assert np.all((result >= 0) & (result <= 1))
@@ -55,7 +65,7 @@ class TestBandPassFilter:
     @pytest.mark.parametrize("shape_is_real_fourier", [True, False])
     def test_call_method(
         self,
-        band_pass_filter: BandPassFilter,
+        band_pass_filter: BandPassReconstructed,
         use_gaussian: bool,
         return_real_fourier: bool,
         shape_is_real_fourier: bool,
@@ -68,22 +78,20 @@ class TestBandPassFilter:
 
         assert isinstance(result, dict)
         assert "data" in result
-        assert "sampling_rate" in result
         assert "is_multiplicative_filter" in result
         assert isinstance(result["data"], type(be.ones((1,))))
         assert result["is_multiplicative_filter"] is True
 
-    def test_default_values(self, band_pass_filter: BandPassFilter):
+    def test_default_values(self, band_pass_filter: BandPassReconstructed):
         assert band_pass_filter.lowpass is None
         assert band_pass_filter.highpass is None
         assert band_pass_filter.sampling_rate == 1
         assert band_pass_filter.use_gaussian is True
         assert band_pass_filter.return_real_fourier is False
-        assert band_pass_filter.shape_is_real_fourier is False
 
     @pytest.mark.parametrize("shape", ((10, 10), (20, 20, 20), (30, 30)))
     def test_return_real_fourier(self, shape: Tuple[int]):
-        bpf = BandPassFilter(return_real_fourier=True)
+        bpf = BandPassReconstructed(return_real_fourier=True)
         result = bpf(shape=shape, lowpass=0.2, highpass=0.8)
         expected_shape = tuple(compute_fourier_shape(shape, False))
         assert result["data"].shape == expected_shape
@@ -146,7 +154,11 @@ class TestLinearWhiteningFilter:
     ):
         data = be.random.random(shape)
         result = LinearWhiteningFilter()(
-            data=data, n_bins=n_bins, batch_dimension=batch_dimension, order=order
+            shape=shape,
+            data=data,
+            n_bins=n_bins,
+            batch_dimension=batch_dimension,
+            order=order,
         )
 
         assert isinstance(result, dict)
@@ -161,7 +173,9 @@ class TestLinearWhiteningFilter:
     def test_call_method_with_data_rfft(self):
         shape = (30, 30, 30)
         data_rfft = be.fft.rfftn(be.random.random(shape))
-        result = LinearWhiteningFilter()(data_rfft=data_rfft)
+        result = LinearWhiteningFilter()(
+            shape=shape, data_rfft=data_rfft, return_real_fourier=True
+        )
 
         assert isinstance(result, dict)
         assert result.get("data", False) is not False
@@ -172,7 +186,7 @@ class TestLinearWhiteningFilter:
     @pytest.mark.parametrize("shape", [(10, 10), (20, 20, 20), (30, 30, 30)])
     def test_filter_mask_range(self, shape: Tuple[int]):
         data = be.random.random(shape)
-        result = LinearWhiteningFilter()(data=data)
+        result = LinearWhiteningFilter()(shape=shape, data=data)
 
         filter_mask = result["data"]
         assert np.all(filter_mask >= 0) and np.all(filter_mask <= 1)

tests/test_analyzer.py

@@ -1,5 +1,3 @@
-from tempfile import mkstemp
-
 import pytest
 import numpy as np
 
@@ -24,6 +22,7 @@ PEAK_CALLER_CHILDREN = [
     PeakCallerScipy,
     PeakClustering,
 ]
+np.random.seed(123)
 
 
 class TestPeakCallers:
@@ -54,21 +53,23 @@ class TestPeakCallers:
     @pytest.mark.parametrize("num_peaks", (1, 100))
     @pytest.mark.parametrize("minimum_score", (None, 0.5))
     def test__call__(self, peak_caller, num_peaks, minimum_score):
-        peak_caller = peak_caller(
-            shape=self.data.shape,
-            num_peaks=num_peaks,
-            min_distance=self.min_distance,
-            min_score=minimum_score,
-        )
-        peak_caller(
+        kwargs = {
+            "shape": self.data.shape,
+            "num_peaks": num_peaks,
+            "min_distance": self.min_distance,
+            "min_score": minimum_score,
+        }
+        peak_caller = peak_caller(**kwargs)
+        state = peak_caller(
+            peak_caller.init_state(),
             self.data.copy(),
             rotation_matrix=self.rotation_matrix,
         )
-        candidates = tuple(peak_caller)
+        state = peak_caller.result(state)
         if minimum_score is None:
-            assert len(candidates[0] <= num_peaks)
+            assert len(state[0] <= num_peaks)
         else:
-            peaks = candidates[0].astype(int)
+            peaks = state[0].astype(int)
             print(self.data[tuple(peaks.T)])
             assert np.all(self.data[tuple(peaks.T)] >= minimum_score)
 
@@ -78,21 +79,20 @@ class TestPeakCallers:
         peak_caller1 = peak_caller(
             shape=self.data.shape, num_peaks=num_peaks, min_distance=self.min_distance
         )
-        peak_caller1(self.data, rotation_matrix=self.rotation_matrix)
+        state1 = peak_caller1.init_state()
+        state1 = peak_caller1(state1, self.data, rotation_matrix=self.rotation_matrix)
 
         peak_caller2 = peak_caller(
             shape=self.data.shape, num_peaks=num_peaks, min_distance=self.min_distance
         )
-        peak_caller2(self.data, rotation_matrix=self.rotation_matrix)
+        state2 = peak_caller2.init_state()
+        state2 = peak_caller2(state2, self.data, rotation_matrix=self.rotation_matrix)
 
-        parameters = [tuple(peak_caller1), tuple(peak_caller2)]
-
-        result = tuple(
-            peak_caller.merge(
-                candidates=parameters,
-                num_peaks=num_peaks,
-                min_distance=self.min_distance,
-            )
+        states = [peak_caller1.result(state1), peak_caller2.result(state2)]
+        result = peak_caller.merge(
+            results=states,
+            num_peaks=num_peaks,
+            min_distance=self.min_distance,
         )
         assert [len(res) == 2 for res in result]
 
@@ -122,7 +122,9 @@ class TestRecursiveMasking:
             rotation_space = self.rotation_space
             rotation_mapping = self.rotation_mapping
 
-        peak_caller(
+        state = peak_caller.init_state()
+        state = peak_caller(
+            state,
             self.data.copy(),
             rotation_matrix=self.rotation_matrix,
             mask=self.mask,
@@ -130,11 +132,10 @@ class TestRecursiveMasking:
             rotation_mapping=rotation_mapping,
         )
 
-        candidates = tuple(peak_caller)
         if minimum_score is None:
-            assert len(candidates[0] <= num_peaks)
+            assert len(state[0] <= num_peaks)
         else:
-            peaks = candidates[0].astype(int)
+            peaks = state[0].astype(int)
             assert np.all(self.data[tuple(peaks.T)] >= minimum_score)
 
 
@@ -157,8 +158,9 @@ class TestMaxScoreOverRotations:
             shape=self.data.shape,
             use_memmap=use_memmap,
         )
-        score_analyzer(self.data, rotation_matrix=self.rotation_matrix)
-        res = tuple(score_analyzer)
+        state = score_analyzer.init_state()
+        state = score_analyzer(state, self.data, rotation_matrix=self.rotation_matrix)
+        res = score_analyzer.result(state)
         assert np.allclose(res[0].shape, self.data.shape)
         assert res[0].dtype == be._float_dtype
         assert res[1].size == self.data.ndim
@@ -174,11 +176,13 @@ class TestMaxScoreOverRotations:
             translation_offset=np.zeros(self.data.ndim, dtype=int),
             use_memmap=use_memmap,
         )
-        score_analyzer(self.data, rotation_matrix=self.rotation_matrix)
+        state = score_analyzer.init_state()
+        state = score_analyzer(state, self.data, rotation_matrix=self.rotation_matrix)
 
         data2 = self.data * 2
-        score_analyzer(data2, rotation_matrix=self.rotation_matrix)
-        scores, translation_offset, rotations, mapping = tuple(score_analyzer)
+        score_analyzer(state, data2, rotation_matrix=self.rotation_matrix)
+        scores, translation_offset, rotations, mapping = score_analyzer.result(state)
+
         assert np.all(scores >= score_threshold)
         max_scores = np.maximum(self.data, data2)
         max_scores = np.maximum(max_scores, score_threshold)
@@ -193,7 +197,8 @@ class TestMaxScoreOverRotations:
             translation_offset=np.zeros(self.data.ndim, dtype=int),
             use_memmap=use_memmap,
         )
-        score_analyzer(self.data, rotation_matrix=self.rotation_matrix)
+        state1 = score_analyzer.init_state()
+        state1, score_analyzer(state1, self.data, rotation_matrix=self.rotation_matrix)
 
         data2 = self.data * 2
         score_analyzer2 = MaxScoreOverRotations(
@@ -202,12 +207,12 @@ class TestMaxScoreOverRotations:
             translation_offset=np.zeros(self.data.ndim, dtype=int),
             use_memmap=use_memmap,
         )
-        score_analyzer2(data2, rotation_matrix=self.rotation_matrix)
-
-        parameters = [tuple(score_analyzer), tuple(score_analyzer2)]
+        state2 = score_analyzer2.init_state()
+        state2 = score_analyzer2(state2, data2, rotation_matrix=self.rotation_matrix)
+        states = [score_analyzer.result(state1), score_analyzer2.result(state2)]
 
         ret = MaxScoreOverRotations.merge(
-            parameters, use_memmap=use_memmap, score_threshold=score_threshold
+            results=states, use_memmap=use_memmap, score_threshold=score_threshold
         )
         scores, translation, rotations, mapping = ret
         assert np.all(scores >= score_threshold)

tests/test_backends.py

@@ -423,6 +423,7 @@ class TestBackends:
             rotation_matrix=rotation_matrix,
             out=out,
             out_mask=out_mask,
+            batched=True,
         )
 
         arr_b = backend.to_backend_array(arr_b)

tests/test_matching_cli.py

@@ -6,26 +6,10 @@ from os import remove, makedirs
 
 import pytest
 import numpy as np
-from tme import Density
+from tme import Density, Orientations
 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...")
-
-
+np.random.seed(42)
 available_backends = (x for x in be.available_backends() if x != "mlx")
 
 
@@ -48,7 +32,6 @@ def argdict_to_command(input_args, executable: str):
 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)
 
@@ -65,6 +48,19 @@ class TestMatchTemplate:
         cls.template_mask_path = tempfile.NamedTemporaryFile(
             delete=False, suffix=".mrc"
         ).name
+        cls.tempdir = tempfile.TemporaryDirectory().name
+        makedirs(cls.tempdir, exist_ok=True)
+
+        orientations = Orientations(
+            translations=((10, 10, 10), (12, 10, 15)),
+            rotations=((0, 0, 0), (45, 12, 90)),
+            scores=(0, 0),
+            details=(-1, -1),
+        )
+        cls.orientations_path = tempfile.NamedTemporaryFile(
+            delete=False, suffix=".star"
+        ).name
+        orientations.to_file(cls.orientations_path)
 
         Density(target, sampling_rate=5).to_file(cls.target_path)
         Density(template, sampling_rate=5).to_file(cls.template_path)
@@ -76,6 +72,8 @@ class TestMatchTemplate:
         cls.try_delete(cls.template_path)
         cls.try_delete(cls.target_mask_path)
         cls.try_delete(cls.template_mask_path)
+        cls.try_delete(cls.orientations_path)
+        cls.try_delete(cls.tempdir)
 
     @staticmethod
     def try_delete(file_path: str):
@@ -88,8 +86,8 @@ class TestMatchTemplate:
         except Exception:
             pass
 
-    @staticmethod
     def run_matching(
+        self,
         use_template_mask: bool,
         test_filter: bool,
         call_peaks: bool,
@@ -99,6 +97,7 @@ class TestMatchTemplate:
         target_mask_path: str,
         use_target_mask: bool = False,
         backend: str = "numpyfftw",
+        test_rejection_sampling: bool = False,
     ):
         output_path = tempfile.NamedTemporaryFile(delete=False, suffix="pickle").name
 
@@ -108,25 +107,24 @@ class TestMatchTemplate:
             "-n": 1,
             "-a": 60,
             "-o": output_path,
-            "--pad_edges": False,
-            "--pad_fourier": False,
+            "--pad-edges": False,
             "--backend": backend,
         }
 
         if use_template_mask:
-            argdict["--template_mask"] = template_mask_path
+            argdict["--template-mask"] = template_mask_path
 
         if use_target_mask:
-            argdict["--target_mask"] = target_mask_path
+            argdict["--target-mask"] = target_mask_path
 
-        if backend in ("cupy", "pytorch") and True in test_gpu:
-            argdict["--use_gpu"] = True
+        if test_rejection_sampling:
+            argdict["--orientations"] = self.orientations_path
 
         if test_filter:
             argdict["--lowpass"] = 30
             argdict["--defocus"] = 3000
-            argdict["--tilt_angles"] = "40,40:10"
-            argdict["--wedge_axes"] = "0,2"
+            argdict["--tilt-angles"] = "40,40"
+            argdict["--wedge-axes"] = "2,0"
             argdict["--whiten"] = True
 
         if call_peaks:
@@ -142,13 +140,18 @@ class TestMatchTemplate:
     @pytest.mark.parametrize("call_peaks", (False, True))
     @pytest.mark.parametrize("use_template_mask", (False, True))
     @pytest.mark.parametrize("test_filter", (False, True))
+    @pytest.mark.parametrize("test_rejection_sampling", (False, True))
     def test_match_template(
         self,
         backend: bool,
         call_peaks: bool,
         use_template_mask: bool,
         test_filter: bool,
+        test_rejection_sampling: bool,
     ):
+        if backend == "jax" and (call_peaks or test_rejection_sampling):
+            return None
+
         self.run_matching(
             use_template_mask=use_template_mask,
             use_target_mask=True,
@@ -159,6 +162,7 @@ class TestMatchTemplate:
             target_path=self.target_path,
             template_mask_path=self.template_mask_path,
             target_mask_path=self.target_mask_path,
+            test_rejection_sampling=test_rejection_sampling,
         )
 
 
@@ -175,20 +179,20 @@ class TestPostprocessing(TestMatchTemplate):
             "target_path": cls.target_path,
             "template_mask_path": cls.template_mask_path,
             "target_mask_path": cls.target_mask_path,
+            "test_rejection_sampling": False,
         }
 
         cls.score_pickle = cls.run_matching(
+            cls,
             call_peaks=False,
             **matching_kwargs,
         )
-        cls.peak_pickle = cls.run_matching(call_peaks=True, **matching_kwargs)
-        cls.tempdir = tempfile.TemporaryDirectory().name
+        cls.peak_pickle = cls.run_matching(cls, call_peaks=True, **matching_kwargs)
 
     @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)))
@@ -203,28 +207,28 @@ class TestPostprocessing(TestMatchTemplate):
         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,
-            "--num_peaks": 3,
+            "--input-file": self.score_pickle,
+            "--output-format": "orientations",
+            "--output-prefix": f"{self.tempdir}/temp",
+            "--peak-oversampling": peak_oversampling,
+            "--num-peaks": 3,
         }
 
         if score_cutoff is not None:
             if len(score_cutoff) == 1:
-                argdict["--n_false_positives"] = 1
+                argdict["--n-false-positives"] = 1
             else:
                 min_score, max_score = score_cutoff
-                argdict["--minimum_score"] = min_score
-                argdict["--maximum_score"] = max_score
+                argdict["--min-score"] = min_score
+                argdict["--max-score"] = max_score
 
         match distance_cutoff_strategy:
             case 1:
-                argdict["--mask_edges"] = True
+                argdict["--mask-edges"] = True
             case 2:
-                argdict["--min_distance"] = 5
+                argdict["--min-distance"] = 5
             case 3:
-                argdict["--min_boundary_distance"] = 5
+                argdict["--min-boundary-distance"] = 5
 
         cmd = argdict_to_command(argdict, executable="postprocess.py")
         ret = subprocess.run(cmd, capture_output=True, shell=True)
@@ -244,14 +248,15 @@ class TestPostprocessing(TestMatchTemplate):
             input_file = self.peak_pickle
 
         argdict = {
-            "--input_file": input_file,
-            "--output_format": output_format,
-            "--output_prefix": f"{self.tempdir}/temp",
-            "--num_peaks": 3,
-            "--peak_caller": "PeakCallerMaximumFilter",
+            "--input-file": input_file,
+            "--output-format": output_format,
+            "--output-prefix": f"{self.tempdir}/temp",
+            "--num-peaks": 3,
+            "--peak-caller": "PeakCallerMaximumFilter",
         }
         cmd = argdict_to_command(argdict, executable="postprocess.py")
         ret = subprocess.run(cmd, capture_output=True, shell=True)
+        print(ret)
 
         match output_format:
             case "orientations":
@@ -264,7 +269,8 @@ class TestPostprocessing(TestMatchTemplate):
                 assert exists(f"{self.tempdir}/temp.star")
             case "relion5":
                 assert exists(f"{self.tempdir}/temp.star")
-
+            case "pickle":
+                assert exists(f"{self.tempdir}/temp.pickle")
         assert ret.returncode == 0
 
     def test_postprocess_score_local_optimization(self):
@@ -272,12 +278,62 @@ class TestPostprocessing(TestMatchTemplate):
         makedirs(self.tempdir, exist_ok=True)
 
         argdict = {
-            "--input_file": self.score_pickle,
-            "--output_format": "orientations",
-            "--output_prefix": f"{self.tempdir}/temp",
-            "--num_peaks": 1,
-            "--local_optimization": True,
+            "--input-file": self.score_pickle,
+            "--output-format": "orientations",
+            "--output-prefix": f"{self.tempdir}/temp",
+            "--num-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
+
+
+class TestEstimateMemoryUsage(TestMatchTemplate):
+    @classmethod
+    def setup_class(cls):
+        super().setup_class()
+
+    @pytest.mark.parametrize("ncores", (1, 4, 8))
+    @pytest.mark.parametrize("pad_edges", (False, True))
+    def test_estimation(self, ncores, pad_edges):
+
+        argdict = {
+            "-m": self.target_path,
+            "-i": self.template_path,
+            "--ncores": ncores,
+            "--pad-edges": pad_edges,
+            "--score": "FLCSphericalMask",
+        }
+
+        cmd = argdict_to_command(argdict, executable="estimate_memory_usage.py")
+        ret = subprocess.run(cmd, capture_output=True, shell=True)
+        assert ret.returncode == 0
+
+
+class TestPreprocess(TestMatchTemplate):
+    @classmethod
+    def setup_class(cls):
+        super().setup_class()
+
+    @pytest.mark.parametrize("backend", available_backends)
+    @pytest.mark.parametrize("align_axis", (False, True))
+    @pytest.mark.parametrize("invert_contrast", (False, True))
+    def test_estimation(self, backend, align_axis, invert_contrast):
+
+        argdict = {
+            "-m": self.target_path,
+            "--backend": backend,
+            "--lowpass": 40,
+            "--sampling-rate": 5,
+            "-o": f"{self.tempdir}/out.mrc",
+        }
+        if align_axis:
+            argdict["--align-axis"] = 2
+
+        if invert_contrast:
+            argdict["--invert-contrast"] = True
+
+        cmd = argdict_to_command(argdict, executable="preprocess.py")
+        ret = subprocess.run(cmd, capture_output=True, shell=True)
+        assert ret.returncode == 0

tests/test_matching_data.py

@@ -7,7 +7,7 @@ from tme.backends import backend as be
 from tme.matching_data import MatchingData
 
 
-class TestDensity:
+class TestMatchingData:
     def setup_method(self):
         target = np.zeros((50, 50, 50))
         target[20:30, 30:40, 12:17] = 1
@@ -87,9 +87,9 @@ class TestDensity:
         matching_data.target_mask = self.target
         matching_data.template_mask = self.template
 
-        ret = matching_data.subset_by_slice()
+        ret, offset = matching_data.subset_by_slice()
 
-        assert type(ret) == type(matching_data)
+        assert isinstance(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)
@@ -107,10 +107,10 @@ class TestDensity:
         template_slice = MatchingData._shape_to_slice(
             shape=np.divide(self.template.shape, 2).astype(int)
         )
-        ret = matching_data.subset_by_slice(
+        ret, offset = matching_data.subset_by_slice(
             target_slice=target_slice, template_slice=template_slice
         )
-        assert type(ret) == type(matching_data)
+        assert isinstance(ret, type(matching_data))
 
         assert np.allclose(
             ret.target.shape, np.divide(self.target.shape, 2).astype(int)

tests/test_matching_exhaustive.py

@@ -7,7 +7,6 @@ from tme.matching_data import MatchingData
 from tme.memory import MATCHING_MEMORY_REGISTRY
 from tme.analyzer import MaxScoreOverRotations, PeakCallerSort
 from tme.matching_exhaustive import (
-    scan,
     scan_subsets,
     MATCHING_EXHAUSTIVE_REGISTER,
     register_matching_exhaustive,
@@ -36,7 +35,7 @@ class TestMatchExhaustive:
         self.coordinates_weights = None
         self.rotations = None
 
-    @pytest.mark.parametrize("evaluate_peak", (False, True))
+    @pytest.mark.parametrize("evaluate_peak", (True,))
     @pytest.mark.parametrize("score", tuple(MATCHING_EXHAUSTIVE_REGISTER.keys()))
     @pytest.mark.parametrize("job_schedule", ((2, 1),))
     @pytest.mark.parametrize("pad_edge", (False, True))

tests/test_matching_optimization.py

@@ -24,6 +24,7 @@ coordinate_to_coordinate = [
     for k, v in MATCHING_OPTIMIZATION_REGISTER.items()
     if issubclass(v, _MatchCoordinatesToCoordinates)
 ]
+np.random.seed(42)
 
 
 class TestMatchDensityToDensity:
@@ -52,9 +53,7 @@ class TestMatchDensityToDensity:
 
     @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)
+        self.test_initialization(method=method, notest=True)()
 
 
 class TestMatchDensityToCoordinates:
@@ -97,9 +96,7 @@ class TestMatchDensityToCoordinates:
 
     @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)
+        self.test_initialization(method=method, notest=True)()
 
 
 class TestMatchCoordinateToCoordinates:
@@ -135,9 +132,7 @@ class TestMatchCoordinateToCoordinates:
 
     @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)
+        self.test_initialization(method=method, notest=True)()
 
 
 class TestOptimizeMatch: