splitting a sparse matrix into two

Question

Question: How can I split 1 sparse matrix into 2, based on the values in a list?

That is, I have a sparse matrix X:

>>print type(X)
<class 'scipy.sparse.csr.csr_matrix'>

that I visualize in my head as a list of lists, to look like this:

>>print X.todense()
[[1,3,4]
 [3,2,2]
 [4,8,1]]

And I have a list y that looks like this:

y = [-1, 
      3, 
     -4]

How can I separate X into two sparse matrices, depending on whether the corresponding value in y is positive or negative? For example, how can I get:

>>print X_pos.todense()
 [[3,2,2]] 
>>print X_neg.todense()
 [[1,3,4]
  [4,8,1]]

The result (X_pos and X_neg) should also be sparse matrices obviously as it's just splitting a sparse matrix to begin with.

Thanks!

Answer 1

Use np.where to generate two arrays of indices for the positive and negative y values, then use those to index into your sparse matrix.

>>> X = csr_matrix([[1,3,4], [3,2,2], [4,8,1]])
>>> y = np.array([-1, 3, -4])
>>> y_pos = np.where(y > 0)[0]
>>> y_neg = np.where(y < 0)[0]
>>> X_pos = X[y_pos]
>>> X_neg = X[y_neg]

You now have to CSR matrices containing the desired elements:

>>> X_pos
<1x3 sparse matrix of type '<type 'numpy.int64'>'
    with 3 stored elements in Compressed Sparse Row format>
>>> X_neg
<2x3 sparse matrix of type '<type 'numpy.int64'>'
    with 6 stored elements in Compressed Sparse Row format>
>>> X_pos.A
array([[3, 2, 2]])
>>> X_neg.A
array([[1, 3, 4],
       [4, 8, 1]])