napari_test

import numpy as np
import xarray as xr
from scipy.interpolate import griddata
import napari
import matplotlib.pyplot as plt
from vispy.color import Colormap
import cmocean.cm as cm

# Make land colormap with transparent ocean:
cmap = plt.cm.gist_earth
cmap_arr = cmap(np.arange(cmap.N))
cmap_arr[:10,-1] = 0
ls_colormap = Colormap(cmap_arr)

#Load Orography
orog_disk= xr.open_dataset("../../Downloads/orography.nc")
orog_arr=orog_disk.to_array()
lsmask=orog_arr.copy()
lsmask.data[lsmask.data<100]=-8000

#Load ERA5 Wind Data
ds_disk = xr.open_dataset("../../Downloads/era_wind_data2.nc")
u_arr,v_arr=ds_disk.to_array()
del v_arr

#Keep only the lower levels
u_arr_interp=u_arr.sel(level=np.arange(50,1001,50))


#We want a lower and upper section of the atmosphere to get the two jets,
#but they both need to be defined on all levels

x=u_arr_interp.data.copy()
u_arr_interp_l=u_arr_interp.copy()
x[:,:9]=0
u_arr_interp_l.data=x

x=u_arr_interp.data.copy()
u_arr_interp_h=u_arr_interp.copy()
x[:,11:]=0
u_arr_interp_h.data=x


#Now the pretty visualisation

default_angles=(5.34,-17.23,-32.51)
default_zoom=1.27
cmap_bounds=[-50,50]


#Add wind field
viewer=napari.view_image(u_arr_interp[:,::-1,::-1],colormap=u_colormap,\
                        ndisplay=3,contrast_limits=cmap_bounds,scale=[1,4,1,1])
#Add land
viewer.add_image(lsmask.data[0,0][::-1],colormap=ls_colormap,opacity=0.4)

#Add isolevels for lower and upper jet
viewer.add_image(u_arr_interp_l[:,::-1,::-1],colormap="gray",\
    contrast_limits=cmap_bounds,scale=[1,4,1,1],rendering="iso",iso_threshold=0.87)
viewer.add_image(u_arr_interp_h[:,::-1,::-1],colormap="bop orange",\
    contrast_limits=cmap_bounds,scale=[1,4,1,1],rendering="iso",iso_threshold=0.886)

#Make it look great:
viewer.text_overlay.visible = True
viewer.layers.selection.active=viewer.layers[0]
viewer.camera.angles=default_angles
viewer.camera.zoom=default_zoom

#This bit means that when you move the sliders, the xarray data
#gets interrogated for the new coordinates
@viewer.dims.events.current_step.connect
def update_text(event):
    points=viewer.dims.point
    data=viewer.layers.selection.active.data
    #Invert coord order to match napari
    coord_names=list(data.dims)
    points
    coord_vals=[data[coord] for coord in coord_names]
    viewer_position=[cv[int(p)] for cv,p in zip(coord_vals,points)]
    #deal with datetimes
    viewer_position=[vp.item() if vp.dtype!="<M8[ns]" else np.datetime_as_string(vp,unit="h") for vp in viewer_position]
    for i in viewer.dims.displayed:
        viewer_position[i]="All"
    annotation="".join([cn+": "+str(vp)+" " for cn,vp in zip(coord_names,viewer_position)])
    viewer.text_overlay.text = annotation

@joshdorrington , I'd love to try making this a reproducible notebook accessing cloud hosted public data, but just so I can understand a bit better would it be possible to get the orography.nc and the era_wind_data2.nc input files? Feel free to email me also at rsignell usgs.gov!