pytme 0.1.8__cp311-cp311-macosx_14_0_arm64.whl → 0.2.0__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
         ----------
@@ -318,7 +317,7 @@ class Density:
             if use_memmap:
                 warnings.warn(
                     f"Cannot open gzipped file {filename} as memmap."
-                    f" Please gunzip {filename} to use memmap functionality."
+                    f" Please run 'gunzip {filename}' to use memmap functionality."
                 )
             use_memmap = False
 
@@ -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
 
@@ -1453,34 +1545,35 @@ class Density:
         --------
         :py:meth:`Density.centered` returns a tuple containing a centered version
         of the current :py:class:`Density` instance, as well as an array with
-        translations. The translation corresponds to the shift that between the
-        center of mass and the center of the internal :py:attr:`Density.data` attribute.
+        translations. The translation corresponds to the shift between the original and
+        current center of mass.
 
         >>> import numpy as np
         >>> from tme import Density
         >>> dens = Density(np.ones((5,5)))
         >>> centered_dens, translation = dens.centered(0)
         >>> translation
-        array([-4.4408921e-16,  4.4408921e-16])
+        array([-0.5, -0.5])
 
         :py:meth:`Density.centered` extended the :py:attr:`Density.data` attribute
         of the current :py:class:`Density` instance and modified
         :py:attr:`Density.origin` accordingly.
 
         >>> centered_dens
-        Origin: (-1.0, -1.0), sampling_rate: (1, 1), Shape: (7, 7)
+        Origin: (-1.0, -1.0), sampling_rate: (1, 1), Shape: (8, 8)
 
-        :py:meth:`Density.centered` achieves centering via zero-padding the
-        internal :py:attr:`Density.data` attribute:
+        :py:meth:`Density.centered` achieves centering via zero-padding and
+        transforming the internal :py:attr:`Density.data` attribute:
 
         >>> centered_dens.data
-        array([[0., 0., 0., 0., 0., 0., 0.],
-               [0., 1., 1., 1., 1., 1., 0.],
-               [0., 1., 1., 1., 1., 1., 0.],
-               [0., 1., 1., 1., 1., 1., 0.],
-               [0., 1., 1., 1., 1., 1., 0.],
-               [0., 1., 1., 1., 1., 1., 0.],
-               [0., 0., 0., 0., 0., 0., 0.]])
+        array([[0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
+               [0.  , 0.25, 0.5 , 0.5 , 0.5 , 0.5 , 0.25, 0.  ],
+               [0.  , 0.5 , 1.  , 1.  , 1.  , 1.  , 0.5 , 0.  ],
+               [0.  , 0.5 , 1.  , 1.  , 1.  , 1.  , 0.5 , 0.  ],
+               [0.  , 0.5 , 1.  , 1.  , 1.  , 1.  , 0.5 , 0.  ],
+               [0.  , 0.5 , 1.  , 1.  , 1.  , 1.  , 0.5 , 0.  ],
+               [0.  , 0.25, 0.5 , 0.5 , 0.5 , 0.5 , 0.25, 0.  ],
+               [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ]])
 
         `centered_dens` is sufficiently large to represent all rotations that
         could be applied to the :py:attr:`Density.data` attribute. Lets look
@@ -1503,10 +1596,11 @@ class Density:
         ret.adjust_box(box)
 
         new_shape = np.maximum(ret.shape, self.shape)
+        new_shape = np.add(new_shape, np.mod(new_shape, 2))
         ret.pad(new_shape)
 
         center = self.center_of_mass(ret.data, cutoff)
-        shift = np.subtract(np.divide(ret.shape, 2), center)
+        shift = np.subtract(np.divide(np.subtract(ret.shape, 1), 2), center)
 
         ret = ret.rigid_transform(
             translation=shift,
@@ -2068,7 +2162,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 +2177,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 +2190,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 +2201,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 +2208,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 +2220,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 +2264,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 +2279,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