Spearman's Correlation

5

Hi

Here is simplified version for computation of Spearman’s rank correlation coefficient. It’s non differentiable, but much faster than scipy.stats.spearmanr version.

def _get_ranks(x: torch.Tensor) -> torch.Tensor:
    tmp = x.argsort()
    ranks = torch.zeros_like(tmp)
    ranks[tmp] = torch.arange(len(x))
    return ranks

def spearman_correlation(x: torch.Tensor, y: torch.Tensor):
    """Compute correlation between 2 1-D vectors
    Args:
        x: Shape (N, )
        y: Shape (N, )
    """
    x_rank = _get_ranks(x)
    y_rank = _get_ranks(y)
    
    n = x.size(0)
    upper = 6 * torch.sum((x_rank - y_rank).pow(2))
    down = n * (n ** 2 - 1.0)
    return 1.0 - (upper / down)

Time comparison

x = torch.rand(1000)
y = torch.rand(1000)
%timeit spearman_correlation(x, y)
206 µs ± 15.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit spearmanr(x_n, y_n)[0]
592 µs ± 2.22 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
2 Likes