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

tests/test_orientations.py

@@ -0,0 +1,173 @@
+from tempfile import mkstemp
+
+import pytest
+import numpy as np
+
+from tme import Orientations
+
+
+class TestDensity:
+    def setup_method(self):
+        self.translations = np.random.rand(100, 3).astype(np.float32)
+        self.rotations = np.random.rand(100, 3).astype(np.float32)
+        self.scores = np.random.rand(100).astype(np.float32)
+        self.details = np.random.rand(100).astype(np.float32)
+
+        self.orientations = Orientations(
+            translations=self.translations,
+            rotations=self.rotations,
+            scores=self.scores,
+            details=self.details,
+        )
+
+    def teardown_method(self):
+        self.translations = None
+        self.rotations = None
+        self.scores = None
+        self.details = None
+        self.orientations = None
+
+    def test_initialization(self):
+        orientations = Orientations(
+            translations=self.translations,
+            rotations=self.rotations,
+            scores=self.scores,
+            details=self.details,
+        )
+
+        assert np.array_equal(self.translations, orientations.translations)
+        assert np.array_equal(self.rotations, orientations.rotations)
+        assert np.array_equal(self.scores, orientations.scores)
+        assert np.array_equal(self.details, orientations.details)
+
+    def test_initialization_type(self):
+        orientations = Orientations(
+            translations=self.translations.astype(int),
+            rotations=self.rotations.astype(int),
+            scores=self.scores.astype(int),
+            details=self.details.astype(int),
+        )
+        assert np.issubdtype(orientations.translations.dtype, np.floating)
+        assert np.issubdtype(orientations.rotations.dtype, np.floating)
+        assert np.issubdtype(orientations.scores.dtype, np.floating)
+        assert np.issubdtype(orientations.details.dtype, np.floating)
+
+    def test_initialization_error(self):
+        with pytest.raises(ValueError):
+            _ = Orientations(
+                translations=self.translations,
+                rotations=np.random.rand(self.translations.shape[0] + 1),
+                scores=self.scores,
+                details=self.details,
+            )
+
+        with pytest.raises(ValueError):
+            _ = Orientations(
+                translations=np.random.rand(self.translations.shape[0]),
+                rotations=np.random.rand(self.translations.shape[0] + 1),
+                scores=self.scores,
+                details=self.details,
+            )
+            _ = Orientations(
+                translations=self.translations,
+                rotations=np.random.rand(self.translations.shape[0]),
+                scores=self.scores,
+                details=self.details,
+            )
+
+        assert True
+
+    @pytest.mark.parametrize("file_format", ("text", "relion", "tbl"))
+    def test_to_file(self, file_format: str):
+        _, output_file = mkstemp(suffix=f".{file_format}")
+        self.orientations.to_file(output_file)
+        assert True
+
+    @pytest.mark.parametrize("file_format", ("text", "star", "tbl"))
+    def test_from_file(self, file_format: str):
+        _, output_file = mkstemp(suffix=f".{file_format}")
+        self.orientations.to_file(output_file)
+        orientations_new = Orientations.from_file(output_file)
+
+        assert np.array_equal(
+            self.orientations.translations, orientations_new.translations
+        )
+
+    @pytest.mark.parametrize("input_format", ("text", "star", "tbl"))
+    @pytest.mark.parametrize("output_format", ("text", "star", "tbl"))
+    def test_file_format_io(self, input_format: str, output_format: str):
+        _, output_file = mkstemp(suffix=f".{input_format}")
+        _, output_file2 = mkstemp(suffix=f".{output_format}")
+
+        self.orientations.to_file(output_file)
+        orientations_new = Orientations.from_file(output_file)
+        orientations_new.to_file(output_file2)
+
+        assert True
+
+    @pytest.mark.parametrize("drop_oob", (True, False))
+    @pytest.mark.parametrize("shape", (10, 40, 80))
+    @pytest.mark.parametrize("odd", (True, False))
+    def test_extraction(self, shape: int, drop_oob: bool, odd: bool):
+        if odd:
+            shape = shape + (1 - shape % 2)
+
+        data = np.random.rand(50, 50, 50)
+        translations = np.array([[25, 25, 25], [15, 25, 35], [35, 25, 15], [0, 15, 49]])
+        orientations = Orientations(
+            translations=translations,
+            rotations=np.random.rand(*translations.shape),
+            scores=np.random.rand(translations.shape[0]),
+            details=np.random.rand(translations.shape[0]),
+        )
+        extraction_shape = np.repeat(np.array(shape), data.ndim)
+        orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
+            target_shape=data.shape,
+            extraction_shape=extraction_shape,
+            drop_out_of_box=drop_oob,
+            return_orientations=True,
+        )
+        assert orientations.translations.shape[0] == len(cand_slices)
+        assert len(cand_slices) == len(obs_slices)
+
+        cand_slices2, obs_slices2 = orientations.get_extraction_slices(
+            target_shape=data.shape,
+            extraction_shape=extraction_shape,
+            drop_out_of_box=drop_oob,
+            return_orientations=False,
+        )
+        assert cand_slices == cand_slices2
+        assert obs_slices == obs_slices2
+
+        # Check whether extraction slices are pasted in center
+        out = np.zeros(extraction_shape, dtype=data.dtype)
+        center = np.divide(extraction_shape, 2) + np.mod(extraction_shape, 2)
+        center = center.astype(int)
+        for index, (cand_slice, obs_slice) in enumerate(zip(cand_slices, obs_slices)):
+            out[cand_slice] = data[obs_slice]
+            assert np.allclose(
+                out[tuple(center)],
+                data[tuple(orientations.translations[index].astype(int))],
+            )
+
+    @pytest.mark.parametrize(
+        "order", (("x", "y", "z"), ("z", "y", "x"), ("y", "x", "z"))
+    )
+    def test_txt_sort(self, order: str):
+        _, output_file = mkstemp(suffix=".tsv")
+        translations = ((50, 30, 20), (10, 5, 30))
+
+        with open(output_file, mode="w", encoding="utf-8") as ofile:
+            ofile.write("\t".join([str(x) for x in order]) + "\n")
+            for translation in translations:
+                ofile.write("\t".join([str(x) for x in translation]) + "\n")
+
+        translations = np.array(translations).astype(np.float32)
+        orientations = Orientations.from_file(output_file)
+
+        out_order = zip(order, range(len(order)))
+        out_order = tuple(
+            x[1] for x in sorted(out_order, key=lambda x: x[0], reverse=True)
+        )
+
+        assert np.array_equal(translations[..., out_order], orientations.translations)

tests/test_packaging.py

@@ -0,0 +1,95 @@
+import pytest
+import numpy as np
+
+from tme.package import GaussianKernel, KernelFitting
+from tme.analyzer import PeakCallerFast
+
+class TestKernelFitting:
+    def setup_method(self):
+        self.number_of_peaks = 100
+        self.min_distance = 5
+        self.data = np.random.rand(100, 100, 100)
+        self.rotation_matrix = np.eye(3)
+
+        self.kernel_box = np.random.choice(np.arange(5, 15), size=self.data.ndim)
+        self.template = np.random.rand(*self.kernel_box)
+
+    @pytest.mark.parametrize(
+        "kernel",
+        [
+            (GaussianKernel),
+        ],
+    )
+    def test_initialization(self, kernel):
+        kernel_params = kernel.estimate_parameters(self.template)
+        _ = KernelFitting(
+            number_of_peaks=self.number_of_peaks,
+            min_distance=self.min_distance,
+            kernel_box=self.kernel_box,
+            kernel_class=kernel,
+            kernel_params=kernel_params,
+            peak_caller=PeakCallerFast,
+        )
+
+    @pytest.mark.parametrize("kernel",[GaussianKernel])
+    def test__call__(self, kernel):
+        kernel_params = kernel.estimate_parameters(self.template)
+        score_analyzer = KernelFitting(
+            number_of_peaks=self.number_of_peaks,
+            min_distance=self.min_distance,
+            kernel_box=self.kernel_box,
+            kernel_class=kernel,
+            kernel_params=kernel_params,
+            peak_caller=PeakCallerFast,
+        )
+        score_analyzer(self.data, rotation_matrix=self.rotation_matrix)
+
+
+class TestKernel:
+    def setup_method(self):
+        self.number_of_peaks = 100
+        self.min_distance = 5
+        self.data = np.random.rand(100, 100, 100)
+        self.rotation_matrix = np.eye(3)
+
+    def test_initialization(self):
+        _ = GaussianKernel()
+
+    @pytest.mark.parametrize(
+        "kernel",
+        [
+            GaussianKernel,
+        ],
+    )
+    def test_fit(self, kernel):
+        params, success, final_error = kernel.fit(self.data)
+        assert isinstance(params, tuple)
+        assert len(params) == 3
+
+    @pytest.mark.parametrize(
+        "kernel",
+        [
+            GaussianKernel,
+        ],
+    )
+    def test_fit_height(self, kernel):
+        params, success, final_error = kernel.fit(self.data)
+        height, mean, cov = params
+        params, succes, final_error = kernel.fit_height(
+            data=self.data, mean=mean, cov=cov
+        )
+        assert np.allclose(params[1], mean)
+        assert np.allclose(params[2], cov)
+
+    @pytest.mark.parametrize(
+        "kernel",
+        [
+            GaussianKernel,
+        ],
+    )
+    def test_score(self, kernel):
+        params, *_ = kernel.fit(self.data)
+        params2, *_ = kernel.fit(self.data)
+
+        score = kernel.score(params, params2)
+        assert score == 1

tests/test_parser.py

@@ -0,0 +1,33 @@
+import pytest
+from importlib_resources import files
+
+from tme.parser import Parser, PDBParser
+
+
+class TestParser:
+    def setup_method(self):
+        self.pdb_file = str(files("tests.data").joinpath("Structures/5khe.pdb"))
+
+    def teardown_method(self):
+        self.pdb_file = None
+
+    def test_initialize_parser_error(self):
+        with pytest.raises(TypeError):
+            _ = Parser(self.pdb_file)
+
+    def test_parser_get(self):
+        parser = PDBParser(self.pdb_file)
+        assert parser.get("NOTPRESENT", None) is None
+        assert parser.get("record_type", None) is not None
+
+    def test_parser_keys(self):
+        parser = PDBParser(self.pdb_file)
+        assert parser.keys() == parser._data.keys()
+
+    def test_parser_values(self):
+        parser = PDBParser(self.pdb_file)
+        assert str(parser.values()) == str(parser._data.values())
+
+    def test_parser_items(self):
+        parser = PDBParser(self.pdb_file)
+        assert parser.items() == parser._data.items()

tests/test_structure.py

@@ -0,0 +1,243 @@
+from os import remove
+from tempfile import mkstemp
+from importlib_resources import files
+
+import pytest
+import numpy as np
+
+from tme import Structure
+from tme.matching_utils import euler_to_rotationmatrix, minimum_enclosing_box
+
+
+STRUCTURE_ATTRIBUTES = [
+    "record_type",
+    "atom_serial_number",
+    "atom_name",
+    "atom_coordinate",
+    "alternate_location_indicator",
+    "residue_name",
+    "chain_identifier",
+    "residue_sequence_number",
+    "code_for_residue_insertion",
+    "occupancy",
+    "temperature_factor",
+    "segment_identifier",
+    "element_symbol",
+    "charge",
+    "metadata",
+]
+
+
+class TestStructure:
+    def setup_method(self):
+        self.structure = Structure.from_file(
+            str(files("tests.data").joinpath("Structures/5khe.cif"))
+        )
+        _, self.path = mkstemp()
+
+    def teardown_method(self):
+        del self.structure
+        remove(self.path)
+
+    def compare_structures(self, structure1, structure2, exclude_attributes=[]):
+        for attribute in STRUCTURE_ATTRIBUTES:
+            if attribute in exclude_attributes:
+                continue
+            value = getattr(structure1, attribute)
+            value_comparison = getattr(structure2, attribute)
+            if isinstance(value, np.ndarray):
+                assert np.all(value_comparison == value)
+            else:
+                assert value == value_comparison
+
+    def test_initialization(self):
+        structure = Structure(
+            record_type=self.structure.record_type,
+            atom_serial_number=self.structure.atom_serial_number,
+            atom_name=self.structure.atom_name,
+            atom_coordinate=self.structure.atom_coordinate,
+            alternate_location_indicator=self.structure.alternate_location_indicator,
+            residue_name=self.structure.residue_name,
+            chain_identifier=self.structure.chain_identifier,
+            residue_sequence_number=self.structure.residue_sequence_number,
+            code_for_residue_insertion=self.structure.code_for_residue_insertion,
+            occupancy=self.structure.occupancy,
+            temperature_factor=self.structure.temperature_factor,
+            segment_identifier=self.structure.segment_identifier,
+            element_symbol=self.structure.element_symbol,
+            charge=self.structure.charge,
+            metadata=self.structure.metadata,
+        )
+
+        for attribute in STRUCTURE_ATTRIBUTES:
+            value = getattr(self.structure, attribute)
+            value_comparison = getattr(structure, attribute)
+            if isinstance(value, np.ndarray):
+                assert np.all(value_comparison == value)
+            else:
+                assert value == value_comparison
+
+    @pytest.mark.parametrize(
+        "modified_attribute",
+        [
+            ("record_type"),
+            ("atom_serial_number"),
+            ("atom_name"),
+            ("atom_coordinate"),
+            ("alternate_location_indicator"),
+            ("residue_name"),
+            ("chain_identifier"),
+            ("residue_sequence_number"),
+            ("code_for_residue_insertion"),
+            ("occupancy"),
+            ("temperature_factor"),
+            ("segment_identifier"),
+            ("element_symbol"),
+        ],
+    )
+    def test_initialization_errors(self, modified_attribute):
+        kwargs = {
+            attribute: getattr(self.structure, attribute)
+            for attribute in STRUCTURE_ATTRIBUTES
+            if attribute != modified_attribute
+        }
+        kwargs[modified_attribute] = getattr(self.structure, modified_attribute)[:1]
+
+        with pytest.raises(ValueError):
+            Structure(**kwargs)
+
+    def test__getitem__(self):
+        ret_single_index = self.structure[1]
+        ret = self.structure[[1]]
+        self.compare_structures(ret_single_index, ret)
+
+        ret = self.structure[self.structure.record_type == "ATOM"]
+        assert np.all(ret.record_type == "ATOM")
+
+        ret = self.structure[self.structure.element_symbol == "C"]
+        assert np.all(ret.element_symbol == "C")
+
+    def test__repr__(self):
+        unique_chains = "-".join(
+            [
+                ",".join([str(x) for x in entity])
+                for entity in self.structure.metadata["unique_chains"]
+            ]
+        )
+
+        min_atom = np.min(self.structure.atom_serial_number)
+        max_atom = np.max(self.structure.atom_serial_number)
+        n_atom = self.structure.atom_serial_number.size
+
+        min_residue = np.min(self.structure.residue_sequence_number)
+        max_residue = np.max(self.structure.residue_sequence_number)
+        n_residue = np.unique(self.structure.residue_sequence_number).size
+
+        repr_str = (
+            f"Structure object at {id(self.structure)}\n"
+            f"Unique Chains: {unique_chains}, "
+            f"Atom Range: {min_atom}-{max_atom} [N = {n_atom}], "
+            f"Residue Range: {min_residue}-{max_residue} [N = {n_residue}]"
+        )
+        assert repr_str == self.structure.__repr__()
+
+    @pytest.mark.parametrize(
+        "path",
+        [
+            str(files("tests.data").joinpath("Structures/5khe.cif")),
+            str(files("tests.data").joinpath("Structures/5khe.pdb")),
+        ],
+    )
+    def test_fromfile(self, path):
+        _ = Structure.from_file(path)
+
+    def test_fromfile_error(self):
+        with pytest.raises(NotImplementedError):
+            _ = Structure.from_file("madeup.extension")
+
+    @pytest.mark.parametrize("file_format", [("cif"), ("pdb")])
+    def test_to_file(self, file_format):
+        _, path = mkstemp()
+        path = f"{path}.{file_format}"
+        self.structure.to_file(path)
+        read = self.structure.from_file(path)
+        comparison = self.structure.copy()
+
+        self.compare_structures(comparison, read, exclude_attributes=["metadata"])
+
+    def test_to_file_error(self):
+        _, path = mkstemp()
+        path = f"{path}.RAISERROR"
+        with pytest.raises(NotImplementedError):
+            self.structure.to_file(path)
+
+    def test_subset_by_chain(self):
+        chain = "A"
+        ret = self.structure.subset_by_chain(chain=chain)
+        assert np.all(ret.chain_identifier == chain)
+
+    def test_subset_by_chain_range(self):
+        chain, start, stop = "A", 0, 20
+        ret = self.structure.subset_by_range(chain=chain, start=start, stop=stop)
+        assert np.all(ret.chain_identifier == chain)
+        assert np.all(
+            np.logical_and(
+                ret.residue_sequence_number >= start,
+                ret.residue_sequence_number <= stop,
+            )
+        )
+
+    def test_center_of_mass(self):
+        center_of_mass = self.structure.center_of_mass()
+        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(
+                self.structure.atom_name, weight_type="RAISEERROR"
+            )
+
+    def test_compare_structures(self):
+        rmsd = Structure.compare_structures(self.structure, self.structure)
+        assert rmsd == 0
+
+        rmsd = Structure.compare_structures(
+            self.structure, self.structure, weighted=True
+        )
+        assert rmsd == 0
+
+        translation = (3, 0, 0)
+        structure_transform = self.structure.rigid_transform(
+            translation=translation,
+            rotation_matrix=np.eye(self.structure.atom_coordinate.shape[1]),
+        )
+        rmsd = Structure.compare_structures(self.structure, structure_transform)
+        assert np.allclose(rmsd, np.linalg.norm(translation))
+
+    def test_comopare_structures_error(self):
+        ret = self.structure[[1, 2, 3, 4, 5]]
+        with pytest.raises(ValueError):
+            Structure.compare_structures(self.structure, ret)
+
+    def test_align_structures(self):
+        rotation_matrix = euler_to_rotationmatrix((20, -10, 45))
+        translation = (10, 0, -15)
+
+        structure_transform = self.structure.rigid_transform(
+            rotation_matrix=rotation_matrix, translation=translation
+        )
+        aligned, final_rmsd = Structure.align_structures(
+            self.structure, structure_transform
+        )
+        assert final_rmsd <= 0.1
+
+        aligned, final_rmsd = Structure.align_structures(
+            self.structure, structure_transform, sampling_rate=1
+        )
+        assert final_rmsd <= 1

tme/__init__.py

@@ -1,4 +1,3 @@
-from . import extensions
 from .__version__ import __version__
 from .density import Density
 from .preprocessor import Preprocessor

tme/__version__.py

@@ -1 +1 @@
-__version__ = "0.2.3"
+__version__ = "0.2.4"

tme/backends/jax_backend.py

@@ -208,13 +208,14 @@ class JaxBackend(NumpyFFTWBackend):
         # Applying the filter leads to more FFTs
         fastt_shape = matching_data._template.shape
         if create_template_filter:
-            _, fastt_shape, _, tshift = matching_data._fourier_padding(
-                target_shape=self.to_numpy_array(matching_data._template.shape),
-                template_shape=self.to_numpy_array(
-                    [1 for _ in matching_data._template.shape]
-                ),
-                pad_fourier=False,
-            )
+            # _, fastt_shape, _, tshift = matching_data._fourier_padding(
+            #     target_shape=self.to_numpy_array(matching_data._template.shape),
+            #     template_shape=self.to_numpy_array(
+            #         [1 for _ in matching_data._template.shape]
+            #     ),
+            #     pad_fourier=False,
+            # )
+            fastt_shape = matching_data._template.shape
 
         ret, template_filter, target_filter = [], 1, 1
         rotation_mapping = {

tme/density.py

@@ -744,9 +744,16 @@ class Density:
         >>> )
 
         :py:meth:`Density.from_structure` supports a variety of methods to convert
-        atoms into densities. In addition to 'atomic_weight', 'atomic_number',
-        and 'van_der_waals_radius', its possible to use experimentally determined
-        scattering factors from various sources:
+        atoms into densities
+
+        >>> density = Density.from_structure(
+        >>>    filename_or_structure = path_to_structure,
+        >>>    weight_type = "gaussian",
+        >>>    weight_type_args={"resolution": "20"}
+        >>> )
+
+        In addition its possible to use experimentally determined scattering factors
+        from various sources:
 
         >>> density = Density.from_structure(
         >>>    filename_or_structure = path_to_structure,
@@ -754,7 +761,7 @@ class Density:
         >>>    weight_type_args={"source": "dt1969"}
         >>> )
 
-        or a lowpass filtered representation introduced in [1]_:
+        and their lowpass filtered representation introduced in [1]_:
 
         >>> density = Density.from_structure(
         >>>    filename_or_structure = path_to_structure,