refactor(evals): vectorize correlation metrics

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 = []