Apply a fonction to multiple vectors stored in a numpy array without loops

Question

I have a function taking two vectors with same size, making some calculations with them, and returning a third vector. Let's now consider that I have a multiple dimension array containing many vectors that I want to pass to my function as the first argument, and a fixed vector that I want to pass as the second argument. Below is an example. Is there a way to simplify the code by removing the loops?

def foo(x, y):
    result = np.zeros(x.shape)
    for i in range(y.size-1):
        result[i+1] = result[i] + (x[i+1] + x[i]) / (y[i+1] + y[i])
    return result

a = np.arange(2*3*4).reshape(2,3,4)
b = np.arange(4)*10
c = np.ones(a.shape)*-9999.
for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        c[i, j, :] = foo(a[i, j, :], b)

Thanks for your help!

EDIT: Below is the real function I'm trying to implement.

def press2alt(press, Tv):
    """
    Convert pressure level to altitude level with hydrostatic atmosphere calculation (hypsometric
    equation). The altitude reference (z = 0 km) is taken for the largest pressure level.
    :param press: pressure level profile [hPa]
    :param Tv: virtual temperature profile [K]
    :return: altitude level profile [km]
    """
    press_c = np.copy(press)
    if press[0] < press[1]:
        press_c = press_c[::-1] # high press first
        Tv = Tv[::-1]

    alt = np.zeros(press_c.size)
    for i in range(alt.size-1):
        alt[i+1] = alt[i] + DRY_AIR_GAS_CST/STD_GRAV_ACC*(Tv[i+1]-Tv[i])* \
                   (np.log(press_c[i])-np.log(press_c[i+1]))/(np.log(Tv[i+1])-np.log(Tv[i]))
    
    if press[0] < press[1]:
        alt = alt[::-1]
        
    return alt


# Get altitude at each pressure level
z = np.ones(tv.shape)*FILL_VALUE_FLOAT
    for i_month in range(tv.shape[0]):
        for i_lat in range(tv.shape[2]):
            for i_lon in range(tv.shape[3]):
                z[i_month, :, i_lat, i_lon] = \
                    press2alt(pressure_level, tv[i_month, :, i_lat, i_lon])

Answer 1

The c from your sample (which you should have shown) is:

In [164]: c
Out[164]: 
array([[[0.        , 0.1       , 0.2       , 0.3       ],
        [0.        , 0.9       , 1.26666667, 1.52666667],
        [0.        , 1.7       , 2.33333333, 2.75333333]],

       [[0.        , 2.5       , 3.4       , 3.98      ],
        [0.        , 3.3       , 4.46666667, 5.20666667],
        [0.        , 4.1       , 5.53333333, 6.43333333]]])

np.vectorize with signature turns out to be easier to use than I first thought:

In [165]: f = np.vectorize(foo, signature="(n),(n)->(n)")
In [166]: f(a, b)
Out[166]: 
array([[[0.        , 0.1       , 0.2       , 0.3       ],
        [0.        , 0.9       , 1.26666667, 1.52666667],
        [0.        , 1.7       , 2.33333333, 2.75333333]],

       [[0.        , 2.5       , 3.4       , 3.98      ],
        [0.        , 3.3       , 4.46666667, 5.20666667],
        [0.        , 4.1       , 5.53333333, 6.43333333]]])

But vectorize does not improve speed:

In [167]: %%timeit
     ...: c = np.ones(a.shape)*-9999.
     ...: for i in range(a.shape[0]):
     ...:     for j in range(a.shape[1]):
     ...:         c[i, j, :] = foo(a[i, j, :], b)
     ...: 
57 µs ± 126 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [168]: timeit f(a, b)
206 µs ± 3.05 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In other recent cases I've found that vectorize does improve in relative performance with larger arrays, but that wasn't with signature.

The function can be rewritten to accept arrays of any size, as long as the iteration on the last dimension is correct. Basically I use ... in the indexing:

def myfoo(a, b):
    result = np.zeros(a.shape)
    for i in range(a.shape[-1] - 1):
        result[..., i + 1] = result[..., i] + (a[..., i + 1] + 
             a[...,i]) / ( b[..., i + 1] + b[..., i])
    return result

In [182]: myfoo(a, b)
Out[182]: 
array([[[0.        , 0.1       , 0.2       , 0.3       ],
        [0.        , 0.9       , 1.26666667, 1.52666667],
        [0.        , 1.7       , 2.33333333, 2.75333333]],

       [[0.        , 2.5       , 3.4       , 3.98      ],
        [0.        , 3.3       , 4.46666667, 5.20666667],
        [0.        , 4.1       , 5.53333333, 6.43333333]]])
In [183]: timeit myfoo(a, b)
65.8 µs ± 483 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

This doesn't help with speed, possibly because the last axis, size 4, is comparable to the 2*3 iterations of the first. I expect it will do better, relatively, if the initial dimensions get much larger.

We may be able to improve speed by replacing the i iteration on:

(a[..., i + 1] + a[...,i]) / ( b[..., i + 1] + b[..., i])

with

(a[...,1:]+a[...,:-1])/(b[...,1:]+b[...,:-1])

edit

In [192]: ab = (a[..., 1:] + a[..., :-1]) / (b[..., 1:] + b[..., :-1])
In [193]: ab
Out[193]: 
array([[[0.1       , 0.1       , 0.1       ],
        [0.9       , 0.36666667, 0.26      ],
        [1.7       , 0.63333333, 0.42      ]],

       [[2.5       , 0.9       , 0.58      ],
        [3.3       , 1.16666667, 0.74      ],
        [4.1       , 1.43333333, 0.9       ]]])
In [194]: ab.cumsum(axis=2)
Out[194]: 
array([[[0.1       , 0.2       , 0.3       ],
        [0.9       , 1.26666667, 1.52666667],
        [1.7       , 2.33333333, 2.75333333]],

       [[2.5       , 3.4       , 3.98      ],
        [3.3       , 4.46666667, 5.20666667],
        [4.1       , 5.53333333, 6.43333333]]])

Those are the values - except for the leading 0's column.

In [195]: timeit((a[..., 1:] + a[..., :-1]) / (b[..., 1:] + b[..., :-1])).cumsum
     ...: (axis=2)
18.8 µs ± 36.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)