Removing square submatrices that are 0 in 2d numpy array

Question

Let's say I have a square 2D array which is a block diagonal matrix size (N,N). Each of the blocks is (say) (N_b,N_b). Some of these blocks are just 0s (eg a (N_b,N_b) matrix of zeros). I would like to squeeze these blocks to be left with a matrix that doesn't have the (N_b,N_b) 0 arrays.

Additionally, I would additionally like to take one of these "squeezed" matrices and expand it to be the original size (with zero padding), with the requirement that all elements are in the same place

Everything I can think of requires loops and a lot of accounting. I am happy to use sparse matrices, but that seems to complicate things even further.

The first part of the problem ("squeezing") could be done in this way:

import numpy as np
import scipy.sparse as sp


A = np.random.randn(3,3)*5
B = A*0.

M = np.array(sp.block_diag ( (A, B, A, A, A, B, A) ).todense())


n_b = A.shape[0]
n_blocks = M.shape[0]/n_b

block_locs = []
nonzero_blocks = []
for this_block in xrange ( n_blocks ):
    if np.all(M[n_b*this_block:(this_block+1)*n_b, n_b*this_block:(this_block+1)*n_b] != 0):
        nonzero_blocks.append (M[n_b*this_block:(this_block+1)*n_b, n_b*this_block:(this_block+1)*n_b] )
        block_locs.append ( this_block )
squeezed_M = sp.block_diag ( nonzero_blocks ).tolil()


blocky = []
this_squeeze_block = 0
for this_block in xrange(n_blocks):
    if this_block in block_locs:
        blocky.append ( squeezed_M[this_squeeze_block*n_b:(this_squeeze_block+1)*n_b,
            this_squeeze_block * n_b:(this_squeeze_block + 1) * n_b])
        this_squeeze_block += 1
    else:
        blocky.append ( np.zeros((n_b, n_b)))
A_big = sp.block_diag(blocky).todense()

print np.allclose(M, A_big)

Here's a plot of the original matrix (with zero blocks, top) and the squeezed version (bottom):

EDIT A more concise way that exploits the symmetry of the problem could be this:

import numpy as np
import scipy.sparse as sp
import matplotlib.pyplot as plt

# First create a big matrix with gaps
n_b = 3
A = np.random.randn(n_b,n_b)*5
B = A*0.
M = sp.block_diag ( (A, B, A, A, A, B, A) )

# Select the rows which are not all zeros
no_zeros = np.any(np.array((M != 0).todense()), axis=1)
n_blocks = no_zeros.sum() 
# where's the meat?
meat = np.outer ( no_zeros, no_zeros )
# We need to use a lil matrix to address by a boolean array
# And then we just reshape
A_squeeze = M.tolil()[meat].reshape((n_blocks, n_blocks))

reco = sp.lil_matrix ( M.shape)
reco[np.outer(no_zeros, no_zeros)]=A_squeeze.todense().ravel()

np.allclose(M.todense(), reco.todense())
plt.subplot(1,2,1)
plt.spy(M)
plt.subplot(1,2,2)
plt.spy(A_squeeze)

Answer 1

You can use a find the diagonal elements that are equal to zero, and then create a 2D boolean mask to select your array.

import numpy as np

a = [[1,1,0,0,0,0,0],
     [1,1,0,0,0,0,0],
     [0,0,0,0,0,0,0],
     [0,0,0,0,0,0,0],
     [0,0,0,0,2,2,2],
     [0,0,0,0,2,2,2],
     [0,0,0,0,2,2,2]]

# find the diagonal elements not equal to zero
b = np.diag(a) != 0

# create a boolean 2d mask
mask = b[np.newaxis] * np.transpose(b[np.newaxis])

> mask  # this is the mask output
# array([[ True,  True, False, False,  True,  True,  True],
#        [ True,  True, False, False,  True,  True,  True],
#        [False, False, False, False, False, False, False],
#        [False, False, False, False, False, False, False],
#        [ True,  True, False, False,  True,  True,  True],
#        [ True,  True, False, False,  True,  True,  True],
#        [ True,  True, False, False,  True,  True,  True]], dtype=bool)

See how the mask will eliminate the unwanted rows and columns? However, because numpy does not allow for jagged arrays, so we will have to reshape a and mask to a 1D arrays, use boolean indexing, then reshape back to square.

a_squished = np.ravel(a)[np.ravel(mask)]

# reshape the squished array to a 2D square
out = a_squished.reshape(                         # reshape the array
          np.sqrt(a_squished.size).astype(int),   # pass to int to stop warning
          -1                                      # fill next axis with leftover
      )

> out
# array([[1, 1, 0, 0, 0],
#        [1, 1, 0, 0, 0],
#        [0, 0, 2, 2, 2],
#        [0, 0, 2, 2, 2],
#        [0, 0, 2, 2, 2]])