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 don't know if there is anything in NumPy to do that directly, but you can always implement it yourself. This is not particularly smart or efficient, as it requires multiple advanced indexing operations, but does what you need:

import numpy as np

def get_with_default(a, indices, default=0):
    # Ensure inputs are arrays
    a = np.asarray(a)
    indices = tuple(np.broadcast_arrays(*indices))
    if len(indices) <= 0 or len(indices) > a.ndim:
        raise ValueError('invalid number of indices.')
    # Make mask of indices out of bounds
    mask = np.zeros(indices[0].shape, np.bool)
    for ind, s in zip(indices, a.shape):
        mask |= (ind < 0) | (ind >= s)
    # Only do masking if necessary
    n_mask = np.count_nonzero(mask)
    # Shortcut for the case where all is masked
    if n_mask == mask.size:
        return np.full_like(a, default)
    if n_mask > 0:
        # Ensure index arrays are contiguous so masking works right
        indices = tuple(map(np.ascontiguousarray, indices))
        for ind in indices:
            # Replace masked indices with zeros
            ind[mask] = 0
    # Get values
    res = a[indices]
    if n_mask > 0:
        # Replace values of masked indices with default value
        res[mask] = default
    return res

# Test
print(get_with_default(
    np.array([[1,2,3],[4,5,6]]),
    (np.array([0, 0, 1, 1, 2, 2]), np.array([1, 2, 2, 3, 3, 5])),
    13
))
# [ 2  3  6 13 13 13]

Answer 2

I also needed a solution to this, but I wanted a solution that worked in N dimensions. I made Markus' solution work for N-dimensions, including selecting from an array with more dimensions than the coordinates point to.

def get_with_defaults(arr, coords, nodata):
    coords, shp = np.array(coords), np.array(arr.shape)
    # Get values from arr, clipping to valid ranges
    res = arr[tuple(np.clip(c, 0, s-1) for c, s in zip(coords, shp))]
    # Set any output where one of the coords was out of range to nodata
    res[np.any(~((0 <= coords) & (coords < shp[:len(coords), None])), axis=0)] = nodata
    return res


import numpy as np

if __name__ == '__main__':
    A = np.array([[1,2,3],[4,5,6]])
    B = np.array([[[1, -9],[2, -8],[3, -7]],[[4, -6],[5, -5],[6, -4]]])
    coords1 = [[0, 0, 1, 1, 2, 2], [1, 2, 2, 3, 3, 5]]
    coords2 = [[0, 0, 1, 1, 2, 2], [1, 2, 2, 3, 3, 5], [1, 1, 1, 1, 1, 1]]

    out1 = get_with_defaults(A, coords1, 13)
    out2 = get_with_defaults(B, coords1, 13)
    out3 = get_with_defaults(B, coords2, 13)

    print(out1)
    # [2, 3, 6, 13, 13, 13]
    print(out2)
    # [[ 2 -8]
    #  [ 3 -7]
    #  [ 6 -4]
    #  [13 13]
    #  [13 13]
    #  [13 13]]
    print(out3)
    # [-8, -7, -4, 13, 13, 13]

Answer 3

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