pytme 0.1.9__cp311-cp311-macosx_14_0_arm64.whl → 0.2.0b0__cp311-cp311-macosx_14_0_arm64.whl

tme/density.py

@@ -1,4 +1,4 @@
-""" Implements class to represent electron density maps.
+""" Representation of N-dimensional densities
 
     Copyright (c) 2023 European Molecular Biology Laboratory
 
@@ -12,6 +12,7 @@ from gzip import open as gzip_open
 from typing import Tuple, Dict, Set
 from os.path import splitext, basename
 
+import h5py
 import mrcfile
 import numpy as np
 import skimage.io as skio
@@ -26,7 +27,6 @@ from scipy.ndimage import (
 )
 from scipy.spatial import ConvexHull
 
-from .matching_optimization import FitRefinement
 from .structure import Structure
 from .matching_utils import (
     minimum_enclosing_box,
@@ -129,8 +129,8 @@ class Density:
         Parameters
         ----------
         filename : str
-            Path to a file in CCP4/MRC, EM or a format supported by skimage.io.imread.
-            The file can be gzip compressed.
+            Path to a file in CCP4/MRC, EM, HDF5 or a format supported by
+            skimage.io.imread. The file can be gzip compressed.
         subset : tuple of slices, optional
             Slices representing the desired subset along each dimension.
         use_memmap : bool, optional
@@ -185,8 +185,9 @@ class Density:
 
         Notes
         -----
-        If ``filename`` ends with ".em" or ".em.gz" the method will parse it as EM file.
-        Otherwise it defaults to the CCP4/MRC format and on failure, switches to
+        If ``filename`` ends with ".em" or ".em.gz" the method will parse it as EM file,
+        if it ends with "h5" or "h5.gz" the method will parse the file as HDF5.
+        Otherwise the method defaults to the CCP4/MRC format and on failure, switches to
         :obj:`skimage.io.imread` regardless of the extension. Currently, the later does not
         extract origin or sampling_rate information from the file.
 
@@ -197,8 +198,10 @@ class Density:
         """
         try:
             func = cls._load_mrc
-            if filename.endswith(".em") or filename.endswith(".em.gz"):
+            if filename.endswith("em") or filename.endswith("em.gz"):
                 func = cls._load_em
+            elif filename.endswith("h5") or filename.endswith("h5.gz"):
+                func = cls._load_hdf5
             data, origin, sampling_rate, meta = func(
                 filename=filename, subset=subset, use_memmap=use_memmap
             )
@@ -213,7 +216,7 @@ class Density:
     @classmethod
     def _load_mrc(
         cls, filename: str, subset: Tuple[int] = None, use_memmap: bool = False
-    ) -> Tuple[NDArray]:
+    ) -> Tuple[NDArray, NDArray, NDArray, Dict]:
         """
         Extracts data from a CCP4/MRC file.
 
@@ -228,12 +231,8 @@ class Density:
 
         Returns
         -------
-        NDArray
-            The data attribute of the CCP4/MRC file.
-        NDArray
-            The coordinate origin of the data.
-        NDArray
-            The sampling rate of the data.
+        Tuple[NDArray, NDArray, NDArray, Dict]
+            File data, coordinate origin, sampling rate array and metadata dictionary.
 
         References
         ----------
@@ -355,7 +354,7 @@ class Density:
     @classmethod
     def _load_em(
         cls, filename: str, subset: Tuple[int] = None, use_memmap: bool = False
-    ) -> Tuple[NDArray]:
+    ) -> Tuple[NDArray, NDArray, NDArray, Dict]:
         """
         Extracts data from a EM file.
 
@@ -370,12 +369,8 @@ class Density:
 
         Returns
         -------
-        NDArray
-            The data attribute of the EM file.
-        NDArray
-            The coordinate origin of the data.
-        NDArray
-            The sampling rate of the data.
+        Tuple[NDArray, NDArray, NDArray, Dict]
+            File data, coordinate origin, sampling rate array and metadata dictionary.
 
         References
         ----------
@@ -383,11 +378,11 @@ class Density:
 
         Warns
         -----
-            Warns if the pixel size is zero.
+            If the sampling rate is zero.
 
         Notes
         -----
-        A pixel size of zero will be treated as missing value and changed to one. This
+        A sampling rate of zero will be treated as missing value and changed to one. This
         function does not yet extract an origin like :py:meth:`Density._load_mrc`.
 
         See Also
@@ -483,10 +478,15 @@ class Density:
                 f"Expected length of slices : {n_dims}, got : {len(slices)}"
             )
 
-        if any([slices[i].stop > shape[i] for i in range(n_dims)]):
+        if any(
+            [
+                slices[i].stop > shape[i] or slices[i].start > shape[i]
+                for i in range(n_dims)
+            ]
+        ):
             raise ValueError(f"Subset exceeds data dimensions ({shape}).")
 
-        if any([slices[i].stop < 0 for i in range(n_dims)]):
+        if any([slices[i].stop < 0 or slices[i].start < 0 for i in range(n_dims)]):
             raise ValueError("Subsets have to be non-negative.")
 
     @classmethod
@@ -560,7 +560,7 @@ class Density:
         return subset_data
 
     @staticmethod
-    def _load_skio(filename: str) -> Tuple[NDArray]:
+    def _load_skio(filename: str) -> Tuple[NDArray, NDArray, NDArray, Dict]:
         """
         Uses :obj:`skimage.io.imread` to extract data from filename [1]_.
 
@@ -571,12 +571,8 @@ class Density:
 
         Returns
         -------
-        NDArray
-            The data attribute of the file.
-        NDArray
-            The coordinate origin of the data.
-        NDArray
-            The sampling rate of the data.
+        Tuple[NDArray, NDArray, NDArray, Dict]
+            File data, coordinate origin, sampling rate array and metadata dictionary.
 
         References
         ----------
@@ -601,6 +597,51 @@ class Density:
         )
         return data, np.zeros(data.ndim), np.ones(data.ndim), {}
 
+    @staticmethod
+    def _load_hdf5(
+        filename: str, subset: Tuple[slice], use_memmap: bool = False, **kwargs
+    ) -> "Density":
+        """
+        Extracts data from an H5 file.
+
+        Parameters
+        ----------
+        filename : str
+            Path to a file in CCP4/MRC format.
+        subset : tuple of slices, optional
+            Slices representing the desired subset along each dimension.
+        use_memmap : bool, optional
+            Whether the Density objects data attribute should be memmory mapped.
+
+        Returns
+        -------
+        Density
+            An instance of the Density class populated with the data from the HDF5 file.
+
+        See Also
+        --------
+        :py:meth:`Density._save_hdf5`
+        """
+        subset = ... if subset is None else subset
+
+        with h5py.File(filename, mode="r") as infile:
+            data = infile["data"]
+            data_attributes = [
+                infile["data"].id.get_offset(),
+                infile["data"].shape,
+                infile["data"].dtype,
+            ]
+            origin = infile["origin"][...].copy()
+            sampling_rate = infile["sampling_rate"][...].copy()
+            metadata = {key: val for key, val in infile.attrs.items()}
+            if not use_memmap:
+                return data[subset], origin, sampling_rate, metadata
+
+        offset, shape, dtype = data_attributes
+        data = np.memmap(filename, dtype=dtype, shape=shape, offset=offset)[subset]
+
+        return data, origin, sampling_rate, metadata
+
     @classmethod
     def from_structure(
         cls,
@@ -798,9 +839,9 @@ class Density:
 
         Notes
         -----
-        If ``filename`` ends with ".em" or ".em.gz", the method will create an EM file.
-        Otherwise, it defaults to the CCP4/MRC format, and on failure, falls back
-        to :obj:`skimage.io.imsave`.
+        If ``filename`` ends with "em" or "em.gz" will create an EM file, "h5" or
+        "h5.gz" will create a HDF5 file. Otherwise, the method defaults to the CCP4/MRC
+        format, and on failure, falls back to :obj:`skimage.io.imsave`.
 
         See Also
         --------
@@ -811,13 +852,15 @@ class Density:
 
         try:
             func = self._save_mrc
-            if filename.endswith(".em") or filename.endswith(".em.gz"):
+            if filename.endswith("em") or filename.endswith("em.gz"):
                 func = self._save_em
+            elif filename.endswith("h5") or filename.endswith("h5.gz"):
+                func = self._save_hdf5
             _ = func(filename=filename, gzip=gzip)
         except ValueError:
             _ = self._save_skio(filename=filename, gzip=gzip)
 
-    def _save_mrc(self, filename: str, gzip: bool) -> None:
+    def _save_mrc(self, filename: str, gzip: bool = False) -> None:
         """
         Writes current class instance to disk as mrc file.
 
@@ -843,20 +886,16 @@ class Density:
             mrc.header["origin"] = tuple(self.origin[::-1])
             mrc.voxel_size = tuple(self.sampling_rate[::-1])
 
-    def _save_em(self, filename: str, gzip: bool) -> None:
+    def _save_em(self, filename: str, gzip: bool = False) -> None:
         """
         Writes data to disk as an .em file.
 
         Parameters
         ----------
         filename : str
-            Path to write the .em file to.
-        data : NDArray
-            Data to be saved.
-        origin : NDArray
-            Coordinate origin of the data.
-        sampling_rate : NDArray
-            Sampling rate of the data.
+            Path to write to.
+        gzip : bool, optional
+            If True, the output will be gzip compressed.
 
         References
         ----------
@@ -886,7 +925,7 @@ class Density:
             f.write(b" " * 256)
             f.write(self.data.tobytes())
 
-    def _save_skio(self, filename: str, gzip: bool) -> None:
+    def _save_skio(self, filename: str, gzip: bool = False) -> None:
         """
         Uses :obj:`skimage.io.imsave` to write data to filename [1]_.
 
@@ -904,12 +943,54 @@ class Density:
         swap, kwargs = filename, {}
         if gzip:
             swap = BytesIO()
-            kwargs["format"] = splitext(basename(filename.replace(".gz", "")))[1]
+            kwargs["format"] = splitext(basename(filename.replace(".gz", "")))[
+                1
+            ].replace(".", "")
         skio.imsave(fname=swap, arr=self.data.astype("float32"), **kwargs)
         if gzip:
             with gzip_open(filename, "wb") as outfile:
                 outfile.write(swap.getvalue())
 
+    def _save_hdf5(self, filename: str, gzip: bool = False) -> None:
+        """
+        Saves the Density instance data to an HDF5 file, with optional compression.
+
+        Parameters
+        ----------
+        filename : str
+            Path to write to.
+        gzip : bool, optional
+            If True, the output will be gzip compressed.
+
+        See Also
+        --------
+        :py:meth:`Density._load_hdf5`
+        """
+        compression = "gzip" if gzip else None
+        with h5py.File(filename, mode="w") as f:
+            f.create_dataset(
+                "data",
+                data=self.data,
+                shape=self.data.shape,
+                dtype=self.data.dtype,
+                compression=compression,
+            )
+            f.create_dataset("origin", data=self.origin)
+            f.create_dataset("sampling_rate", data=self.sampling_rate)
+
+            self.metadata["mean"] = self.metadata.get("mean", 0)
+            self.metadata["std"] = self.metadata.get("std", 0)
+            self.metadata["min"] = self.metadata.get("min", 0)
+            self.metadata["max"] = self.metadata.get("max", 0)
+            if type(self.data) != np.memmap:
+                self.metadata["mean"] = self.data.mean()
+                self.metadata["std"] = self.data.std()
+                self.metadata["min"] = self.data.min()
+                self.metadata["max"] = self.data.max()
+
+            for key, val in self.metadata.items():
+                f.attrs[key] = val
+
     @property
     def empty(self) -> "Density":
         """
@@ -1098,7 +1179,13 @@ class Density:
     @property
     def sampling_rate(self) -> NDArray:
         """
-        Returns sampling rate along data axis.
+        Returns the value of the current instance's :py:attr:`Density.sampling_rate`
+        attribute.
+
+        Returns
+        -------
+        NDArray
+            Sampling rate along axis.
         """
         return self._sampling_rate
 
@@ -1114,7 +1201,12 @@ class Density:
     @property
     def metadata(self) -> Dict:
         """
-        Returns dictionary with metadata information, empty by default.
+        Returns the current instance's :py:attr:`Density.metadata` dictionary attribute.
+
+        Returns
+        -------
+        Dict
+            Metadata dictionary. Empty by default.
         """
         return self._metadata
 
@@ -2068,7 +2160,7 @@ class Density:
 
         If voxel sizes of target and template dont match coordinates are scaled
         to the numerically smaller voxel size. Instances are prealigned based on their
-        center of mass. Finally :py:class:`tme.matching_optimization.FitRefinement` is
+        center of mass. Finally :py:meth:`tme.matching_optimization.optimize_match` is
         used to determine translation and rotation to map template to target.
 
         Parameters
@@ -2083,7 +2175,7 @@ class Density:
             The cutoff value for the template map, by default 0.
         scoring_method : str, optional
             The scoring method to use for alignment. See
-            :py:class:`tme.matching_optimization.FitRefinement` for available methods,
+            :py:class:`tme.matching_optimization.create_score_object` for available methods,
             by default "NormalizedCrossCorrelation".
 
         Returns
@@ -2096,6 +2188,8 @@ class Density:
         -----
         No densities below cutoff_template are present in the returned Density object.
         """
+        from .matching_optimization import optimize_match, create_score_object
+
         target_sampling_rate = np.array(target.sampling_rate)
         template_sampling_rate = np.array(template.sampling_rate)
 
@@ -2105,7 +2199,6 @@ class Density:
         template_sampling_rate = np.repeat(
             template_sampling_rate, template.data.ndim // template_sampling_rate.size
         )
-
         if not np.allclose(target_sampling_rate, template_sampling_rate):
             print(
                 "Voxel size of target and template do not match. "
@@ -2113,7 +2206,6 @@ class Density:
             )
 
         target_coordinates = target.to_pointcloud(cutoff_target)
-        target_weights = target.data[tuple(target_coordinates)]
 
         template_coordinates = template.to_pointcloud(cutoff_template)
         template_weights = template.data[tuple(template_coordinates)]
@@ -2126,23 +2218,24 @@ class Density:
         target_coordinates = target_coordinates * target_scaling[:, None]
         template_coordinates = template_coordinates * template_scaling[:, None]
 
-        target_mass_center = cls.center_of_mass(target.data, cutoff_target)
-        template_mass_center = cls.center_of_mass(template.data, cutoff_template)
         mass_center_difference = np.subtract(
-            target_mass_center, template_mass_center
+            cls.center_of_mass(target.data, cutoff_target),
+            cls.center_of_mass(template.data, cutoff_template),
         ).astype(int)
         template_coordinates += mass_center_difference[:, None]
 
-        matcher = FitRefinement()
-        translation, rotation_matrix, score = matcher.refine(
-            target_coordinates=target_coordinates,
+        score_object = create_score_object(
+            score=scoring_method,
+            target=target.data,
             template_coordinates=template_coordinates,
-            target_weights=target_weights,
             template_weights=template_weights,
-            scoring_class=scoring_method,
             sampling_rate=np.ones(template.data.ndim),
         )
 
+        translation, rotation_matrix, score = optimize_match(
+            score_object=score_object, optimization_method="basinhopping"
+        )
+
         translation += mass_center_difference
         translation = np.divide(translation, template_scaling)
 
@@ -2169,7 +2262,7 @@ class Density:
 
         If voxel sizes of target and template dont match coordinates are scaled
         to the numerically smaller voxel size. Prealignment is done by center's
-        of mass. Finally :py:class:`tme.matching_optimization.FitRefinement` is used to
+        of mass. Finally :py:class:`tme.matching_optimization.optimize_match` is used to
         determine translation and rotation to match a template to target.
 
         Parameters
@@ -2184,7 +2277,7 @@ class Density:
             The cutoff value for the template map, by default 0.
         scoring_method : str, optional
             The scoring method to use for template matching. See
-            :py:class:`tme.matching_optimization.FitRefinement` for available methods,
+            :py:class:`tme.matching_optimization.create_score_object` for available methods,
             by default "NormalizedCrossCorrelation".
 
         Returns

tme/matching_constrained.py

@@ -0,0 +1,195 @@
+import numpy as np
+from typing import Tuple, Dict
+
+from scipy.ndimage import map_coordinates
+
+from tme.types import ArrayLike
+from tme.backends import backend
+from tme.matching_data import MatchingData
+from tme.matching_exhaustive import _normalize_under_mask
+
+
+class MatchDensityToDensity:
+    def __init__(
+        self,
+        matching_data: "MatchingData",
+        pad_target_edges: bool = False,
+        pad_fourier: bool = False,
+        rotate_mask: bool = True,
+        interpolation_order: int = 1,
+        negate_score: bool = False,
+    ):
+        self.rotate_mask = rotate_mask
+        self.interpolation_order = interpolation_order
+
+        target_pad = matching_data.target_padding(pad_target=pad_target_edges)
+        matching_data = matching_data.subset_by_slice(target_pad=target_pad)
+
+        fast_shape, fast_ft_shape, fourier_shift = matching_data.fourier_padding(
+            pad_fourier=pad_fourier
+        )
+
+        self.target = backend.topleft_pad(matching_data.target, fast_shape)
+        self.target_mask = matching_data.target_mask
+
+        self.template = matching_data.template
+        self.template_rot = backend.preallocate_array(
+            fast_shape, backend._default_dtype
+        )
+
+        self.template_mask, self.template_mask_rot = 1, 1
+        rotate_mask = False if matching_data.template_mask is None else rotate_mask
+        if matching_data.template_mask is not None:
+            self.template_mask = matching_data.template_mask
+            self.template_mask_rot = backend.topleft_pad(
+                matching_data.template_mask, fast_shape
+            )
+
+        self.score_sign = -1 if negate_score else 1
+
+    @staticmethod
+    def rigid_transform(
+        arr,
+        rotation_matrix,
+        translation,
+        arr_mask=None,
+        out=None,
+        out_mask=None,
+        order: int = 1,
+        use_geometric_center: bool = False,
+    ):
+        rotate_mask = arr_mask is not None
+        return_type = (out is None) + 2 * rotate_mask * (out_mask is None)
+        translation = np.zeros(arr.ndim) if translation is None else translation
+
+        center = np.floor(np.array(arr.shape) / 2)[:, None]
+        grid = np.indices(arr.shape, dtype=np.float32).reshape(arr.ndim, -1)
+        np.subtract(grid, center, out=grid)
+        np.matmul(rotation_matrix.T, grid, out=grid)
+        np.add(grid, center, out=grid)
+
+        if out is None:
+            out = np.zeros_like(arr)
+
+        map_coordinates(arr, grid, order=order, output=out.ravel())
+
+        if out_mask is None and arr_mask is not None:
+            out_mask = np.zeros_like(arr_mask)
+
+        if arr_mask is not None:
+            map_coordinates(arr_mask, grid, order=order, output=out_mask.ravel())
+
+        match return_type:
+            case 0:
+                return None
+            case 1:
+                return out
+            case 2:
+                return out_mask
+            case 3:
+                return out, out_mask
+
+    @staticmethod
+    def angles_to_rotationmatrix(angles: Tuple[float]) -> ArrayLike:
+        angles = backend.to_numpy_array(angles)
+        rotation_matrix = euler_to_rotationmatrix(angles)
+        return backend.to_backend_array(rotation_matrix)
+
+    def format_translation(self, translation: Tuple[float] = None) -> ArrayLike:
+        if translation is None:
+            return backend.zeros(self.template.ndim, backend._default_dtype)
+
+        return backend.to_backend_array(translation)
+
+    def score_translation(self, x: Tuple[float]) -> float:
+        translation = self.format_translation(x)
+        rotation_matrix = self.angles_to_rotationmatrix((0, 0, 0))
+
+        return self(translation=translation, rotation_matrix=rotation_matrix)
+
+    def score_angles(self, x: Tuple[float]) -> float:
+        translation = self.format_translation(None)
+        rotation_matrix = self.angles_to_rotationmatrix(x)
+
+        return self(translation=translation, rotation_matrix=rotation_matrix)
+
+    def score(self, x: Tuple[float]) -> float:
+        split = len(x) // 2
+        translation, angles = x[:split], x[split:]
+
+        translation = self.format_translation(translation)
+        rotation_matrix = self.angles_to_rotationmatrix(angles)
+
+        return self(translation=translation, rotation_matrix=rotation_matrix)
+
+
+class FLC(MatchDensityToDensity):
+    def __init__(self, **kwargs: Dict):
+        super().__init__(**kwargs)
+
+        if self.target_mask is not None:
+            backend.multiply(self.target, self.target_mask, out=self.target)
+
+        self.target_square = backend.square(self.target)
+
+        _normalize_under_mask(
+            template=self.template,
+            mask=self.template_mask,
+            mask_intensity=backend.sum(self.template_mask),
+        )
+
+        self.template = backend.reverse(self.template)
+        self.template_mask = backend.reverse(self.template_mask)
+
+    def __call__(self, translation: ArrayLike, rotation_matrix: ArrayLike) -> float:
+        if self.rotate_mask:
+            self.rigid_transform(
+                arr=self.template,
+                arr_mask=self.template_mask,
+                rotation_matrix=rotation_matrix,
+                translation=translation,
+                out=self.template_rot,
+                out_mask=self.template_mask_rot,
+                use_geometric_center=False,
+                order=self.interpolation_order,
+            )
+        else:
+            self.rigid_transform(
+                arr=self.template,
+                rotation_matrix=rotation_matrix,
+                translation=translation,
+                out=self.template_rot,
+                use_geometric_center=False,
+                order=self.interpolation_order,
+            )
+        n_observations = backend.sum(self.template_mask_rot)
+
+        _normalize_under_mask(
+            template=self.template_rot,
+            mask=self.template_mask_rot,
+            mask_intensity=n_observations,
+        )
+
+        ex2 = backend.sum(
+            backend.divide(
+                backend.sum(
+                    backend.multiply(self.target_square, self.template_mask_rot),
+                ),
+                n_observations,
+            )
+        )
+        e2x = backend.square(
+            backend.divide(
+                backend.sum(backend.multiply(self.target, self.template_mask_rot)),
+                n_observations,
+            )
+        )
+
+        denominator = backend.maximum(backend.subtract(ex2, e2x), 0.0)
+        denominator = backend.sqrt(denominator)
+        denominator = backend.multiply(denominator, n_observations)
+
+        overlap = backend.sum(backend.multiply(self.template_rot, self.target))
+
+        score = backend.divide(overlap, denominator) * self.score_sign
+        return score