Skip to content

Instantly share code, notes, and snippets.

@dwf
Created February 15, 2011 19:50
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 2 You must be signed in to fork a gist
  • Save dwf/828099 to your computer and use it in GitHub Desktop.
Save dwf/828099 to your computer and use it in GitHub Desktop.
Pull out a submatrix from a COO SciPy sparse matrix.
import scipy as S
def coo_submatrix_pull(matr, rows, cols):
"""
Pulls out an arbitrary i.e. non-contiguous submatrix out of
a sparse.coo_matrix.
"""
if type(matr) != S.sparse.coo_matrix:
raise TypeError('Matrix must be sparse COOrdinate format')
gr = -1 * N.ones(matr.shape[0])
gc = -1 * N.ones(matr.shape[1])
lr = len(rows)
lc = len(cols)
ar = N.arange(0, lr)
ac = N.arange(0, lc)
gr[rows[ar]] = ar
gc[cols[ac]] = ac
mrow = matr.row
mcol = matr.col
newelem = (gr[mrow] > -1) & (gc[mcol] > -1)
newrows = mrow[newelem]
newcols = mcol[newelem]
return S.sparse.coo_matrix((matr.data[newelem], N.array([gr[newrows],
gc[newcols]])),(lr, lc))
@amueller
Copy link

Very helpful :)
you forgot to import numpy, though ;)

@dwf
Copy link
Author

dwf commented Mar 30, 2012

That I most certainly did!

@bhttchr6
Copy link

bhttchr6 commented May 7, 2020

Thanks!, super useful

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment