Numpy: given the nonzero indices of a matrix how to extract the elements into a submatrix

Question

I have very sparse matrices, so I want to extract the smallest rectangular region of a matrix that has non-zero values. I know that numpy.nonzero(a) gives you the indices of the elements that are non-zero, but how can I use this to extract a submatrix that contains the elements of the matrix at those indices.

To give an example, this is what I am aiming for:

>>> test
array([[0, 0, 0, 0, 0, 0],
       [0, 1, 1, 1, 1, 0],
       [0, 0, 1, 1, 0, 0]])
>>> np.nonzero(test)
(array([1, 1, 1, 1, 2, 2]), array([1, 2, 3, 4, 2, 3]))
>>> submatrix(test)
array([[1, 1, 1, 1],
       [0, 1, 1, 0]])

Does anyone know a simple way to do this in numpy? Thanks.

Answer 1

It seems like you're looking to find the smallest region of your matrix that contains all the nonzero elements. If that's true, here's a method:

import numpy as np

def submatrix(arr):
    x, y = np.nonzero(arr)
    # Using the smallest and largest x and y indices of nonzero elements, 
    # we can find the desired rectangular bounds.  
    # And don't forget to add 1 to the top bound to avoid the fencepost problem.
    return arr[x.min():x.max()+1, y.min():y.max()+1]

test = np.array([[0, 0, 0, 0, 0, 0],
                 [0, 1, 1, 1, 1, 0],
                 [0, 0, 1, 1, 0, 0]])

print submatrix(test)

# Result:  
# [[1 1 1 1]
#  [0 1 1 0]]