How to vectorized version of this code using np.meshgrid and np.linspace

Question

How can I write a vectorized version of this code cell below? The code should be fully vectorized, with no for-loops, using np.meshgrid and np.linspace?

def eval_on_grid_unvectorized(func, extent, numsteps):


     """Evaluates func(x1, x2) for each combination in a 2D grid.

    func: callable - function to evaluate for each grid element

    extent: tuple - grid extent as (x1min, x1max, x2min, x2max)

    numsteps: int - number of grid steps (same for each 
    dimension)

    """
    x1min, x1max, x2min, x2max = extent

    x1 = np.empty((numsteps, numsteps))
    x2 = np.empty((numsteps, numsteps))
    y  = np.empty((numsteps, numsteps))
    for i in range(numsteps):
        for j in range(numsteps):
           x1[i,j] = x1min + j*(x1max-x1min)/(numsteps-1)
           x2[i,j] = x2min + i*(x2max-x2min)/(numsteps-1)
           y[i,j] = func(x1[i,j], x2[i,j])
    return x1, x2, y

Answer 1

Your function, without the y; produces

In [57]: eval_on_grid_unvectorized(None,(0,5,0,6),6)
Out[57]: 
(array([[0., 1., 2., 3., 4., 5.],
        [0., 1., 2., 3., 4., 5.],
        [0., 1., 2., 3., 4., 5.],
        [0., 1., 2., 3., 4., 5.],
        [0., 1., 2., 3., 4., 5.],
        [0., 1., 2., 3., 4., 5.]]),
 array([[0. , 0. , 0. , 0. , 0. , 0. ],
        [1.2, 1.2, 1.2, 1.2, 1.2, 1.2],
        [2.4, 2.4, 2.4, 2.4, 2.4, 2.4],
        [3.6, 3.6, 3.6, 3.6, 3.6, 3.6],
        [4.8, 4.8, 4.8, 4.8, 4.8, 4.8],
        [6. , 6. , 6. , 6. , 6. , 6. ]]))

Meshgrid and linspace can do the same:

In [59]: np.meshgrid(np.linspace(0,5,6), np.linspace(0,6,6),indexing='xy')
Out[59]: 
[array([[0., 1., 2., 3., 4., 5.],
        [0., 1., 2., 3., 4., 5.],
        [0., 1., 2., 3., 4., 5.],
        [0., 1., 2., 3., 4., 5.],
        [0., 1., 2., 3., 4., 5.],
        [0., 1., 2., 3., 4., 5.]]),
 array([[0. , 0. , 0. , 0. , 0. , 0. ],
        [1.2, 1.2, 1.2, 1.2, 1.2, 1.2],
        [2.4, 2.4, 2.4, 2.4, 2.4, 2.4],
        [3.6, 3.6, 3.6, 3.6, 3.6, 3.6],
        [4.8, 4.8, 4.8, 4.8, 4.8, 4.8],
        [6. , 6. , 6. , 6. , 6. , 6. ]])]

But as Jérôme points out, as long as func can only work with scalar values, you can't "vectorize", in the sense of doing fast whole array calculations.

If the func is something like x+y, then we can simply pass the arrays to it:

In [60]: _[0]+_[1]
Out[60]: 
array([[ 0. ,  1. ,  2. ,  3. ,  4. ,  5. ],
       [ 1.2,  2.2,  3.2,  4.2,  5.2,  6.2],
       [ 2.4,  3.4,  4.4,  5.4,  6.4,  7.4],
       [ 3.6,  4.6,  5.6,  6.6,  7.6,  8.6],
       [ 4.8,  5.8,  6.8,  7.8,  8.8,  9.8],
       [ 6. ,  7. ,  8. ,  9. , 10. , 11. ]])

The key to what we normally call vectorization in numpy is to write the calculation in a way that uses the compiled numpy methods, and operators. It requires real knowledge of numpy; there isn't a magic short cut to let you jump from scalar Python calculations to efficient numpy ones.