0

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
Joy
  • 75
  • 1
  • 6
  • This is not possible to efficiently vectorize arbitrary `func` function with Numpy. Numpy provides function to vectorize functions from the user point of view but they are inefficient since they just do a naive loop internally. Thus, you need to provide the code of `func`. If `func` is a generic function and you cannot specialize this code nor the one of `func`, then there is no way to vectorize the function (at least with Numpy). – Jérôme Richard Jan 23 '22 at 16:47
  • Thanks Richard! func is a lambda function. See below where eval_on_grid_unvectorized is called : x1, x2, y = eval_on_grid_unvectorized(lambda x1, x2: x1 + x2, (-1, 1, 0, 2), 3) – Joy Jan 23 '22 at 16:50

1 Answers1

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.

hpaulj
  • 221,503
  • 14
  • 230
  • 353