From 773cb9957445c15a69c7f6b73ffe052ed7a18637 Mon Sep 17 00:00:00 2001 From: Kenya Otsuka Date: Thu, 17 Sep 2026 14:24:42 +0000 Subject: [PATCH 1/2] test(evals): cover NaN and degenerate inputs in metrics Add tests for behavior that is currently only exercised indirectly through the pickled 2-d fixtures: undefined correlations on zero-variance columns and rows, NaN column removal (and its opt-out), the mean/std arguments of pattern_correlation, and pairwise_identification against a definition-based oracle built from cdist, including a non-correlation metric. --- tests/evals/test_metrics.py | 173 ++++++++++++++++++++++++++++++++++++ 1 file changed, 173 insertions(+) diff --git a/tests/evals/test_metrics.py b/tests/evals/test_metrics.py index 1d3db0cf..cdea6c9b 100644 --- a/tests/evals/test_metrics.py +++ b/tests/evals/test_metrics.py @@ -1,12 +1,33 @@ import unittest import pickle +import warnings import numpy as np +from scipy.spatial.distance import cdist from bdpy.evals.metrics import profile_correlation, pattern_correlation, pairwise_identification +def identification_accuracy(pred, true, metric='correlation'): + '''Pair-wise identification accuracy, computed straight from its definition. + + For each sample, count how many of the other candidates are closer to it + than its own counterpart, using the same similarity `1 - distance` as + `pairwise_identification`. + ''' + d = 1 - cdist( + pred.reshape(pred.shape[0], -1), + true.reshape(true.shape[0], -1), + metric=metric + ) + n_sample = d.shape[0] + return np.array([ + np.sum(np.delete(d[i], i) < d[i, i]) / (n_sample - 1) + for i in range(n_sample) + ]) + + class TestMetrics(unittest.TestCase): def test_profile_correlation(self): # 2-d array @@ -97,5 +118,157 @@ def test_2d_nan(self): d['ident_acc'], rtol=1e-12, atol=1e-12 ) + +class TestMetricsNanAndDegenerateInputs(unittest.TestCase): + """Behavior around NaN and zero-variance inputs.""" + + def test_profile_correlation_constant_column(self): + # A column with zero variance has an undefined correlation. + n_sample, n_feat = 10, 8 + rand = np.random.RandomState(0) + x = rand.rand(n_sample, n_feat) + y = rand.rand(n_sample, n_feat) + x[:, 3] = 7.0 + + r = profile_correlation(x, y).ravel() + + self.assertTrue(np.isnan(r[3])) + np.testing.assert_allclose( + np.delete(r, 3), + [ + np.corrcoef(x[:, j], y[:, j])[0, 1] + for j in range(n_feat) if j != 3 + ], + rtol=1e-12, atol=1e-12 + ) + + def test_pattern_correlation_constant_row(self): + # A sample with zero variance has an undefined correlation. + n_sample, n_feat = 6, 20 + rand = np.random.RandomState(1) + x = rand.rand(n_sample, n_feat) + y = rand.rand(n_sample, n_feat) + x[2, :] = 5.0 + + r = pattern_correlation(x, y) + + self.assertTrue(np.isnan(r[2])) + np.testing.assert_allclose( + np.delete(r, 2), + [ + np.corrcoef(x[i, :], y[i, :])[0, 1] + for i in range(n_sample) if i != 2 + ], + rtol=1e-12, atol=1e-12 + ) + + def test_pattern_correlation_removes_nan_columns(self): + # A single NaN drops the whole column for every sample. + n_sample, n_feat = 6, 20 + rand = np.random.RandomState(2) + x = rand.rand(n_sample, n_feat) + y = rand.rand(n_sample, n_feat) + x[3, 7] = np.nan + + with self.assertWarns(UserWarning): + r = pattern_correlation(x, y) + + xd = np.delete(x, 7, axis=1) + yd = np.delete(y, 7, axis=1) + np.testing.assert_allclose( + r, + [np.corrcoef(xd[i, :], yd[i, :])[0, 1] for i in range(n_sample)], + rtol=1e-12, atol=1e-12 + ) + + def test_pattern_correlation_keeps_nan_columns_when_disabled(self): + n_sample, n_feat = 6, 20 + rand = np.random.RandomState(3) + x = rand.rand(n_sample, n_feat) + y = rand.rand(n_sample, n_feat) + x[3, 7] = np.nan + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter('always') + r = pattern_correlation(x, y, remove_nan=False) + + self.assertFalse(any( + 'NaN column removed' in str(w.message) for w in caught + )) + self.assertTrue(np.isnan(r[3])) + self.assertFalse(np.isnan(np.delete(r, 3)).any()) + + def test_pattern_correlation_with_mean_and_std(self): + n_sample, n_feat = 6, 20 + rand = np.random.RandomState(4) + x = rand.rand(n_sample, n_feat) + y = rand.rand(n_sample, n_feat) + mean = rand.rand(n_feat) + std = rand.rand(n_feat) + 0.5 + + r = pattern_correlation(x, y, mean=mean, std=std) + + xs = (x - mean) / std + ys = (y - mean) / std + np.testing.assert_allclose( + r, + [np.corrcoef(xs[i, :], ys[i, :])[0, 1] for i in range(n_sample)], + rtol=1e-12, atol=1e-12 + ) + + def test_pairwise_identification_matches_definition(self): + # Mix well-predicted and poorly-predicted samples so that the accuracy + # is not degenerate, and include an exactly correlated and an exactly + # anti-correlated sample. + n_sample, n_feat = 12, 40 + rand = np.random.RandomState(5) + true = rand.rand(n_sample, n_feat) + pred = true + 5.0 * rand.rand(n_sample, n_feat) + pred[0] = true[0] + pred[1] = -true[1] + + cr = pairwise_identification(pred, true) + + self.assertEqual(cr.shape, (n_sample,)) + # Guard against a degenerate all-correct case. + self.assertTrue(np.any((cr[2:] > 0) & (cr[2:] < 1))) + np.testing.assert_allclose( + cr, identification_accuracy(pred, true), rtol=1e-12, atol=1e-12 + ) + + def test_pairwise_identification_constant_row(self): + # A sample with zero variance gives an undefined correlation distance; + # those comparisons are dropped instead of failing. + n_sample, n_feat = 10, 30 + rand = np.random.RandomState(6) + true = rand.rand(n_sample, n_feat) + pred = true + 5.0 * rand.rand(n_sample, n_feat) + pred[4, :] = 3.0 + + with self.assertWarns(UserWarning): + cr = pairwise_identification(pred, true) + + self.assertEqual(cr.shape, (n_sample,)) + np.testing.assert_allclose( + np.delete(cr, 4), + np.delete(identification_accuracy(pred, true), 4), + rtol=1e-12, atol=1e-12 + ) + + def test_pairwise_identification_non_correlation_metric(self): + n_sample, n_feat = 10, 30 + rand = np.random.RandomState(7) + true = rand.rand(n_sample, n_feat) + pred = true + 5.0 * rand.rand(n_sample, n_feat) + + cr = pairwise_identification(pred, true, metric='euclidean') + + self.assertTrue(np.any((cr > 0) & (cr < 1))) + np.testing.assert_allclose( + cr, identification_accuracy(pred, true, metric='euclidean'), + rtol=1e-12, atol=1e-12 + ) + + if __name__ == '__main__': unittest.main() \ No newline at end of file From 037d2e980db4be52696ce8583270ca85415e9542 Mon Sep 17 00:00:00 2001 From: Kenya Otsuka Date: Thu, 17 Sep 2026 14:40:21 +0000 Subject: [PATCH 2/2] refactor(evals): vectorize correlation metrics profile_correlation and pattern_correlation called np.corrcoef once per unit and per sample, and pairwise_identification went through cdist, whose correlation metric walks every pair in a Python-level C loop. On the feature sizes produced by recent vision encoders (hundreds of thousands of units) this dominates the evaluation runtime. Compute the same quantities with array operations instead: - profile_correlation and pattern_correlation accumulate centered dot products over blocks of the axis they do not reduce over, which keeps the two-pass formula of np.corrcoef, including its clipping to [-1, 1] and the NaN it yields for a zero-variance column or row. - pairwise_identification derives 1 - cdist(p, t, 'correlation') from a single matrix product when the metric is 'correlation'. Other metrics and the single-trial branch still go through cdist, and the NaN handling that follows is untouched, so the emitted warnings are unchanged. scipy clips the cosine it builds the correlation distance from, so the matrix product is clipped the same way. Blocking the work also bounds peak memory, which the previous standardization and NaN scan did not: they materialized full-size copies. Measured with random arrays of the given (n_samples, n_units), in seconds, best of three runs: 500x5000 50x1000000 1500x40000 before after before after before after profile_correlation 0.17 0.01 28.12 0.30 2.75 0.34 pattern_correlation 0.03 0.01 0.39 0.35 0.45 0.38 pairwise_identification 0.45 0.03 3.26 0.64 89.26 1.43 profile_correlation gains as the unit count grows, since it used to make one np.corrcoef call per unit. pairwise_identification gains as the sample count grows, since its cost is pairs times units. pattern_correlation gains least, and at some shapes is only break-even: it already looped over the sample axis, the small one, so the gain there is in memory. Peak RSS of the 1500x40000 pairwise_identification run drops from 931 MiB to 149 MiB. --- bdpy/evals/metrics.py | 236 +++++++++++++++++++++++++++++++----- tests/evals/test_metrics.py | 40 ++++++ 2 files changed, 249 insertions(+), 27 deletions(-) diff --git a/bdpy/evals/metrics.py b/bdpy/evals/metrics.py index cc436cbc..07f8d8d7 100644 --- a/bdpy/evals/metrics.py +++ b/bdpy/evals/metrics.py @@ -1,10 +1,208 @@ """bdpy.evals.metrics""" import warnings +from typing import Iterator, Optional import numpy as np from scipy.spatial.distance import cdist +# Upper bound on the number of array elements held by a single working block +# (about 32 MiB in float64). The correlation metrics below work on blocks of +# whichever axis they do not reduce over, so that peak memory stays bounded +# even when the feature dimension reaches several hundred thousand. +_BLOCK_ELEMENTS = 2 ** 22 + + +def _index_blocks(length: int, other_length: int) -> Iterator[slice]: + """Split an axis into slices small enough to be processed at once. + + Parameters + ---------- + length : int + Number of indices to split. + other_length : int + Length of the other axis of the arrays to be processed. + + Yields + ------ + slice + Index range holding at most ``_BLOCK_ELEMENTS`` elements. + """ + step = max(1, _BLOCK_ELEMENTS // max(other_length, 1)) + for start in range(0, length, step): + yield slice(start, min(start + step, length)) + + +def _nan_column_mask( + x: np.ndarray, + y: np.ndarray, + mean: Optional[np.ndarray], + std: Optional[np.ndarray], +) -> np.ndarray: + """Flag the columns holding a NaN in either array. + + Parameters + ---------- + x, y : numpy.ndarray + Arrays of shape (n_samples, n_units). + mean, std : numpy.ndarray or None + Flattened standardization parameters, applied before the check when + both are given. + + Returns + ------- + numpy.ndarray + Boolean mask of shape (n_units,). + """ + n_sample, n_feat = x.shape + nan_cols = np.empty(n_feat, dtype=bool) + + for cols in _index_blocks(n_feat, n_sample): + xb = x[:, cols] + yb = y[:, cols] + if mean is not None and std is not None: + xb = (xb - mean[cols]) / std[cols] + yb = (yb - mean[cols]) / std[cols] + nan_cols[cols] = np.isnan(xb).any(axis=0) | np.isnan(yb).any(axis=0) + + return nan_cols + + +def _column_correlation(x: np.ndarray, y: np.ndarray) -> np.ndarray: + """Correlate `x` and `y` column by column, over whole columns at once. + + Parameters + ---------- + x, y : numpy.ndarray + Arrays of shape (n_samples, n_units). + + Returns + ------- + numpy.ndarray + Correlation of each column, of shape (n_units,). A column with zero + variance gives NaN, as ``numpy.corrcoef`` does. + """ + n_sample, n_feat = x.shape + r = np.empty(n_feat, dtype=np.float64) + + for cols in _index_blocks(n_feat, n_sample): + xb = np.array(x[:, cols], dtype=np.float64) + yb = np.array(y[:, cols], dtype=np.float64) + xb -= xb.mean(axis=0) + yb -= yb.mean(axis=0) + + with np.errstate(invalid='ignore', divide='ignore'): + rb = np.einsum('ij,ij->j', xb, yb) / np.sqrt(np.einsum('ij,ij->j', xb, xb)) + rb /= np.sqrt(np.einsum('ij,ij->j', yb, yb)) + + r[cols] = rb + + # numpy.corrcoef clips its result, so that a rounding error cannot push a + # perfect correlation past 1. + return np.clip(r, -1.0, 1.0) + + +def _row_correlation( + x: np.ndarray, + y: np.ndarray, + mean: Optional[np.ndarray] = None, + std: Optional[np.ndarray] = None, + keep_cols: Optional[np.ndarray] = None, +) -> np.ndarray: + """Correlate `x` and `y` row by row, over whole rows at once. + + Parameters + ---------- + x, y : numpy.ndarray + Arrays of shape (n_samples, n_units). + mean, std : numpy.ndarray or None + Flattened standardization parameters, applied when both are given. + keep_cols : numpy.ndarray or None + Boolean mask of the columns to keep, applied after standardization. + + Returns + ------- + numpy.ndarray + Correlation of each row, of shape (n_samples,). A row with zero + variance gives NaN, as ``numpy.corrcoef`` does. + """ + n_sample, n_feat = x.shape + r = np.empty(n_sample, dtype=np.float64) + + for rows in _index_blocks(n_sample, n_feat): + xb = np.array(x[rows], dtype=np.float64) + yb = np.array(y[rows], dtype=np.float64) + + if mean is not None and std is not None: + xb -= mean + xb /= std + yb -= mean + yb /= std + + if keep_cols is not None: + xb = xb[:, keep_cols] + yb = yb[:, keep_cols] + + xb -= xb.mean(axis=1, keepdims=True) + yb -= yb.mean(axis=1, keepdims=True) + + with np.errstate(invalid='ignore', divide='ignore'): + rb = np.einsum('ij,ij->i', xb, yb) / np.sqrt(np.einsum('ij,ij->i', xb, xb)) + rb /= np.sqrt(np.einsum('ij,ij->i', yb, yb)) + + r[rows] = rb + + # numpy.corrcoef clips its result, so that a rounding error cannot push a + # perfect correlation past 1. + return np.clip(r, -1.0, 1.0) + + +def _correlation_similarity(p: np.ndarray, t: np.ndarray) -> np.ndarray: + """Correlate every row of `p` with every row of `t`. + + Equivalent to ``1 - cdist(p, t, metric='correlation')``, but computed as + matrix products instead of one pass per pair. + + Parameters + ---------- + p, t : numpy.ndarray + Arrays of shape (n_samples, n_units). + + Returns + ------- + numpy.ndarray + Correlation matrix of shape (n_samples of `p`, n_samples of `t`). A row + with zero variance gives NaN, as the correlation distance does. + """ + n_p, n_feat = p.shape + n_t = t.shape[0] + n_row = max(n_p, n_t) + + sum_p = np.zeros(n_p) + sum_t = np.zeros(n_t) + for cols in _index_blocks(n_feat, n_row): + sum_p += p[:, cols].sum(axis=1, dtype=np.float64) + sum_t += t[:, cols].sum(axis=1, dtype=np.float64) + mean_p = sum_p / n_feat + mean_t = sum_t / n_feat + + d = np.zeros((n_p, n_t)) + sq_p = np.zeros(n_p) + sq_t = np.zeros(n_t) + for cols in _index_blocks(n_feat, n_row): + pb = np.array(p[:, cols], dtype=np.float64) - mean_p[:, np.newaxis] + tb = np.array(t[:, cols], dtype=np.float64) - mean_t[:, np.newaxis] + d += pb @ tb.T + sq_p += np.einsum('ij,ij->i', pb, pb) + sq_t += np.einsum('ij,ij->i', tb, tb) + + with np.errstate(invalid='ignore', divide='ignore'): + d /= np.sqrt(sq_p)[:, np.newaxis] * np.sqrt(sq_t)[np.newaxis, :] + + # scipy clips the cosine it derives the correlation distance from, so that + # a rounding error cannot push a perfect correlation past 1. + return np.clip(d, -1.0, 1.0) + def profile_correlation(x, y): """Profile correlation.""" @@ -16,17 +214,7 @@ def profile_correlation(x, y): _x = x.reshape(n_sample, -1) _y = y.reshape(n_sample, -1) - n_feat = _y.shape[1] - - r = np.array( - [ - np.corrcoef( - _x[:, j].ravel(), - _y[:, j].ravel() - )[0, 1] - for j in range(n_feat) - ] - ) + r = _column_correlation(_x, _y) r = r.reshape((1,) + orig_shape[1:]) @@ -43,30 +231,21 @@ def pattern_correlation(x, y, mean=None, std=None, remove_nan=True): _x = x.reshape(n_sample, -1) _y = y.reshape(n_sample, -1) + m = None + s = None if mean is not None and std is not None: m = mean.reshape(-1) s = std.reshape(-1) - _x = (_x - m) / s - _y = (_y - m) / s - + keep_cols = None if remove_nan: # Remove nan columns based on the decoded features - nan_cols = np.isnan(_x).any(axis=0) | np.isnan(_y).any(axis=0) + nan_cols = _nan_column_mask(_x, _y, m, s) if nan_cols.any(): warnings.warn('NaN column removed ({})'.format(np.sum(nan_cols))) - _x = _x[:, ~nan_cols] - _y = _y[:, ~nan_cols] + keep_cols = ~nan_cols - r = np.array( - [ - np.corrcoef( - _x[i, :].ravel(), - _y[i, :].ravel() - )[0, 1] - for i in range(n_sample) - ] - ) + r = _row_correlation(_x, _y, m, s, keep_cols) return r @@ -137,7 +316,10 @@ def pairwise_identification(pred, true, metric='correlation', remove_nan=True, r cr.append(ac) cr = np.asarray(cr) else: - d = 1 - cdist(p, t, metric=metric) + if metric == 'correlation': + d = _correlation_similarity(p, t) + else: + d = 1 - cdist(p, t, metric=metric) if remove_nan_dist: cr = [] diff --git a/tests/evals/test_metrics.py b/tests/evals/test_metrics.py index cdea6c9b..701041bd 100644 --- a/tests/evals/test_metrics.py +++ b/tests/evals/test_metrics.py @@ -2,10 +2,12 @@ import pickle import warnings +from unittest.mock import patch import numpy as np from scipy.spatial.distance import cdist +from bdpy.evals import metrics from bdpy.evals.metrics import profile_correlation, pattern_correlation, pairwise_identification @@ -269,6 +271,44 @@ def test_pairwise_identification_non_correlation_metric(self): rtol=1e-12, atol=1e-12 ) + def test_results_do_not_depend_on_the_block_size(self): + # The metrics process the units in blocks to bound memory usage. Shrink + # a block so that several of them are needed, with the NaN column + # falling inside a block rather than on its boundary. + n_sample, n_feat, nan_col = 8, 37, 17 + rand = np.random.RandomState(8) + x = rand.rand(n_sample, n_feat) + y = rand.rand(n_sample, n_feat) + x[5, nan_col] = np.nan + + with patch.object(metrics, '_BLOCK_ELEMENTS', 3 * n_sample): + r_prof = profile_correlation(x, y).ravel() + with self.assertWarns(UserWarning): + r_patt = pattern_correlation(x, y) + with self.assertWarns(UserWarning): + cr = pairwise_identification(x, y) + + self.assertTrue(np.isnan(r_prof[nan_col])) + np.testing.assert_allclose( + np.delete(r_prof, nan_col), + [ + np.corrcoef(x[:, j], y[:, j])[0, 1] + for j in range(n_feat) if j != nan_col + ], + rtol=1e-12, atol=1e-12 + ) + + xd = np.delete(x, nan_col, axis=1) + yd = np.delete(y, nan_col, axis=1) + np.testing.assert_allclose( + r_patt, + [np.corrcoef(xd[i, :], yd[i, :])[0, 1] for i in range(n_sample)], + rtol=1e-12, atol=1e-12 + ) + np.testing.assert_allclose( + cr, identification_accuracy(xd, yd), rtol=1e-12, atol=1e-12 + ) + if __name__ == '__main__': unittest.main() \ No newline at end of file