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

{pytme-0.2.4.dist-info → pytme-0.2.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytme
-Version: 0.2.4
+Version: 0.2.5
 Summary: Python Template Matching Engine
 Author: Valentin Maurer
 Author-email: Valentin Maurer <valentin.maurer@embl-hamburg.de>

{pytme-0.2.4.dist-info → pytme-0.2.5.dist-info}/RECORD

@@ -1,8 +1,8 @@
-pytme-0.2.4.data/scripts/estimate_ram_usage.py,sha256=R1NDpFajcF-MonJ4a43SfDlA-nxBYwK7D2quzCdsVFM,2767
-pytme-0.2.4.data/scripts/match_template.py,sha256=fDxH0yYudh4bWimmum1hCtjasG1EVJp6mKZ8a6zDt0Q,39852
-pytme-0.2.4.data/scripts/postprocess.py,sha256=50PwDfOWe2Fdws4J5K-k2SgM55fARlAWCnIsv-l0i-4,24414
-pytme-0.2.4.data/scripts/preprocess.py,sha256=A2nQlNr2fvrZ6C89jGsscgWk85KuDQIPKloQGBhExeE,4380
-pytme-0.2.4.data/scripts/preprocessor_gui.py,sha256=AHgL8j7nVCH3srsyGYWU7i3mCxeu00H-mR2qObR90GA,39071
+pytme-0.2.5.data/scripts/estimate_ram_usage.py,sha256=R1NDpFajcF-MonJ4a43SfDlA-nxBYwK7D2quzCdsVFM,2767
+pytme-0.2.5.data/scripts/match_template.py,sha256=fDxH0yYudh4bWimmum1hCtjasG1EVJp6mKZ8a6zDt0Q,39852
+pytme-0.2.5.data/scripts/postprocess.py,sha256=50PwDfOWe2Fdws4J5K-k2SgM55fARlAWCnIsv-l0i-4,24414
+pytme-0.2.5.data/scripts/preprocess.py,sha256=A2nQlNr2fvrZ6C89jGsscgWk85KuDQIPKloQGBhExeE,4380
+pytme-0.2.5.data/scripts/preprocessor_gui.py,sha256=AHgL8j7nVCH3srsyGYWU7i3mCxeu00H-mR2qObR90GA,39071
 scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 scripts/estimate_ram_usage.py,sha256=rN7haobnHg3YcgGJIp81FNiCzy8-saJGeEurQlmQmNQ,2768
 scripts/eval.py,sha256=ebJVLxbRlB6TI5YHNr0VavZ4lmaRdf8QVafyiDhh_oU,2528
@@ -22,7 +22,7 @@ tests/test_matching_cli.py,sha256=9qLrUM3nuIkY_LaKuzxtTjOqtgC9jUCMZXTWhUxYBGw,93
 tests/test_matching_data.py,sha256=TyvnSJPzdLPiXYWdz9coQ-m4H1tUS_cbD0WaBdvrevg,6062
 tests/test_matching_exhaustive.py,sha256=-Xm8jro1YJ3uPcPUCNazWX9Y7CxsTwPeKoy2Vr8TVH8,5531
 tests/test_matching_memory.py,sha256=jypztjDwTxvixHKteZzcvTzEVyOyJzVNK5JxlzInLcE,1145
-tests/test_matching_optimization.py,sha256=nG2R03eWGsZmdLPI0O0yBVKeD2k0C65OMJ7YKV5HLsc,9823
+tests/test_matching_optimization.py,sha256=BOPeI37zK8DWxprXgH5MIDqPWYyirIGoNZaFLvFm2BE,7946
 tests/test_matching_utils.py,sha256=jdQc8L5RhFFo5T_EeiIrDs6z8G43yAI-iBe5FKr5KZc,11401
 tests/test_orientations.py,sha256=4ngFOUskcsumaaVCYIxapjdjX0A7KUM6VGkjLV-y9_Y,6776
 tests/test_packaging.py,sha256=fZhTXqa_1ONSZSVW581QRZaFPw13l0Wpc-p8AgQ4b3E,2783
@@ -57,25 +57,25 @@ tests/preprocessing/test_frequency_filters.py,sha256=zXBRBvuFjiKWiT8yTXWYfYF3kKG
 tests/preprocessing/test_preprocessor.py,sha256=XyR4xc4YM76PUKuTIiive76Q85DdcyDAvbXNGcoKL8w,4820
 tests/preprocessing/test_utils.py,sha256=M6rmFl7a3JaBdONvPHhkbkzoDjjOAwrztPTHXqsbN6o,3037
 tme/__init__.py,sha256=R0cxXFmTvL3p7y6D0zX_rfjChbXNU_-tYI4FTWZ16Ns,177
-tme/__version__.py,sha256=SBl2EPFW-ltPvQ7vbVWItyAsz3aKYIpjO7vcfr84GkU,22
+tme/__version__.py,sha256=Xsa3ayOMVkhUWm4t06YeyHE0apjpZefxLH4ylp0CDtU,22
 tme/analyzer.py,sha256=kYxtzBCp02OJrlMME8Njh-9wIZAvsuhTHNAveNplssY,50466
-tme/density.py,sha256=APQgpZ8ILUsGCZTQDfAgcR_d-tm0T5m_AsPWIPr5pcg,84357
+tme/density.py,sha256=SUlLrEeKY7wl-ePhXb6b5yA0HVSEZic-5ICRbaPFPck,84311
 tme/extensions.cpython-311-darwin.so,sha256=KM0UcTYdq0Gib8rC6yVt7KnXrVvPNcWS-cvuYW2aPso,392496
-tme/matching_data.py,sha256=UJnDk5CTDyoytD1CamUSmQR82WvtUUoFqUEadTHTIQ8,25394
-tme/matching_exhaustive.py,sha256=g6znkqHcrgPMvKw2wYaHT_N9B_pR6bj1xbtTByJ-QW0,19593
-tme/matching_optimization.py,sha256=Y8HfecXiOvAHXM1viBaQ_aXljqqTnGwlOlFe0MJpDRQ,45082
+tme/matching_data.py,sha256=SWcW6qCfO3vImOOlIxxWfyzXkNSdyJox7JKVKHa_PT8,25450
+tme/matching_exhaustive.py,sha256=E_XPl7Z26loDfL9hP8T5t8ID4DxpLUWFN9FhxNbG3fs,19601
+tme/matching_optimization.py,sha256=S7Mmzln-VUbEClcbETxYCVahUWLRhq97hrnVVtH9T_I,45673
 tme/matching_scores.py,sha256=CECxl2Lh0TMLfZYnoCJXy3euGf8i9J0eHsAD7sqvWGU,30962
-tme/matching_utils.py,sha256=C4x4lxJq0_e1R-c0-JkGYM2MoqECAgJY5D1w4qaac5k,40046
+tme/matching_utils.py,sha256=dGAfwA3VMkaOwTG8_36oPxSNg2ihZJyQ42LORKJbUnQ,40044
 tme/memory.py,sha256=6xeIMAncQkgYDi6w-PIYgFEWRTUPu0_OTCeRO0p9r9Q,11029
 tme/orientations.py,sha256=KsYXJuLRLYXRHsDjP9_Tn1jXxIVPSaYkw1wRrWH3nUQ,26027
 tme/parser.py,sha256=fNiCAdsWI4ql92F1Ob4suiVzpjUOBlh2lad1iNY_FP8,13772
 tme/preprocessor.py,sha256=8UgPuNb0GwZ7JQoBZQisgp0r-wFKwvo0Jxb0u9kb2fg,40412
-tme/structure.py,sha256=9cG9I4muinstujpj79ZJgVQEABly8OEt9Uha26FXJLM,65800
+tme/structure.py,sha256=9kuOwLgOXkPLo4JCBFPB5fUAsR5JqOb7uFXGyURuyiM,65849
 tme/types.py,sha256=NAY7C4qxE6yz-DXVtClMvFfoOV-spWGLNfpLATZ1LcU,442
 tme/backends/__init__.py,sha256=4S68W2WJNZ9t33QSrRs6aL3OIyEVFo_zVsqXjS1iWYA,5185
 tme/backends/_jax_utils.py,sha256=YuNJHCYnSqOESMV-9LPr-ZxBg6Zvax2euBjsZM-j-64,5906
 tme/backends/cupy_backend.py,sha256=1nnCJ4nT7tJsXu1mrJGCy7x0Yg1wWVRg4SdzsQ2qiiw,9284
-tme/backends/jax_backend.py,sha256=femPIcppQVPKMyVIowgoOFFWOArvAX16aLvGho8qBNQ,10414
+tme/backends/jax_backend.py,sha256=xGVsdEjvuNYACuAxncJF07yo4-dZOEmeomGEeHek8b4,10002
 tme/backends/matching_backend.py,sha256=KfCOKD_rA9el3Y7BeH17KJ1apCUIIhhvn-vmbkb3CB0,33750
 tme/backends/mlx_backend.py,sha256=FJhqmCzgjXAjWGX1HhHFrCy_We4YwQQBkKFNG05ctzM,7788
 tme/backends/npfftw_backend.py,sha256=JDTc_1QcVi9jU3yLQF7jkgwQz_Po60OhkKuV2V3g5v8,16997
@@ -110,10 +110,10 @@ tme/preprocessing/_utils.py,sha256=1K8xPquM0v1MASwsMpIc3ZWxxpUFt69LezVZY5QcJnY,6
 tme/preprocessing/composable_filter.py,sha256=zmXN_NcuvvtstFdU6yYQ09z-XJFE4Y-kkMCL4vHy-jc,778
 tme/preprocessing/compose.py,sha256=NFB6oUQOwn8foy82i3Lm5DeZUd_5dmcKdhuwX8E6wpo,1454
 tme/preprocessing/frequency_filters.py,sha256=XPG6zRF_VSPH4CWFj1BLICm3_jNrzmiHaln0JZR7CrU,12755
-tme/preprocessing/tilt_series.py,sha256=6OptAfqISxzZOtHIx5MdSaJf7VGFeDntz2jWekpZMus,37307
-pytme-0.2.4.dist-info/LICENSE,sha256=K1IUNSVAz8BXbpH5EA8y5FpaHdvFXnAF2zeK95Lr2bY,18467
-pytme-0.2.4.dist-info/METADATA,sha256=fq2oJEwpG-N_8XYERPQwaBD-IbXySwuJ75BG8NP0WtQ,5278
-pytme-0.2.4.dist-info/WHEEL,sha256=uY16WuvBs6SVLr1w0jr9fTUdSkt0n_9cWxlDSGwcm3o,109
-pytme-0.2.4.dist-info/entry_points.txt,sha256=ff3LQL3FCWfCYOwFiP9zatm7laUbnwCkuPELkQVyUO4,241
-pytme-0.2.4.dist-info/top_level.txt,sha256=ovCUR7UXXouH3zYt_fJLoqr_vtjp1wudFgjVAnztQLE,18
-pytme-0.2.4.dist-info/RECORD,,
+tme/preprocessing/tilt_series.py,sha256=cCbUHufhPs_fniqXPm_JkqqeADAhgw3hxB0AX1Vhwz0,37273
+pytme-0.2.5.dist-info/LICENSE,sha256=K1IUNSVAz8BXbpH5EA8y5FpaHdvFXnAF2zeK95Lr2bY,18467
+pytme-0.2.5.dist-info/METADATA,sha256=h5QwF5Hkrt1YakNH9pPvzy3eB_wjbLZQApqykRJPboQ,5278
+pytme-0.2.5.dist-info/WHEEL,sha256=iyHOmzQZtL_L_KJL9-PDLQa8YEVkAYN5-On75avjx-8,109
+pytme-0.2.5.dist-info/entry_points.txt,sha256=ff3LQL3FCWfCYOwFiP9zatm7laUbnwCkuPELkQVyUO4,241
+pytme-0.2.5.dist-info/top_level.txt,sha256=ovCUR7UXXouH3zYt_fJLoqr_vtjp1wudFgjVAnztQLE,18
+pytme-0.2.5.dist-info/RECORD,,

{pytme-0.2.4.dist-info → pytme-0.2.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (73.0.1)
+Generator: setuptools (74.1.2)
 Root-Is-Purelib: false
 Tag: cp311-cp311-macosx_14_0_arm64
 

tests/test_matching_optimization.py

@@ -101,56 +101,6 @@ class TestMatchDensityToCoordinates:
         score = instance()
         assert isinstance(score, float)
 
-    def test_map_coordinates_to_array(self):
-        ret = _MatchCoordinatesToDensity.map_coordinates_to_array(
-            coordinates=self.coordinates.astype(np.float32),
-            array_shape=self.target.shape,
-            array_origin=np.zeros(self.target.ndim),
-            sampling_rate=np.ones(self.target.ndim),
-        )
-        assert len(ret) == 2
-
-        in_vol, in_vol_mask = ret
-
-        assert in_vol_mask is None
-        assert np.allclose(in_vol.shape, self.coordinates.shape[1])
-
-    def test_map_coordinates_to_array_mask(self):
-        ret = _MatchCoordinatesToDensity.map_coordinates_to_array(
-            coordinates=self.coordinates.astype(np.float32),
-            array_shape=self.target.shape,
-            array_origin=self.origin,
-            sampling_rate=self.sampling_rate,
-            coordinates_mask=self.coordinates.astype(np.float32),
-        )
-        assert len(ret) == 2
-
-        in_vol, in_vol_mask = ret
-        assert np.allclose(in_vol, in_vol_mask)
-
-    def test_array_from_coordinates(self):
-        ret = _MatchCoordinatesToDensity.array_from_coordinates(
-            coordinates=self.coordinates,
-            weights=self.coordinates_weights,
-            sampling_rate=self.sampling_rate,
-        )
-        assert len(ret) == 3
-        arr, positions, origin = ret
-        assert arr.ndim == self.coordinates.shape[0]
-        assert positions.shape == self.coordinates.shape
-        assert origin.shape == (self.coordinates.shape[0],)
-
-        assert np.allclose(origin, self.coordinates.min(axis=1))
-
-        ret = _MatchCoordinatesToDensity.array_from_coordinates(
-            coordinates=self.coordinates,
-            weights=self.coordinates_weights,
-            sampling_rate=self.sampling_rate,
-            origin=self.origin,
-        )
-        arr, positions, origin = ret
-        assert np.allclose(origin, self.origin)
-
 
 class TestMatchCoordinateToCoordinates:
     def setup_method(self):

tme/__version__.py

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

tme/backends/jax_backend.py

@@ -208,13 +208,6 @@ 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 = matching_data._template.shape
 
         ret, template_filter, target_filter = [], 1, 1
@@ -260,8 +253,8 @@ class JaxBackend(NumpyFFTWBackend):
             base, targets = None, self._array_backend.stack(targets)
             scores, rotations = scan_inner(
                 targets,
-                self.topleft_pad(matching_data.template, fastt_shape),
-                self.topleft_pad(matching_data.template_mask, fastt_shape),
+                matching_data.template,
+                matching_data.template_mask,
                 matching_data.rotations,
                 template_filter,
                 target_filter,

tme/density.py

@@ -697,8 +697,7 @@ class Density:
         References
         ----------
         .. [1]  Sorzano, Carlos et al (Mar. 2015). Fast and accurate conversion
-            of atomic models into electron density maps. AIMS Biophysics
-            2, 8–20.
+            of atomic models into electron density maps. AIMS Biophysics 2, 8–20.
 
         Examples
         --------
@@ -743,17 +742,16 @@ class Density:
         >>>    filter_by_residues = {"SER", "THR", "CYS", "ASN", "GLN", "TYR"}
         >>> )
 
-        :py:meth:`Density.from_structure` supports a variety of methods to convert
-        atoms into densities
+        In addtion, :py:meth:`Density.from_structure` supports a variety of methods
+        to convert atoms into densities, such as Gaussians
 
         >>> density = Density.from_structure(
         >>>    filename_or_structure = path_to_structure,
         >>>    weight_type = "gaussian",
-        >>>    weight_type_args={"resolution": "20"}
+        >>>    weight_type_args={"resolution": 20}
         >>> )
 
-        In addition its possible to use experimentally determined scattering factors
-        from various sources:
+        experimentally determined scattering factors
 
         >>> density = Density.from_structure(
         >>>    filename_or_structure = path_to_structure,
@@ -1688,11 +1686,10 @@ class Density:
             Resampling method to use, defaults to `spline`. Availabe options are:
 
             +---------+----------------------------------------------------------+
-            | spline  | Smooth spline interpolation via :obj:`scipy.ndimage.zoom`|
+            | spline  | Spline interpolation using :obj:`scipy.ndimage.zoom`     |
             +---------+----------------------------------------------------------+
-            | fourier | Frequency preserving Fourier cropping                    |
+            | fourier | Fourier cropping                                         |
             +---------+----------------------------------------------------------+
-
         order : int, optional
             Order of spline used for interpolation, by default 1. Ignored when
             ``method`` is `fourier`.

tme/matching_data.py

@@ -478,7 +478,8 @@ class MatchingData:
         shape_diff = np.multiply(
             np.subtract(target_shape, template_shape), 1 - batch_mask
         )
-        if np.sum(shape_diff < 0):
+        shape_mask = shape_diff < 0
+        if np.sum(shape_mask):
             shape_shift = np.divide(shape_diff, 2)
             offset = np.mod(shape_diff, 2)
             if pad_fourier:
@@ -491,8 +492,7 @@ class MatchingData:
                     "Template is larger than target and padding is turned off. Consider "
                     "swapping them or activate padding. Correcting the shift for now."
                 )
-
-            shape_shift = np.add(shape_shift, offset)
+            shape_shift = np.multiply(np.add(shape_shift, offset), shape_mask)
             fourier_shift = np.subtract(fourier_shift, shape_shift).astype(int)
 
         fourier_shift = tuple(fourier_shift.astype(int))

tme/matching_exhaustive.py

@@ -90,12 +90,12 @@ def _setup_template_filter_apply_target_filter(
         #     pad_fourier=False,
         # )
         fastt_shape = matching_data._template.shape
-        matching_data.template = be.reverse(
-            be.topleft_pad(matching_data.template, fastt_shape)
-        )
-        matching_data.template_mask = be.reverse(
-            be.topleft_pad(matching_data.template_mask, fastt_shape)
-        )
+        # matching_data.template = be.reverse(
+        #     be.topleft_pad(matching_data.template, fastt_shape)
+        # )
+        # matching_data.template_mask = be.reverse(
+        #     be.topleft_pad(matching_data.template_mask, fastt_shape)
+        # )
         matching_data._set_matching_dimension(
             target_dims=matching_data._target_dims,
             template_dims=matching_data._template_dims,
@@ -208,7 +208,7 @@ def scan(
 
     Examples
     --------
-    Schematically, using :py:meth:`scan` is similar to :py:meth:`scan_subsets`,
+    Schematically, :py:meth:`scan` is identical to :py:meth:`scan_subsets`,
     with the distinction that the objects contained in ``matching_data`` are not
     split and the search is only parallelized over angles.
     Assuming you have followed the example in :py:meth:`scan_subsets`, :py:meth:`scan`

tme/matching_optimization.py

@@ -1,17 +1,16 @@
-""" Implements various methods for non-exhaustive template matching
-    based on numerical optimization.
+""" Implements methods for non-exhaustive template matching.
 
     Copyright (c) 2023 European Molecular Biology Laboratory
 
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
-
-from typing import Tuple, Dict
+import warnings
+from typing import Tuple, List, Dict
 from abc import ABC, abstractmethod
 
 import numpy as np
 from scipy.spatial import KDTree
-from scipy.ndimage import laplace, map_coordinates
+from scipy.ndimage import laplace, map_coordinates, sobel
 from scipy.optimize import (
     minimize,
     basinhopping,
@@ -157,15 +156,13 @@ class _MatchDensityToDensity(ABC):
         if out is None:
             out = np.zeros_like(arr)
 
-        map_coordinates(arr, self.grid_out, order=order, output=out.ravel())
+        self._interpolate(arr, self.grid_out, order=order, out=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, self.grid_out, order=order, output=out_mask.ravel()
-            )
+            self._interpolate(arr_mask, self.grid_out, order=order, out=out.ravel())
 
         match return_type:
             case 0:
@@ -177,6 +174,12 @@ class _MatchDensityToDensity(ABC):
             case 3:
                 return out, out_mask
 
+    @staticmethod
+    def _interpolate(data, positions, order: int = 1, out=None):
+        return map_coordinates(
+            data, positions, order=order, mode="constant", output=out
+        )
+
     def score_translation(self, x: Tuple[float]) -> float:
         """
         Computes the score after a given translation.
@@ -291,6 +294,8 @@ class _MatchCoordinatesToDensity(_MatchDensityToDensity):
         A d-dimensional mask to be applied to the target.
     negate_score : bool, optional
         Whether the final score should be multiplied by negative one. Default is True.
+    return_gradient : bool, optional
+        Invoking __call_ returns a tuple of score and parameter gradient. Default is False.
     **kwargs : Dict, optional
         Keyword arguments propagated to downstream functions.
     """
@@ -303,40 +308,43 @@ class _MatchCoordinatesToDensity(_MatchDensityToDensity):
         template_mask_coordinates: NDArray = None,
         target_mask: NDArray = None,
         negate_score: bool = True,
+        return_gradient: bool = False,
+        interpolation_order: int = 1,
         **kwargs: Dict,
     ):
-        self.eps = be.eps(target.dtype)
-        self.target_density = target
-        self.target_mask_density = target_mask
+        self.target = target.astype(np.float32)
+        self.target_mask = None
+        if target_mask is not None:
+            self.target_mask = target_mask.astype(np.float32)
 
-        self.template_weights = template_weights
-        self.template_coordinates = template_coordinates
-        self.template_coordinates_rotated = np.copy(self.template_coordinates).astype(
-            np.float32
+        self.eps = be.eps(self.target.dtype)
+
+        self.target_grad = np.stack(
+            [sobel(self.target, axis=i) for i in range(self.target.ndim)]
         )
-        if template_mask_coordinates is None:
-            template_mask_coordinates = template_coordinates.copy()
 
-        self.template_mask_coordinates = template_mask_coordinates
-        self.template_mask_coordinates_rotated = template_mask_coordinates
+        self.n_points = template_coordinates.shape[1]
+        self.template = template_coordinates.astype(np.float32)
+        self.template_rotated = np.zeros_like(self.template)
+        self.template_weights = template_weights.astype(np.float32)
+        self.template_center = np.mean(self.template, axis=1)[:, None]
+
+        self.template_mask, self.template_mask_rotated = None, None
         if template_mask_coordinates is not None:
-            self.template_mask_coordinates_rotated = np.copy(
-                self.template_mask_coordinates
-            ).astype(np.float32)
+            self.template_mask = template_mask_coordinates.astype(np.float32)
+            self.template_mask_rotated = np.empty_like(self.template_mask)
 
         self.denominator = 1
         self.score_sign = -1 if negate_score else 1
+        self.interpolation_order = interpolation_order
 
-        self.in_volume, self.in_volume_mask = self.map_coordinates_to_array(
-            coordinates=self.template_coordinates_rotated,
-            coordinates_mask=self.template_mask_coordinates_rotated,
-            array_origin=be.zeros(target.ndim),
-            array_shape=self.target_density.shape,
-            sampling_rate=be.full(target.ndim, fill_value=1),
+        self._target_values = self._interpolate(
+            self.target, self.template, order=self.interpolation_order
         )
 
-        if hasattr(self, "_post_init"):
-            self._post_init(**kwargs)
+        if return_gradient and not hasattr(self, "grad"):
+            raise NotImplementedError(f"{type(self)} does not have grad method.")
+        self.return_gradient = return_gradient
 
     def score(self, x: Tuple[float]):
         """
@@ -356,119 +364,39 @@ class _MatchCoordinatesToDensity(_MatchDensityToDensity):
         translation, rotation_matrix = _format_rigid_transform(x)
 
         rigid_transform(
-            coordinates=self.template_coordinates,
-            coordinates_mask=self.template_mask_coordinates,
+            coordinates=self.template,
+            coordinates_mask=self.template_mask,
             rotation_matrix=rotation_matrix,
             translation=translation,
-            out=self.template_coordinates_rotated,
-            out_mask=self.template_mask_coordinates_rotated,
+            out=self.template_rotated,
+            out_mask=self.template_mask_rotated,
             use_geometric_center=False,
         )
 
-        self.in_volume, self.in_volume_mask = self.map_coordinates_to_array(
-            coordinates=self.template_coordinates_rotated,
-            coordinates_mask=self.template_mask_coordinates_rotated,
-            array_origin=be.zeros(rotation_matrix.shape[0]),
-            array_shape=self.target_density.shape,
-            sampling_rate=be.full(rotation_matrix.shape[0], fill_value=1),
+        self._target_values = self._interpolate(
+            self.target, self.template_rotated, order=self.interpolation_order
         )
 
-        return self()
+        score = self()
+        if not self.return_gradient:
+            return score
 
-    @staticmethod
-    def array_from_coordinates(
-        coordinates: NDArray,
-        weights: NDArray,
-        sampling_rate: NDArray,
-        origin: NDArray = None,
-        shape: NDArray = None,
-    ) -> Tuple[NDArray, NDArray, NDArray]:
-        """
-        Create a volume from coordinates, using given weights and voxel size.
+        return score, self.grad()
 
-        Parameters
-        ----------
-        coordinates : NDArray
-            An array representing the coordinates [d x N].
-        weights : NDArray
-            An array representing the weights for each coordinate [N].
-        sampling_rate : NDArray
-            The size of a voxel in the volume.
-        origin : NDArray, optional
-            The origin of the volume.
-        shape : NDArray, optional
-            The shape of the volume.
+    def _interpolate_gradient(self, positions):
+        ret = be.zeros(positions.shape, dtype=positions.dtype)
 
-        Returns
-        -------
-        tuple
-            Returns the generated volume, positions of coordinates, and origin.
-        """
-        if origin is None:
-            origin = coordinates.min(axis=1)
-
-        positions = np.divide(coordinates - origin[:, None], sampling_rate[:, None])
-        positions = positions.astype(int)
-
-        if shape is None:
-            shape = positions.max(axis=1) + 1
+        for k in range(self.target_grad.shape[0]):
+            ret[k, :] = self._interpolate(
+                self.target_grad[k], positions, order=self.interpolation_order
+            )
 
-        arr = np.zeros(shape, dtype=np.float32)
-        np.add.at(arr, tuple(positions), weights)
-        return arr, positions, origin
+        return ret
 
     @staticmethod
-    def map_coordinates_to_array(
-        coordinates: NDArray,
-        array_shape: NDArray,
-        array_origin: NDArray,
-        sampling_rate: NDArray,
-        coordinates_mask: NDArray = None,
-    ) -> Tuple[NDArray, NDArray]:
-        """
-        Map coordinates to a volume based on given voxel size and origin.
-
-        Parameters
-        ----------
-        coordinates : NDArray
-            An array representing the coordinates to be mapped [d x N].
-        array_shape : NDArray
-            The shape of the array to which the coordinates are mapped.
-        array_origin : NDArray
-            The origin of the array to which the coordinates are mapped.
-        sampling_rate : NDArray
-            The size of a voxel in the array.
-        coordinates_mask : NDArray, optional
-            An array representing the mask for the coordinates [d x T].
-
-        Returns
-        -------
-        tuple
-            Returns transformed coordinates, transformed coordinates mask,
-            mask for in_volume points, and mask for in_volume points in mask.
-        """
-        np.divide(
-            coordinates - array_origin[:, None], sampling_rate[:, None], out=coordinates
-        )
-
-        in_volume = np.logical_and(
-            coordinates < np.array(array_shape)[:, None],
-            coordinates >= 0,
-        ).min(axis=0)
-
-        in_volume_mask = None
-        if coordinates_mask is not None:
-            np.divide(
-                coordinates_mask - array_origin[:, None],
-                sampling_rate[:, None],
-                out=coordinates_mask,
-            )
-            in_volume_mask = np.logical_and(
-                coordinates_mask < np.array(array_shape)[:, None],
-                coordinates_mask >= 0,
-            ).min(axis=0)
-
-        return in_volume, in_volume_mask
+    def _torques(positions, center, gradients):
+        positions_center = (positions - center).T
+        return be.cross(positions_center, gradients.T).T
 
 
 class _MatchCoordinatesToCoordinates(_MatchDensityToDensity):
@@ -635,14 +563,43 @@ class CrossCorrelation(_MatchCoordinatesToDensity):
 
     def __call__(self) -> float:
         """Returns the score of the current configuration."""
-        score = np.dot(
-            self.target_density[
-                tuple(self.template_coordinates_rotated[:, self.in_volume].astype(int))
-            ],
-            self.template_weights[self.in_volume],
+        score = be.dot(self._target_values, self.template_weights)
+        score /= self.denominator * self.score_sign
+        return score
+
+    def grad(self):
+        """
+        Calculate the gradient of the cost function w.r.t. translation and rotation.
+
+        .. math::
+
+            \\nabla f = -\\frac{1}{N} \\begin{bmatrix}
+            \\sum_i w_i \\nabla v(x_i) \\\\
+            \\sum_i w_i (r_i \\times \\nabla v(x_i))
+            \\end{bmatrix}
+
+        where :math:`N` is the number of points, :math:`w_i` are weights,
+        :math:`x_i` are rotated template positions, and :math:`r_i` are
+        positions relative to the template center.
+
+        Returns
+        -------
+        np.ndarray
+            Negative gradient of the cost function: [dx, dy, dz, dRx, dRy, dRz].
+
+        """
+        grad = self._interpolate_gradient(positions=self.template_rotated)
+        torque = self._torques(
+            positions=self.template_rotated, gradients=grad, center=self.template_center
         )
-        score /= self.denominator
-        return score * self.score_sign
+
+        translation_grad = be.sum(grad * self.template_weights, axis=1)
+        torque_grad = be.sum(torque * self.template_weights, axis=1)
+
+        # <u, dv/dx> / <u, r x dv/dx>
+        total_grad = be.concatenate([translation_grad, torque_grad])
+        total_grad = be.divide(total_grad, self.n_points, out=total_grad)
+        return -total_grad
 
 
 class LaplaceCrossCorrelation(CrossCorrelation):
@@ -658,15 +615,18 @@ class LaplaceCrossCorrelation(CrossCorrelation):
 
     __doc__ += _MatchCoordinatesToDensity.__doc__
 
-    def _post_init(self, **kwargs):
-        self.target_density = laplace(self.target_density)
+    def __init__(self, **kwargs):
+        kwargs["target"] = laplace(kwargs["target"])
 
-        arr, positions, _ = self.array_from_coordinates(
-            self.template_coordinates,
-            self.template_weights,
-            np.ones(self.template_coordinates.shape[0]),
-        )
-        self.template_weights = laplace(arr)[tuple(positions)]
+        coordinates = kwargs["template_coordinates"]
+        origin = coordinates.min(axis=1)
+        positions = (coordinates - origin[:, None]).astype(int)
+        shape = positions.max(axis=1) + 1
+        arr = np.zeros(shape, dtype=np.float32)
+        np.add.at(arr, tuple(positions), kwargs["template_weights"])
+
+        kwargs["template_weights"] = laplace(arr)[tuple(positions)]
+        super().__init__(**kwargs)
 
 
 class NormalizedCrossCorrelation(CrossCorrelation):
@@ -696,24 +656,76 @@ class NormalizedCrossCorrelation(CrossCorrelation):
     __doc__ += _MatchCoordinatesToDensity.__doc__
 
     def __call__(self) -> float:
-        n_observations = be.sum(self.in_volume_mask)
-        target_coordinates = be.astype(
-            self.template_mask_coordinates_rotated[:, self.in_volume_mask], int
+        denominator = be.multiply(
+            np.linalg.norm(self.template_weights), np.linalg.norm(self._target_values)
         )
-        target_weight = self.target_density[tuple(target_coordinates)]
-        ex2 = be.divide(be.sum(be.square(target_weight)), n_observations)
-        e2x = be.square(be.divide(be.sum(target_weight), n_observations))
-
-        denominator = be.maximum(be.subtract(ex2, e2x), 0.0)
-        denominator = be.sqrt(denominator)
-        denominator = be.multiply(denominator, n_observations)
 
-        if denominator <= self.eps:
+        if denominator <= 0:
             return 0.0
 
         self.denominator = denominator
         return super().__call__()
 
+    def grad(self):
+        """
+        Calculate the normalized gradient of the cost function w.r.t. translation and rotation.
+
+        .. math::
+
+            \\nabla f = -\\frac{1}{N|w||v|^3} \\begin{bmatrix}
+            (\\sum_i w_i \\nabla v(x_i))|v|^2 - (\\sum_i v(x_i)
+                \\nabla v(x_i))(w \\cdot v) \\\\
+            (\\sum_i w_i (r_i \\times \\nabla v(x_i)))|v|^2 - (\\sum_i v(x_i)
+            (r_i \\times \\nabla v(x_i)))(w \\cdot v)
+            \\end{bmatrix}
+
+        where :math:`N` is the number of points, :math:`w` are weights,
+        :math:`v` are target values, :math:`x_i` are rotated template positions,
+        and :math:`r_i` are positions relative to the template center.
+
+        Returns
+        -------
+        np.ndarray
+            Negative normalized gradient: [dx, dy, dz, dRx, dRy, dRz].
+
+        """
+        grad = self._interpolate_gradient(positions=self.template_rotated)
+        torque = self._torques(
+            positions=self.template_rotated, gradients=grad, center=self.template_center
+        )
+
+        norm = be.multiply(
+            be.power(be.sqrt(be.sum(be.square(self._target_values))), 3),
+            be.sqrt(be.sum(be.square(self.template_weights))),
+        )
+
+        # (<u,dv/dx> * |v|**2 - <u,v> * <v,dv/dx>)/(|w|*|v|**3)
+        translation_grad = be.multiply(
+            be.sum(be.multiply(grad, self.template_weights), axis=1),
+            be.sum(be.square(self._target_values)),
+        )
+        translation_grad -= be.multiply(
+            be.sum(be.multiply(grad, self._target_values), axis=1),
+            be.sum(be.multiply(self._target_values, self.template_weights)),
+        )
+
+        # (<u,r x dv/dx> * |v|**2 - <u,v> * <v,r x dv/dx>)/(|w|*|v|**3)
+        torque_grad = be.multiply(
+            be.sum(be.multiply(torque, self.template_weights), axis=1),
+            be.sum(be.square(self._target_values)),
+        )
+        torque_grad -= be.multiply(
+            be.sum(be.multiply(torque, self._target_values), axis=1),
+            be.sum(be.multiply(self._target_values, self.template_weights)),
+        )
+
+        total_grad = be.concatenate([translation_grad, torque_grad])
+        if norm > 0:
+            total_grad = be.divide(total_grad, norm, out=total_grad)
+
+        total_grad = be.divide(total_grad, self.n_points, out=total_grad)
+        return -total_grad
+
 
 class NormalizedCrossCorrelationMean(NormalizedCrossCorrelation):
     """
@@ -802,33 +814,33 @@ class MaskedCrossCorrelation(_MatchCoordinatesToDensity):
 
     def __call__(self) -> float:
         """Returns the score of the current configuration."""
+
+        in_volume = np.logical_and(
+            self.template_rotated < np.array(self.target.shape)[:, None],
+            self.template_rotated >= 0,
+        ).min(axis=0)
+        in_volume_mask = np.logical_and(
+            self.template_mask_rotated < np.array(self.target.shape)[:, None],
+            self.template_mask_rotated >= 0,
+        ).min(axis=0)
+
         mask_overlap = np.sum(
-            self.target_mask_density[
-                tuple(
-                    self.template_mask_coordinates_rotated[
-                        :, self.in_volume_mask
-                    ].astype(int)
-                )
+            self.target_mask[
+                tuple(self.template_mask_rotated[:, in_volume_mask].astype(int))
             ],
         )
         mask_overlap = np.fmax(mask_overlap, np.finfo(float).eps)
 
-        mask_target = self.target_density[
-            tuple(
-                self.template_mask_coordinates_rotated[:, self.in_volume_mask].astype(
-                    int
-                )
-            )
+        mask_target = self.target[
+            tuple(self.template_mask_rotated[:, in_volume_mask].astype(int))
         ]
         denominator1 = np.subtract(
             np.sum(mask_target**2),
             np.divide(np.square(np.sum(mask_target)), mask_overlap),
         )
         mask_template = np.multiply(
-            self.template_weights[self.in_volume],
-            self.target_mask_density[
-                tuple(self.template_coordinates_rotated[:, self.in_volume].astype(int))
-            ],
+            self.template_weights[in_volume],
+            self.target_mask[tuple(self.template_rotated[:, in_volume].astype(int))],
         )
         denominator2 = np.subtract(
             np.sum(mask_template**2),
@@ -840,10 +852,8 @@ class MaskedCrossCorrelation(_MatchCoordinatesToDensity):
         denominator = np.sqrt(np.multiply(denominator1, denominator2))
 
         numerator = np.dot(
-            self.target_density[
-                tuple(self.template_coordinates_rotated[:, self.in_volume].astype(int))
-            ],
-            self.template_weights[self.in_volume],
+            self.target[tuple(self.template_rotated[:, in_volume].astype(int))],
+            self.template_weights[in_volume],
         )
 
         numerator -= np.divide(
@@ -877,21 +887,9 @@ class PartialLeastSquareDifference(_MatchCoordinatesToDensity):
 
     def __call__(self) -> float:
         """Returns the score of the current configuration."""
-        score = np.sum(
-            np.square(
-                np.subtract(
-                    self.target_density[
-                        tuple(
-                            self.template_coordinates_rotated[:, self.in_volume].astype(
-                                int
-                            )
-                        )
-                    ],
-                    self.template_weights[self.in_volume],
-                )
-            )
+        score = be.sum(
+            be.square(be.subtract(self._target_values, self.template_weights))
         )
-        score += np.sum(np.square(self.template_weights[np.invert(self.in_volume)]))
         return score * self.score_sign
 
 
@@ -917,10 +915,7 @@ class MutualInformation(_MatchCoordinatesToDensity):
     def __call__(self) -> float:
         """Returns the score of the current configuration."""
         p_xy, target, template = np.histogram2d(
-            self.target_density[
-                tuple(self.template_coordinates_rotated[:, self.in_volume].astype(int))
-            ],
-            self.template_weights[self.in_volume],
+            self._target_values, self.template_weights
         )
         p_x, p_y = np.sum(p_xy, axis=1), np.sum(p_xy, axis=0)
 
@@ -947,7 +942,7 @@ class Envelope(_MatchCoordinatesToDensity):
     References
     ----------
     .. [1]  Daven Vasishtan and Maya Topf, "Scoring functions for cryoEM density
-            fitting", Journal of Structural Biology, vol. 174, no. 2,
+            fitting", Journal of Structural Biology, vol. 1174, no. 2,
             pp. 333--343, 2011. DOI: https://doi.org/10.1016/j.jsb.2011.01.012
     """
 
@@ -956,22 +951,22 @@ class Envelope(_MatchCoordinatesToDensity):
     def __init__(self, target_threshold: float = None, **kwargs):
         super().__init__(**kwargs)
         if target_threshold is None:
-            target_threshold = np.mean(self.target_density)
-        self.target_density = np.where(self.target_density > target_threshold, -1, 1)
-        self.target_density_present = np.sum(self.target_density == -1)
-        self.target_density_absent = np.sum(self.target_density == 1)
+            target_threshold = np.mean(self.target)
+        self.target = np.where(self.target > target_threshold, -1, 1)
+        self.target_present = np.sum(self.target == -1)
+        self.target_absent = np.sum(self.target == 1)
         self.template_weights = np.ones_like(self.template_weights)
 
     def __call__(self) -> float:
         """Returns the score of the current configuration."""
-        score = self.target_density[
-            tuple(self.template_coordinates_rotated[:, self.in_volume].astype(int))
-        ]
-        unassigned_density = self.target_density_present - (score == -1).sum()
+        score = self._target_values
+        unassigned_density = self.target_present - (score == -1).sum()
 
-        score = score.sum() - unassigned_density - 2 * np.sum(np.invert(self.in_volume))
-        min_score = -self.target_density_present - 2 * self.target_density_absent
-        score = (score - 2 * min_score) / (2 * self.target_density_present - min_score)
+        # Out of volume values will be set to 0
+        score = score.sum() - unassigned_density
+        score -= 2 * np.sum(np.invert(np.abs(self._target_values) > 0))
+        min_score = -self.target_present - 2 * self.target_absent
+        score = (score - 2 * min_score) / (2 * self.target_present - min_score)
 
         return score * self.score_sign
 
@@ -1271,7 +1266,15 @@ def optimize_match(
 
     x0 = np.zeros(2 * ndim) if x0 is None else x0
 
+    return_gradient = getattr(score_object, "return_gradient", False)
+    if optimization_method != "minimize" and return_gradient:
+        warnings.warn("Gradient only considered for optimization_method='minimize'.")
+        score_object.return_gradient = False
+
     initial_score = score_object.score(x=x0)
+    if isinstance(initial_score, (List, Tuple)):
+        initial_score = initial_score[0]
+
     if optimization_method == "basinhopping":
         result = basinhopping(
             x0=x0,
@@ -1287,10 +1290,14 @@ def optimize_match(
             maxiter=maxiter,
         )
     elif optimization_method == "minimize":
-        print(maxiter)
+        if hasattr(score_object, "grad") and not return_gradient:
+            warnings.warn(
+                "Consider initializing score object with return_gradient=True."
+            )
         result = minimize(
             x0=x0,
             fun=score_object.score,
+            jac=return_gradient,
             bounds=bounds,
             constraints=linear_constraint,
             options={"maxiter": maxiter},

tme/matching_utils.py

@@ -645,7 +645,7 @@ def get_rotation_matrices(
         dets = np.linalg.det(ret)
         neg_dets = dets < 0
         ret[neg_dets, :, -1] *= -1
-        ret[0] = np.eye(dim, dtype = ret.dtype)
+        ret[0] = np.eye(dim, dtype=ret.dtype)
     return ret
 
 

tme/preprocessing/tilt_series.py

@@ -91,9 +91,9 @@ def create_reconstruction_filter(
         if tilt_angles is False:
             raise ValueError("'ramp' filter requires specifying tilt angles.")
         size = filter_shape[0]
-        ret = fftfreqn((size,), sampling_rate = 1, compute_euclidean_norm = True)
+        ret = fftfreqn((size,), sampling_rate=1, compute_euclidean_norm=True)
         min_increment = np.radians(np.min(np.abs(np.diff(np.sort(tilt_angles)))))
-        ret *= (min_increment * size)
+        ret *= min_increment * size
         np.fmin(ret, 1, out=ret)
 
         ret = np.tile(ret[:, np.newaxis], (1, filter_shape[1]))
@@ -193,7 +193,7 @@ class ReconstructFromTilt:
         volume_temp_rotated = be.zeros(shape, dtype=be._float_dtype)
         volume = be.zeros(shape, dtype=be._float_dtype)
 
-        slices = tuple(slice(a//2, (a//2) + 1) for a in shape)
+        slices = tuple(slice(a // 2, (a // 2) + 1) for a in shape)
         subset = tuple(
             slice(None) if i != opening_axis else slices[opening_axis]
             for i in range(len(shape))
@@ -423,12 +423,9 @@ class Wedge:
 
         return wedges
 
-    def weight_relion(self,
-        shape: Tuple[int],
-        opening_axis: int,
-        tilt_axis: int,
-        **kwargs
-        ) -> NDArray:
+    def weight_relion(
+        self, shape: Tuple[int], opening_axis: int, tilt_axis: int, **kwargs
+    ) -> NDArray:
         """
         Generate weighted wedges based on the RELION 1.4 formalism, weighting each
         angle using the cosine of the angle and a Gaussian lowpass filter computed
@@ -545,7 +542,7 @@ class WedgeReconstructed:
         angles: Tuple[float] = None,
         opening_axis: int = 0,
         tilt_axis: int = 2,
-        weights : Tuple[float] = None,
+        weights: Tuple[float] = None,
         weight_wedge: bool = False,
         create_continuous_wedge: bool = False,
         frequency_cutoff: float = 0.5,
@@ -675,7 +672,7 @@ class WedgeReconstructed:
         opening_axis: int,
         tilt_axis: int,
         weights: Tuple[float] = None,
-        reconstruction_filter : str = None,
+        reconstruction_filter: str = None,
         **kwargs: Dict,
     ) -> NDArray:
         """
@@ -715,7 +712,7 @@ class WedgeReconstructed:
         weights = np.repeat(weights, angles.size // weights.size)
         plane = np.zeros(
             (shape[opening_axis], shape[tilt_axis] + (1 - shape[tilt_axis] % 2)),
-            dtype=np.float32
+            dtype=np.float32,
         )
 
         # plane = np.zeros((shape[opening_axis], int(2 * np.max(shape)) + 1), dtype=np.float32)
@@ -723,7 +720,7 @@ class WedgeReconstructed:
         rec_filter = 1
         if reconstruction_filter is not None:
             rec_filter = create_reconstruction_filter(
-                plane.shape[::-1], filter_type = reconstruction_filter, tilt_angles = angles
+                plane.shape[::-1], filter_type=reconstruction_filter, tilt_angles=angles
             ).T
 
         subset = tuple(

tme/structure.py

@@ -1117,16 +1117,16 @@ class Structure:
         ----------
         positions : NDArray
             Array containing atomic positions in z,y,x format (n,d).
-        weights : [float]
+        weights : tuple of float
             The weights to use for the entries in positions.
-        resolution : float
+        resolution : float, optional
             The product of resolution and sigma_factor gives the sigma used to
             compute the discretized Gaussian.
-        sigma_factor : float
+        sigma_factor : float, optional
             The factor used with resolution to compute sigma. Default is 1 / (π√2).
-        cutoff_value : float
+        cutoff_value : float, optional
             The cutoff value for the Gaussian kernel. Default is 4.0.
-        sampling_rate : float
+        sampling_rate : float, optional
             Sampling rate along each dimension. One third of resolution by default.
 
         References
@@ -1160,8 +1160,8 @@ class Structure:
         positions = positions[:, ::-1]
         origin = positions.min(axis=0) - pad * sampling_rate
         positions = np.rint(np.divide((positions - origin), sampling_rate)).astype(int)
-        shape = positions.max(axis=0).astype(int) + pad + 1
 
+        shape = positions.max(axis=0).astype(int) + pad + 1
         out = np.zeros(shape, dtype=np.float32)
         np.add.at(out, tuple(positions.T), weights)
 
@@ -1299,10 +1299,10 @@ class Structure:
             )
 
         temp = self.subset_by_chain(chain=chain)
-        positions, atoms, shape, sampling_rate, origin = temp._coordinate_to_position(
+        positions, atoms, _shape, sampling_rate, origin = temp._coordinate_to_position(
             shape=shape, sampling_rate=sampling_rate, origin=origin
         )
-        volume = np.zeros(shape, dtype=np.float32)
+        volume = np.zeros(_shape, dtype=np.float32)
         if weight_type in ("atomic_weight", "atomic_number"):
             weights = temp._get_atom_weights(atoms=atoms, weight_type=weight_type)
             np.add.at(volume, tuple(positions.T), weights)