<jupyter_start><jupyter_text><jupyter_code># Fast spearman correlations based on https://saket-choudhary.me/blog/2022/03/09/sparsespearman/
import scipy.sparse as sparse
import scipy.stats as stats
from scipy.sparse import csc_matrix
import numpy as np

def cor_sparse_pearson(X, Y=None, cov=False):
    if not sparse.issparse(X):
      # should be CSC
      X = csc_matrix(X)
    nrows = X.shape[0]
    meanX = X.mean(axis=0)
    if Y is None:
        covmat = X.T * X - nrows * meanX.T * meanX
        covmat = covmat / (nrows - 1)
        var = np.outer(np.diag(covmat), np.diag(covmat))
        cormat = np.divide(covmat, np.sqrt(var))
    else:
        if X.shape[0] != Y.shape[0]:
            raise RuntimeError("X and Y have different # of columns")
        if not sparse.issparse(Y):
          Y = csc_matrix(Y)
        meanY = Y.mean(axis=0)
        covmat = X.T * Y - nrows * meanX.T * meanY
        covmat = covmat / (nrows - 1)        
        meanX2 = np.power(meanX, 2)
        meanY2 = np.power(meanY, 2)
        sdX = np.sqrt((X.power(2).sum(axis=0) - nrows * meanX2) / (nrows - 1))
        sdY = np.sqrt((Y.power(2).sum(axis=0) - nrows * meanY2) / (nrows - 1))
        cormat = covmat / (sdX.T * sdY)
    if cov:
        return covmat
    return cormat


def sparsify_ranks(X):
    if not sparse.issparse(X):
        # should be CSC
        X = csc_matrix(X)

    non_zeros_per_col = X.getnnz(axis=0)
    n_zeros_per_col = X.shape[0] - non_zeros_per_col
    offsets = (n_zeros_per_col - 1) / 2

    x = X.data
    nzdata = np.split(x, np.cumsum(non_zeros_per_col))[:-1]
    for index, data in enumerate(nzdata):
        nzdata[index] = offsets[index] + stats.rankdata(data)
    X.data = np.concatenate(nzdata).ravel()
    return X


def cor_sparse_spearman(X, Y=None, cov=None):
    rankX = sparsify_ranks(X)
    if Y is None:
        # Calculate pearson correlation on rank matrices
        return cor_sparse_pearson(X=rankX, cov=cov)
    rankY = sparsify_ranks(Y)
    return cor_sparse_pearson(X=rankX, Y=rankY, cov=cov)


def cor_dense_pearson(X, Y):
  X = (X - X.mean(axis=0)) / X.std(axis=0)
  Y = (Y - Y.mean(axis=0)) / Y.std(axis=0)
  pearson_r = np.dot(X.T, Y) / X.shape[0]
  return pearson_r
X = sparse.random(1000, 10000, density=0.1, format='csc')
Y = sparse.random(1000, 5000, density=0.125, format='csc')
%%timeit
corXY_dense_pearson = cor_dense_pearson(X.todense(), Y.todense())
%%timeit
corXY_sparse_pearson = cor_sparse_pearson(X, Y)
# check
corXY_dense_pearson = cor_dense_pearson(X.todense(), Y.todense())
corXY_sparse_pearson = cor_sparse_pearson(X, Y)

np.allclose(corXY_sparse_pearson, corXY_dense_pearson)
%%timeit
X_ranks = np.apply_along_axis(stats.rankdata, axis=0, arr=X.todense())
Y_ranks = np.apply_along_axis(stats.rankdata, axis=0, arr=Y.todense())
corXY_dense_spearman =  cor_dense_pearson(X_ranks, Y_ranks)
%%timeit
corXY_sparse_spearman = cor_sparse_spearman(X, Y)
corXY_sparse_spearman.shape
# check 
X_ranks = np.apply_along_axis(stats.rankdata, axis=0, arr=X.todense())
Y_ranks = np.apply_along_axis(stats.rankdata, axis=0, arr=Y.todense())

corXY_sparse_spearman = cor_sparse_spearman(X, Y)
corXY_dense_spearman =  cor_dense_pearson(X_ranks, Y_ranks)

np.allclose(corXY_sparse_spearman, corXY_dense_spearman)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1330498177399757, 'fraction_numerical': 0.019441069258809233, 'mean_word_length': 3.443994601889339, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '264942', 'license': 'BSD-2-Clause', 'n_tokens_mistral': 1449, 'n_tokens_neox': 1351, 'n_words': 255, 'path': '2022-03-09/SparseSpearman_Python.ipynb', 'repo_name': 'saketkc/blog', 'size': '3292'}