Numpy: Apply Function To Two Numpy Arrays And Return Two Numpy Arrays

I have two input numpy arrays with, respectively, latitude and longitude coordinates of a set of points: lats and lons. I have inherited a function which converts each (lat,lon) pa

Solution 1:

Let's define a trivial conversion

def convert(lat, lon):
    return lat*np.pi/180, lon*np.pi/180

frompyfunc is a useful way of applying 'scalar' function to arrays; We can even have it take in 2 arrays, and return 2 arrays (in a tuple)

In [233]: f = np.frompyfunc(convert,2,2)
In [234]: lats=np.linspace(-45,45,5)
In [235]: lons=np.linspace(0,100,5)
In [236]: out = f(lats, lons)
In [237]: out
Out[237]: 
(array([-0.7853981633974483, -0.39269908169872414, 0.0, 0.39269908169872414,
        0.7853981633974483], dtype=object),
 array([0.0, 0.4363323129985824, 0.8726646259971648, 1.3089969389957472,
        1.7453292519943295], dtype=object))

One feature is that it returns an object array, while you probably want a float array:

In [238]: out[0].astype(float)
Out[238]: array([-0.78539816, -0.39269908,  0.        ,  0.39269908,  0.78539816])

Or with unpacking:

In [239]: rlat, rlon = f(lats, lons)
In [240]: rlat.astype(float)
Out[240]: array([-0.78539816, -0.39269908,  0.        ,  0.39269908,  0.78539816])

frompyfunc does iterate through the inputs. In other tests it tends to be 2x faster than more explicit loops. And in this case, since it returns a tuple you don't have to call it twice to get 2 results.

As written, this convert works just as well with arrays as with scalars, so

In [241]: convert(lats, lons)
Out[241]: 
(array([-0.78539816, -0.39269908,  0.        ,  0.39269908,  0.78539816]),
 array([ 0.        ,  0.43633231,  0.87266463,  1.30899694,  1.74532925]))

which will be much faster than any version that loops in Python.

So for real speed you want convert to work directly with arrays. But if it can't do that, then frompyfunc is a modest improvement over do-it-yourself loops.

Another advantage to frompyfunc - it applies array broadcasting, as in

 f( lats[:,None], lons[None,:])