Numpy Индексирование 2D-массива с индексами за пределами

Question

Я обнаружил существенное узкое место в следующем коде:

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

Учитывая 2D-матрицу и индекс, обращающийся к матрице, он проверяет индексы за пределами границ и возвращает значение по заданному индексу . В противном случае он возвращает DEFAULT_VAL в случае индексов, выходящих за границы.

Этот метод вызывается много раз (даже миллионы вызовов), что является медленным. Итак, я пытаюсь векторизовать его, используя numpy. К сожалению, я не могу найти способ сделать это.

Если бы мне не приходилось заботиться о значениях, выходящих за пределы, я бы сделал следующее:

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

I Я думал о способе использовать numpy для выполнения этой задачи, но пока не нашел способа.

Есть ли способ выполнить sh это?

Answer 1

Показанный вами код

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

- лучшее, что вы можете сделать, учитывая, что indices - это кортеж с двумя значениями.

Лучше посмотрите на оптимальный способ вызовите указанный выше метод. Я предлагаю, если вы можете, вместо того, чтобы вызывать get_values с одним кортежем, вызывать с возможностью большого количества таких кортежей. Тогда вы можете хотя бы попробовать написать векторизованную версию get_values. С одним кортежем здесь нет ничего, что можно векторизовать. индексы и 2 соответствует двум измерениям, тогда вы можете использовать

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