Zero padding slice past end of array in numpy

Question

In numpy, is there a way to zero pad entries if I'm slicing past the end of the array, such that I get something that is the size of the desired slice?

For example,

>>> x = np.ones((3,3,))
>>> x
array([[ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.]])
>>> x[1:4, 1:4] # would behave as x[1:3, 1:3] by default
array([[ 1.,  1.,  0.],
       [ 1.,  1.,  0.],
       [ 0.,  0.,  0.]])
>>> x[-1:2, -1:2]
 array([[ 0.,  0.,  0.],
       [ 0.,  1.,  1.],
       [ 0.,  1.,  1.]])

Visually, I'd like the out-of-bounds areas to be zero padded:

I'm dealing with images and would like to zero pad to signify moving off the image for my application.

My current plan is to use np.pad to make the entire array larger prior to slicing, but indexing seems to be a bit tricky. Is there a potentially easier way?

Answer 1

As far as I know there is no numpy solution (nor in any package I know) for such a problem. You could do it yourself but it will be a really, really complicated one even if you only want basic slicing. I would suggest you manually np.pad your array and simply offset your start/stop/step before you actually slice it.

However if all you need to support are integers and slices without step I have some "working code" for this:

import numpy as np

class FunArray(np.ndarray):
    def __getitem__(self, item):

        all_in_slices = []
        pad = []
        for dim in range(self.ndim):
            # If the slice has no length then it's a single argument.
            # If it's just an integer then we just return, this is
            # needed for the representation to work properly
            # If it's not then create a list containing None-slices
            # for dim>=1 and continue down the loop
            try:
                len(item)
            except TypeError:
                if isinstance(item, int):
                    return super().__getitem__(item)
                newitem = [slice(None)]*self.ndim
                newitem[0] = item
                item = newitem
            # We're out of items, just append noop slices
            if dim >= len(item):
                all_in_slices.append(slice(0, self.shape[dim]))
                pad.append((0, 0))
            # We're dealing with an integer (no padding even if it's
            # out of bounds)
            if isinstance(item[dim], int):
                all_in_slices.append(slice(item[dim], item[dim]+1))
                pad.append((0, 0))
            # Dealing with a slice, here it get's complicated, we need
            # to correctly deal with None start/stop as well as with
            # out-of-bound values and correct padding
            elif isinstance(item[dim], slice):
                # Placeholders for values
                start, stop = 0, self.shape[dim]
                this_pad = [0, 0]
                if item[dim].start is None:
                    start = 0
                else:
                    if item[dim].start < 0:
                        this_pad[0] = -item[dim].start
                        start = 0
                    else:
                        start = item[dim].start
                if item[dim].stop is None:
                    stop = self.shape[dim]
                else:
                    if item[dim].stop > self.shape[dim]:
                        this_pad[1] = item[dim].stop - self.shape[dim]
                        stop = self.shape[dim]
                    else:
                        stop = item[dim].stop
                all_in_slices.append(slice(start, stop))
                pad.append(tuple(this_pad))

        # Let numpy deal with slicing
        ret = super().__getitem__(tuple(all_in_slices))
        # and padding
        ret = np.pad(ret, tuple(pad), mode='constant', constant_values=0)

        return ret

This can be used as follows:

>>> x = np.arange(9).reshape(3, 3)
>>> x = x.view(FunArray)
>>> x[0:2]
array([[0, 1, 2],
       [3, 4, 5]])
>>> x[-3:2]
array([[0, 0, 0],
       [0, 0, 0],
       [0, 0, 0],
       [0, 1, 2],
       [3, 4, 5]])
>>> x[-3:2, 2]
array([[0],
       [0],
       [0],
       [2],
       [5]])
>>> x[-1:4, -1:4]
array([[0, 0, 0, 0, 0],
       [0, 0, 1, 2, 0],
       [0, 3, 4, 5, 0],
       [0, 6, 7, 8, 0],
       [0, 0, 0, 0, 0]])

Note that this may be contain Bugs and "not cleanly coded" parts, I've never used this except in trivial cases.