使用索引作为参数沿 numpy 数组应用 function

Question

I am trying to find a way to run a function through a whole ndarray using the indexes of every value as parameters. A regular loop is quite slow and I can't find a way to make it work using numpy built-in functions. The next example code summarizes what i'm trying to do. Thanks in advance.

import numpy as np

arr = np.arange(10).reshape(2, 5)

print(arr)

def example_func(row, col, value):
    return(row + col + value)

for row in range(arr.shape[0]):
    for col in range(arr.shape[1]):
        arr[row, col] = example_func(row, col, arr[row, col])

print(arr)

[[0 1 2 3 4]
 [5 6 7 8 9]]

[[ 0  2  4  6  8]
 [ 6  8 10 12 14]]

Answer 1

What you try to do can be done with meshgrid .

Return coordinate matrices from coordinate vectors.

n_rows, n_cols = arr.shape
col_matrix, row_matrix = np.meshgrid(np.arange(n_cols), np.arange(n_rows))
result = arr + col_matrix + row_matrix
print(result)

This returns

[[ 0  2  4  6  8]
 [ 6  8 10 12 14]]