Python numpy grid transformation using universal functions

Question

Here is my problem : I manipulate 432*46*136*136 grids representing time*(space) encompassed in numpy arrays with numpy and python. I have one array alt, which encompasses the altitudes of the grid points, and another array temp which stores the temperature of the grid points.

It is problematic for a comparison : if T1 and T2 are two results, T1[t0,z0,x0,y0] and T2[t0,z0,x0,y0] represent the temperature at H1[t0,z0,x0,y0] and H2[t0,z0,x0,y0] meters, respectively. But I want to compare the temperature of points at the same altitude, not at the same grid point.

Hence I want to modify the z-axis of my matrices to represent the altitude and not the grid point. I create a function conv(alt[t,z,x,y]) which attributes a number between -20 and 200 to each altitude. Here is my code :

def interpolation_extended(self,temp,alt):
    [t,z,x,y]=temp.shape
    new=np.zeros([t,220,x,y])
    for l in range(0,t):
       for j in range(0,z):
          for lat in range(0,x):
             for lon in range(0,y):
                new[l,conv(alt[l,j,lat,lon]),lat,lon]=temp[l,j,lat,lon]
    return new

But this takes definitely too much time, I can't work this it. I tried to write it using universal functions with numpy :

def interpolation_extended(self,temp,alt):
    [t,z,x,y]=temp.shape
    new=np.zeros([t,220,x,y])
    for j in range(0,z):
       new[:,conv(alt[:,j,:,:]),:,:]=temp[:,j,:,:]
    return new

But that does not work. Do you have any idea of doing this in python/numpy without using 4 nested loops ?

Thank you

Answer 1

I can't really try the code since I don't have your matrices, but something like this should do the job.

First, instead of declaring conv as a function, get the whole altitude projection for all your data:

conv = np.round(alt / 500.).astype(int)

Using np.round, the numpys version of round, it rounds all the elements of the matrix by vectorizing operations in C, and thus, you get a new array very quickly (at C speed). The following line aligns the altitudes to start in 0, by shifting all the array by its minimum value (in your case, -20):

conv -= conv.min()

the line above would transform your altitude matrix from [-20, 200] to [0, 220] (better for indexing).

With that, interpolation can be done easily by getting multidimensional indices:

t, z, y, x = np.indices(temp.shape)

the vectors above contain all the indices needed to index your original matrix. You can then create the new matrix by doing:

new_matrix[t, conv[t, z, y, x], y, x] = temp[t, z, y, x]

without any loop at all.

Let me know if it works. It might give you some erros since is hard for me to test it without data, but it should do the job.

---

The following toy example works fine:

A = np.random.randn(3,4,5) # Random 3x4x5 matrix -- your temp matrix
B = np.random.randint(0, 10, 3*4*5).reshape(3,4,5) # your conv matrix with altitudes from 0 to 9
C = np.zeros((3,10,5)) # your new matrix

z, y, x = np.indices(A.shape)
C[z, B[z, y, x], x] = A[z, y, x]

C contains your results by altitude.