Search code examples
pythonnumpymatplotlibscipyscientific-computing

Dynamic 3d array slicing and sorting


I'm a bit confused on how to accurately slice and sort 3D array in numpy. There seem to be many ways to do this manually but I need to do this using numpy.where(). For example if lo360 are 2D longitude values, lat2d latitude values in 2D, yi is a 1D array of longitude values, and xi is a 1D array of latitude values.

xi and yi dynamically change to represent a small geographical region while lo360 and lat2d are static latitudes and longitudes of the planet of the type (-90,90) and (0,360). xi of similar form as lo360 but yi is descending instead of ascending. So if I have a 3D array representing A(levels,lat,lon) and I want to extract a region:

slice2d = np.where( (lo360 <= xi.max()) &                                                    
                    (lo360 >= xi.min()) &                                                    
                    (lat2d <= yi.max()) &                                                    
                    (lat2d >= yi.min()) )
lon_old = lo360[slice2d]; print lon_old.shape
(441,)

This returns a 1D array when I wanted a 2D slice. The data is correct though so this is not my problem.

Then when I tried to slice the 3D array A[i][slice2d] I get a 1D array that is not easy to verify dynamically. I used griddata to the 3D array to xi and yi resolution but I change the yi lat to ascending: yi = yi[::-1]:

for i in np.arange(4):                                                          
    nvals[i] = matplotlib.mlab.griddata(lat_old,lon_old,                                     
                                            mvals[i][slice2d],                                   
                                            yi,xi)

Here is where I think the problem starts, I need the results to have descending lats so I do this to nvals: nvals = nvals[:,::-1,:]. But the data is all screwed up. I suspect some error in the indexing but since python returned no errors then I'm doing something with the indexing thinking one thing but getting another.

Perhaps one of you experts can spot something screwy or maybe suggest a better way. I'll attach the image when I figure out how to attach files.


Solution

  • It seems like the result returned by griddata is transposed - it gives the x-axis along the columns and y-axis along the rows:

    import numpy as np
    
    #lats x lons
    a2d=np.arange(20).reshape( (4,5) )
    print a2d
    
    lats=np.arange(4)
    lats2d=np.ones(5)*lats[:,None]
    yi=[1,3]
    nlats=np.sum(np.bitwise_and(lats>=np.min(yi),lats<=np.max(yi)))
    
    lons=np.arange(5)
    lons2d=np.ones(4)[:,None]*lons
    xi=[1,2]
    nlons=np.sum(np.bitwise_and(lons>=np.min(xi),lons<=np.max(xi)))
    
    
    #slice= lats2d>=1 & lats2d<=2 & lons2d>=1 & lons2d<=2
    s1=np.bitwise_and(lats2d>=np.min(yi),lats2d<=np.max(yi))
    s2=np.bitwise_and(lons2d>=np.min(xi),lons2d<=np.max(xi))
    slice=np.bitwise_and(s1,s2)
    print slice
    slice=np.where(slice)
    print a2d[slice].reshape( (nlats,nlons) )
    
    import matplotlib.mlab as mlab
    
    print mlab.griddata(lats2d[slice],lons2d[slice],a2d[slice],
    #              np.array([1.3,2.1,2.9]),np.array([1.1,1.9]))
                np.array([1,2,3]),np.array([1,2]))