Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Compute large, sparse correlation matrices in parallel using dask.
import dask
import dask.array as da
import dask.dataframe as dd
import sparse
@dask.delayed(pure=True)
def corr_on_chunked(chunk1, chunk2, corr_thresh=0.9):
return sparse.COO.from_numpy((np.dot(chunk1, chunk2.T) > corr_thresh))
def chunked_corr_sparse_dask(data, chunksize=5000, corr_thresh=0.9):
# Gets the correlation of a large DataFrame, chunking the computation
# Returns a sparse directed adjancy matrix (old->young)
# Adapted from https://stackoverflow.com/questions/24717513/python-numpy-corrcoef-memory-error
numrows = data.shape[0]
data -= np.mean(data, axis=1)[:,None] # subtract means form the input data
data /= np.sqrt(np.sum(data**2, axis=1))[:,None] # normalize the data
rows = []
for r in range(0, numrows, chunksize):
cols = []
for c in range(0, numrows, chunksize):
r1 = r + chunksize
c1 = c + chunksize
chunk1 = data[r:r1]
chunk2 = data[c:c1]
delayed_array = corr_on_chunked(chunk1, chunk2, corr_thresh=corr_thresh)
cols.append(da.from_delayed(
delayed_array,
dtype='bool',
shape=(chunksize, chunksize),
))
rows.append(da.hstack(cols))
res = da.vstack(rows).compute()
res = sparse.triu(res, k=1)
return res.tocsr()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment