How to crop zero edges of a numpy array?

Question

I have this ugly, un-pythonic beast:

def crop(dat, clp=True):
    '''Crops zero-edges of an array and (optionally) clips it to [0,1].

    Example:
    >>> crop( np.array(
    ...       [[0,0,0,0,0,0],
    ...        [0,0,0,0,0,0],
    ...        [0,1,0,2,9,0],
    ...        [0,0,0,0,0,0],
    ...        [0,7,4,1,0,0],
    ...        [0,0,0,0,0,0]]
    ...     ))
    array([[1, 0, 1, 1],
           [0, 0, 0, 0],
           [1, 1, 1, 0]])
    '''
    if clp: np.clip( dat, 0, 1, out=dat )
    while np.all( dat[0,:]==0 ):
        dat = dat[1:,:]
    while np.all( dat[:,0]==0 ):
        dat = dat[:,1:]
    while np.all( dat[-1,:]==0 ):
        dat = dat[:-1,:]
    while np.all( dat[:,-1]==0 ):
        dat = dat[:,:-1]
    return dat
    # Below gets rid of zero-lines/columns in the middle
    #+so not usable.
    #dat = dat[~np.all(dat==0, axis=1)]      
    #dat = dat[:, ~np.all(dat == 0, axis=0)]

How do I tame it, and make it beautiful?

Answer 1

Try incorporating something like this:

# argwhere will give you the coordinates of every non-zero point
true_points = np.argwhere(dat)
# take the smallest points and use them as the top left of your crop
top_left = true_points.min(axis=0)
# take the largest points and use them as the bottom right of your crop
bottom_right = true_points.max(axis=0)
out = dat[top_left[0]:bottom_right[0]+1,  # plus 1 because slice isn't
          top_left[1]:bottom_right[1]+1]  # inclusive

This could be expanded without reasonable difficulty for the general n-d case.

Answer 2

This should work in any number of dimensions. I believe it is also quite efficient because swapping axes and slicing create only views on the array, not copies (which rules out functions such as take() or compress() which one might be tempted to use) or any temporaries. However it is not significantly 'nicer' than your own solution.

def crop2(dat, clp=True):
    if clp: np.clip( dat, 0, 1, out=dat )
    for i in range(dat.ndim):
        dat = np.swapaxes(dat, 0, i)  # send i-th axis to front
        while np.all( dat[0]==0 ):
            dat = dat[1:]
        while np.all( dat[-1]==0 ):
            dat = dat[:-1]
        dat = np.swapaxes(dat, 0, i)  # send i-th axis to its original position
    return dat

Answer 3

Definitely not the prettiest approach but wanted to try something else.

def _fill_gap(a):
    """
    a = 1D array of `True`s and `False`s.
    Fill the gap between first and last `True` with `True`s.

    Doesn't do a copy of `a` but in this case it isn't really needed.
    """
    a[slice(*a.nonzero()[0].take([0,-1]))] = True
    return a

def crop3(d, clip=True):
    dat = np.array(d)
    if clip: np.clip(dat, 0, 1, out=dat)
    dat = np.compress(_fill_gap(dat.any(axis=0)), dat, axis=1)
    dat = np.compress(_fill_gap(dat.any(axis=1)), dat, axis=0)
    return dat

But it works.

In [639]: crop3(np.array(
     ...:   [[0,0,0,0,0,0],
     ...:    [0,0,0,0,0,0],
     ...:    [0,1,0,2,9,0],
     ...:    [0,0,0,0,0,0],
     ...:    [0,7,4,1,0,0],
     ...:    [0,0,0,0,0,0]]))
Out[639]:
array([[1, 0, 1, 1],
       [0, 0, 0, 0],
       [1, 1, 1, 0]])