N-D indexing with defaults in NumPy

Question

Can I index NumPy N-D array with fallback to default values for out-of-bounds indexes? Example code below for some imaginary np.get_with_default(a, indexes, default):

import numpy as np
print(np.get_with_default(
    np.array([[1,2,3],[4,5,6]]), # N-D array
    [(np.array([0, 0, 1, 1, 2, 2]), np.array([1, 2, 2, 3, 3, 5]))], # N-tuple of indexes along each axis
    13, # Default for out-of-bounds fallback
))

should print

[2 3 6 13 13 13]

I'm looking for some built-in function for this. If such not exists then at least some short and efficient implementation to do that.

Answer 1

I arrived at this question because I was looking for exactly the same. I came up with the following function, which does what you ask for 2 dimension. It could likely be generalised to N dimensions.

def get_with_defaults(a, xx, yy, nodata):
   # get values from a, clipping the index values to valid ranges
   res = a[np.clip(yy, 0, a.shape[0] - 1), np.clip(xx, 0, a.shape[1] - 1)]
   # compute a mask for both x and y, where all invalid index values are set to true
   myy = np.ma.masked_outside(yy, 0, a.shape[0] - 1).mask
   mxx = np.ma.masked_outside(xx, 0, a.shape[1] - 1).mask
   # replace all values in res with NODATA, where either the x or y index are invalid
   np.choose(myy + mxx, [res, nodata], out=res)
   return res

xx and yy are the index array, a is indexed by (y,x).

This gives:

>>> a=np.zeros((3,2),dtype=int)
>>> get_with_defaults(a, (-1, 1000, 0, 1, 2), (0, -1, 0, 1, 2), -1)
array([-1, -1,  0,  0, -1])

As an alternative, the following implementation achieves the same and is more concise:

def get_with_default(a, xx, yy, nodata):
   # get values from a, clipping the index values to valid ranges
   res = a[np.clip(yy, 0, a.shape[0] - 1), np.clip(xx, 0, a.shape[1] - 1)]
   # replace all values in res with NODATA (gets broadcasted to the result array), where
   # either the x or y index are invalid
   res[(yy < 0) | (yy >= a.shape[0]) | (xx < 0) | (xx >= a.shape[1])] = nodata
   return res