Generate Streamplot With 2D Raster Arrays-CodePudding

I am attempting to generate a streamplot in Matplotlib using u-component & v-component winds from a gridded analysis, and am unsure how to compile these arrays into the Streamplot function. These data are two-dimensional, in GeoTIFF format, and is read into Python using xarray/rioxarray as follows:

import rioxarray as rxr
import numpy as np
ugrd = '/path/to/file/rtma2p5_ucomp.tif'
vgrd = '/path/to/file/rtma2p5_vcomp.tif'
        
ucomp = rxr.open_rasterio(ugrd)
ux, uy = np.meshgrid(ucomp['x'], ucomp['y'])
vcomp = rxr.open_rasterio(vgrd)
vx, vy = np.meshgrid(vcomp['x'], vcomp['y'])
     
X, Y = ucomp['x'], ucomp['y']
        
# Uncertainty on U,V
stream = plt.streamplot(X,Y,U,V)

From the Streamplot documentation, I can leverage rioxarray to pass the 1-dimensional grid arrays as the X & Y arguments for the function. However, I am uncertain how to approach the U & V components, given both are 2-dimensional meshed into x and y components after executing numpy meshgrid.

What additional steps can I take to pass each two-dimensional wind component into the function?

CodePudding user response：

The simplest way to do this is definitely to use xarray.Dataset.plot.streamplot:

ds = xr.Dataset({'u': ucomp, 'v': 'vcomp'})
ds.plot.streamplot(x="x", y="y", u="u", v="v")

If you want to instead with matplotlib, you're nearly there! As you mention, the arguments do need to be 1D arrays. np.meshgrid is already producing the x and y 2D coordinate labels (you've assigned them to ux, uy and also vx, vy - I assume these are the same). So for each (x, y) point indexing ucomp and vcomp, there is a correponding set of coordinates (ux[i, j], uy[i, j]). This allows you to reshape all four arrays to get the inputs you need:

plt.streamplot(
    ux.ravel(),
    uy.ravel(),
    ucomp.values.ravel(),
    vcomp.values.ravel(),
)