Selecting rows in scipy sparse matrix

Is there a way to select the rows corresponding to some given indices in a scipy sparse matrix? The dummy approach does not work:

sparse.eye(3)[:2, :]

returns

TypeError: 'dia_matrix' object is not subscriptable

Solution

When asking a question like this, you should say more than 'returns error'. What error? That matters.

But I'll do that work for you:

In [143]: m =sparse.eye(3)
In [144]: m
Out[144]: 
<3x3 sparse matrix of type '<class 'numpy.float64'>'
    with 3 stored elements (1 diagonals) in DIAgonal format>
In [145]: m[:2,:]
...
TypeError: 'dia_matrix' object is not subscriptable

The error is significant. It tells us that this particular sparse format has not implemented indexing. We'd get the same error with the common coo format. But with csr (or lil) format, indexing works:

In [146]: m =sparse.eye(3, format='csr')
In [147]: m
Out[147]: 
<3x3 sparse matrix of type '<class 'numpy.float64'>'
    with 3 stored elements in Compressed Sparse Row format>
In [148]: m[:2,:]
Out[148]: 
<2x3 sparse matrix of type '<class 'numpy.float64'>'
    with 2 stored elements in Compressed Sparse Row format>
In [149]: _.A
Out[149]: 
array([[1., 0., 0.],
       [0., 1., 0.]])

When generating a sparse matrix I like to display it's repr, which tells me the format and the size. print(m) (str) shows the values in coo style.

sparse.eye produces dia format as the default, since the nonzero values are all on one diagonal. Other sparse functions produce different default formats.

dia page shows a getrow method:

In [153]: sparse.eye(3).getrow(1)
Out[153]: 
<1x3 sparse matrix of type '<class 'numpy.float64'>'
    with 1 stored elements in Compressed Sparse Row format>

But notice the format of the returned matrix - csr, not dia. The formats have methods that readily convert among each other. Some operations do the necessary conversion for us.