From c48e4dece0068f948d500b4f46974e7c45b80a5c Mon Sep 17 00:00:00 2001 From: Jan Kosinski Date: Sat, 27 Jun 2026 12:13:02 +0200 Subject: [PATCH 1/5] Fix NormalizedCrossCorrelationMean: centre by the matched mean, not the global array mean MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The docstring and the score name define the Pearson correlation about the mean over the matched points: subtract mean(target_weights) and mean(template_weights). The implementation instead subtracted the *global* target-array mean in __init__ (target - target.mean()) — a constant offset that does not centre the overlapping values, so the score did not match its own documentation. When the template covers a small high-signal region of a largely empty map (a subunit, or any contour/envelope-restricted point cloud) the global mean ≈ 0 while the matched mean is large, so the results differ substantially. Fix: compute the correlation in __call__ by centring the matched target values (self._target_values) and the template weights by their own means — the documented Pearson r over the overlap. It now equals Chimera "measure correlation ... envelope true" (the cam reported by efitter/Assembline) when given a contour-thresholded template point cloud. Note: the inherited analytic gradient is not updated for the mean-subtracted form; use a gradient-free optimizer (optimize_match's default basinhopping). --- tme/matching_optimization.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/tme/matching_optimization.py b/tme/matching_optimization.py index b67c4c64..8a16a00d 100644 --- a/tme/matching_optimization.py +++ b/tme/matching_optimization.py @@ -700,12 +700,19 @@ 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() - ) - super().__init__(**kwargs) + def __call__(self) -> float: + """Returns the score of the current configuration.""" + # Centre by the mean of the *matched* values, as the docstring and the score name + # ("...Mean") specify. The previous implementation subtracted the global target-array + # mean in __init__ — a constant offset that does not centre the overlapping values and + # so did not compute the documented Pearson-about-the-mean (it ignored the local mean + # of the target over the template footprint). + target_values = self._target_values - self._target_values.mean() + template_weights = self.template_weights - self.template_weights.mean() + denominator = float(np.linalg.norm(target_values) * np.linalg.norm(template_weights)) + if denominator < self.eps: + return 0.0 + return float(np.dot(target_values, template_weights) / denominator) * self.score_sign class MaskedCrossCorrelation(_MatchCoordinatesToDensity): From dd86e6c8cd0d410f02e214ba5d9a4177586ba2f5 Mon Sep 17 00:00:00 2001 From: Jan Kosinski Date: Sun, 28 Jun 2026 12:54:26 +0200 Subject: [PATCH 2/5] NCCMean: add the matched-mean analytic gradient Update the inherited gradient for the mean-subtracted form. NormalizedCrossCorrelationMean inherited NormalizedCrossCorrelation.grad (the gradient of the NON-centred score), inconsistent with the mean-subtracted __call__. Add NormalizedCrossCorrelationMean.grad = the same formula on the centred values v-mean(v), w-mean(w); the mean-subtraction adds no terms because sum(v-mean(v)) = sum(w-mean(w)) = 0. Verified numerically: cosine +0.999 with finite differences and the same scale factor as NormalizedCrossCorrelation.grad. (cherry picked from commit 6cf65790df5c844313b0d1d72a282eedd43825ee) --- tme/matching_optimization.py | 45 ++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tme/matching_optimization.py b/tme/matching_optimization.py index 8a16a00d..a382f3f6 100644 --- a/tme/matching_optimization.py +++ b/tme/matching_optimization.py @@ -714,6 +714,51 @@ def __call__(self) -> float: return 0.0 return float(np.dot(target_values, template_weights) / denominator) * self.score_sign + def grad(self): + """Gradient of the matched-mean score w.r.t. translation and rotation. + + Identical in form to :py:meth:`NormalizedCrossCorrelation.grad`, but evaluated on the + *centered* values ``v - mean(v)`` and weights ``w - mean(w)``. The mean-subtraction adds no + extra terms: because ``sum(v - mean(v)) = sum(w - mean(w)) = 0``, the derivatives of the two + means cancel against the centered weights/values (``d/dp = `` and + ``d||v-v̄||/dp = /||v-v̄||``). This makes the inherited gradient consistent with + the (mean-subtracted) ``__call__``; the previous version inherited the non-centered gradient. + """ + 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 - self._target_values.mean() + weights = self.template_weights - self.template_weights.mean() + + 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 -total_grad + class MaskedCrossCorrelation(_MatchCoordinatesToDensity): """ From 90f71f6b2b8f4241092dcb7f631d84294ac69000 Mon Sep 17 00:00:00 2001 From: Jan Kosinski Date: Sun, 28 Jun 2026 14:19:15 +0200 Subject: [PATCH 3/5] NCCMean: score_sign-aware gradient, ASCII docstring, regression tests (review) - grad() returns score_sign * total_grad, so it is correct for both negate_score settings (the parent CrossCorrelation/NormalizedCrossCorrelation grads hardcode -total_grad, valid only for negate_score=True). - Replace non-ASCII symbols (combining diacritics, em dash) in the comment/docstring with ASCII to avoid the bidirectional-Unicode source warning. - Add TestNormalizedCrossCorrelationMean: score equals the matched Pearson correlation; differs from the old global-mean centring under a strong local offset; finite-difference direction of grad(); and grad sign follows negate_score. (cherry picked from commit 21d49db9b2266f5b063a5fd55a89aebdffca244c) --- tests/test_matching_optimization.py | 73 +++++++++++++++++++++++++++-- tme/matching_optimization.py | 20 ++++---- 2 files changed, 80 insertions(+), 13 deletions(-) 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 a382f3f6..a9e27328 100644 --- a/tme/matching_optimization.py +++ b/tme/matching_optimization.py @@ -704,7 +704,7 @@ def __call__(self) -> float: """Returns the score of the current configuration.""" # Centre by the mean of the *matched* values, as the docstring and the score name # ("...Mean") specify. The previous implementation subtracted the global target-array - # mean in __init__ — a constant offset that does not centre the overlapping values and + # mean in __init__ -- a constant offset that does not centre the overlapping values and # so did not compute the documented Pearson-about-the-mean (it ignored the local mean # of the target over the template footprint). target_values = self._target_values - self._target_values.mean() @@ -717,12 +717,16 @@ def __call__(self) -> float: def grad(self): """Gradient of the matched-mean score w.r.t. translation and rotation. - Identical in form to :py:meth:`NormalizedCrossCorrelation.grad`, but evaluated on the - *centered* values ``v - mean(v)`` and weights ``w - mean(w)``. The mean-subtraction adds no - extra terms: because ``sum(v - mean(v)) = sum(w - mean(w)) = 0``, the derivatives of the two - means cancel against the centered weights/values (``d/dp = `` and - ``d||v-v̄||/dp = /||v-v̄||``). This makes the inherited gradient consistent with - the (mean-subtracted) ``__call__``; the previous version inherited the non-centered gradient. + Same form as :py:meth:`NormalizedCrossCorrelation.grad`, but evaluated on the *centred* + values ``vc = v - mean(v)`` and weights ``wc = w - mean(w)``. The mean-subtraction adds no + extra terms: because ``sum(vc) = sum(wc) = 0``, the derivatives of the two means cancel + against the centred weights/values (``d/dp = `` and + ``d||vc||/dp = / ||vc||``), so this is the non-centred formula with v, w replaced + by vc, wc. The previous version inherited the non-centred gradient, inconsistent with the + (mean-subtracted) ``__call__``. + + Returns ``score_sign * d(score)/dp`` so it is correct for both ``negate_score`` settings + (the parent classes hardcode ``-total_grad``, which only holds for ``negate_score=True``). """ grad = self._interpolate_gradient(positions=self.template_rotated) torque = self._torques( @@ -757,7 +761,7 @@ def grad(self): 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 + return self.score_sign * total_grad class MaskedCrossCorrelation(_MatchCoordinatesToDensity): From f47480b7662cf3a935e088d61686d6aa510a9768 Mon Sep 17 00:00:00 2001 From: Jan Kosinski Date: Sun, 28 Jun 2026 14:35:36 +0200 Subject: [PATCH 4/5] NCCMean: tidy docstrings to final form (rationale moved to the PR description) (cherry picked from commit 8f3136070cb6796c78beae09880283725207c7b4) --- tme/matching_optimization.py | 24 +++++++----------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/tme/matching_optimization.py b/tme/matching_optimization.py index a9e27328..6530f380 100644 --- a/tme/matching_optimization.py +++ b/tme/matching_optimization.py @@ -702,11 +702,7 @@ class NormalizedCrossCorrelationMean(NormalizedCrossCorrelation): def __call__(self) -> float: """Returns the score of the current configuration.""" - # Centre by the mean of the *matched* values, as the docstring and the score name - # ("...Mean") specify. The previous implementation subtracted the global target-array - # mean in __init__ -- a constant offset that does not centre the overlapping values and - # so did not compute the documented Pearson-about-the-mean (it ignored the local mean - # of the target over the template footprint). + # Centre by the mean of the matched values (the target sampled over the template footprint). target_values = self._target_values - self._target_values.mean() template_weights = self.template_weights - self.template_weights.mean() denominator = float(np.linalg.norm(target_values) * np.linalg.norm(template_weights)) @@ -715,18 +711,12 @@ def __call__(self) -> float: return float(np.dot(target_values, template_weights) / denominator) * self.score_sign def grad(self): - """Gradient of the matched-mean score w.r.t. translation and rotation. - - Same form as :py:meth:`NormalizedCrossCorrelation.grad`, but evaluated on the *centred* - values ``vc = v - mean(v)`` and weights ``wc = w - mean(w)``. The mean-subtraction adds no - extra terms: because ``sum(vc) = sum(wc) = 0``, the derivatives of the two means cancel - against the centred weights/values (``d/dp = `` and - ``d||vc||/dp = / ||vc||``), so this is the non-centred formula with v, w replaced - by vc, wc. The previous version inherited the non-centred gradient, inconsistent with the - (mean-subtracted) ``__call__``. - - Returns ``score_sign * d(score)/dp`` so it is correct for both ``negate_score`` settings - (the parent classes hardcode ``-total_grad``, which only holds for ``negate_score=True``). + """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( From 9a913d4116ee81a56f87d2b14c8863781509051e Mon Sep 17 00:00:00 2001 From: Jan Kosinski Date: Sun, 28 Jun 2026 15:11:58 +0200 Subject: [PATCH 5/5] NCCMean: use the backend (be) for the score and gradient compute Use be.mean / be.dot / be.sqrt(be.sum(be.square(...))) instead of numpy in NormalizedCrossCorrelationMean.__call__ and grad, so they run on any backend (numpy/cupy/jax), consistent with the rest of matching_optimization. Numerically identical (tests unchanged). (cherry picked from commit 4cf16848b68bfbdcdb705618439b3f4c9ad55118) --- tme/matching_optimization.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tme/matching_optimization.py b/tme/matching_optimization.py index 6530f380..7fae1bf8 100644 --- a/tme/matching_optimization.py +++ b/tme/matching_optimization.py @@ -703,12 +703,14 @@ class NormalizedCrossCorrelationMean(NormalizedCrossCorrelation): 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 - self._target_values.mean() - template_weights = self.template_weights - self.template_weights.mean() - denominator = float(np.linalg.norm(target_values) * np.linalg.norm(template_weights)) + 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)) + ) if denominator < self.eps: return 0.0 - return float(np.dot(target_values, template_weights) / denominator) * self.score_sign + 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). @@ -723,8 +725,8 @@ def grad(self): positions=self.template_rotated, gradients=grad, center=self.template_center ) - values = self._target_values - self._target_values.mean() - weights = self.template_weights - self.template_weights.mean() + 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),