Избегайте Numpy Index For цикла

Question

Есть ли способ избежать использования второго for цикл для такой операции?

for x in range(Size_1):
    for y in range(Size_2):
        k[x,y] = np.sqrt(x+y) - y

Или есть лучший способ оптимизировать это? Прямо сейчас это невероятно медленно для больших размеров.

Answer 1

Вот векторизованное решение с broadcasting -

X,Y = np.ogrid[:Size_1,:Size_2]
k_out = np.sqrt(X+Y) - Y

Answer 2

Дополнение решения Дивакара: если Y а также X не новые диапазоны, но некоторые ранее существующие векторы чисел, используйте np.ix_:

Y, X = np.array([[1.3, 3.5, 2], [2.0, -1, 1]])
Y, X = np.ix_(Y, X) # does the same as Y = Y[:, None]; X = X[None, :]
out = np.sqrt(Y+X) - X