pytme 0.3b0__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1__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,
@@ -72,17 +82,16 @@ class TestBandPassFilter:
         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

tests/preprocessing/test_utils.py

@@ -51,6 +51,24 @@ class TestPreprocessUtils:
         assert fgrid.shape == tuple(tilt_shape)
         assert fgrid.max() <= np.sqrt(1 / sampling_rate * len(shape))
 
+    @pytest.mark.parametrize("shape", ((15, 15, 15), (31, 31, 31), (64, 64, 64)))
+    @pytest.mark.parametrize("sampling_rate", (0.5, 1, 2))
+    @pytest.mark.parametrize("angle", (-5, 0, 5))
+    def test_freqgrid_comparison(self, shape, sampling_rate, angle):
+        grid = frequency_grid_at_angle(
+            shape=shape,
+            angle=angle,
+            sampling_rate=sampling_rate,
+            opening_axis=2,
+            tilt_axis=0,
+        )
+        grid2 = fftfreqn(
+            shape=shape[1:], sampling_rate=sampling_rate, compute_euclidean_norm=True
+        )
+
+        # These should be equal for cubical input shapes
+        assert np.allclose(grid, grid2)
+
     @pytest.mark.parametrize("n", [10, 100, 1000])
     @pytest.mark.parametrize("sampling_rate", range(1, 4))
     def test_fftfreqn(self, n, sampling_rate):

tests/test_analyzer.py

@@ -97,125 +97,125 @@ class TestPeakCallers:
         assert [len(res) == 2 for res in result]
 
 
-# class TestRecursiveMasking:
-#     def setup_method(self):
-#         self.num_peaks = 100
-#         self.min_distance = 5
-#         self.data = np.random.rand(100, 100, 100)
-#         self.rotation_matrix = np.eye(3)
-#         self.mask = np.random.rand(20, 20, 20)
-#         self.rotation_space = np.zeros_like(self.data)
-#         self.rotation_mapping = {0: (0, 0, 0)}
-
-#     @pytest.mark.parametrize("num_peaks", (1, 100))
-#     @pytest.mark.parametrize("compute_rotation", (True, False))
-#     @pytest.mark.parametrize("minimum_score", (None, 0.5))
-#     def test__call__(self, num_peaks, compute_rotation, minimum_score):
-#         peak_caller = PeakCallerRecursiveMasking(
-#             shape=self.data.shape,
-#             num_peaks=num_peaks,
-#             min_distance=self.min_distance,
-#             min_score=minimum_score,
-#         )
-#         rotation_space, rotation_mapping = None, None
-#         if compute_rotation:
-#             rotation_space = self.rotation_space
-#             rotation_mapping = self.rotation_mapping
-
-#         state = peak_caller.init_state()
-#         state = peak_caller(
-#             state,
-#             self.data.copy(),
-#             rotation_matrix=self.rotation_matrix,
-#             mask=self.mask,
-#             rotation_space=rotation_space,
-#             rotation_mapping=rotation_mapping,
-#         )
-
-#         if minimum_score is None:
-#             assert len(state[0] <= num_peaks)
-#         else:
-#             peaks = state[0].astype(int)
-#             assert np.all(self.data[tuple(peaks.T)] >= minimum_score)
-
-
-# class TestMaxScoreOverRotations:
-#     def setup_method(self):
-#         self.num_peaks = 100
-#         self.min_distance = 5
-#         self.data = np.random.rand(100, 100, 100)
-#         self.rotation_matrix = np.eye(3)
-
-#     def test_initialization(self):
-#         _ = MaxScoreOverRotations(
-#             shape=self.data.shape,
-#             translation_offset=np.zeros(self.data.ndim, dtype=int),
-#         )
-
-#     @pytest.mark.parametrize("use_memmap", [False, True])
-#     def test__iter__(self, use_memmap: bool):
-#         score_analyzer = MaxScoreOverRotations(
-#             shape=self.data.shape,
-#             use_memmap=use_memmap,
-#         )
-#         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
-#         assert np.allclose(res[2].shape, self.data.shape)
-#         assert len(res) == 4
-
-#     @pytest.mark.parametrize("use_memmap", [False, True])
-#     @pytest.mark.parametrize("score_threshold", [0, 1e10, -1e10])
-#     def test__call__(self, use_memmap: bool, score_threshold: float):
-#         score_analyzer = MaxScoreOverRotations(
-#             shape=self.data.shape,
-#             score_threshold=score_threshold,
-#             translation_offset=np.zeros(self.data.ndim, dtype=int),
-#             use_memmap=use_memmap,
-#         )
-#         state = score_analyzer.init_state()
-#         state = score_analyzer(state, self.data, rotation_matrix=self.rotation_matrix)
-
-#         data2 = self.data * 2
-#         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)
-#         assert np.allclose(scores, max_scores)
-
-#     @pytest.mark.parametrize("use_memmap", [False, True])
-#     @pytest.mark.parametrize("score_threshold", [0, 1e10, -1e10])
-#     def test_merge(self, use_memmap: bool, score_threshold: float):
-#         score_analyzer = MaxScoreOverRotations(
-#             shape=self.data.shape,
-#             score_threshold=score_threshold,
-#             translation_offset=np.zeros(self.data.ndim, dtype=int),
-#             use_memmap=use_memmap,
-#         )
-#         state1 = score_analyzer.init_state()
-#         state1, score_analyzer(state1, self.data, rotation_matrix=self.rotation_matrix)
-
-#         data2 = self.data * 2
-#         score_analyzer2 = MaxScoreOverRotations(
-#             shape=self.data.shape,
-#             score_threshold=score_threshold,
-#             translation_offset=np.zeros(self.data.ndim, dtype=int),
-#             use_memmap=use_memmap,
-#         )
-#         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(
-#             results=states, use_memmap=use_memmap, score_threshold=score_threshold
-#         )
-#         scores, translation, rotations, mapping = ret
-#         assert np.all(scores >= score_threshold)
-#         max_scores = np.maximum(self.data, data2)
-#         max_scores = np.maximum(max_scores, score_threshold)
-#         assert np.allclose(scores, max_scores)
+class TestRecursiveMasking:
+    def setup_method(self):
+        self.num_peaks = 100
+        self.min_distance = 5
+        self.data = np.random.rand(100, 100, 100)
+        self.rotation_matrix = np.eye(3)
+        self.mask = np.random.rand(20, 20, 20)
+        self.rotation_space = np.zeros_like(self.data)
+        self.rotation_mapping = {0: (0, 0, 0)}
+
+    @pytest.mark.parametrize("num_peaks", (1, 100))
+    @pytest.mark.parametrize("compute_rotation", (True, False))
+    @pytest.mark.parametrize("minimum_score", (None, 0.5))
+    def test__call__(self, num_peaks, compute_rotation, minimum_score):
+        peak_caller = PeakCallerRecursiveMasking(
+            shape=self.data.shape,
+            num_peaks=num_peaks,
+            min_distance=self.min_distance,
+            min_score=minimum_score,
+        )
+        rotation_space, rotation_mapping = None, None
+        if compute_rotation:
+            rotation_space = self.rotation_space
+            rotation_mapping = self.rotation_mapping
+
+        state = peak_caller.init_state()
+        state = peak_caller(
+            state,
+            self.data.copy(),
+            rotation_matrix=self.rotation_matrix,
+            mask=self.mask,
+            rotation_space=rotation_space,
+            rotation_mapping=rotation_mapping,
+        )
+
+        if minimum_score is None:
+            assert len(state[0] <= num_peaks)
+        else:
+            peaks = state[0].astype(int)
+            assert np.all(self.data[tuple(peaks.T)] >= minimum_score)
+
+
+class TestMaxScoreOverRotations:
+    def setup_method(self):
+        self.num_peaks = 100
+        self.min_distance = 5
+        self.data = np.random.rand(100, 100, 100)
+        self.rotation_matrix = np.eye(3)
+
+    def test_initialization(self):
+        _ = MaxScoreOverRotations(
+            shape=self.data.shape,
+            translation_offset=np.zeros(self.data.ndim, dtype=int),
+        )
+
+    @pytest.mark.parametrize("use_memmap", [False, True])
+    def test__iter__(self, use_memmap: bool):
+        score_analyzer = MaxScoreOverRotations(
+            shape=self.data.shape,
+            use_memmap=use_memmap,
+        )
+        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
+        assert np.allclose(res[2].shape, self.data.shape)
+        assert len(res) == 4
+
+    @pytest.mark.parametrize("use_memmap", [False, True])
+    @pytest.mark.parametrize("score_threshold", [0, 1e10, -1e10])
+    def test__call__(self, use_memmap: bool, score_threshold: float):
+        score_analyzer = MaxScoreOverRotations(
+            shape=self.data.shape,
+            score_threshold=score_threshold,
+            translation_offset=np.zeros(self.data.ndim, dtype=int),
+            use_memmap=use_memmap,
+        )
+        state = score_analyzer.init_state()
+        state = score_analyzer(state, self.data, rotation_matrix=self.rotation_matrix)
+
+        data2 = self.data * 2
+        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)
+        assert np.allclose(scores, max_scores)
+
+    @pytest.mark.parametrize("use_memmap", [False, True])
+    @pytest.mark.parametrize("score_threshold", [0, 1e10, -1e10])
+    def test_merge(self, use_memmap: bool, score_threshold: float):
+        score_analyzer = MaxScoreOverRotations(
+            shape=self.data.shape,
+            score_threshold=score_threshold,
+            translation_offset=np.zeros(self.data.ndim, dtype=int),
+            use_memmap=use_memmap,
+        )
+        state1 = score_analyzer.init_state()
+        state1, score_analyzer(state1, self.data, rotation_matrix=self.rotation_matrix)
+
+        data2 = self.data * 2
+        score_analyzer2 = MaxScoreOverRotations(
+            shape=self.data.shape,
+            score_threshold=score_threshold,
+            translation_offset=np.zeros(self.data.ndim, dtype=int),
+            use_memmap=use_memmap,
+        )
+        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(
+            results=states, use_memmap=use_memmap, score_threshold=score_threshold
+        )
+        scores, translation, rotations, mapping = ret
+        assert np.all(scores >= score_threshold)
+        max_scores = np.maximum(self.data, data2)
+        max_scores = np.maximum(max_scores, score_threshold)
+        assert np.allclose(scores, max_scores)

tests/test_backends.py

@@ -292,16 +292,10 @@ class TestBackends:
 
     @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
     @pytest.mark.parametrize("fast_shape", ((10, 15, 100), (55, 23, 17)))
-    def test_build_fft(self, backend, fast_shape):
+    def test_fft(self, backend, fast_shape):
         _, fast_shape, fast_ft_shape = backend.compute_convolution_shapes(
             fast_shape, (1 for _ in range(len(fast_shape)))
         )
-        rfftn, irfftn = backend.build_fft(
-            fwd_shape=fast_shape,
-            inv_shape=fast_ft_shape,
-            real_dtype=backend._float_dtype,
-            cmpl_dtype=backend._complex_dtype,
-        )
         arr = np.random.rand(*fast_shape)
         out = np.zeros(fast_ft_shape)
 
@@ -310,11 +304,11 @@ class TestBackends:
             backend.to_backend_array(out), backend._complex_dtype
         )
 
-        rfftn(
+        backend.rfftn(
             backend.astype(backend.to_backend_array(arr), backend._float_dtype),
             complex_arr,
         )
-        irfftn(complex_arr, real_arr)
+        backend.irfftn(complex_arr, real_arr)
         assert np.allclose(arr, backend.to_numpy_array(real_arr), rtol=0.3)
 
     @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
@@ -423,6 +417,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_density.py

@@ -98,7 +98,6 @@ class TestDensity:
         assert np.allclose(density.data, self.density.data)
         assert np.allclose(density.sampling_rate, self.density.sampling_rate)
         assert np.allclose(density.origin, self.density.origin)
-        assert density.metadata == self.density.metadata
 
     def test_from_file_baseline(self):
         self.test_to_file(gzip=False)

tests/test_matching_cli.py

@@ -107,15 +107,15 @@ class TestMatchTemplate:
             "-n": 1,
             "-a": 60,
             "-o": output_path,
-            "--pad_edges": 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 test_rejection_sampling:
             argdict["--orientations"] = self.orientations_path
@@ -123,8 +123,8 @@ class TestMatchTemplate:
         if test_filter:
             argdict["--lowpass"] = 30
             argdict["--defocus"] = 3000
-            argdict["--tilt_angles"] = "40,40"
-            argdict["--wedge_axes"] = "2,0"
+            argdict["--tilt-angles"] = "40,40"
+            argdict["--wedge-axes"] = "2,0"
             argdict["--whiten"] = True
 
         if call_peaks:
@@ -207,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)
@@ -248,11 +248,11 @@ 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)
@@ -278,11 +278,11 @@ 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)
@@ -302,7 +302,7 @@ class TestEstimateMemoryUsage(TestMatchTemplate):
             "-m": self.target_path,
             "-i": self.template_path,
             "--ncores": ncores,
-            "--pad_edges": pad_edges,
+            "--pad-edges": pad_edges,
             "--score": "FLCSphericalMask",
         }
 
@@ -325,14 +325,14 @@ class TestPreprocess(TestMatchTemplate):
             "-m": self.target_path,
             "--backend": backend,
             "--lowpass": 40,
-            "--sampling_rate": 5,
+            "--sampling-rate": 5,
             "-o": f"{self.tempdir}/out.mrc",
         }
         if align_axis:
-            argdict["--align_axis"] = 2
+            argdict["--align-axis"] = 2
 
         if invert_contrast:
-            argdict["--invert_contrast"] = True
+            argdict["--invert-contrast"] = True
 
         cmd = argdict_to_command(argdict, executable="preprocess.py")
         ret = subprocess.run(cmd, capture_output=True, shell=True)

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)