numpy对两个二维数组进行迭代

Question

Say I have two matrices:

X, Y = np.meshgrid(np.arange(0, 2, 0.1), np.arange(3, 5, 0.1))

And a function, something like:

def func(x) :
    return x[0]**2 + x[1]**2

How can I fill a matrix Z (of size np.shape(X)), where each entry is formed by calling func on the two corresponding values of X and Y , ie:

Z[i, j] = func([X[i, j], Y[i, j]])

Is there a way without using a double nested for-loop?

Answer 1

For given numpy arrays X and Y , you could just do -

Zout = X**2 + Y**2

If you are actually constructing X and Y like that, there is a direct way to get Z with broadcasting and thus avoid np.meshgrid , like so -

Zout = np.arange(0, 2, 0.1)**2 + np.arange(3, 5, 0.1)[:,None]**2

Answer 2

This is also works as a vectorized form of function evaluation:

import numpy as np
X, Y = np.meshgrid(np.arange(0, 2, 0.1), np.arange(3, 5, 0.1))
def func(x) :
    return x[0]**2 + x[1]**2

Z = func([X,Y])