Numpy 2D array indexing with indices out of bounds

Question

I found a substantial bottleneck in the following code:

def get_value(matrix, index):
    if (index[0] >= 0 and index[1] >= 0 and
        index[0] < matrix.shape[0] and
        index[1] < matrix.shape[1]):
        return matrix[index[0], index[1]]
    return DEFAULT_VAL

Given a 2D matrix and an index accessing the matrix, it checks for out-of-bounds indices and returns the value at the given index. Otherwise, it returns a DEFAULT_VAL in case of out-of-bounds indices.

This method is called many times (even millions of calls), which is slow. So, I am trying to vectorize it using numpy. Unfortunately, I cannot find a way to do it.

If I didn't have to care about out-of-bounds values, I'd do the following:

def get_values(matrix, indices):
    return matrix[indices[:,0], indices[:,1]]

I've been thinking of a way to utilize numpy to do this task, but I haven't found a way yet.

Is there a way to accomplish this?

Answer 1

The code you have shown

def get_values(matrix, indices):
    return matrix[indices[:,0], indices[:,1]]

is the best you can do given that indices is a tuple with two values.

You should rather look at the optimal way to call the above method. I suggest if you can, rather then calling get_values with a single tuple, call with a possibility large number of such tuples. Then you can atleast try to write a vectorized version of get_values. With single tuple there is nothing you can vectorize here.

Vectorized method

Assuming that your indices is a numpy array of size n X 2 where n is the number of indices and 2 corresponds to two dimensions then you can use

index = np.random.randint(0,500, size=(10000,2))
matrix = np.random.randn(1000,1000)

def get_value(matrix, index, default_value=-1):
  result = np.zeros(len(index))+default_value
  mask = (index[:,0] < matrix.shape[0]) & (index[:,1] < matrix.shape[1])
  valid = index[mask]
  result[mask] = matrix[valid[:, 0], valid[:, 1]]
  return result


assert np.all(get_value(matrix, np.array(([0,1001],[1001,1001]))) == -1)

%timeit get_value(matrix, index, -1): 1 loop, best of 3: 264 ms per loop