# print('saved')
# except:
# pass
# %%
#ws_ccmp1=xr.open_dataset('datasets/CCMP_windspeed.nc')
#wu=xr.open_dataset('datasets/uwnd.10m.mon.mean.nc').sel(level=10).uwnd
#wv=xr.open_dataset('datasets/vwnd.10m.mon.mean.nc').sel(level=10).vwnd
ws_ccmp=xr.open_dataset('processed/CCMP_ws_1deg_global.nc')
wu=ws_ccmp.uwnd
wv=ws_ccmp.vwnd
# %% Test Horizontal Divergence
w = VectorWind(wu, wv)
#spd = w.magnitude()
divergence = w.divergence().sel(lat=slice(20,-20),lon=slice(120,290),time=slice('1997-07-01','2020-01-01'))
#div.mean(dim='time').plot()
wu=wu.sel(lat=slice(-20,20),lon=slice(120,290),time=slice('1997-07-01','2020-01-01'))
wv=wv.sel(lat=slice(-20,20),lon=slice(120,290),time=slice('1997-07-01','2020-01-01'))
# %% Prepare Figure
lanina=pd.read_csv('processed/indexes/la_nina_events.csv')
cp_nino=pd.read_csv('processed/indexes/cp_events.csv')
ep_nino=pd.read_csv('processed/indexes/ep_events.csv')
fp='processed/combined_dataset/month_data_exports.nc'
info=xr.open_mfdataset(fp).sel(Mooring=195).to_dataframe()
mpl.rcParams['mathtext.default'] = 'regular'
# Read zonal and meridional wind components from file using the xarray module.
# The components are in separate files.
ds = xr.open_mfdataset(
[example_data_path(f) for f in ('uwnd_mean.nc', 'vwnd_mean.nc')])
uwnd = ds['uwnd']
vwnd = ds['vwnd']
# Create a VectorWind instance to handle the computations.
w = VectorWind(uwnd, vwnd)
# Compute components of rossby wave source: absolute vorticity, divergence,
# irrotational (divergent) wind components, gradients of absolute vorticity.
eta = w.absolutevorticity()
div = w.divergence()
uchi, vchi = w.irrotationalcomponent()
etax, etay = w.gradient(eta)
etax.attrs['units'] = 'm**-1 s**-1'
etay.attrs['units'] = 'm**-1 s**-1'
# Combine the components to form the Rossby wave source term.
S = eta * -1. * div - (uchi * etax + vchi * etay)
# Pick out the field for December at 200 hPa.
S_dec = S[S['time.month'] == 12]
# Plot Rossby wave source.
clevs = [-30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30]
ax = plt.subplot(111, projection=ccrs.PlateCarree(central_longitude=180))
S_dec *= 1e11
# Read zonal and meridional wind components from file using the xarray module.
# The components are in separate files.
ds = xr.open_mfdataset([example_data_path(f)
for f in ('uwnd_mean.nc', 'vwnd_mean.nc')])
uwnd = ds['uwnd']
vwnd = ds['vwnd']
# Create a VectorWind instance to handle the computations.
w = VectorWind(uwnd, vwnd)
# Compute components of rossby wave source: absolute vorticity, divergence,
# irrotational (divergent) wind components, gradients of absolute vorticity.
eta = w.absolutevorticity()
div = w.divergence()
uchi, vchi = w.irrotationalcomponent()
etax, etay = w.gradient(eta)
etax.attrs['units'] = 'm**-1 s**-1'
etay.attrs['units'] = 'm**-1 s**-1'
# Combine the components to form the Rossby wave source term.
S = eta * -1. * div - (uchi * etax + vchi * etay)
# Pick out the field for December at 200 hPa.
S_dec = S[S['time.month'] == 12]
# Plot Rossby wave source.
clevs = [-30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30]
ax = plt.subplot(111, projection=ccrs.PlateCarree(central_longitude=180))
S_dec *= 1e11
