Python: speed up matrix coordinate mapping through iteration

Question

I'm trying to get this code to run as fast as possible and at the moment is very inefficient.

I have a 4D matrix of scalar data. The 4 dimensions correspond to latitude, longitude, altitude and time. The data is stored in a numpy array and its shape is (5,5,30,2).

In 4 different lists I am keeping the "map" for each axis, storing what value corresponds to each index. For example, the map arrays could look like:

mapLatitude = [45.,45.2,45.4,45.6,45.8]
mapLongitude = [-10.8,-10.6,-10.4,-10.2,-10.]
mapAltitude = [0,50,100,150,...,1450]
mapTime = [1345673,1345674]

This means that in the data matrix, the data point at location 0,1,3,0 corresponds to
Lat = 45, Lon = -10.6, Alt = 150, Time = 1345673.

Now, I need to generate a new array containing the coordinates of each point in my data matrix.

So far, this is what I've written:

import numpy as np

# data = np.array([<all data>])
coordinateMatrix = [ 
   (mapLatitude[index[0]],
    mapLongitude[index[1]],
    mapAltitude[index[2]],
    mapTime[index[3]] ) for index in numpy.ndindex(data.shape) ]

This works, but takes quite a long time, especially when the data matrix increases in size (I need to use this with matrices with a shape like (100,100,150,30) ).

If it helps, I need to generate this coordinateMatrix to feed it to scipy.interpolate.NearestNDInterpolator .

Any suggestions on how to speed this up?
Thank you very much!

Answer 1

If you turn your lists into ndarray's you can use broadcasting as follows:

coords = np.zeros((5, 5, 30, 2, 4))
coords[..., 0] = np.array(mapLatitude).reshape(5, 1, 1, 1)
coords[..., 1] = np.array(mapLongitude).reshape(1, 5, 1, 1)
coords[..., 2] = np.array(mapAltitude).reshape(1, 1, 30, 1)
coords[..., 3] = np.array(mapTime).reshape(1, 1, 1, 2)

For more general inputs something like this should work:

def makeCoordinateMatrix(*coords) :
    dims = len(coords)
    coords = [np.array(a) for a in coords]
    shapes = tuple([len(a) for a in coords])
    ret = np.zeros(shapes + (dims,))
    for j, a in enumerate(coords) :
        ret[..., j] = a.reshape((len(a),) + (1,) * (dims - j - 1))
    return ret

coordinateMatrix = makeCoordinateMatrix(mapLatitude, mapLongitude,
                                        mapAltitude, mapTime)