numpy iterate over two 2d arrays

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, i.e.:

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

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

It's best to avoid fractional ranges (like np.arange(0, 2, 0.1)) as floating point inaccuracies can cause unpredictable end-points. Use linspace instead. — Alex Riley
– Alex Riley, Commented Jan 2, 2016 at 18:58

NKN · Accepted Answer · 2016-01-02 18:59:02Z

3

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])

answered Jan 2, 2016 at 18:59

NKN

6,4526 gold badges37 silver badges54 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Divakar · Accepted Answer · 2016-01-02 18:54:46Z

3

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

answered Jan 2, 2016 at 18:54

Divakar

222k19 gold badges273 silver badges374 bronze badges

2 Comments

Tom Over a year ago

I did mean for an arbitrary function, but your answer is still really informative. Could you possibly explain what the [:,None] does please?

Divakar Over a year ago

@Tom That one is to extend the dimension by introducing a new axis with None. So, in this case, we convert a 1D to a 2D with it. And the idea with extending dimension(s) here is to bring in broadcasting.

Collectives™ on Stack Overflow

numpy iterate over two 2d arrays

2 Answers 2

Comments

2 Comments

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

2 Comments

Related