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

tests/test_backends.py

@@ -331,16 +331,6 @@ class TestBackends:
         real_arr = backend.irfftn(complex_arr)
         assert np.allclose(arr, backend.to_numpy_array(real_arr), rtol=0.3)
 
-    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
-    def test_extract_center(self, backend):
-        new_shape = np.divide(self.x1.shape, 2).astype(int)
-        base = self.backend.extract_center(arr=self.x1, newshape=new_shape)
-        other = backend.extract_center(
-            arr=backend.to_backend_array(self.x1), newshape=new_shape
-        )
-
-        assert np.allclose(base, backend.to_numpy_array(other), rtol=0.01)
-
     @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
     def test_compute_convolution_shapes(self, backend):
         base = self.backend.compute_convolution_shapes(self.x1.shape, self.x2.shape)
@@ -359,32 +349,69 @@ class TestBackends:
         elif dim == 3:
             arr[20:25, 21:26, 26:31] = 1
 
-        rotation_matrix = np.eye(dim)
-        rotation_matrix[0, 0] = -1
+        from tme.rotations import get_rotation_matrices
+
+        np.random.seed(42)
+        rotation_matrix = get_rotation_matrices(
+            dim=dim, angular_sampling=10, use_optimized_set=False
+        )[-1]
 
         out = np.zeros_like(arr)
+        out.setflags(write=True)
 
         arr_mask, out_mask = None, None
         if create_mask:
             arr_mask = np.multiply(np.random.rand(*arr.shape) > 0.5, 1.0)
             out_mask = np.zeros_like(arr_mask)
+            out_mask.setflags(write=True)
+
+        out, _ = NumpyFFTWBackend().rigid_transform(
+            arr=arr,
+            arr_mask=arr_mask,
+            rotation_matrix=rotation_matrix,
+            out=out,
+            out_mask=out_mask,
+            order=1,
+            use_geometric_center=True,
+        )
+
+        arr = backend.to_backend_array(arr.copy())
+        out_be = backend.to_backend_array(out.copy())
+        if create_mask:
             arr_mask = backend.to_backend_array(arr_mask)
             out_mask = backend.to_backend_array(out_mask)
 
-        arr = backend.to_backend_array(arr)
-        out = backend.to_backend_array(arr)
-
         rotation_matrix = backend.to_backend_array(rotation_matrix)
 
-        backend.rigid_transform(
+        out_be, _ = backend.rigid_transform(
             arr=arr,
             arr_mask=arr_mask,
             rotation_matrix=rotation_matrix,
-            out=out,
+            out=out_be,
             out_mask=out_mask,
+            order=1,
+            use_geometric_center=True,
         )
+        out_be = backend.to_numpy_array(out_be)
+        assert np.allclose(out, out_be, atol=0.3)
 
-        assert np.round(arr.sum(), 3) == np.round(out.sum(), 3)
+    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
+    @pytest.mark.parametrize("create_mask", (False, True))
+    def test_rigid_transform_identity(self, backend, create_mask):
+        dim = 3
+        shape = tuple(50 for _ in range(dim))
+
+        arr = np.zeros(shape)
+        arr[20:25, 21:26, 26:31] = 1
+
+        rotation_matrix = backend.to_backend_array(np.eye(dim))
+        out, _ = backend.rigid_transform(
+            arr=backend.to_backend_array(arr),
+            rotation_matrix=backend.to_backend_array(rotation_matrix),
+            order=1,
+            use_geometric_center=True,
+        )
+        assert np.allclose(out, arr, atol=0.01)
 
     @pytest.mark.parametrize("dim", (2, 3))
     @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
@@ -418,6 +445,7 @@ class TestBackends:
             out=out,
             out_mask=out_mask,
             batched=True,
+            order=1,
         )
 
         arr_b = backend.to_backend_array(arr_b)
@@ -430,6 +458,7 @@ class TestBackends:
                 rotation_matrix=rotation_matrix,
                 out=out_b[i],
                 out_mask=out_mask if out_mask is None else out_mask[i],
+                order=1,
             )
 
         assert np.allclose(arr, arr_b)

tests/test_density.py

@@ -101,9 +101,10 @@ class TestDensity:
 
     def test_from_file_baseline(self):
         self.test_to_file(gzip=False)
-        density = Density.from_file(str(BASEPATH.joinpath("Maps/emd_8621.mrc.gz")))
-        assert np.allclose(density.origin, (4.35, 2.90, -1.45), rtol=0.1)
-        assert np.allclose(density.sampling_rate, (1.45), rtol=0.3)
+        density = Density.from_file(str(BASEPATH.joinpath("Raw/em_map.map")))
+        assert density.shape == (19, 14, 20)
+        assert np.allclose(density.origin, (-52.8, -10.56, -52.8), rtol=0.1)
+        assert np.allclose(density.sampling_rate, (5.28), rtol=0.1)
 
     @pytest.mark.parametrize("extension", ("mrc", "em", "tiff", "h5"))
     @pytest.mark.parametrize("gzip", (True, False))
@@ -315,7 +316,7 @@ class TestDensity:
         "cutoff", [DEFAULT_DATA.min() - 1, 0, DEFAULT_DATA.max() - 0.1]
     )
     def test_centered(self, cutoff):
-        centered_density, translation = self.density.centered(cutoff=cutoff)
+        centered_density = self.density.centered(cutoff=cutoff)
         com = centered_density.center_of_mass(centered_density.data, 0)
 
         difference = np.abs(
@@ -334,7 +335,7 @@ class TestDensity:
             box = temp.minimum_enclosing_box(cutoff=0, use_geometric_center=True)
             temp.adjust_box(box)
         else:
-            temp, translation = temp.centered()
+            temp = temp.centered()
 
         swaps = set(permutations([0, 1, 2]))
         temp_matrix = np.eye(temp.data.ndim).astype(np.float32)
@@ -436,7 +437,7 @@ class TestDensity:
         arr = np.linalg.norm(np.indices(shape) - position, axis=0)
         arr = (arr <= radius).astype(np.float32)
 
-        center_of_mass = Density.center_of_mass(arr)
+        center_of_mass = Density(arr).center_of_mass()
         assert np.allclose(center, center_of_mass)
 
     @pytest.mark.parametrize(
@@ -451,7 +452,7 @@ class TestDensity:
         target[5:10, 15:22, 10:13] = 1
 
         target = Density(target, sampling_rate=(1, 1, 1), origin=(0, 0, 0))
-        target, translation = target.centered(cutoff=0)
+        target = target.centered(cutoff=0)
 
         template = target.copy()
 
@@ -474,21 +475,79 @@ class TestDensity:
         assert np.allclose(np.linalg.inv(rotation), initial_rotation, atol=0.2)
 
     def test_match_structure_to_density(self):
-        density = Density.from_file("tests/data/Maps/emd_8621.mrc.gz")
-        density = density.resample(density.sampling_rate * 4)
-        structure = Structure.from_file(
-            "tests/data/Structures/5uz4.cif", filter_by_residues=None
+        mask = create_mask(
+            mask_type="ellipse",
+            center=(20, 20, 20),
+            radius=(10, 5, 10),
+            shape=(50, 50, 50),
         )
 
+        coordinates = np.array(np.where(mask > 0), dtype=np.float32).T
+        structure = Structure(
+            record_type=[
+                "ATOM",
+            ]
+            * coordinates.shape[0],
+            atom_serial_number=list(range(coordinates.shape[0])),
+            atom_name=[
+                "C",
+            ]
+            * coordinates.shape[0],
+            atom_coordinate=coordinates,
+            alternate_location_indicator=[
+                ".",
+            ]
+            * coordinates.shape[0],
+            residue_name=[
+                "GLY",
+            ]
+            * coordinates.shape[0],
+            chain_identifier=[
+                "A",
+            ]
+            * coordinates.shape[0],
+            residue_sequence_number=[
+                0,
+            ]
+            * coordinates.shape[0],
+            code_for_residue_insertion=[
+                "?",
+            ]
+            * coordinates.shape[0],
+            occupancy=[
+                0,
+            ]
+            * coordinates.shape[0],
+            temperature_factor=[
+                0,
+            ]
+            * coordinates.shape[0],
+            segment_identifier=[
+                "1",
+            ]
+            * coordinates.shape[0],
+            element_symbol=[
+                "C",
+            ]
+            * coordinates.shape[0],
+            charge=[
+                "?",
+            ]
+            * coordinates.shape[0],
+            metadata={},
+        )
+        density = Density(mask, sampling_rate=1.0)
+        density = density.resample(density.sampling_rate * 2)
+
         initial_translation = np.array([-1, 0, 5])
         initial_rotation = euler_to_rotationmatrix((-10, 2, 5))
-        structure.rigid_transform(
+        structure_mod = structure.rigid_transform(
             translation=initial_translation, rotation_matrix=initial_rotation
         )
         np.random.seed(12)
         ret = Density.match_structure_to_density(
             target=density,
-            template=structure,
+            template=structure_mod,
             cutoff_target=0,
             scoring_method="CrossCorrelation",
             maxiter=10,

tests/test_extensions.py

@@ -53,6 +53,7 @@ class TestExtensions:
     @pytest.mark.parametrize("min_distance", [0, 5, 10])
     def test_find_candidate_indices(self, dimension, dtype, min_distance):
         coordinates = COORDINATES[dimension].astype(dtype)
+        print(coordinates.shape)
 
         min_distance = np.array([min_distance]).astype(dtype)[0]
 

tests/test_matching_cli.py

@@ -10,7 +10,7 @@ from tme import Density, Orientations
 from tme.backends import backend as be
 
 np.random.seed(42)
-available_backends = (x for x in be.available_backends() if x != "mlx")
+available_backends = tuple(x for x in be.available_backends() if x != "mlx")
 
 
 def argdict_to_command(input_args, executable: str):
@@ -29,7 +29,7 @@ def argdict_to_command(input_args, executable: str):
     return " ".join(ret)
 
 
-class TestMatchTemplate:
+class TestSetup:
     @classmethod
     def setup_class(cls):
         target = np.random.rand(20, 20, 20)
@@ -98,6 +98,7 @@ class TestMatchTemplate:
         use_target_mask: bool = False,
         backend: str = "numpyfftw",
         test_rejection_sampling: bool = False,
+        background_correction: str = None,
     ):
         output_path = tempfile.NamedTemporaryFile(delete=False, suffix="pickle").name
 
@@ -120,6 +121,9 @@ class TestMatchTemplate:
         if test_rejection_sampling:
             argdict["--orientations"] = self.orientations_path
 
+        if background_correction is not None:
+            argdict["--background-correction"] = background_correction
+
         if test_filter:
             argdict["--lowpass"] = 30
             argdict["--defocus"] = 3000
@@ -136,11 +140,15 @@ class TestMatchTemplate:
         assert ret.returncode == 0
         return output_path
 
+
+class TestMatchTemplate(TestSetup):
+
     @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))
     @pytest.mark.parametrize("test_rejection_sampling", (False, True))
+    @pytest.mark.parametrize("background_correction", (None, "phase-scrambling"))
     def test_match_template(
         self,
         backend: bool,
@@ -148,8 +156,14 @@ class TestMatchTemplate:
         use_template_mask: bool,
         test_filter: bool,
         test_rejection_sampling: bool,
+        background_correction: str,
     ):
-        if backend == "jax" and (call_peaks or test_rejection_sampling):
+        # Jax does not support peak calling yet
+        if backend == "jax" and call_peaks:
+            return None
+
+        # These use different analyzers raising an error in the interface
+        if call_peaks and test_rejection_sampling:
             return None
 
         self.run_matching(
@@ -163,10 +177,11 @@ class TestMatchTemplate:
             template_mask_path=self.template_mask_path,
             target_mask_path=self.target_mask_path,
             test_rejection_sampling=test_rejection_sampling,
+            background_correction=background_correction,
         )
 
 
-class TestPostprocessing(TestMatchTemplate):
+class TestPostprocessing(TestSetup):
     @classmethod
     def setup_class(cls):
         super().setup_class()
@@ -256,7 +271,6 @@ class TestPostprocessing(TestMatchTemplate):
         }
         cmd = argdict_to_command(argdict, executable="postprocess.py")
         ret = subprocess.run(cmd, capture_output=True, shell=True)
-        print(ret)
 
         match output_format:
             case "orientations":
@@ -289,14 +303,14 @@ class TestPostprocessing(TestMatchTemplate):
         assert ret.returncode == 0
 
 
-class TestEstimateMemoryUsage(TestMatchTemplate):
+class TestEstimateMemoryUsage(TestSetup):
     @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):
+    def test_estimation_cli(self, ncores, pad_edges):
 
         argdict = {
             "-m": self.target_path,
@@ -311,7 +325,7 @@ class TestEstimateMemoryUsage(TestMatchTemplate):
         assert ret.returncode == 0
 
 
-class TestPreprocess(TestMatchTemplate):
+class TestPreprocess(TestSetup):
     @classmethod
     def setup_class(cls):
         super().setup_class()
@@ -319,7 +333,7 @@ class TestPreprocess(TestMatchTemplate):
     @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):
+    def test_preprocess_cli(self, backend, align_axis, invert_contrast):
 
         argdict = {
             "-m": self.target_path,

tests/test_matching_exhaustive.py

@@ -1,5 +1,5 @@
-import numpy as np
 import pytest
+import numpy as np
 
 from scipy.ndimage import laplace
 
@@ -7,7 +7,7 @@ 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_subsets,
+    match_exhaustive,
     MATCHING_EXHAUSTIVE_REGISTER,
     register_matching_exhaustive,
 )
@@ -35,11 +35,11 @@ class TestMatchExhaustive:
         self.coordinates_weights = None
         self.rotations = None
 
-    @pytest.mark.parametrize("evaluate_peak", (True,))
+    @pytest.mark.parametrize("evaluate_peak", (True, False))
     @pytest.mark.parametrize("score", tuple(MATCHING_EXHAUSTIVE_REGISTER.keys()))
     @pytest.mark.parametrize("job_schedule", ((2, 1),))
     @pytest.mark.parametrize("pad_edge", (False, True))
-    def test_scan_subset(
+    def test_match_exhaustive(
         self,
         score: str,
         job_schedule: int,
@@ -64,7 +64,7 @@ class TestMatchExhaustive:
         if evaluate_peak:
             callback_class = PeakCallerSort
 
-        ret = scan_subsets(
+        ret = match_exhaustive(
             matching_data=matching_data,
             matching_setup=setup,
             matching_score=process,

tests/test_matching_utils.py

@@ -15,7 +15,7 @@ from tme.matching_utils import (
     apply_convolution_mode,
     write_pickle,
     load_pickle,
-    _normalize_template_overflow_safe,
+    _standardize_safe,
 )
 
 BASEPATH = files("tests.data")
@@ -127,12 +127,12 @@ class TestMatchingUtils:
         loaded_data = load_pickle(filename)
         assert np.array_equal(loaded_data, data)
 
-    def test_normalize_template_overflow_safe(self):
+    def test_standardize_safe(self):
         template = be.random.random((10, 10)).astype(be.float32)
         mask = be.ones_like(template)
         n_observations = 100.0
 
-        result = _normalize_template_overflow_safe(template, mask, n_observations)
+        result = _standardize_safe(template, mask, n_observations)
         assert result.shape == template.shape
         assert result.dtype == template.dtype
         assert np.allclose(result.mean(), 0, atol=0.1)

tests/test_rotations.py

@@ -35,24 +35,19 @@ class TestRotations:
         )
 
     @pytest.mark.parametrize(
-        "initial_vector, target_vector, convention",
+        "initial_vector, target_vector",
         [
-            ([1, 0, 0], [0, 1, 0], None),
-            ([0, 1, 0], [0, 0, 1], "zyx"),
-            ([1, 1, 1], [1, 0, 0], "xyz"),
+            ([1, 0, 0], [0, 1, 0]),
+            ([0, 1, 0], [0, 0, 1]),
+            ([1, 1, 1], [1, 0, 0]),
         ],
     )
-    def test_align_vectors(self, initial_vector, target_vector, convention):
-        result = align_vectors(initial_vector, target_vector, convention)
+    def test_align_vectors(self, initial_vector, target_vector):
+        result = align_vectors(initial_vector, target_vector)
 
         assert isinstance(result, np.ndarray)
-        if convention is None:
-            assert result.shape == (3, 3)
-            assert np.allclose(np.dot(result, result.T), np.eye(3), atol=1e-6)
-        else:
-            assert len(result) == 3
-            result = Rotation.from_euler(convention, result, degrees=True).as_matrix()
-            assert np.allclose(np.dot(result, result.T), np.eye(3), atol=1e-6)
+        assert result.shape == (3, 3)
+        assert np.allclose(np.dot(result, result.T), np.eye(3), atol=1e-6)
 
         rotated = np.dot(Rotation.from_matrix(result).as_matrix(), initial_vector)
         assert np.allclose(
@@ -62,11 +57,11 @@ class TestRotations:
         )
 
     @pytest.mark.parametrize(
-        "cone_angle, cone_sampling, axis_angle, axis_sampling, vector, n_symmetry, convention",
+        "cone_angle, cone_sampling, axis_angle, axis_sampling, vector, n_symmetry",
         [
-            (30, 5, 360, None, (1, 0, 0), 1, None),
-            (45, 10, 180, 15, (0, 1, 0), 2, "zyx"),
-            (60, 15, 90, 30, (0, 0, 1), 4, "xyz"),
+            (30, 5, 360, None, (1, 0, 0), 1),
+            (45, 10, 180, 15, (0, 1, 0), 2),
+            (60, 15, 90, 30, (0, 0, 1), 4),
         ],
     )
     def test_get_cone_rotations(
@@ -77,7 +72,6 @@ class TestRotations:
         axis_sampling,
         vector,
         n_symmetry,
-        convention,
     ):
         result = get_cone_rotations(
             cone_angle=cone_angle,
@@ -86,14 +80,10 @@ class TestRotations:
             axis_sampling=axis_sampling,
             reference=vector,
             n_symmetry=n_symmetry,
-            seq=convention,
         )
 
         assert isinstance(result, np.ndarray)
-        if convention is None:
-            assert result.shape[1:] == (3, 3)
-        else:
-            assert result.shape[1] == 3
+        assert result.shape[1:] == (3, 3)
 
     def test_euler_conversion(self):
         rotation_matrix_initial = np.array(

tests/test_structure.py

@@ -6,7 +6,6 @@ import pytest
 import numpy as np
 
 from tme import Structure
-from tme.matching_utils import minimum_enclosing_box
 from tme.rotations import euler_to_rotationmatrix
 
 
@@ -196,11 +195,6 @@ class TestStructure:
         assert center_of_mass.shape[0] == self.structure.atom_coordinate.shape[1]
         assert np.allclose(center_of_mass, [-0.89391639, 29.94908928, -2.64736741])
 
-    def test_centered(self):
-        ret, translation = self.structure.centered()
-        box = minimum_enclosing_box(coordinates=self.structure.atom_coordinate.T)
-        assert np.allclose(ret.center_of_mass(), np.divide(box, 2), atol=1)
-
     def test__get_atom_weights_error(self):
         with pytest.raises(NotImplementedError):
             self.structure._get_atom_weights(
@@ -242,6 +236,6 @@ class TestStructure:
         assert final_rmsd <= 0.1
 
         aligned, final_rmsd = Structure.align_structures(
-            self.structure, structure_transform, sampling_rate=1
+            self.structure, structure_transform
         )
         assert final_rmsd <= 1

tme/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.1"
+__version__ = "0.3.2"