subclassing numpy ndarray to simulate padding with zeros on out of bounds indices

Question

I want to have a 2D numpy array (NxN) which will return zeros when an index is either negative or out of bounds (ie I want to suppress the usual wrap-around indexing that occurs when an index is -1 or the IndexError when your index is N). I could do this literally, simply adding a border of zeros around the array and treating it as a 1-based array instead of a 0-based array, but this seems inelegant.

I stumbled across a few answers here relating to subclassing the numpy ndarray class and defining your own __getitem__ attribute. My first attempt looked like this:

import numpy as np

class zeroPaddedArray(np.ndarray):
    def __getitem__(self, index):
        x,y = index
        if x < 0 or y < 0 or x >= self.shape[0] or y >= self.shape[1]:
            return 0
        return super(zeroPaddedArray, self).__getitem__(index)

This sort of works, but only allows you to access array elements as arr[x,y] , and throws an error when you try arr[x][y] . It also completely breaks a lot of other functions such as print . print arr gives TypeError: 'int' object is not iterable .

My next attempt was to check if a tuple was given for the index, and to default to the old behavior if not.

import numpy as np

class zeroPaddedArray(np.ndarray):
    def __getitem__(self, index):
        if type(index) is tuple:
            x,y = index
            if x < 0 or y < 0 or x >= self.shape[0] or y >= self.shape[1]:
                return 0
        return super(zeroPaddedArray, self).__getitem__(index)
    else:
        return super(zeroPaddedArray, self).__getitem__(index)

This gives me the requires zero-padding behavior when indexing as a tuple ( arr[-1,-1] correctly gives 0 ), but allows other functions to work properly. However, now I get different result depending on the way I index things. For example:

a = np.ones((5,5))
b = a.view(zeroPaddedArray)
print b[-1][-1]
print b[-1,-1]

gives

>>>1.0
>>>0

I think this is probably usable as-is for my purposes, but I am unsatisfied. How can I tweak this to give the desired zero-padding behavior regardless of the syntax I use for indexing, without breaking all of the other ndarray functionality?

Answer 1

Here's how take could be used:

In [34]: a=np.zeros((5,5),int)
In [35]: a[1:4,1:4].flat=np.arange(9)
In [36]: a
Out[36]: 
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]])
In [37]: np.take(a, np.arange(-1,6),1,mode="clip")
Out[37]: 
array([[0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 1, 2, 0, 0],
       [0, 0, 3, 4, 5, 0, 0],
       [0, 0, 6, 7, 8, 0, 0],
       [0, 0, 0, 0, 0, 0, 0]])
In [38]: np.take(a, np.arange(-1,6),0,mode="clip")
Out[38]: 
array([[0, 0, 0, 0, 0],
       [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],
       [0, 0, 0, 0, 0]])

I'm using the clip mode to expand the 0 boundary.

There is also an np.pad function, though with all of its generality it is quite long, and no speed solution (it ends up doing 2 concatenates per dimension).

np.lib.index_tricks.py has some good examples of using custom classes to playing indexing tricks.

Before getting too deep into subclassing ndarray I'd suggest writing functions or trick classes to test your ideas.