Search code examples
pythonmatplotlibmultidimensional-arrayplotimshow

Plotting a 2d contour plot in python with sparse data


I have some output data from an ocean circulation model (MITgcm). It is an idealised channel (Cartesian) so the geometry is not confusing, luckily.

I'd like to plot some of the fields (velocity, temperature etc.) in the y-z plane. The simulation domain involves 30 vertical layers where each layer is an 800x400 y-x grid. I have each of the fields stored in numpy arrays with shape (30, 800, 400) going z,y,x respectively.

I can easily plot x-y plane slices for the 30 vertical levels. I can do this using matplotlib's contourf or imshow and changing the extent to the correct physical values in km.

The problem is that the vertical layers are unevenly spaced. I have the grid data for Z which tells me what physical depth (in metres) each of the layers corresponds to.

Z is: [-5. -15. -25. -36. -49. -65. -84. -105.5 -130.5 -159.5 -192.5 -230. -273. -322.5 -379. -443. -515. -596. -688. -792. -909.5 -1042.5 -1192.5 -1362. -1553.5 -1770. -2015. -2285. -2565. -2845.]

I tried to get round this by creating an empty matrix with 2985 (as the full domain depth is 2985m) 'vertical' layers, and inputting the y-data at the corresponding positions for the 30 layers as given by Z above (here yz_zonal is a (30,800) matrix of data values):

yz_matrix = np.empty((2985, 800)) #empty matrix for yz-plane data, vertical extent is 2985 (m)

for i in range(len(Z)):
     yz_matrix[round(-Z[i])] = yz_zonal[i] #set matrix values to correct depths

Then if I try to plot yz_matrix using matplotlib's imshow, by doing:

fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_xlabel('y (km)')
ax.set_ylabel('z (m)')
yzplot = ax.imshow(yz_matrix, aspect='auto', interpolation='gaussian', cmap='inferno', extent=[0,2000,-2985,0])
plt.colorbar(yzplot)

I just get this figure:BAD y-z plot of velocity data

There are 30 strips of data values at the correct physical z positions, but there's a whole load of zeros in between them. I only want to interpolate the data in between the 30 strips and ignore all of the other points.

It would be brilliant if anyone could sort this out for me. Thanks in advance!

Peter


Solution

  • You may directly plot the yz_matrix as a pcolormesh, giving a meshgrid of the z and y data as coordinates. This would lead to different sized cells which extent up to next value in z. See left picture below.

    You may also interpolate your data on a new finer grid. To this end, scipy.interpolate.griddata may be used.

    enter image description here

    import numpy as np
    import matplotlib.pyplot as plt
    from scipy.interpolate import griddata
    
    z = np.array([-5.,-15.,-25.,-36.,-49.,-65.,-84.,-105.5,-130.5,-159.5,-192.5,
                  -230.,-273.,-322.5,-379.,-443.,-515.,-596.,-688.,-792.,-909.5,
                  -1042.5,-1192.5,-1362.,-1553.5,-1770.,-2015.,-2285.,-2565.,-2845.])
    y = np.arange(0,100)
    yz_matrix = np.cumsum(np.random.rand(len(z), len(y)), axis=0)
    
    fig, (ax, ax2) = plt.subplots(ncols=2)
    
    # plot raw data as pcolormesh
    Y,Z = np.meshgrid(y,z[::-1])
    ax.pcolormesh(Y,Z, yz_matrix, cmap='inferno')
    ax.set_title("pcolormesh data")
    
    # now interpolate data to new grid 
    zi = np.arange(-2845,-5)
    YI,ZI = np.meshgrid(y,zi)
    points = np.c_[Y.flatten(),Z.flatten()]
    
    interp = griddata(points, yz_matrix.flatten(), (YI,ZI), method='linear')
    
    ax2.pcolormesh(YI,ZI, interp, cmap='inferno')
    ax2.set_title("pcolormesh interpolated")
    
    plt.show()