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

tests/preprocessing/test_compose.py

@@ -0,0 +1,76 @@
+import pytest
+
+from tme.preprocessing import Compose
+from tme.backends import backend as be
+
+
+def mock_transform1(**kwargs):
+    return {"data": be.ones((10, 10)), "is_multiplicative_filter": True}
+
+
+def mock_transform2(**kwargs):
+    return {"data": be.ones((10, 10)) * 2, "is_multiplicative_filter": True}
+
+
+def mock_transform3(**kwargs):
+    return {"extra_info": "test"}
+
+
+class TestCompose:
+    @pytest.fixture
+    def compose_instance(self):
+        return Compose((mock_transform1, mock_transform2, mock_transform3))
+
+    def test_init(self):
+        transforms = (mock_transform1, mock_transform2)
+        compose = Compose(transforms)
+        assert compose.transforms == transforms
+
+    def test_call_empty_transforms(self):
+        compose = Compose(())
+        result = compose()
+        assert result == {}
+
+    def test_call_single_transform(self):
+        compose = Compose((mock_transform1,))
+        result = compose()
+        assert "data" in result
+        assert result.get("is_multiplicative_filter", False)
+        assert be.allclose(result["data"], be.ones((10, 10)))
+
+    def test_call_multiple_transforms(self, compose_instance):
+        result = compose_instance()
+        assert "data" in result
+        assert "extra_info" not in result
+        assert be.allclose(result["data"], be.ones((10, 10)) * 2)
+
+    def test_multiplicative_filter_composition(self):
+        compose = Compose((mock_transform1, mock_transform2))
+        result = compose()
+        assert "data" in result
+        assert be.allclose(result["data"], be.ones((10, 10)) * 2)
+
+    @pytest.mark.parametrize(
+        "kwargs", [{}, {"extra_param": "test"}, {"data": be.zeros((5, 5))}]
+    )
+    def test_call_with_kwargs(self, compose_instance, kwargs):
+        result = compose_instance(**kwargs)
+        assert "data" in result
+        assert "extra_info" not in result
+
+    def test_non_multiplicative_filter(self):
+        def non_mult_transform(**kwargs):
+            return {"data": be.ones((10, 10)) * 3, "is_multiplicative_filter": False}
+
+        compose = Compose((mock_transform1, non_mult_transform))
+        result = compose()
+        assert "data" in result
+        assert be.allclose(result["data"], be.ones((10, 10)) * 3)
+
+    def test_error_handling(self):
+        def error_transform(**kwargs):
+            raise ValueError("Test error")
+
+        compose = Compose((mock_transform1, error_transform))
+        with pytest.raises(ValueError, match="Test error"):
+            compose()

tests/preprocessing/test_frequency_filters.py

@@ -0,0 +1,178 @@
+import pytest
+import numpy as np
+from typing import Tuple
+
+from tme.backends import backend as be
+from tme.preprocessing._utils import compute_fourier_shape
+from tme.preprocessing.frequency_filters import BandPassFilter, LinearWhiteningFilter
+
+
+class TestBandPassFilter:
+    @pytest.fixture
+    def band_pass_filter(self):
+        return BandPassFilter()
+
+    @pytest.mark.parametrize(
+        "shape, lowpass, highpass, sampling_rate",
+        [
+            ((10, 10), 0.2, 0.8, 1),
+            ((20, 20, 20), 0.1, 0.9, 2),
+            ((30, 30), None, 0.5, 1),
+            ((40, 40), 0.3, None, 0.5),
+        ],
+    )
+    def test_discrete_bandpass(
+        self, shape: Tuple[int], lowpass: float, highpass: float, sampling_rate: float
+    ):
+        result = BandPassFilter.discrete_bandpass(
+            shape, lowpass, highpass, sampling_rate
+        )
+        assert isinstance(result, type(be.ones((1,))))
+        assert result.shape == shape
+        assert np.all((result >= 0) & (result <= 1))
+
+    @pytest.mark.parametrize(
+        "shape, lowpass, highpass, sampling_rate",
+        [
+            ((10, 10), 0.2, 0.8, 1),
+            ((20, 20, 20), 0.1, 0.9, 2),
+            ((30, 30), None, 0.5, 1),
+            ((40, 40), 0.3, None, 0.5),
+        ],
+    )
+    def test_gaussian_bandpass(
+        self, shape: Tuple[int], lowpass: float, highpass: float, sampling_rate: float
+    ):
+        result = BandPassFilter.gaussian_bandpass(
+            shape, lowpass, highpass, sampling_rate
+        )
+        assert isinstance(result, type(be.ones((1,))))
+        assert result.shape == shape
+        assert np.all((result >= 0) & (result <= 1))
+
+    @pytest.mark.parametrize("use_gaussian", [True, False])
+    @pytest.mark.parametrize("return_real_fourier", [True, False])
+    @pytest.mark.parametrize("shape_is_real_fourier", [True, False])
+    def test_call_method(
+        self,
+        band_pass_filter: BandPassFilter,
+        use_gaussian: bool,
+        return_real_fourier: bool,
+        shape_is_real_fourier: bool,
+    ):
+        band_pass_filter.use_gaussian = use_gaussian
+        band_pass_filter.return_real_fourier = return_real_fourier
+        band_pass_filter.shape_is_real_fourier = shape_is_real_fourier
+
+        result = band_pass_filter(shape=(10, 10), lowpass=0.2, highpass=0.8)
+
+        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):
+        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)
+        result = bpf(shape=shape, lowpass=0.2, highpass=0.8)
+        expected_shape = tuple(compute_fourier_shape(shape, False))
+        assert result["data"].shape == expected_shape
+
+
+class TestLinearWhiteningFilter:
+    @pytest.mark.parametrize(
+        "shape, n_bins, batch_dimension",
+        [
+            ((10, 10), None, None),
+            ((20, 20, 20), 15, 0),
+            ((30, 30, 30), 20, 1),
+            ((40, 40, 40, 40), 25, 2),
+        ],
+    )
+    def test_compute_spectrum(
+        self, shape: Tuple[int], n_bins: int, batch_dimension: int
+    ):
+        data_rfft = be.fft.rfftn(be.random.random(shape))
+        bins, radial_averages = LinearWhiteningFilter._compute_spectrum(
+            data_rfft, n_bins, batch_dimension
+        )
+        data_shape = tuple(
+            int(x) for i, x in enumerate(data_rfft.shape) if i != batch_dimension
+        )
+
+        assert isinstance(bins, np.ndarray)
+        assert isinstance(radial_averages, np.ndarray)
+        assert bins.shape == data_shape
+        assert radial_averages.ndim == 1
+        assert np.all(radial_averages >= 0) and np.all(radial_averages <= 1)
+
+    @pytest.mark.parametrize("shape", ((10, 10), (21, 20, 31)))
+    @pytest.mark.parametrize("shape_is_real_fourier", (False, True))
+    @pytest.mark.parametrize("order", (1, 3))
+    def test_interpolate_spectrum(
+        self, shape: Tuple[int], shape_is_real_fourier: bool, order: int
+    ):
+        spectrum = be.random.random(100)
+        result = LinearWhiteningFilter()._interpolate_spectrum(
+            spectrum, shape, shape_is_real_fourier, order
+        )
+        assert result.shape == tuple(shape)
+        assert isinstance(result, np.ndarray)
+
+    @pytest.mark.parametrize(
+        "shape, n_bins, batch_dimension, order",
+        [
+            ((10, 10), None, None, 1),
+            ((20, 20, 20), 15, 0, 2),
+            ((30, 30, 30), 20, 1, None),
+        ],
+    )
+    def test_call_method(
+        self,
+        shape: Tuple[int],
+        n_bins: int,
+        batch_dimension: int,
+        order: int,
+    ):
+        data = be.random.random(shape)
+        result = LinearWhiteningFilter()(
+            data=data, n_bins=n_bins, batch_dimension=batch_dimension, order=order
+        )
+
+        assert isinstance(result, dict)
+        assert result.get("data", False) is not False
+        assert result.get("is_multiplicative_filter", False)
+        assert isinstance(result["data"], type(be.ones((1,))))
+        data_shape = tuple(
+            int(x) for i, x in enumerate(data.shape) if i != batch_dimension
+        )
+        assert result["data"].shape == tuple(compute_fourier_shape(data_shape, False))
+
+    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)
+
+        assert isinstance(result, dict)
+        assert result.get("data", False) is not False
+        assert result.get("is_multiplicative_filter", False)
+        assert isinstance(result["data"], type(be.ones((1,))))
+        assert result["data"].shape == data_rfft.shape
+
+    @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)
+
+        filter_mask = result["data"]
+        assert np.all(filter_mask >= 0) and np.all(filter_mask <= 1)

tests/preprocessing/test_preprocessor.py

@@ -0,0 +1,136 @@
+import pytest
+import numpy as np
+
+from tme import Density, Structure, Preprocessor
+
+
+class TestPreprocessor:
+    def setup_method(self):
+        self.density = Density.from_file(filename="tests/data/Raw/em_map.map")
+        self.structure = Structure.from_file("tests/data/Structures/5khe.cif")
+        self.structure_density = Density.from_structure(
+            filename_or_structure="tests/data/Structures/5khe.cif",
+            origin=self.density.origin,
+            shape=self.density.shape,
+            sampling_rate=self.density.sampling_rate,
+        )
+        self.preprocessor = Preprocessor()
+
+    def teardown_method(self):
+        self.density = None
+        self.structure_density = None
+
+    def test_initialization(self):
+        _ = Preprocessor()
+
+    def test_apply_method_error(self):
+        with pytest.raises(TypeError):
+            self.preprocessor.apply_method(method=None, parameters={})
+
+        with pytest.raises(NotImplementedError):
+            self.preprocessor.apply_method(method="None", parameters={})
+
+    def test_method_to_id_error(self):
+        with pytest.raises(TypeError):
+            self.preprocessor.method_to_id(method=None, parameters={})
+
+        with pytest.raises(NotImplementedError):
+            self.preprocessor.method_to_id(method="None", parameters={})
+
+    def test_method_to_id(self):
+        ret = self.preprocessor.method_to_id(method="gaussian_filter", parameters={})
+        assert isinstance(ret, str)
+
+    @pytest.mark.parametrize("low_sigma,high_sigma", [(0, 1), (3, 5)])
+    def test_difference_of_gaussian_filter(self, low_sigma, high_sigma):
+        _ = self.preprocessor.difference_of_gaussian_filter(
+            template=self.structure_density.data,
+            low_sigma=low_sigma,
+            high_sigma=high_sigma,
+        )
+
+    @pytest.mark.parametrize("smallest_size,largest_size", [(1, 10), (2, 20)])
+    def test_bandpass_filter(self, smallest_size, largest_size):
+        _ = self.preprocessor.bandpass_filter(
+            template=self.structure_density.data,
+            lowpass=smallest_size,
+            highpass=largest_size,
+            sampling_rate=1,
+        )
+
+    @pytest.mark.parametrize("lbd,sigma_range", [(1, (2, 4)), (20, (1, 6))])
+    def test_local_gaussian_alignment_filter(self, lbd, sigma_range):
+        _ = self.preprocessor.local_gaussian_alignment_filter(
+            template=self.structure_density.data,
+            target=self.density.data,
+            lbd=lbd,
+            sigma_range=sigma_range,
+        )
+
+    @pytest.mark.parametrize(
+        "lbd,sigma_range,gaussian_sigma", [(1, (2, 4), 1), (20, (1, 6), 3)]
+    )
+    def test_local_gaussian_filter(self, lbd, sigma_range, gaussian_sigma):
+        _ = self.preprocessor.local_gaussian_filter(
+            template=self.structure_density.data,
+            lbd=lbd,
+            sigma_range=sigma_range,
+            gaussian_sigma=gaussian_sigma,
+        )
+
+    @pytest.mark.parametrize(
+        "edge_algorithm",
+        ["sobel", "prewitt", "laplace", "gaussian", "gaussian_laplace"],
+    )
+    @pytest.mark.parametrize("reverse", [(True), (False)])
+    def test_edge_gaussian_filter(self, edge_algorithm, reverse):
+        _ = self.preprocessor.edge_gaussian_filter(
+            template=self.structure_density.data,
+            edge_algorithm=edge_algorithm,
+            reverse=reverse,
+            sigma=3,
+        )
+
+    @pytest.mark.parametrize("width", range(1, 9, 3))
+    def test_mean_filter(self, width):
+        _ = self.preprocessor.mean_filter(
+            template=self.structure_density.data,
+            width=width,
+        )
+
+    @pytest.mark.parametrize("width", range(1, 9, 3))
+    def test_kaiserb_filter(self, width):
+        _ = self.preprocessor.kaiserb_filter(
+            template=self.structure_density.data,
+            width=width,
+        )
+
+    @pytest.mark.parametrize("width", range(1, 9, 3))
+    def test_blob_filter(self, width):
+        _ = self.preprocessor.blob_filter(
+            template=self.structure_density.data,
+            width=width,
+        )
+
+    @pytest.mark.parametrize("width", range(1, 9, 3))
+    def test_hamming_filter(self, width):
+        _ = self.preprocessor.hamming_filter(
+            template=self.structure_density.data,
+            width=width,
+        )
+
+    @pytest.mark.parametrize("rank", range(1, 9, 3))
+    def test_rank_filter(self, rank):
+        _ = self.preprocessor.rank_filter(
+            template=self.structure_density.data,
+            rank=rank,
+        )
+
+    @pytest.mark.parametrize("infinite_plane", [False, True])
+    def test_continuous_wedge_mask(self, infinite_plane):
+        _ = self.preprocessor.continuous_wedge_mask(
+            start_tilt=50,
+            stop_tilt=-40,
+            shape=(50, 50, 50),
+            infinite_plane=infinite_plane,
+        )

tests/preprocessing/test_utils.py

@@ -0,0 +1,79 @@
+import pytest
+import numpy as np
+
+from tme.preprocessing._utils import (
+    fftfreqn,
+    centered_grid,
+    shift_fourier,
+    compute_fourier_shape,
+    crop_real_fourier,
+    compute_tilt_shape,
+    frequency_grid_at_angle,
+)
+
+
+class TestPreprocessUtils:
+    @pytest.mark.parametrize("reduce_dim", (False, True))
+    @pytest.mark.parametrize("shape", ((10,), (10, 15), (10, 15, 30)))
+    def test_compute_tilt_shape(self, shape, reduce_dim):
+        tilt_shape = compute_tilt_shape(
+            shape=shape, opening_axis=0, reduce_dim=reduce_dim
+        )
+        if reduce_dim:
+            assert len(tilt_shape) == len(shape) - 1
+        else:
+            assert len(tilt_shape) == len(shape)
+            assert tilt_shape[0] == 1
+
+    @pytest.mark.parametrize("shape", ((10,), (10, 15), (10, 15, 30)))
+    def test_centered_grid(self, shape):
+        grid = centered_grid(shape=shape)
+        assert grid.shape[0] == len(shape)
+        center = tuple(int(x) // 2 for x in shape)
+        for i in range(grid.shape[0]):
+            assert grid[i][center] == 0
+            assert np.max(grid[i]) <= center[i]
+
+    @pytest.mark.parametrize("shape", ((10, 15, 30),))
+    @pytest.mark.parametrize("sampling_rate", (0.5, 1, 2))
+    @pytest.mark.parametrize("angle", (-5, 0, 5))
+    @pytest.mark.parametrize("wedge", ((0, 1), (1, 0)))
+    def test_frequency_grid_at_angle(self, shape, sampling_rate, angle, wedge):
+        opening, tilt = wedge
+        fgrid = frequency_grid_at_angle(
+            shape=shape,
+            angle=angle,
+            sampling_rate=sampling_rate,
+            opening_axis=opening,
+            tilt_axis=tilt,
+        )
+        tilt_shape = compute_tilt_shape(shape, opening_axis=opening, reduce_dim=True)
+        assert fgrid.shape == tuple(tilt_shape)
+        assert fgrid.max() <= np.sqrt(1 / sampling_rate * len(shape))
+
+    @pytest.mark.parametrize("n", [10, 100, 1000])
+    @pytest.mark.parametrize("sampling_rate", range(1, 4))
+    def test_fftfreqn(self, n, sampling_rate):
+        assert np.allclose(
+            fftfreqn(
+                shape=(n,), sampling_rate=sampling_rate, compute_euclidean_norm=True
+            ),
+            np.abs(np.fft.ifftshift(np.fft.fftfreq(n=n, d=sampling_rate))),
+        )
+
+    @pytest.mark.parametrize("shape", ((10,), (10, 15), (10, 15, 30)))
+    def test_crop_real_fourier(self, shape):
+        data = np.random.rand(*shape)
+        data_crop = crop_real_fourier(data)
+        assert data_crop.shape == tuple(compute_fourier_shape(data.shape, False))
+
+    @pytest.mark.parametrize("real", (False, True))
+    @pytest.mark.parametrize("shape", ((10,), (10, 15), (10, 15, 30)))
+    def test_compute_fourier_shape(self, shape, real: bool):
+        data = np.random.rand(*shape)
+        func = np.fft.rfftn if real else np.fft.fftn
+        assert func(data).shape == tuple(compute_fourier_shape(data.shape, not real))
+
+    def test_shift_fourier(self):
+        data = np.random.rand(10)
+        assert np.allclose(shift_fourier(data, False), np.fft.ifftshift(data))