CODEWITHSUNDEEP
pythonarraysfunctionloopsnumpy
FaCoffee
FaCoffee

Reputation: 7909

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) pair into an (E,N) pair:

def convert(lat,lon): #takes two floats as arguments (unit: degrees)
   ...
   computation #Actual function is too long to post
   ...
   return N,E #returns two floats (unit: meters)

My question: how could I efficiently apply the same function simultaneously to both input numpy arrays?

I was thinking of modifying the function so that it returns a list:

return [N,E]

in such a way that:

rows = int(lat.shape[0]) #lat and lon have the same shape
cols = int(lat.shape[1])
easting=numpy.zeros(shape=(rows,cols))
northing=numpy.zeros(shape=(rows,cols))
for i in range(0, rows):
    for j in range(0, cols):
        northing=convert(lon[i][j])[0] #first element of the returned list
        easting=convert(lat[i][j])[1] #second element of the returned list

I have not yet tested this, but by looking at it I don't feel very comfortable this will work. Any insights will be much appreciated.

Upvotes: 3

Views: 3726

Answers (1)

hpaulj
hpaulj

Reputation: 231375

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

Upvotes: 4

Related Questions