pytme 0.1.5__cp311-cp311-macosx_14_0_arm64.whl

tme/matching_optimization.py

@@ -0,0 +1,1123 @@
+""" Implements various methods for non-exhaustive template matching
+    based on numerical optimization.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Tuple, Dict
+from abc import ABC, abstractmethod
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.optimize import (
+    differential_evolution,
+    LinearConstraint,
+    basinhopping,
+)
+from scipy.ndimage import laplace
+from scipy.spatial import KDTree
+
+from .matching_utils import rigid_transform, euler_to_rotationmatrix
+
+
+class MatchCoordinatesToDensity(ABC):
+    """
+    A class to template match coordinate sets.
+
+    Parameters
+    ----------
+    target_coordinates : NDArray
+        The coordinates of the target.
+    template_coordinates : NDArray
+        The coordinates of the template.
+    target_weights : NDArray
+        The weights of the target.
+    template_weights : NDArray
+        The weights of the template.
+    sampling_rate : NDArray
+        The size of the voxel.
+    template_mask_coordinates : NDArray, optional
+        The coordinates of the template mask. Default is None.
+    target_mask_coordinates : NDArray, optional
+        The coordinates of the target mask. Default is None.
+    **kwargs : dict, optional
+        Other keyword arguments.
+    """
+
+    def __init__(
+        self,
+        target_coordinates: NDArray,
+        template_coordinates: NDArray,
+        target_weights: NDArray,
+        template_weights: NDArray,
+        sampling_rate: NDArray,
+        template_mask_coordinates: NDArray = None,
+        target_mask_coordinates: NDArray = None,
+        **kwargs,
+    ):
+        target, _, origin = FitRefinement.array_from_coordinates(
+            target_coordinates, target_weights, sampling_rate
+        )
+        self.target_density = target
+        self.target_origin = origin
+        self.sampling_rate = sampling_rate
+
+        self.template_weights = template_weights
+        self.template_coordinates = template_coordinates
+        self.template_coordinates_rotated = np.empty(
+            self.template_coordinates.shape, dtype=np.float32
+        )
+
+        self.target_mask_density = None
+        if target_mask_coordinates is not None:
+            target_mask, *_ = FitRefinement.array_from_coordinates(
+                coordinates=target_mask_coordinates.astype(np.float32),
+                weights=np.ones(target_mask_coordinates.shape[1]),
+                shape=self.target_density.shape,
+                origin=self.target_origin,
+                sampling_rate=self.sampling_rate,
+            )
+            self.target_mask_density = target_mask
+
+        self.template_mask_coordinates = None
+        self.template_mask_coordinates_rotated = None
+        if template_mask_coordinates is not None:
+            self.template_mask_coordinates = template_mask_coordinates
+            self.template_mask_coordinates_rotated = np.empty(
+                self.template_mask_coordinates.shape, dtype=np.float32
+            )
+
+    def __call__(self, x: NDArray):
+        """
+        Return the score for a given transformation.
+
+        Parameters
+        ----------
+        x : NDArray
+            The input transformation parameters.
+
+        Returns
+        -------
+        float
+            The negative score from the scoring function.
+        """
+        translation, rotation = x[:3], x[3:]
+        rotation_matrix = euler_to_rotationmatrix(rotation)
+
+        rigid_transform(
+            coordinates=self.template_coordinates,
+            coordinates_mask=self.template_mask_coordinates,
+            rotation_matrix=rotation_matrix,
+            translation=translation,
+            out=self.template_coordinates_rotated,
+            out_mask=self.template_mask_coordinates_rotated,
+            use_geometric_center=False,
+        )
+
+        mapping = FitRefinement.map_coordinates_to_array(
+            coordinates=self.template_coordinates_rotated,
+            coordinates_mask=self.template_mask_coordinates_rotated,
+            array_origin=self.target_origin,
+            array_shape=self.target_density.shape,
+            sampling_rate=self.sampling_rate,
+        )
+
+        return -self.scoring_function(
+            transformed_coordinates=mapping[0],
+            transformed_coordinates_mask=mapping[1],
+            in_volume=mapping[2],
+            in_volume_mask=mapping[3],
+        )
+
+    @abstractmethod
+    def scoring_function(*args, **kwargs):
+        """
+        Computes a scoring metric for a given set of coordinates.
+
+        This function is not intended to be called directly, but should rather be
+        defined by classes inheriting from :py:class:`MatchCoordinatesToDensity`
+        to parse a given file format.
+        """
+
+
+class MatchCoordinatesToCoordinates(ABC):
+    """
+    A class to template match coordinate sets.
+
+    Parameters
+    ----------
+    target_coordinates : NDArray
+        The coordinates of the target.
+    template_coordinates : NDArray
+        The coordinates of the template.
+    target_weights : NDArray
+        The weights of the target.
+    template_weights : NDArray
+        The weights of the template.
+    sampling_rate : NDArray
+        The size of the voxel.
+    template_mask_coordinates : NDArray, optional
+        The coordinates of the template mask. Default is None.
+    target_mask_coordinates : NDArray, optional
+        The coordinates of the target mask. Default is None.
+    **kwargs : dict, optional
+        Other keyword arguments.
+    """
+
+    def __init__(
+        self,
+        target_coordinates: NDArray,
+        template_coordinates: NDArray,
+        target_weights: NDArray,
+        template_weights: NDArray,
+        template_mask_coordinates: NDArray = None,
+        target_mask_coordinates: NDArray = None,
+        **kwargs,
+    ):
+        self.target_weights = target_weights
+        self.target_coordinates = target_coordinates
+
+        self.template_weights = template_weights
+        self.template_coordinates = template_coordinates
+        self.template_coordinates_rotated = np.empty(
+            self.template_coordinates.shape, dtype=np.float32
+        )
+        self.target_mask_coordinates = target_mask_coordinates
+
+        self.template_mask_coordinates = None
+        self.template_mask_coordinates_rotated = None
+        if template_mask_coordinates is not None:
+            self.template_mask_coordinates = template_mask_coordinates
+            self.template_mask_coordinates_rotated = np.empty(
+                self.template_mask_coordinates.shape, dtype=np.float32
+            )
+
+    def __call__(self, x: NDArray):
+        """
+        Return the score for a given transformation.
+
+        Parameters
+        ----------
+        x : NDArray
+            The input transformation parameters.
+
+        Returns
+        -------
+        float
+            The negative score from the scoring function.
+        """
+        translation, rotation = x[:3], x[3:]
+        rotation_matrix = euler_to_rotationmatrix(rotation)
+
+        rigid_transform(
+            coordinates=self.template_coordinates,
+            coordinates_mask=self.template_mask_coordinates,
+            rotation_matrix=rotation_matrix,
+            translation=translation,
+            out=self.template_coordinates_rotated,
+            out_mask=self.template_mask_coordinates_rotated,
+            use_geometric_center=False,
+        )
+
+        return -self.scoring_function(
+            transformed_coordinates=self.template_coordinates_rotated,
+            transformed_coordinates_mask=self.template_mask_coordinates_rotated,
+        )
+
+    @abstractmethod
+    def scoring_function(*args, **kwargs):
+        """
+        Computes a scoring metric for a given set of coordinates.
+
+        This function is not intended to be called directly, but should rather be
+        defined by classes inheriting from :py:class:`MatchCoordinatesToDensity`
+        to parse a given file format.
+        """
+
+
+class CrossCorrelation(MatchCoordinatesToDensity):
+    """
+    Class representing the Cross-Correlation matching score.
+
+    Cross-Correlation score formula:
+
+    .. math::
+
+        \\text{score} = \\text{target_weights} \\cdot \\text{template_weights}
+
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.denominator = 1
+
+    def scoring_function(
+        self,
+        transformed_coordinates: NDArray,
+        transformed_coordinates_mask: NDArray,
+        in_volume: NDArray,
+        in_volume_mask: NDArray,
+    ) -> float:
+        """
+        Compute the Cross-Correlation score.
+
+        Parameters
+        ----------
+        transformed_coordinates : NDArray
+            Transformed coordinates.
+        transformed_coordinates_mask : NDArray
+            Mask for the transformed coordinates.
+        in_volume : NDArray
+            Binary mask indicating which ``transformed_coordinates`` are in the
+            target volume.
+        in_volume_mask : NDArray
+            Binary mask indicating which ``transformed_coordinates`` are in the
+            target mask volume.
+
+        Returns
+        -------
+        float
+            The Cross-Correlation score.
+        """
+        score = np.dot(
+            self.target_density[tuple(transformed_coordinates[:, in_volume])],
+            self.template_weights[in_volume],
+        )
+        score /= self.denominator
+        return score
+
+
+class LaplaceCrossCorrelation(CrossCorrelation):
+    """
+    Class representing the Laplace Cross-Correlation matching score.
+
+    The score is computed like CrossCorrelation, but with Laplace filtered
+    weights, indicated by the Laplace operator :math:`\\nabla^{2}`.
+
+    .. math::
+
+        \\text{score} = \\nabla^{2} \\text{target_weights} \\cdot
+                        \\nabla^{2} \\text{template_weights}
+
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.target_density = laplace(self.target_density)
+
+        arr, positions, _ = FitRefinement.array_from_coordinates(
+            self.template_coordinates, self.template_weights, self.sampling_rate
+        )
+        self.template_weights = laplace(arr)[tuple(positions)]
+
+
+class NormalizedCrossCorrelation(CrossCorrelation):
+    """
+    Class representing the Normalized Cross-Correlation matching score.
+
+    The score is computed by normalizing the dot product of `target_weights` and
+    `template_weights` with the product of their norms. This normalization ensures
+    the score lies between -1 and 1, providing a measure of similarity that's invariant
+    to scale.
+
+    .. math::
+
+        \\text{score} = \\frac{\\text{target_weights} \\cdot \\text{template_weights}}
+                        {\\text{max(target_norm} \\times \\text{template_norm, eps)}}
+
+    Where:
+
+    .. math::
+
+        \\text{target_norm} = ||\\text{target_weights}||
+
+    .. math::
+
+        \\text{template_norm} = ||\\text{template_weights}||
+
+    Here, :math:`||.||` denotes the L2 (Euclidean) norm.
+
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        target_norm = np.linalg.norm(self.target_density[self.target_density != 0])
+        template_norm = np.linalg.norm(self.template_weights)
+        self.denominator = np.fmax(target_norm * template_norm, np.finfo(float).eps)
+
+
+class NormalizedCrossCorrelationMean(NormalizedCrossCorrelation):
+    """
+    Class representing the Mean Normalized Cross-Correlation matching score.
+
+    This class extends the Normalized Cross-Correlation by computing the score
+    after subtracting the mean from both `target_weights` and `template_weights`.
+    This modification enhances the matching score's sensitivity to patterns
+    over flat regions in the data.
+
+    Mathematically, the Mean Normalized Cross-Correlation score is computed as:
+
+    .. math::
+
+        \\text{score} = \\frac{(\\text{target_weights} - \\text{mean(target_weights)})
+                        \\cdot (\\text{template_weights} -
+                        \\text{mean(template_weights)})}
+                        {\\text{max(target_norm} \\times \\text{template_norm, eps)}}
+
+    Where:
+
+    .. math::
+
+        \\text{target_norm} = ||\\text{target_weights} - \\text{mean(target_weights)}||
+
+    .. math::
+
+        \\text{template_norm} = ||\\text{template_weights} -
+        \\text{mean(template_weights)}||
+
+    Here, :math:`||.||` denotes the L2 (Euclidean) norm, and :math:`\\text{mean(.)}`
+    computes the mean of the respective weights.
+    """
+
+    def __init__(self, **kwargs):
+        print(kwargs["target_weights"].mean())
+        kwargs["target_weights"] -= kwargs["target_weights"].mean()
+        kwargs["template_weights"] -= kwargs["template_weights"].mean()
+        super().__init__(**kwargs)
+
+
+class MaskedCrossCorrelation(MatchCoordinatesToDensity):
+    """
+    Class representing the Masked Cross-Correlation matching score.
+
+    The Masked Cross-Correlation computes the similarity between `target_weights`
+    and `template_weights` under respective masks. The score is normalized and lies
+    between -1 and 1, providing a measure of similarity even in the presence of
+    missing or masked data.
+
+    The formula for the Masked Cross-Correlation is:
+
+    .. math::
+        \\text{numerator} = \\text{dot}(\\text{target_weights},
+                            \\text{template_weights}) -
+                            \\frac{\\text{sum}(\\text{mask_target}) \\times
+                            \\text{sum}(\\text{mask_template})}
+                                  {\\text{mask_overlap}}
+
+    .. math::
+        \\text{denominator1} = \\text{sum}(\\text{mask_target}^2) -
+                               \\frac{\\text{sum}(\\text{mask_target})^2}
+                                     {\\text{mask_overlap}}
+
+    .. math::
+        \\text{denominator2} = \\text{sum}(\\text{mask_template}^2) -
+                               \\frac{\\text{sum}(\\text{mask_template})^2}
+                                     {\\text{mask_overlap}}
+
+    .. math::
+        \\text{denominator} = \\sqrt{\\text{denominator1} \\times \\text{denominator2}}
+
+    .. math::
+        \\text{score} = \\frac{\\text{numerator}}{\\text{denominator}}
+                        \\text{ if denominator } \\neq 0
+                        \\text{ else } 0
+
+    Where:
+
+    -   mask_target and mask_template are binary masks for the target_weights
+        and template_weights respectively.
+
+    -   mask_overlap represents the number of overlapping non-zero elements in
+        the masks.
+
+    References
+    ----------
+    .. [1]  Masked FFT registration, Dirk Padfield, CVPR 2010 conference
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def scoring_function(
+        self,
+        transformed_coordinates: NDArray,
+        transformed_coordinates_mask: NDArray,
+        in_volume: NDArray,
+        in_volume_mask: NDArray,
+    ) -> float:
+        """
+        Compute the Masked Cross-Correlation score.
+
+        Parameters
+        ----------
+        transformed_coordinates : NDArray
+            Transformed coordinates.
+        transformed_coordinates_mask : NDArray
+            Mask for the transformed coordinates.
+        in_volume : NDArray
+            Binary mask indicating which ``transformed_coordinates`` are in the
+            target volume.
+        in_volume_mask : NDArray
+            Binary mask indicating which ``transformed_coordinates`` are in the
+            target mask volume.
+
+        Returns
+        -------
+        float
+            The Masked Cross-Correlation score.
+        """
+        mask_overlap = np.sum(
+            self.target_mask_density[
+                tuple(transformed_coordinates_mask[:, in_volume_mask])
+            ],
+        )
+        mask_overlap = np.fmax(mask_overlap, np.finfo(float).eps)
+
+        mask_target = self.target_density[
+            tuple(transformed_coordinates_mask[:, in_volume_mask])
+        ]
+        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[in_volume],
+            self.target_mask_density[tuple(transformed_coordinates[:, in_volume])],
+        )
+        denominator2 = np.subtract(
+            np.sum(mask_template**2),
+            np.divide(np.square(np.sum(mask_template)), mask_overlap),
+        )
+
+        denominator1 = np.fmax(denominator1, 0.0)
+        denominator2 = np.fmax(denominator2, 0.0)
+        denominator = np.sqrt(np.multiply(denominator1, denominator2))
+
+        numerator = np.dot(
+            self.target_density[tuple(transformed_coordinates[:, in_volume])],
+            self.template_weights[in_volume],
+        )
+
+        numerator -= np.divide(
+            np.multiply(np.sum(mask_target), np.sum(mask_template)), mask_overlap
+        )
+
+        if denominator == 0:
+            return 0
+
+        score = numerator / denominator
+        return score
+
+
+class PartialLeastSquareDifference(MatchCoordinatesToDensity):
+    """
+    Class representing the Partial Least Square Difference matching score.
+
+    The Partial Least Square Difference (PLSQ) between the target :math:`f` and the
+    template :math:`g` is calculated as:
+
+    .. math::
+
+        \\text{d(f,g)} = \\sum_{i=1}^{n} \\| f(\\mathbf{p}_i) - g(\\mathbf{q}_i) \\|^2
+
+    References
+    ----------
+    .. [1]  Daven Vasishtan and Maya Topf, "Scoring functions for cryoEM density
+            fitting", Journal of Structural Biology, vol. 174, no. 2,
+            pp. 333--343, 2011. DOI: https://doi.org/10.1016/j.jsb.2011.01.012
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def scoring_function(
+        self,
+        transformed_coordinates: NDArray,
+        transformed_coordinates_mask: NDArray,
+        in_volume: NDArray,
+        in_volume_mask: NDArray,
+    ) -> float:
+        """
+        Compute the Partial Least Square Difference score.
+
+        Given the transformed coordinates and their associated mask, this function
+        computes the difference between target and template densities.
+
+        Parameters
+        ----------
+        transformed_coordinates : NDArray
+            Transformed coordinates.
+        transformed_coordinates_mask : NDArray
+            Mask for the transformed coordinates.
+        in_volume : NDArray
+            Binary mask indicating which ``transformed_coordinates`` are in the
+            target volume.
+        in_volume_mask : NDArray
+            Binary mask indicating which ``transformed_coordinates`` are in the
+            target mask volume.
+
+        Returns
+        -------
+        float
+            The negative of the Partial Least Square Difference score.
+        """
+        score = np.sum(
+            np.square(
+                np.subtract(
+                    self.target_density[tuple(transformed_coordinates[:, in_volume])],
+                    self.template_weights[in_volume],
+                )
+            )
+        )
+        score += np.sum(np.square(self.template_weights[np.invert(in_volume)]))
+
+        return -score
+
+
+class Chamfer(MatchCoordinatesToCoordinates):
+    """
+    Class representing the Chamfer matching score.
+
+    The Chamfer distance is computed as:
+
+    .. math::
+
+        \\text{d(f,g)} = \\frac{1}{|X|} \\sum_{\\mathbf{f}_i \\in X}
+        \\inf_{\\mathbf{g} \\in Y} ||\\mathbf{f}_i - \\mathbf{g}||_2
+
+    References
+    ----------
+    .. [1]  Daven Vasishtan and Maya Topf, "Scoring functions for cryoEM density
+            fitting", Journal of Structural Biology, vol. 174, no. 2,
+            pp. 333--343, 2011. DOI: https://doi.org/10.1016/j.jsb.2011.01.012
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.target_tree = KDTree(self.target_coordinates.T)
+
+    def scoring_function(
+        self,
+        transformed_coordinates: NDArray,
+        transformed_coordinates_mask: NDArray,
+        **kwargs,
+    ) -> float:
+        """
+        Compute the Chamfer distance score.
+
+        Given the transformed coordinates and their associated mask, this function
+        calculates the average distance between the rotated template coordinates
+        and the nearest target coordinates.
+
+        Parameters
+        ----------
+        transformed_coordinates : NDArray
+            Transformed coordinates.
+
+        Returns
+        -------
+        float
+            The negative of the Chamfer distance score.
+
+        """
+        dist, _ = self.target_tree.query(self.template_coordinates_rotated.T)
+        score = np.mean(dist)
+        return -score
+
+
+class MutualInformation(MatchCoordinatesToDensity):
+    """
+    Class representing the Mutual Information matching score.
+
+    The Mutual Information (MI) score is calculated as:
+
+    .. math::
+
+        \\text{d(f,g)} = \\sum_{f,g} p(f,g) \\log \\frac{p(f,g)}{p(f)p(g)}
+
+    References
+    ----------
+    .. [1]  Daven Vasishtan and Maya Topf, "Scoring functions for cryoEM density
+            fitting", Journal of Structural Biology, vol. 174, no. 2,
+            pp. 333--343, 2011. DOI: https://doi.org/10.1016/j.jsb.2011.01.012
+
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def scoring_function(
+        self,
+        transformed_coordinates: NDArray,
+        transformed_coordinates_mask: NDArray,
+        in_volume: NDArray,
+        in_volume_mask: NDArray,
+    ) -> float:
+        """
+        Compute the Mutual Information score.
+
+        Given the transformed coordinates and their associated mask, this function
+        computes the mutual information between the target and template densities.
+
+        Parameters
+        ----------
+        transformed_coordinates : NDArray
+            Transformed coordinates.
+        transformed_coordinates_mask : NDArray
+            Mask for the transformed coordinates.
+        in_volume : NDArray
+            Binary mask indicating which ``transformed_coordinates`` are in the
+            target volume.
+        in_volume_mask : NDArray
+            Binary mask indicating which ``transformed_coordinates`` are in the
+            target mask volume.
+
+        Returns
+        -------
+        float
+            The Mutual Information score.
+        """
+        p_xy, target, template = np.histogram2d(
+            self.target_density[tuple(transformed_coordinates[:, in_volume])],
+            self.template_weights[in_volume],
+        )
+        p_x, p_y = np.sum(p_xy, axis=1), np.sum(p_xy, axis=0)
+
+        p_xy /= p_xy.sum()
+        p_x /= p_x.sum()
+        p_y /= p_y.sum()
+
+        logprob = np.divide(p_xy, p_x[:, None] * p_y[None, :] + np.finfo(float).eps)
+        score = np.nansum(p_xy * logprob)
+
+        return score
+
+
+class Envelope(MatchCoordinatesToDensity):
+    """
+    Class representing the Envelope matching score.
+
+    The Envelope score (ENV) is calculated as:
+
+    .. math::
+
+        \\text{d(f,g)} =    \\sum_{\\mathbf{p} \\in P} f'(\\mathbf{p})
+                            \\cdot g'(\\mathbf{p})
+
+    References
+    ----------
+    .. [1]  Daven Vasishtan and Maya Topf, "Scoring functions for cryoEM density
+            fitting", Journal of Structural Biology, vol. 174, no. 2,
+            pp. 333--343, 2011. DOI: https://doi.org/10.1016/j.jsb.2011.01.012
+    """
+
+    def __init__(self, target_threshold: float, **kwargs):
+        super().__init__(**kwargs)
+        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)
+        self.template_weights = np.ones_like(self.template_weights)
+
+    def scoring_function(
+        self,
+        transformed_coordinates: NDArray,
+        transformed_coordinates_mask: NDArray,
+        in_volume: NDArray,
+        in_volume_mask: NDArray,
+    ) -> float:
+        """
+        Compute the Envelope score.
+
+        Given the transformed coordinates and their associated mask, this function
+        computes the envelope score based on target density thresholds.
+
+        Parameters
+        ----------
+        transformed_coordinates : NDArray
+            Transformed coordinates.
+        transformed_coordinates_mask : NDArray
+            Mask for the transformed coordinates.
+        in_volume : NDArray
+            Binary mask indicating which ``transformed_coordinates`` are in the
+            target volume.
+        in_volume_mask : NDArray
+            Binary mask indicating which ``transformed_coordinates`` are in the
+            target mask volume.
+
+        Returns
+        -------
+        float
+            The Envelope score.
+        """
+        score = self.target_density[tuple(transformed_coordinates[:, in_volume])]
+        unassigned_density = self.target_density_present - (score == -1).sum()
+
+        score = score.sum() - unassigned_density - 2 * np.sum(np.invert(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)
+
+        return score
+
+
+class NormalVectorScore(MatchCoordinatesToCoordinates):
+    """
+    Class representing the Normal Vector matching score.
+
+    The Normal Vector Score (NVS) is calculated as:
+
+    .. math::
+
+        \\text{d(f,g)} = \\frac{1}{N} \\sum_{i=1}^{N}
+        \\frac{
+            {\\vec{f}_i} \\cdot {\\vec{g}_i}
+        }{
+            ||\\vec{f}_i|| \\, ||\\vec{g}_i||
+        }
+
+    References
+    ----------
+    .. [1]  Daven Vasishtan and Maya Topf, "Scoring functions for cryoEM density
+            fitting", Journal of Structural Biology, vol. 174, no. 2,
+            pp. 333--343, 2011. DOI: https://doi.org/10.1016/j.jsb.2011.01.012
+
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def scoring_function(
+        self,
+        transformed_coordinates: NDArray,
+        transformed_coordinates_mask: NDArray,
+        **kwargs,
+    ) -> float:
+        """
+        Compute the Normal Vector Score.
+
+        Given the template and target vectors, this function computes the average
+        cosine similarity between the two sets of vectors.
+
+        Parameters
+        ----------
+        template_vectors : NDArray
+            Normal vectors derived from the template.
+        target_vectors : NDArray
+            Normal vectors derived from the target.
+
+        Returns
+        -------
+        float
+            The Normal Vector Score.
+        """
+        numerator = np.multiply(transformed_coordinates, self.target_coordinates)
+        denominator = np.linalg.norm(transformed_coordinates)
+        denominator *= np.linalg.norm(self.target_coordinates)
+        score = np.mean(numerator / denominator)
+        return score
+
+
+MATCHING_OPTIMIZATION_REGISTER = {
+    "CrossCorrelation": CrossCorrelation,
+    "LaplaceCrossCorrelation": LaplaceCrossCorrelation,
+    "NormalizedCrossCorrelationMean": NormalizedCrossCorrelationMean,
+    "NormalizedCrossCorrelation": NormalizedCrossCorrelation,
+    "MaskedCrossCorrelation": MaskedCrossCorrelation,
+    "PartialLeastSquareDifference": PartialLeastSquareDifference,
+    "Envelope": Envelope,
+    "Chamfer": Chamfer,
+    "MutualInformation": MutualInformation,
+    "NormalVectorScore": NormalVectorScore,
+}
+
+
+def register_matching_optimization(match_name: str, match_class: type):
+    """
+    Registers a class to be used by :py:class:`FitRefinement`.
+
+    Parameters
+    ----------
+    match_name : str
+        Name of the matching instance.
+    match_class : type
+        Class pointer.
+
+    Raises
+    ------
+    ValueError
+        If any of the required methods is not defined.
+    """
+    methods_to_check = ["__init__", "__call__", "scoring_function"]
+
+    for method in methods_to_check:
+        if not hasattr(match_class, method):
+            raise ValueError(
+                f"Method '{method}' is not defined in the provided class or object."
+            )
+    MATCHING_OPTIMIZATION_REGISTER[match_name] = match_class
+
+
+class FitRefinement:
+    """
+    A class to refine the fit between target and template coordinates.
+
+    Notes
+    -----
+    By default scipy.optimize.differential_evolution or scipy.optimize.basinhopping
+    are used which can be unreliable if the initial alignment is very poor. Other
+    optimizers can be implemented by subclassing :py:class:`FitRefinement` and
+    overwriting the :py:meth:`FitRefinement.refine` function.
+
+    """
+
+    @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.
+        """
+        coordinates = coordinates.astype(sampling_rate.dtype)
+        np.divide(
+            coordinates - array_origin[:, None], sampling_rate[:, None], out=coordinates
+        )
+        transformed_coordinates = coordinates.astype(int)
+        in_volume = np.logical_and(
+            transformed_coordinates < np.array(array_shape)[:, None],
+            transformed_coordinates >= 0,
+        ).min(axis=0)
+
+        transformed_coordinates_mask, in_volume_mask = None, None
+
+        if coordinates_mask is not None:
+            coordinates_mask = coordinates_mask.astype(sampling_rate.dtype)
+            np.divide(
+                coordinates_mask - array_origin[:, None],
+                sampling_rate[:, None],
+                out=coordinates_mask,
+            )
+            transformed_coordinates_mask = coordinates_mask.astype(int)
+            in_volume_mask = np.logical_and(
+                transformed_coordinates_mask < np.array(array_shape)[:, None],
+                transformed_coordinates_mask >= 0,
+            ).min(axis=0)
+
+        return (
+            transformed_coordinates,
+            transformed_coordinates_mask,
+            in_volume,
+            in_volume_mask,
+        )
+
+    @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.
+
+        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.
+
+        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
+
+        arr = np.zeros(shape, dtype=np.float32)
+        np.add.at(arr, tuple(positions), weights)
+        return arr, positions, origin
+
+    def refine(
+        self,
+        target_coordinates: NDArray,
+        target_weights: NDArray,
+        template_coordinates: NDArray,
+        template_weights: NDArray,
+        sampling_rate: float = None,
+        translational_uncertainty: Tuple[float] = None,
+        rotational_uncertainty: Tuple[float] = None,
+        scoring_class: str = "CrossCorrelation",
+        scoring_class_parameters: Dict = dict(),
+        local_optimization: bool = True,
+        maxiter: int = 100,
+    ) -> (NDArray, NDArray):
+        """
+        Refines the alignment of template coordinates to target coordinates.
+
+        Parameters
+        ----------
+        target_coordinates : NDArray
+            The coordinates of the target.
+
+        target_weights : NDArray
+            The weights of the target.
+
+        template_coordinates : NDArray
+            The coordinates of the template.
+
+        template_weights : NDArray
+            The weights of the template.
+
+        sampling_rate : float, optional
+            The size of the voxel. Default is None.
+
+        translational_uncertainty : (float,), optional
+            The translational uncertainty. Default is None.
+
+        rotational_uncertainty : (float,), optional
+            The rotational uncertainty. Default is None.
+
+        scoring_class : str, optional
+            The scoring class to be used. Default is "CC".
+
+        scoring_class_parameters : dict, optional
+            The parameters for the scoring class. Default is an empty dictionary.
+
+        local_optimization : bool, optional
+            Whether to use local optimization. Default is True.
+
+        maxiter : int, optional
+            The maximum number of iterations. Default is 100.
+
+        Returns
+        -------
+        tuple
+            A tuple containing the translation and rotation matrix of the refinement,
+            as well as the score of the refinement.
+
+        Raises
+        ------
+        NotNotImplementedError
+            If scoring class is not a part of `MATCHING_OPTIMIZATION_REGISTER`.
+            Individual scores can be added via
+            :py:meth:`register_matching_optimization`.
+
+        See Also
+        --------
+        :py:meth:`register_matching_optimization`
+        """
+        if scoring_class not in MATCHING_OPTIMIZATION_REGISTER:
+            raise NotImplementedError(
+                f"Parameter score has to be one of "
+                f"{', '.join(MATCHING_OPTIMIZATION_REGISTER.keys())}."
+            )
+        scoring_class = MATCHING_OPTIMIZATION_REGISTER.get(scoring_class, None)
+
+        if sampling_rate is None:
+            sampling_rate = np.ones(1)
+        sampling_rate = np.repeat(
+            sampling_rate, target_coordinates.shape[0] // sampling_rate.size
+        )
+
+        score = scoring_class(
+            target_coordinates=target_coordinates,
+            template_coordinates=template_coordinates,
+            target_weights=target_weights,
+            template_weights=template_weights,
+            sampling_rate=sampling_rate,
+            **scoring_class_parameters,
+        )
+
+        initial_score = score(np.zeros(6))
+
+        mass_center_target = np.dot(target_coordinates, target_weights)
+        mass_center_target /= target_weights.sum()
+        mass_center_template = np.dot(template_coordinates, template_weights)
+        mass_center_template /= template_weights.sum()
+
+        if translational_uncertainty is None:
+            mass_center_difference = np.ceil(
+                np.subtract(mass_center_target, mass_center_template)
+            ).astype(int)
+            target_range = np.ceil(
+                np.divide(
+                    np.subtract(
+                        target_coordinates.max(axis=1), target_coordinates.min(axis=1)
+                    ),
+                    2,
+                )
+            ).astype(int)
+            translational_uncertainty = tuple(
+                (center - start, center + start)
+                for center, start in zip(mass_center_difference, target_range)
+            )
+        if rotational_uncertainty is None:
+            rotational_uncertainty = tuple(
+                (-90, 90) for _ in range(target_coordinates.shape[0])
+            )
+
+        uncertainty = (*translational_uncertainty, *rotational_uncertainty)
+        bounds = [bound if bound != (0, 0) else (-1e-9, 1e-9) for bound in uncertainty]
+        linear_constraint = LinearConstraint(
+            np.eye(len(bounds)), np.min(bounds, axis=1), np.max(bounds, axis=1)
+        )
+
+        if local_optimization:
+            result = basinhopping(
+                x0=np.zeros(6),
+                func=score,
+                niter=maxiter,
+                minimizer_kwargs={"method": "COBYLA", "constraints": linear_constraint},
+            )
+        else:
+            result = differential_evolution(
+                func=score,
+                bounds=bounds,
+                constraints=linear_constraint,
+                maxiter=maxiter,
+            )
+
+        print(f"Initial score: {-initial_score} - Refined score: {-result.fun}")
+        if initial_score < result.fun:
+            result.x = np.zeros_like(result.x)
+        translation, rotation = result.x[:3], result.x[3:]
+        rotation_matrix = euler_to_rotationmatrix(rotation)
+        return translation, rotation_matrix, -result.fun