reanalysis_clim = []
reanalysis_ls = []
c = []
for i in list(np.arange(24,yearmax,1)):
di = d.where(d.MY == i, drop=True)
print(i)
di["Ls"] = di.Ls.sel(lat=di.lat[0]).sel(lon=di.lon[0])
di = di.transpose("lat","lon","time")
di = di.sortby("lat", ascending = True)
di = di.sortby("lon", ascending = True)
# Lait scale PV
theta = ilev
laitPV = funcs.lait(di.PV, theta, theta0, kappa = kappa)
di["scaled_PV"] = laitPV
Zi = funcs.calc_eddy_enstr(di.scaled_PV) * 10**8
if i != 28:
reanalysis_clim.append(Zi)
Ls = di.Ls
if i != 28:
reanalysis_ls.append(Ls)
Zi = Zi.chunk({'time':'auto'})
Zi = Zi.rolling(time=smooth,center=True)
Zi = Zi.mean()
ds = xr.open_mfdataset(inpath + 'OpenMARS_Ls270-300_zonal.nc',
decode_times=False)
ds = ds.where(Lsmin <= ds.Ls, drop=True)
ds = ds.where(ds.Ls <= Lsmax, drop=True)
ds["plev"] = ds.plev / 100
ds = ds.where(ds.lat >= 0, drop=True)
t1 = ds.tmp.mean(dim='time')
u1 = ds.uwnd.mean(dim='time')
theta1 = ds.theta.mean(dim='time')
lait1 = funcs.lait(ds.PV, ds.theta, theta0, kappa=kappa)
lait1 = lait1.mean(dim='time') * 10**4
p = ds.plev
lat = ds.lat
#marr = np.ma.masked_array(lait1, np.isnan(lait1),fill_value=-999)
#arr = lait1.where(lait1.plev < 5.5, drop=True)
arr = lait1.load()
lats_max = []
for ilev in range(len(arr.plev)):
marr = arr.sel(plev=arr.plev[ilev])
marr_max = marr.max().values
marr = marr.where(marr >= marr_max, drop=True)
lat_max = marr.lat.values
# reduce dataset
d = data_reduct(d, Lsmin, Lsmax, latmin, latmax)
dp = data_reduct(dp, Lsmin, Lsmax, latmin, latmax)
if jetmax == False:
d = d.mean(dim="lat")
dp = dp.mean(dim="lat")
d = d.sel(ilev=ilev, method='nearest').squeeze()
dp = dp.sel(pfull=pfull / 100, method='nearest').squeeze()
# Lait scale PV
theta = d.ilev
print("Scaling PV")
laitPV = funcs.lait(d.PV, theta, theta0, kappa=kappa)
d["scaled_PV"] = laitPV
x, index = funcs.assign_MY(d)
dsr, N, n = funcs.make_coord_MY(x, index)
Ls = dsr.mars_solar_long.mean(dim='MY')
year_mean = dsr.mean(dim='MY')
year_mean = year_mean.chunk({'new_time': 'auto'})
#year_mean = year_mean.rolling(new_time = 4, center = True)
#year_mean = year_mean.mean()
mean_PV = year_mean.scaled_PV * 10**4
mean_wind = year_mean.uwnd
if jetmax == True:
ax.text(-0.05, 1.05, string.ascii_lowercase[i], transform=ax.transAxes,
size=22, weight='bold')
ax.yaxis.set_major_formatter(ticker.ScalarFormatter())
if i < 3:
my = i+24
else:
my = i+25
print(my)
d_O = d_O1.where(d_O1.MY==my, drop=True)
d_O = d_O.where(Lsmin <= d_O.Ls, drop = True)
d_O = d_O.where(d_O.Ls <= Lsmax, drop = True)
d_O = d_O.where(d_O.lat >= 0, drop = True)
lait_O = funcs.lait(d_O.PV, d_O.theta, theta_0, kappa=kappa)
pv_o = lait_O.mean(dim='time').mean(dim='lon') *10**4
t_o = d_O.theta.mean(dim='time').mean(dim='lon')
u_o = d_O.uwnd.mean(dim='time').mean(dim='lon')
lats_max = []
arr = pv_o.load()
for jlev in range(len(arr.plev)):
marr = arr.sel(plev=arr.plev[jlev])
marr_max = marr.max().values
marr = marr.where(marr >= marr_max,drop=True)
lat_max = marr.lat.values
#lat_max, max_mag = calc_PV_max(marr[:,ilev], lait1.lat)
lats_max.append(lat_max)