pytme 0.2.0__cp311-cp311-macosx_14_0_arm64.whl → 0.2.1__cp311-cp311-macosx_14_0_arm64.whl

tme/density.py

@@ -40,12 +40,12 @@ from .backends import NumpyFFTWBackend
 
 class Density:
     """
-    Contains electron density data and implements operations on it.
+    Abstract representation of N-dimensional densities.
 
     Parameters
     ----------
     data : NDArray
-        Electron density data.
+        Array of data values.
     origin : NDArray, optional
         Origin of the coordinate system. Defaults to zero.
     sampling_rate : NDArray, optional
@@ -267,9 +267,9 @@ class Density:
             # nx := column; ny := row; nz := section
             start = np.array(
                 [
-                    int(mrc.header["nxstart"]),
-                    int(mrc.header["nystart"]),
-                    int(mrc.header["nzstart"]),
+                    mrc.header["nzstart"],
+                    mrc.header["nystart"],
+                    mrc.header["nxstart"],
                 ]
             )
 
@@ -291,17 +291,9 @@ class Density:
             sampling_rate = mrc.voxel_size.astype(
                 [("x", "<f4"), ("y", "<f4"), ("z", "<f4")]
             ).view(("<f4", 3))
-            sampling_rate = sampling_rate[::-1]
-            sampling_rate = np.array(sampling_rate)
+            sampling_rate = np.array(sampling_rate)[::-1]
 
-            if np.all(origin == start):
-                pass
-            elif np.all(origin == 0) and not np.all(start == 0):
-                origin = np.multiply(start, sampling_rate)
-            elif np.all(
-                np.abs(origin.astype(int))
-                != np.abs((start * sampling_rate).astype(int))
-            ) and not np.all(start == 0):
+            if np.allclose(origin, 0) and not np.allclose(start, 0):
                 origin = np.multiply(start, sampling_rate)
 
             extended_header = mrc.header.nsymbt
@@ -878,7 +870,7 @@ class Density:
         compression = "gzip" if gzip else None
         with mrcfile.new(filename, overwrite=True, compression=compression) as mrc:
             mrc.set_data(self.data.astype("float32"))
-            mrc.header.nzstart, mrc.header.nystart, mrc.header.nxstart = np.ceil(
+            mrc.header.nzstart, mrc.header.nystart, mrc.header.nxstart = np.rint(
                 np.divide(self.origin, self.sampling_rate)
             )
             # mrcfile library expects origin to be in xyz format
@@ -1529,11 +1521,10 @@ class Density:
 
         Returns
         -------
-        Density
-            A copy of the class instance whose data center of mass is in the
-            center of the data array.
+        :py:class:`Density`
+            A centered copy of the current class instance.
         NDArray
-            The coordinate translation.
+            The offset between array center and center of mass.
 
         See Also
         --------
@@ -1550,44 +1541,32 @@ class Density:
 
         >>> import numpy as np
         >>> from tme import Density
-        >>> dens = Density(np.ones((5,5)))
+        >>> dens = Density(np.ones((5,5,5)))
         >>> centered_dens, translation = dens.centered(0)
         >>> translation
-        array([-0.5, -0.5])
+        array([0., 0., 0.])
 
         :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: (8, 8)
+        Origin: (-2.0, -2.0, -2.0), sampling_rate: (1, 1, 1), Shape: (9, 9, 9)
 
         :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.  ],
-               [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
-        at a random rotation obtained from
+        rigid-transform of the internal :py:attr:`Density.data` attribute.
+        `centered_dens` is sufficiently large to represent all rotations of the
+        :py:attr:`Density.data` attribute, such as ones obtained from
         :py:meth:`tme.matching_utils.get_rotation_matrices`.
 
         >>> from tme.matching_utils import get_rotation_matrices
-        >>> rotation_matrix = get_rotation_matrices(dim = 2 ,angular_sampling = 10)[0]
+        >>> rotation_matrix = get_rotation_matrices(dim = 3 ,angular_sampling = 10)[0]
         >>> rotated_centered_dens = centered_dens.rigid_transform(
         >>>     rotation_matrix = rotation_matrix,
         >>>     order = None
         >>> )
         >>> print(centered_dens.data.sum(), rotated_centered_dens.data.sum())
-        25.000000000000007 25.000000000000007
+        125.0 125.0
 
         """
         ret = self.copy()
@@ -1596,7 +1575,7 @@ 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))
+        new_shape = np.add(new_shape, 1 - np.mod(new_shape, 2))
         ret.pad(new_shape)
 
         center = self.center_of_mass(ret.data, cutoff)
@@ -1608,9 +1587,9 @@ class Density:
             use_geometric_center=False,
             order=1,
         )
-        offset = np.subtract(center, self.center_of_mass(ret.data, cutoff))
 
-        return ret, offset
+        shift = np.subtract(center, self.center_of_mass(ret.data, cutoff))
+        return ret, shift
 
     @classmethod
     def rotate_array(
@@ -1733,31 +1712,45 @@ class Density:
         Parameters
         ----------
         rotation_matrix : NDArray
-            Rotation matrix to apply to the `Density` instance.
+            Rotation matrix to apply.
         translation : NDArray
-            Translation to apply to the `Density` instance.
+            Translation to apply.
         order : int, optional
-            Order of spline interpolation.
+            Interpolation order to use. Default is 3, has to be in range 0-5.
         use_geometric_center : bool, optional
-            Whether to use geometric or coordinate center. If False,
-            class instance should be centered using :py:meth:`Density.centered`.
+            Use geometric or mass center as rotation center.
 
         Returns
         -------
         Density
-            The transformed instance of :py:class:`tme.density.Density`.
+            The transformed instance of :py:class:`Density`.
 
         Examples
         --------
+        Define the :py:class:`Density` instance
+
         >>> import numpy as np
-        >>> rotation_matrix = np.eye(3)
-        >>> rotation_matrix[0] = -1
-        >>> density.rotate(rotation_matrix = rotation_matrix)
+        >>> from tme import Density
+        >>> dens = Density(np.arange(9).reshape(3,3).astype(np.float32))
+        >>> dens, translation = dens.centered(0)
+
+        and apply the rotation, in this case a mirror around the z-axis
+
+        >>> rotation_matrix = np.eye(dens.data.ndim)
+        >>> rotation_matrix[0, 0] = -1
+        >>> dens_transform = dens.rigid_transform(rotation_matrix = rotation_matrix)
+        >>> dens_transform.data
+        array([[0.        , 0.        , 0.        , 0.        , 0.        ],
+               [0.5       , 3.0833333 , 3.5833333 , 3.3333333 , 0.        ],
+               [0.75      , 4.6666665 , 5.6666665 , 5.4166665 , 0.        ],
+               [0.25      , 1.6666666 , 2.6666667 , 2.9166667 , 0.        ],
+               [0.        , 0.08333334, 0.5833333 , 0.8333333 , 0.        ]],
+              dtype=float32)
 
         Notes
         -----
-        :py:meth:`Density.rigid_transform` that the internal data array is
-        sufficiently sized to accomodate the transform.
+        This function assumes the internal :py:attr:`Density.data` attribute is
+        sufficiently sized to hold the transformation.
 
         See Also
         --------
@@ -2155,6 +2148,7 @@ class Density:
         cutoff_target: float = 0,
         cutoff_template: float = 0,
         scoring_method: str = "NormalizedCrossCorrelation",
+        **kwargs,
     ) -> Tuple["Density", NDArray, NDArray, NDArray]:
         """
         Aligns two :py:class:`Density` instances target and template and returns
@@ -2179,6 +2173,9 @@ class Density:
             The scoring method to use for alignment. See
             :py:class:`tme.matching_optimization.create_score_object` for available methods,
             by default "NormalizedCrossCorrelation".
+        kwargs : dict, optional
+            Optional keyword arguments passed to
+            :py:meth:`tme.matching_optimization.optimize_match`.
 
         Returns
         -------
@@ -2190,8 +2187,18 @@ class Density:
         -----
         No densities below cutoff_template are present in the returned Density object.
         """
+        from .matching_exhaustive import normalize_under_mask
         from .matching_optimization import optimize_match, create_score_object
 
+        template_mask = template.empty
+        template_mask.data[:] = 1
+
+        normalize_under_mask(
+            template=template.data,
+            mask=template_mask.data,
+            mask_intensity=template_mask.data.sum(),
+        )
+
         target_sampling_rate = np.array(target.sampling_rate)
         template_sampling_rate = np.array(template.sampling_rate)
 
@@ -2226,16 +2233,21 @@ class Density:
         ).astype(int)
         template_coordinates += mass_center_difference[:, None]
 
+        coordinates_mask = template_mask.to_pointcloud()
+        coordinates_mask = coordinates_mask * template_scaling[:, None]
+        coordinates_mask += mass_center_difference[:, None]
+
         score_object = create_score_object(
             score=scoring_method,
             target=target.data,
             template_coordinates=template_coordinates,
+            template_mask_coordinates=coordinates_mask,
             template_weights=template_weights,
             sampling_rate=np.ones(template.data.ndim),
         )
 
         translation, rotation_matrix, score = optimize_match(
-            score_object=score_object, optimization_method="basinhopping"
+            score_object=score_object, **kwargs
         )
 
         translation += mass_center_difference
@@ -2257,6 +2269,8 @@ class Density:
         template: "Structure",
         cutoff_target: float = 0,
         scoring_method: str = "NormalizedCrossCorrelation",
+        optimization_method: str = "basinhopping",
+        maxiter: int = 500,
     ) -> Tuple["Structure", NDArray, NDArray]:
         """
         Aligns a :py:class:`tme.structure.Structure` template to :py:class:`Density`
@@ -2281,6 +2295,12 @@ class Density:
             The scoring method to use for template matching. See
             :py:class:`tme.matching_optimization.create_score_object` for available methods,
             by default "NormalizedCrossCorrelation".
+        optimization_method : str, optional
+            Optimizer that is used.
+            See :py:meth:`tme.matching_optimization.optimize_match`.
+        maxiter : int, optional
+            Maximum number of iterations for the optimizer.
+            See :py:meth:`tme.matching_optimization.optimize_match`.
 
         Returns
         -------
@@ -2304,18 +2324,17 @@ class Density:
             cutoff_target=cutoff_target,
             cutoff_template=0,
             scoring_method=scoring_method,
+            optimization_method=optimization_method,
+            maxiter=maxiter,
         )
         out = template.copy()
-        final_translation = np.add(
-            -template_density.origin,
-            np.multiply(translation, template_density.sampling_rate),
-        )
+        final_translation = np.subtract(ret.origin, template_density.origin)
 
         # Atom coordinates are in xyz
         final_translation = final_translation[::-1]
         rotation_matrix = rotation_matrix[::-1, ::-1]
 
-        out.rigid_transform(
+        out = out.rigid_transform(
             translation=final_translation, rotation_matrix=rotation_matrix
         )
 

tme/matching_data.py

@@ -22,20 +22,44 @@ class MatchingData:
 
     Parameters
     ----------
-    target : np.ndarray or Density
-        Target data array for template matching.
-    template : np.ndarray or Density
-        Template data array for template matching.
+    target : np.ndarray or :py:class:`tme.density.Density`
+        Target data.
+    template : np.ndarray or :py:class:`tme.density.Density`
+        Template data.
+    target_mask : np.ndarray or :py:class:`tme.density.Density`, optional
+        Target mask data.
+    template_mask : np.ndarray or :py:class:`tme.density.Density`, optional
+        Template mask data.
+    invert_target : bool, optional
+        Whether to invert and rescale the target before template matching..
+    rotations: np.ndarray, optional
+        Template rotations to sample. Can be a single (d x d) or a stack (n x d x d)
+        of rotation matrices where d is the dimension of the template.
+
+    Examples
+    --------
+    The following achieves the minimal definition of a :py:class:`MatchingData` instance.
+
+    >>> import numpy as np
+    >>> from tme.matching_data import MatchingData
+    >>> target = np.random.rand(50,40,60)
+    >>> template = target[15:25, 10:20, 30:40]
+    >>> matching_data = MatchingData(target=target, template=template)
 
     """
 
-    def __init__(self, target: NDArray, template: NDArray):
-        self._default_dtype = np.float32
-        self._complex_dtype = np.complex64
-
+    def __init__(
+        self,
+        target: NDArray,
+        template: NDArray,
+        template_mask: NDArray = None,
+        target_mask: NDArray = None,
+        invert_target: bool = False,
+        rotations: NDArray = None,
+    ):
         self._target = target
-        self._target_mask = None
-        self._template_mask = None
+        self._target_mask = target_mask
+        self._template_mask = template_mask
         self._translation_offset = np.zeros(len(target.shape), dtype=int)
 
         self.template = template
@@ -46,8 +70,11 @@ class MatchingData:
         self.template_filter = {}
         self.target_filter = {}
 
-        self._invert_target = False
-        self._rotations = None
+        self._invert_target = invert_target
+
+        self._rotations = rotations
+        if rotations is not None:
+            self.rotations = rotations
 
         self._set_batch_dimension()
 
@@ -257,7 +284,7 @@ class MatchingData:
         template_offset[(template_offset.size - len(template_slice)) :] = [
             x.start for x in template_slice
         ]
-        ret._translation_offset = np.add(target_offset, template_offset)
+        ret._translation_offset = target_offset
 
         ret.template_filter = self.template_filter
         ret.target_filter = self.target_filter
@@ -266,7 +293,7 @@ class MatchingData:
         ret._invert_target = self._invert_target
 
         if self._target_mask is not None:
-            ret.target_mask = self.subset_array(
+            ret._target_mask = self.subset_array(
                 arr=self._target_mask, arr_slice=target_slice, padding=target_pad
             )
         if self._template_mask is not None:
@@ -289,19 +316,29 @@ class MatchingData:
 
     def to_backend(self) -> None:
         """
-        Transfer the class instance's numpy arrays to the current backend.
+        Transfer and convert types of class instance's data arrays to the current backend
         """
-        backend_arr = type(backend.zeros((1), dtype=backend._default_dtype))
+        backend_arr = type(backend.zeros((1), dtype=backend._float_dtype))
         for attr_name, attr_value in vars(self).items():
+            converted_array = None
             if isinstance(attr_value, np.ndarray):
                 converted_array = backend.to_backend_array(attr_value.copy())
-                setattr(self, attr_name, converted_array)
             elif isinstance(attr_value, backend_arr):
                 converted_array = backend.to_backend_array(attr_value)
-                setattr(self, attr_name, converted_array)
+            else:
+                continue
+
+            current_dtype = backend.get_fundamental_dtype(converted_array)
+            target_dtype = backend._fundamental_dtypes[current_dtype]
+
+            # Optional, but scores are float so we avoid casting and potential issues
+            if attr_name in ("_template", "_template_mask", "_target", "_target_mask"):
+                target_dtype = backend._float_dtype
 
-        self._default_dtype = backend._default_dtype
-        self._complex_dtype = backend._complex_dtype
+            if target_dtype != current_dtype:
+                converted_array = backend.astype(converted_array, target_dtype)
+
+            setattr(self, attr_name, converted_array)
 
     def _set_batch_dimension(
         self, target_dims: Tuple[int] = None, template_dims: Tuple[int] = None
@@ -350,12 +387,14 @@ class MatchingData:
         matching_dims = target_measurement_dims + batch_dims
 
         target_shape = backend.full(
-            shape=(matching_dims,), fill_value=1, dtype=backend._default_dtype_int
+            shape=(matching_dims,), fill_value=1, dtype=backend._int_dtype
         )
         template_shape = backend.full(
-            shape=(matching_dims,), fill_value=1, dtype=backend._default_dtype_int
+            shape=(matching_dims,), fill_value=1, dtype=backend._int_dtype
+        )
+        batch_mask = backend.full(
+            shape=(matching_dims,), fill_value=1, dtype=backend._int_dtype
         )
-        batch_mask = backend.full(shape=(matching_dims,), fill_value=1, dtype=bool)
 
         target_index, template_index = 0, 0
         for k in range(matching_dims):
@@ -440,7 +479,7 @@ class MatchingData:
             An array indicating the padding for each dimension of the target.
         """
         target_padding = backend.zeros(
-            len(self._output_target_shape), dtype=backend._default_dtype_int
+            len(self._output_target_shape), dtype=backend._int_dtype
         )
 
         if pad_target:
@@ -491,13 +530,14 @@ class MatchingData:
             fourier_pad = backend.full(
                 shape=(len(fourier_pad),),
                 fill_value=1,
-                dtype=backend._default_dtype_int,
+                dtype=backend._int_dtype,
             )
 
         fourier_pad = backend.to_backend_array(fourier_pad)
         if hasattr(self, "_batch_mask"):
             batch_mask = backend.to_backend_array(self._batch_mask)
-            fourier_pad[batch_mask] = 1
+            backend.multiply(fourier_pad, 1 - batch_mask, out=fourier_pad)
+            backend.add(fourier_pad, batch_mask, out=fourier_pad)
 
         pad_shape = backend.maximum(target_shape, template_shape)
         ret = backend.compute_convolution_shapes(pad_shape, fourier_pad)
@@ -510,11 +550,11 @@ class MatchingData:
 
             if hasattr(self, "_batch_mask"):
                 batch_mask = backend.to_backend_array(self._batch_mask)
-                shape_diff[batch_mask] = 0
+                backend.multiply(shape_diff, 1 - batch_mask, out=shape_diff)
 
             backend.add(fourier_shift, shape_diff, out=fourier_shift)
 
-        fourier_shift = backend.astype(fourier_shift, backend._default_dtype_int)
+        fourier_shift = backend.astype(fourier_shift, backend._int_dtype)
 
         return fast_shape, fast_ft_shape, fourier_shift
 
@@ -542,27 +582,27 @@ class MatchingData:
             pass
         else:
             raise ValueError("Rotations have to be a rank 2 or 3 array.")
-        self._rotations = rotations.astype(self._default_dtype)
+        self._rotations = rotations.astype(np.float32)
 
     @property
     def target(self):
-        """Returns the target NDArray."""
+        """Returns the target."""
+        target = self._target
         if isinstance(self._target, Density):
             target = self._target.data
-        else:
-            target = self._target
-        out_shape = backend.to_numpy_array(self._output_target_shape)
-        return target.reshape(tuple(int(x) for x in out_shape))
+
+        out_shape = tuple(int(x) for x in self._output_target_shape)
+        return target.reshape(out_shape)
 
     @property
     def template(self):
-        """Returns the reversed template NDArray."""
+        """Returns the reversed template."""
         template = self._template
         if isinstance(self._template, Density):
             template = self._template.data
         template = backend.reverse(template)
-        out_shape = backend.to_numpy_array(self._output_template_shape)
-        return template.reshape(tuple(int(x) for x in out_shape))
+        out_shape = tuple(int(x) for x in self._output_template_shape)
+        return template.reshape(out_shape)
 
     @template.setter
     def template(self, template: NDArray):
@@ -582,9 +622,7 @@ class MatchingData:
                 shape=template.shape, dtype=float, fill_value=1
             )
 
-        if type(template) == Density:
-            template = template.data
-        self._template = template.astype(self._default_dtype, copy=False)
+        self._template = template
 
     @property
     def target_mask(self):
@@ -594,8 +632,8 @@ class MatchingData:
             target_mask = self._target_mask.data
 
         if target_mask is not None:
-            out_shape = backend.to_numpy_array(self._output_target_shape)
-            target_mask = target_mask.reshape(tuple(int(x) for x in out_shape))
+            out_shape = tuple(int(x) for x in self._output_target_shape)
+            target_mask = target_mask.reshape(out_shape)
 
         return target_mask
 
@@ -605,14 +643,7 @@ class MatchingData:
         if not np.all(self.target.shape == mask.shape):
             raise ValueError("Target and its mask have to have the same shape.")
 
-        if type(mask) == Density:
-            mask.data = mask.data.astype(self._default_dtype, copy=False)
-            self._target_mask = mask
-            self._targetmaskshape = self._target_mask.shape[::-1]
-            return None
-
-        self._target_mask = mask.astype(self._default_dtype, copy=False)
-        self._targetmaskshape = self._target_mask.shape
+        self._target_mask = mask
 
     @property
     def template_mask(self):
@@ -630,24 +661,17 @@ class MatchingData:
 
         if mask is not None:
             mask = backend.reverse(mask)
-            out_shape = backend.to_numpy_array(self._output_template_shape)
-            mask = mask.reshape(tuple(int(x) for x in out_shape))
+            out_shape = tuple(int(x) for x in self._output_template_shape)
+            mask = mask.reshape(out_shape)
         return mask
 
     @template_mask.setter
     def template_mask(self, mask: NDArray):
         """Returns the reversed template mask NDArray."""
         if not np.all(self._templateshape[::-1] == mask.shape):
-            raise ValueError("Target and its mask have to have the same shape.")
-
-        if type(mask) == Density:
-            mask.data = mask.data.astype(self._default_dtype, copy=False)
-            self._template_mask = mask
-            self._templatemaskshape = self._template_mask.shape[::-1]
-            return None
+            raise ValueError("Template and its mask have to have the same shape.")
 
-        self._template_mask = mask.astype(self._default_dtype, copy=False)
-        self._templatemaskshape = self._template_mask.shape[::-1]
+        self._template_mask = mask
 
     def _split_rotations_on_jobs(self, n_jobs: int) -> List[NDArray]:
         """