Plot streamlines on a matplotlib contourf plot

Question

I have a matplotlib contourf plot of longitudes and pressure levels in the vertical. I am trying to plot streamlines on this using the plt.streamplot function in matplotlib and using U and V wind data. If I plot only the streamplot, it works fine. But I cannot get the streamlines to overlay on the contour plot. Here is my code:-

fig, axes = plt.subplots(nrows, ncols, sharex=True, sharey=True)
if (nrows==1 and ncols==1):
    axes=[axes]
else:
    axes=axes.flat
for i, ax in enumerate(axes):
    X,Y = np.meshgrid(x[i],y[i])
    levels=np.arange(vmin,vmax,step)
    h = ax.contourf(X,Y,z[i],cmap=cmap,levels=levels,extend='both')
    w = ax.streamplot(X, Y, W[i], Z[i], linewidth=0.2, color='gray')

And this is the plot I get:

The following is the streamline plot, not sure why the y axis is from 0-120 instead of 0 to 1000:

Answer 1

You use curvilinear coordinate system for contour plot (lat-p). You have to convert u,v to coordinate system of contour something like here (this is example for lat-lon you have to modify it to use pressure levels):

def myStreamPlot(lon,lat,u,v,color='k',density=2.5):
        from scipy.interpolate import griddata

        n,m = u.shape[1],u.shape[0]
        x = np.linspace(np.nanmin(lon), np.nanmax(lon), n)
        y = np.linspace(np.nanmin(lat), np.nanmax(lat), m)
        xi, yi = np.meshgrid(x,y)

        lon = lon.ravel()
        lat = lat.ravel()
        u   = u.ravel()
        v   = v.ravel()

        gu = griddata(zip(lon,lat), u, (xi,yi))
        gv = griddata(zip(lon,lat), v, (xi,yi))
        gspd = np.sqrt(gu**2 + gv**2)
        SL = plt.streamplot(x,y,gu,gv,linewidth=1.,color=color,density=density)

This code use griddata function of scipy.interpolate: https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.griddata.html