diff --git a/tests/test_matching_optimization.py b/tests/test_matching_optimization.py index 63a599fa..02b23f21 100644 --- a/tests/test_matching_optimization.py +++ b/tests/test_matching_optimization.py @@ -1,15 +1,15 @@ import numpy as np import pytest - -from tme.rotations import euler_from_rotationmatrix from tme.matching_optimization import ( MATCHING_OPTIMIZATION_REGISTER, - register_matching_optimization, - _MatchCoordinatesToDensity, + NormalizedCrossCorrelationMean, _MatchCoordinatesToCoordinates, - optimize_match, + _MatchCoordinatesToDensity, create_score_object, + optimize_match, + register_matching_optimization, ) +from tme.rotations import euler_from_rotationmatrix density_to_density = ["FLC"] @@ -219,3 +219,66 @@ def test_register_matching_optimization(self): def test_register_matching_optimization_error(self): with pytest.raises(ValueError): register_matching_optimization(match_name="new_score", match_class=None) + + +class TestNormalizedCrossCorrelationMean: + """Regression tests for the matched-mean fix and its analytic gradient.""" + + def setup_method(self): + rng = np.random.default_rng(0) + n, center, sigma = 40, 20.0, 5.0 + xx, yy, zz = np.mgrid[0:n, 0:n, 0:n] + self.target = np.exp( + -((xx - center) ** 2 + (yy - center) ** 2 + (zz - center) ** 2) / (2 * sigma ** 2) + ).astype(np.float32) + pts = (center + rng.normal(0, 3, (3, 120))).astype(np.float32) + self.coordinates = pts + self.weights = np.exp( + -((pts[0] - center) ** 2 + (pts[1] - center) ** 2 + (pts[2] - center) ** 2) / (2 * sigma ** 2) + ).astype(np.float32) + + def _obj(self, **kwargs): + return NormalizedCrossCorrelationMean( + target=self.target, + template_coordinates=self.coordinates.copy(), + template_weights=self.weights.copy(), + **kwargs, + ) + + def test_score_equals_matched_pearson(self): + # The score is the Pearson correlation of the matched target values and template weights. + obj = self._obj(negate_score=False) + v, w = obj._target_values, obj.template_weights + assert np.isclose(obj(), np.corrcoef(v, w)[0, 1], atol=1e-4) + + def test_differs_from_global_mean_centring(self): + # With a strong local offset (the matched footprint sits in a dense region), centring by the + # matched mean differs from the old global-target-mean centring. + obj = self._obj(negate_score=False) + v, w = obj._target_values, obj.template_weights + global_mean = float(self.target.mean()) + wc = w - w.mean() + old = float(np.dot(v - global_mean, wc) / (np.linalg.norm(v - global_mean) * np.linalg.norm(wc))) + assert v.mean() > global_mean + 0.1 # strong local offset + assert abs(obj() - old) > 1e-2 # matched-mean != global-mean + + def test_grad_matches_finite_difference_direction(self): + # The analytic gradient points along the finite-difference gradient (its magnitude is scaled + # by the shared Sobel factor, as for NormalizedCrossCorrelation.grad). + x0 = np.array([0.6, -0.4, 0.5, 0.05, -0.03, 0.04]) + _, g = self._obj(negate_score=True, return_gradient=True).score(x0) + ref, eps = self._obj(negate_score=True), 1e-4 + fd = np.array([ + (ref.score(x0 + np.eye(6)[i] * eps) - ref.score(x0 - np.eye(6)[i] * eps)) / (2 * eps) + for i in range(6) + ]) + g, fd = np.asarray(g)[:3], fd[:3] # translation block + assert np.dot(g, fd) / (np.linalg.norm(g) * np.linalg.norm(fd)) > 0.99 + + def test_grad_sign_follows_negate_score(self): + # grad() returns score_sign * d(score)/dp, so flipping negate_score flips the gradient + # (consistent with __call__); a hardcoded -total_grad would not. + x0 = np.array([0.6, -0.4, 0.5, 0.05, -0.03, 0.04]) + g_neg = np.asarray(self._obj(negate_score=True, return_gradient=True).score(x0)[1]) + g_pos = np.asarray(self._obj(negate_score=False, return_gradient=True).score(x0)[1]) + assert np.allclose(g_pos, -g_neg) diff --git a/tme/matching_optimization.py b/tme/matching_optimization.py index b67c4c64..7fae1bf8 100644 --- a/tme/matching_optimization.py +++ b/tme/matching_optimization.py @@ -700,12 +700,60 @@ class NormalizedCrossCorrelationMean(NormalizedCrossCorrelation): __doc__ += _MatchCoordinatesToDensity.__doc__ - def __init__(self, **kwargs): - kwargs["target"] = np.subtract(kwargs["target"], kwargs["target"].mean()) - kwargs["template_weights"] = np.subtract( - kwargs["template_weights"], kwargs["template_weights"].mean() + def __call__(self) -> float: + """Returns the score of the current configuration.""" + # Centre by the mean of the matched values (the target sampled over the template footprint). + target_values = self._target_values - be.mean(self._target_values) + template_weights = self.template_weights - be.mean(self.template_weights) + denominator = be.sqrt(be.sum(be.square(target_values))) * be.sqrt( + be.sum(be.square(template_weights)) ) - super().__init__(**kwargs) + if denominator < self.eps: + return 0.0 + return float(be.dot(target_values, template_weights) / denominator) * self.score_sign + + def grad(self): + """Gradient of the score w.r.t. translation (first three) and rotation (last three). + + Has the same form as :py:meth:`NormalizedCrossCorrelation.grad` evaluated on the centred + values ``vc = v - mean(v)`` and weights ``wc = w - mean(w)``; the mean-subtraction + contributes no extra terms because ``sum(vc) = sum(wc) = 0`` (so ``d/dp = `` + and ``d||vc||/dp = / ||vc||``). Returns ``score_sign * d(score)/dp``. + """ + grad = self._interpolate_gradient(positions=self.template_rotated) + torque = self._torques( + positions=self.template_rotated, gradients=grad, center=self.template_center + ) + + values = self._target_values - be.mean(self._target_values) + weights = self.template_weights - be.mean(self.template_weights) + + norm = be.multiply( + be.power(be.sqrt(be.sum(be.square(values))), 3), + be.sqrt(be.sum(be.square(weights))), + ) + values_sq = be.sum(be.square(values)) + values_weights = be.sum(be.multiply(values, weights)) + + translation_grad = be.multiply( + be.sum(be.multiply(grad, weights), axis=1), values_sq + ) + translation_grad -= be.multiply( + be.sum(be.multiply(grad, values), axis=1), values_weights + ) + + torque_grad = be.multiply( + be.sum(be.multiply(torque, weights), axis=1), values_sq + ) + torque_grad -= be.multiply( + be.sum(be.multiply(torque, values), axis=1), values_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 self.score_sign * total_grad class MaskedCrossCorrelation(_MatchCoordinatesToDensity):