pytme 0.3.1.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.2__cp311-cp311-macosx_15_0_arm64.whl

scripts/pytme_runner.py

@@ -8,7 +8,7 @@ import subprocess
 from abc import ABC, abstractmethod
 
 from pathlib import Path
-from dataclasses import dataclass
+from dataclasses import dataclass, field, fields
 from typing import Dict, List, Optional, Any
 
 from tme.backends import backend as be
@@ -149,7 +149,7 @@ class AnalysisDatasetDiscovery(DatasetDiscovery):
 
     #: Glob pattern for TM pickle files, e.g., "/data/results/*.pickle"
     input_patterns: List[str]
-    #: List of glob patterns for background files, e.g., ["/data/bg1/*.pickle", "/data/bg2/*.pickle"]
+    #: List of glob patterns for background files, e.g., ["bg1/*.pickle", "bg2/*."]
     background_patterns: List[str] = None
     #: Target masks, e.g., "/data/masks/*.mrc"
     mask_patterns: Optional[str] = None
@@ -230,46 +230,47 @@ class TMParameters:
     axis_sampling: Optional[float] = None
     axis_symmetry: int = 1
     cone_axis: int = 2
-    invert_cone: bool = False
-    no_use_optimized_set: bool = False
+    invert_cone: bool = field(default=False, metadata={"flag": True})
+    no_use_optimized_set: bool = field(default=False, metadata={"flag": True})
 
     # Microscope parameters
     acceleration_voltage: float = 300.0  # kV
     spherical_aberration: float = 2.7e7  # Å
     amplitude_contrast: float = 0.07
     defocus: Optional[float] = None  # Å
-    phase_shift: float = 0.0  # Dg
+    phase_shift: float = 0.0  # degrees
 
     # Processing options
     lowpass: Optional[float] = None  # Å
     highpass: Optional[float] = None  # Å
     pass_format: str = "sampling_rate"  # "sampling_rate", "voxel", "frequency"
-    no_pass_smooth: bool = True
+    no_pass_smooth: bool = field(default=True, metadata={"flag": False})
     interpolation_order: int = 3
     score_threshold: float = 0.0
     score: str = "FLCSphericalMask"
+    background_correction: Optional[str] = None
 
     # Weighting and correction
     tilt_weighting: Optional[str] = None  # "angle", "relion", "grigorieff"
     wedge_axes: str = "2,0"
-    whiten_spectrum: bool = False
-    scramble_phases: bool = False
-    invert_target_contrast: bool = False
+    whiten_spectrum: bool = field(default=False, metadata={"flag": True})
+    scramble_phases: bool = field(default=False, metadata={"flag": True})
+    invert_target_contrast: bool = field(default=False, metadata={"flag": True})
 
     # CTF parameters
     ctf_file: Optional[Path] = None
-    no_flip_phase: bool = True
-    correct_defocus_gradient: bool = False
+    no_flip_phase: bool = field(default=True, metadata={"flag": False})
+    correct_defocus_gradient: bool = field(default=False, metadata={"flag": True})
 
     # Performance options
-    centering: bool = False
-    pad_edges: bool = False
-    pad_filter: bool = False
-    use_mixed_precision: bool = False
-    use_memmap: bool = False
+    centering: bool = field(default=False, metadata={"flag": True})
+    pad_edges: bool = field(default=False, metadata={"flag": True})
+    pad_filter: bool = field(default=False, metadata={"flag": True})
+    use_mixed_precision: bool = field(default=False, metadata={"flag": True})
+    use_memmap: bool = field(default=False, metadata={"flag": True})
 
     # Analysis options
-    peak_calling: bool = False
+    peak_calling: bool = field(default=False, metadata={"flag": True})
     num_peaks: int = 1000
 
     # Backend selection
@@ -279,7 +280,7 @@ class TMParameters:
     # Reconstruction
     reconstruction_filter: str = "ramp"
     reconstruction_interpolation_order: int = 1
-    no_filter_target: bool = False
+    no_filter_target: bool = field(default=False, metadata={"flag": True})
 
     def __post_init__(self):
         """Validate parameters and convert units."""
@@ -337,6 +338,9 @@ class TMParameters:
             args["ctf-file"] = str(files.metadata)
             args["tilt-angles"] = str(files.metadata)
 
+        if self.background_correction:
+            args["background-correction"] = self.background_correction
+
         # Optional parameters
         if self.lowpass:
             args["lowpass"] = self.lowpass
@@ -378,43 +382,26 @@ class TMParameters:
         return {k: v for k, v in args.items() if v is not None}
 
     def get_flags(self) -> List[str]:
-        """Get boolean flags for pyTME command."""
         flags = []
-        if self.whiten_spectrum:
-            flags.append("whiten-spectrum")
-        if self.scramble_phases:
-            flags.append("scramble-phases")
-        if self.invert_target_contrast:
-            flags.append("invert-target-contrast")
-        if self.centering:
-            flags.append("centering")
-        if self.pad_edges:
-            flags.append("pad-edges")
-        if self.pad_filter:
-            flags.append("pad-filter")
-        if not self.no_pass_smooth:
-            flags.append("no-pass-smooth")
-        if self.use_mixed_precision:
-            flags.append("use-mixed-precision")
-        if self.use_memmap:
-            flags.append("use-memmap")
-        if self.peak_calling:
-            flags.append("peak-calling")
-        if not self.no_flip_phase:
-            flags.append("no-flip-phase")
-        if self.correct_defocus_gradient:
-            flags.append("correct-defocus-gradient")
-        if self.invert_cone:
-            flags.append("invert-cone")
-        if self.no_use_optimized_set:
-            flags.append("no-use-optimized-set")
-        if self.no_filter_target:
-            flags.append("no-filter-target")
+
+        for field_info in fields(self):
+            flag_meta = field_info.metadata.get("flag")
+            if flag_meta is None:
+                continue
+
+            value = getattr(self, field_info.name)
+            if not isinstance(value, bool):
+                continue
+
+            flag_name = field_info.name.replace("_", "-")
+            if (flag_meta is True and value) or (flag_meta is False and not value):
+                flags.append(flag_name)
+
         return flags
 
 
 @dataclass
-class AnalysisParameters:
+class AnalysisParameters(TMParameters):
     """Parameters for template matching analysis and peak calling."""
 
     # Peak calling
@@ -424,13 +411,12 @@ class AnalysisParameters:
     max_score: Optional[float] = None
     min_distance: int = 5
     min_boundary_distance: int = 0
-    mask_edges: bool = False
+    mask_edges: bool = field(default=False, metadata={"flag": True})
     n_false_positives: Optional[int] = None
 
     # Output format
     output_format: str = "relion4"
     output_directory: Optional[str] = None
-    angles_clockwise: bool = False
 
     # Advanced options
     extraction_box_size: Optional[int] = None
@@ -468,15 +454,6 @@ class AnalysisParameters:
 
         return {k: v for k, v in args.items() if v is not None}
 
-    def get_flags(self) -> List[str]:
-        """Get boolean flags for analyze_template_matching command."""
-        flags = []
-        if self.mask_edges:
-            flags.append("mask-edges")
-        if self.angles_clockwise:
-            flags.append("angles-clockwise")
-        return flags
-
 
 @dataclass
 class ComputeResources:
@@ -592,12 +569,10 @@ class ExecutionBackend(ABC):
     @abstractmethod
     def submit_job(self, task) -> str:
         """Submit a single job and return job ID or status."""
-        pass
 
     @abstractmethod
     def submit_jobs(self, tasks: List) -> List[str]:
         """Submit multiple jobs and return list of job IDs."""
-        pass
 
 
 class SlurmBackend(ExecutionBackend):
@@ -841,6 +816,13 @@ def parse_args():
         default="FLCSphericalMask",
         help="Template matching scoring function. Use FLC if mask is not spherical.",
     )
+    tm_group.add_argument(
+        "--background-correction",
+        choices=["phase-scrambling"],
+        required=False,
+        help="Transform cross-correlation into SNR-like values using a given method: "
+        "'phase-scrambling' uses a phase-scrambled template as background",
+    )
     tm_group.add_argument(
         "--score-threshold", type=float, default=0.0, help="Minimum score threshold"
     )
@@ -1006,11 +988,6 @@ def parse_args():
         default="relion4",
         help="Output format for analysis results",
     )
-    output_group.add_argument(
-        "--angles-clockwise",
-        action="store_true",
-        help="Report Euler angles in clockwise format expected by RELION",
-    )
 
     advanced_group = analysis_parser.add_argument_group("Advanced Options")
     advanced_group.add_argument(
@@ -1077,6 +1054,7 @@ def run_matching(args, resources):
         backend=args.backend,
         whiten_spectrum=args.whiten_spectrum,
         scramble_phases=args.scramble_phases,
+        background_correction=args.background_correction,
     )
     print_params = params.to_command_args(files[0], "")
     _ = print_params.pop("target")
@@ -1132,7 +1110,6 @@ def run_analysis(args, resources):
         mask_edges=args.mask_edges,
         n_false_positives=args.n_false_positives,
         output_format=args.output_format,
-        angles_clockwise=args.angles_clockwise,
         extraction_box_size=args.extraction_box_size,
     )
     print_params = params.to_command_args(files[0], Path(""))

scripts/refine_matches.py

@@ -10,11 +10,9 @@ import subprocess
 from sys import exit
 from os import unlink
 from time import time
-from os.path import join
 from typing import Tuple, List, Dict
 
 import numpy as np
-from scipy import optimize
 from sklearn.metrics import roc_auc_score
 
 from tme import Orientations, Density
@@ -66,7 +64,6 @@ def parse_args():
     matching_group.add_argument(
         "-i",
         "--template",
-        dest="template",
         type=str,
         required=True,
         help="Path to a template in PDB/MMCIF or other supported formats (see target).",
@@ -102,7 +99,7 @@ def parse_args():
     )
     matching_group.add_argument(
         "-s",
-        dest="score",
+        "--score",
         type=str,
         default="batchFLCSphericalMask",
         choices=list(MATCHING_EXHAUSTIVE_REGISTER.keys()),
@@ -197,6 +194,7 @@ def create_matching_argdict(args) -> Dict:
         "-n": args.cores,
         "--ctf-file": args.ctf_file,
         "--invert-target-contrast": args.invert_target_contrast,
+        "--backend" : args.backend,
     }
     return arg_dict
 
@@ -252,7 +250,7 @@ class DeepMatcher:
         if args.lowpass_range:
             self.filter_parameters["--lowpass"] = 0
         if args.highpass_range:
-            self.filter_parameters["--highpass"] = 200
+            self.filter_parameters["--highpass"] = 0
 
         self.postprocess_args = create_postprocessing_argdict(args)
         self.log_file = f"{args.output_prefix}_optimization_log.txt"
@@ -309,14 +307,14 @@ class DeepMatcher:
 
         match_template = argdict_to_command(
             self.match_template_args,
-            executable="match_template.py",
+            executable="match_template",
         )
         run_command(match_template)
 
         # Assume we get a new peak for each input in the same order
         postprocess = argdict_to_command(
             self.postprocess_args,
-            executable="postprocess.py",
+            executable="postprocess",
         )
         run_command(postprocess)
 

tests/preprocessing/test_compose.py

@@ -1,28 +1,38 @@
 import pytest
 
-from tme.filters import Compose
+from tme.filters import Compose, ComposableFilter
 from tme.backends import backend as be
 
 
-def mock_transform1(**kwargs):
-    return {"data": be.ones((10, 10)), "is_multiplicative_filter": True}
+class MockFilter(ComposableFilter):
+    def _evaluate(self, *args, **kwargs):
+        return {"data": be.ones((10, 10)), "shape": (10, 10)}
+
+
+class MockFilterNoMult(ComposableFilter):
+    def _evaluate(self, *args, **kwargs):
+        return {
+            "data": be.ones((10, 10)) * 3,
+            "shape": (10, 10),
+            "is_multiplicative_filter": False,
+        }
 
 
-def mock_transform2(**kwargs):
-    return {"data": be.ones((10, 10)) * 2, "is_multiplicative_filter": True}
+mock_transform = MockFilter()
+mock_transform_nomult = MockFilterNoMult()
 
 
-def mock_transform3(**kwargs):
-    return {"extra_info": "test"}
+def mock_transform_error(**kwargs):
+    return {"data": be.ones((10, 10)), "is_multiplicative_filter": True}
 
 
 class TestCompose:
     @pytest.fixture
     def compose_instance(self):
-        return Compose((mock_transform1, mock_transform2, mock_transform3))
+        return Compose((mock_transform, mock_transform, mock_transform))
 
     def test_init(self):
-        transforms = (mock_transform1, mock_transform2)
+        transforms = (mock_transform, mock_transform)
         compose = Compose(transforms)
         assert compose.transforms == transforms
 
@@ -32,45 +42,36 @@ class TestCompose:
         assert result == {}
 
     def test_call_single_transform(self):
-        compose = Compose((mock_transform1,))
-        result = compose()
+        compose = Compose((mock_transform,))
+        result = compose(return_real_fourier=False)
         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()
+        result = compose_instance(return_real_fourier=False)
         assert "data" in result
-        assert "extra_info" not in result
-        assert be.allclose(result["data"], be.ones((10, 10)) * 2)
+        assert be.allclose(result["data"], be.ones((10, 10)))
 
     def test_multiplicative_filter_composition(self):
-        compose = Compose((mock_transform1, mock_transform2))
-        result = compose()
+        compose = Compose((mock_transform, mock_transform))
+        result = compose(return_real_fourier=False)
         assert "data" in result
-        assert be.allclose(result["data"], be.ones((10, 10)) * 2)
+        assert be.allclose(result["data"], be.ones((10, 10)))
 
     @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)
+        result = compose_instance(**kwargs, return_real_fourier=False)
         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()
+        compose = Compose((mock_transform, mock_transform_nomult))
+        result = compose(return_real_fourier=False)
         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()
+        with pytest.raises(ValueError):
+            Compose((mock_transform, mock_transform_error))()

tests/preprocessing/test_frequency_filters.py

@@ -71,52 +71,43 @@ class TestBandPassFilter:
         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 "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: 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
 
     @pytest.mark.parametrize("shape", ((10, 10), (20, 20, 20), (30, 30)))
     def test_return_real_fourier(self, shape: Tuple[int]):
-        bpf = BandPassReconstructed(return_real_fourier=True)
+        bpf = BandPassReconstructed()
         result = bpf(shape=shape, lowpass=0.2, highpass=0.8)
-        expected_shape = tuple(compute_fourier_shape(shape, False))
-        assert result["data"].shape == expected_shape
+        assert result["data"].shape == shape
 
 
 class TestLinearWhiteningFilter:
     @pytest.mark.parametrize(
-        "shape, n_bins, batch_dimension",
+        "shape, n_bins",
         [
-            ((10, 10), None, None),
-            ((20, 20, 20), 15, 0),
-            ((30, 30, 30), 20, 1),
-            ((40, 40, 40, 40), 25, 2),
+            ((10, 10), None),
+            ((20, 20, 20), 15),
+            ((30, 30, 30), 20),
+            ((40, 40, 40, 40), 25),
         ],
     )
-    def test_compute_spectrum(
-        self, shape: Tuple[int], n_bins: int, batch_dimension: int
-    ):
+    def test_compute_spectrum(self, shape: Tuple[int], n_bins: 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
+            data_rfft, n_bins
         )
+        data_shape = tuple(int(x) for i, x in enumerate(data_rfft.shape))
 
         assert isinstance(bins, np.ndarray)
         assert isinstance(radial_averages, np.ndarray)
@@ -140,9 +131,8 @@ class TestLinearWhiteningFilter:
     @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),
+            ((10, 10), None, (), 1),
+            ((20, 20, 20), 15, 0, 1),
         ],
     )
     def test_call_method(
@@ -154,21 +144,17 @@ class TestLinearWhiteningFilter:
     ):
         data = be.random.random(shape)
         result = LinearWhiteningFilter()(
-            shape=shape,
-            data=data,
+            shape=tuple(x for i, x in enumerate(shape) if i != batch_dimension),
+            data_rfft=np.fft.rfftn(data),
             n_bins=n_bins,
-            batch_dimension=batch_dimension,
+            axes=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))
+        assert result["data"].shape == shape
 
     def test_call_method_with_data_rfft(self):
         shape = (30, 30, 30)
@@ -179,14 +165,13 @@ class TestLinearWhiteningFilter:
 
         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()(shape=shape, data=data)
+        result = LinearWhiteningFilter()(shape=shape, data_rfft=np.fft.rfftn(data))
 
         filter_mask = result["data"]
         assert np.all(filter_mask >= 0) and np.all(filter_mask <= 1)

tests/preprocessing/test_preprocessor.py

@@ -1,5 +1,4 @@
 import pytest
-import numpy as np
 
 from tme import Density, Structure, Preprocessor
 
@@ -49,15 +48,6 @@ class TestPreprocessor:
             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(
@@ -125,12 +115,3 @@ class TestPreprocessor:
             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

@@ -74,10 +74,22 @@ class TestPreprocessUtils:
     def test_fftfreqn(self, n, sampling_rate):
         assert np.allclose(
             fftfreqn(
-                shape=(n,), sampling_rate=sampling_rate, compute_euclidean_norm=True
+                shape=(n,),
+                sampling_rate=sampling_rate,
+                compute_euclidean_norm=True,
+                fftshift=True,
             ),
             np.abs(np.fft.ifftshift(np.fft.fftfreq(n=n, d=sampling_rate))),
         )
+        assert np.allclose(
+            fftfreqn(
+                shape=(n,),
+                sampling_rate=sampling_rate,
+                compute_euclidean_norm=True,
+                fftshift=False,
+            ),
+            np.abs(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):

tests/test_analyzer.py

@@ -109,13 +109,11 @@ class TestRecursiveMasking:
 
     @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):
+    def test__call__(self, num_peaks, compute_rotation):
         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:
@@ -131,13 +129,7 @@ class TestRecursiveMasking:
             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)
-
+        assert len(state[0] <= num_peaks)
 
 class TestMaxScoreOverRotations:
     def setup_method(self):