Apply a function of the indices to all elements of a numpy array-CodePudding

I am looking for a way to apply a function to all elements of a numpy array. The function receives not only the value but also the indices of the element as arguments. The goal is perfomance on large 2- or 3-dim. arrays.

(I know there are several ways to do that with a function that receives the value of an element only)

The code to be replaced is

def foo(x, i, j)
    return (i*x)**2 - (j*x)**3  # or some other fancy stuff

...

arr = np.zeros((nx, ny))

...

# nested loop to be replaced, e.g. via vectorization
for i in range(nx):
    for j in range(ny):
        arr[i,j] = foo(arr[i,j], i, j)

CodePudding user response：

You can do this with simple broadcasting rules, by using suitably generated indices with the proper shapes so that standard broadcasting rules match the shape of the input.

This can be generated either manually (e.g. with a combination of np.arange() and np.reshape()) or more concisely with np.ogrid().

import numpy as np


import numpy as np


np.random.seed(0)


def foo(x):
    n, m = arr.shape
    i, j = np.ogrid[:n, :m]
    return (i * x) ** 2 - (j * x) ** 3


n, m = 2, 3
arr = np.random.random((n, m))


foo(arr)
# array([[ 0.        , -0.36581638, -1.7519857 ],
#        [ 0.29689768,  0.10344439, -1.73844954]])

This approach would require potentially large temporary memory arrays for the intermediate results.

A more efficient approach can be obtained by keeping explicit looping to be accelerated with a JIT compiler like numba:

import numba as nb


@nb.njit
def foo_nb(arr):
    n, m = arr.shape
    out = np.empty((n, m), dtype=arr.dtype)
    for i in range(n):
        for j in range(m):
            x = arr[i, j]
            out[i, j] = (i * x) ** 2 - (j * x) ** 3
    return out


foo_nb(arr)
# array([[ 0.        , -0.36581638, -1.7519857 ],
#        [ 0.29689768,  0.10344439, -1.73844954]])