#!/usr/bin/env python # Wenchang Yang (wenchang@princeton.edu) # Mon Mar 28 19:43:11 EDT 2022 if __name__ == '__main__': import sys from misc.timer import Timer tt = Timer('start ' + ' '.join(sys.argv)) import sys, os.path, os, glob, datetime import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt #more imports from scipy.stats import ttest_ind import xlinregress # if __name__ == '__main__': tt.check('end import') # #start from here #season = ['JJA', 'SON', 'MAM', 'DJF', 'annual'][-1]# 'JJA' season = 'annua' sel_season = lambda x: x if 'JJA' in sys.argv: season = 'JJA' sel_season = lambda da: da.isel(time=(da.time.dt.month>=6)&(da.time.dt.month<=8)) elif 'DJF' in sys.argv: season = 'DJF' def sel_season(da): time_new = da.indexes['time'].shift(1, 'MS') da['time'] = time_new da = da.isel(time=(da.time.dt.month>=1)&(da.time.dt.month<=3)) return da elif 'MAM' in sys.argv: season = 'MAM' sel_season = lambda da: da.isel(time=(da.time.dt.month>=3)&(da.time.dt.month<=5)) elif 'SON' in sys.argv: season = 'SON' sel_season = lambda da: da.isel(time=(da.time.dt.month>=9)&(da.time.dt.month<=11)) else: sel_season = lambda da: da #years = slice('1901', '2014') years = slice('1901', '2020') if len(sys.argv)>1: if '1990-2014' in sys.argv: years = slice('1990', '2014') elif '1901-2014' in sys.argv: years = slice('1901', '2014') elif '1999-2014' in sys.argv: years = slice('1999', '2014') elif '1976-1998' in sys.argv: years = slice('1976', '1998') elif '1901-2020' in sys.argv: years = slice('1901', '2020') elif '1990-2016' in sys.argv: years = slice('1990', '2016') elif '1981-2020' in sys.argv: years = slice('1981', '2020') #else: # years = slice('1901', '2014') region = 'NA' if len(sys.argv)>1: if 'global' in sys.argv: region = 'global' if region == 'NA': sel_region = lambda da: da.sel(lon=slice(360-170, 360-50), lat=slice(10, 80)) elif region == 'global': sel_region = lambda da: da case = 'urbanGlobal' if len(sys.argv)>1: if 'urbanNA' in sys.argv: case = 'urbanNA' elif 'urbanAS' in sys.argv: case = 'urbanAS' #func_mean = lambda da: da.mean('en').mean('year') func_lintrend = lambda da: da.mean('en').linregress.on(da.year, dim='year')['slope'] new_ens = True ten_ens = True# and False label = f'{case} - noUrban, {season}, {years.start}-{years.stop}, {region}' if case == 'urbanGlobal' and new_ens: label += ', new ens' if ten_ens: label = label.replace('new ens', '10 ens') else: label = label.replace('new ens', '5 ens') tag = label.replace(',', '_').replace(' ', '') pnwbox = 360-123, 360-119, 45, 52 ofile = __file__.replace('.py', f'_{tag}.nc') #ofile_pvalue = ofile.replace('.nc', '_pvalue.nc') if os.path.exists(ofile): daa = xr.open_dataarray(ofile) print('[loaded]:', ofile) #pvalue = xr.open_dataarray(ofile_pvalue) #print('[loaded]:', ofile_pvalue) else: ifile = '../data_precip_AM4urban_wasteCool_0urban_amip_5ens_1870-2020.nc' da0 = xr.open_dataarray(ifile).load().pipe(sel_season).sel(time=years).pipe(sel_region) \ .groupby('time.year').mean('time') #more ens if ten_ens: #ifile = './data_precip_AM4urban_wasteCool_0urban_amip_new5ens_1870-2014.nc' ifile = './data_precip_AM4urban_wasteCool_0urban_amip_new5ens_1870-2020.nc' da0_more = xr.open_dataarray(ifile).load().pipe(sel_season).sel(time=years).pipe(sel_region) \ .groupby('time.year').mean('time') da0 = xr.concat([da0, da0_more], dim='en') if case == 'urbanGlobal': ifile = '../data_precip_AM4urban_wasteCool_amip_5ens_1870-2020.nc' da = xr.open_dataarray(ifile).load().pipe(sel_season).sel(time=years).pipe(sel_region) \ .groupby('time.year').mean('time') if new_ens: da = da.sel(en=[0, 21, 22, 23, 24]) else: da = da.sel(ens=[0, 11, 12, 13, 14]) #more ens if ten_ens: #ifile = './data_precip_AM4urban_wasteCool_amip_new5ens_1870-2014.nc' ifile = './data_precip_AM4urban_wasteCool_amip_new5ens_1870-2020.nc' da_more = xr.open_dataarray(ifile).load().pipe(sel_season).sel(time=years).pipe(sel_region) \ .groupby('time.year').mean('time') da = xr.concat([da, da_more], dim='en') elif case == 'urbanNA': ifile = '../data_precip_AM4urban_wasteCool_urbanNAonly_amip_5ens_1870-2020.nc' da = xr.open_dataarray(ifile).load().pipe(sel_season).sel(time=years).pipe(sel_region) \ .groupby('time.year').mean('time') elif case == 'urbanAS': ifile = '../data_precip_AM4urban_wasteCool_urbanASonly_amip_5ens_1870-2020.nc' da = xr.open_dataarray(ifile).load().pipe(sel_season).sel(time=years).pipe(sel_region) \ .groupby('time.year').mean('time') #daa = da.pipe(func_mean) - da0.pipe(func_mean) #daa.to_dataset(name=da0.name).to_netcdf(ofile) #daa = da.pipe(func_lintrend) - da0.pipe(func_lintrend) #daa = da.pipe(func_lintrend) - da0.pipe(func_lintrend) daa1 = da.pipe(func_lintrend) daa0 = da0.pipe(func_lintrend) daa = xr.concat([daa1, daa0], dim='icase') daa = daa.pipe(lambda x: x*100).assign_attrs(units='mm/day/century', long_name='linear trend') daa.to_dataset(name='slope').to_netcdf(ofile) print('[saved]:', ofile) """ #pvalue tvalue, pvalue = ttest_ind(da.stack(s=['en', 'year']), da0.stack(s=['en', 'year']), axis=-1) pvalue = xr.DataArray(pvalue, dims=da.stack(s=['en', 'year']).dims[:-1]) for dim in pvalue.dims: pvalue = pvalue.assign_coords({dim: da[dim].values}) pvalue.to_dataset(name='pvalue').to_netcdf(ofile_pvalue) print('[saved]:', ofile_pvalue) """ if __name__ == '__main__': from wyconfig import * #my plot settings from geoplots import mapplot vmax = 0.4 if region == 'NA' else 1 robust = False levels = 21 #if region == 'NA' else 11 fig, ax = plt.subplots() (daa.isel(icase=0) - daa.isel(icase=1)).assign_attrs(long_name='linear trend', units=daa.attrs['units']) \ .plot(vmax=vmax, center=0, levels=levels, cmap='BrBG', extend='both', robust=robust) mapplot() lon0, lon1, lat0, lat1 = pnwbox plt.plot([lon0, lon1, lon1, lon0, lon0], [lat0, lat0, lat1, lat1, lat0], color='gray') plt.title(label) #savefig if len(sys.argv)>1 and 'savefig' in sys.argv[1:]: figname = __file__.replace('.py', f'_{tag}.png') if 'overwritefig' in sys.argv[1:]: wysavefig(figname, overwritefig=True) else: wysavefig(figname) fig, ax = plt.subplots() daa.isel(icase=0).assign_attrs(long_name='linear trend') \ .plot(vmax=vmax, center=0, levels=levels, cmap='BrBG', extend='both', robust=robust) mapplot() lon0, lon1, lat0, lat1 = pnwbox plt.plot([lon0, lon1, lon1, lon0, lon0], [lat0, lat0, lat1, lat1, lat0], color='gray') plt.title(label.replace(' - noUrban', '')) #savefig if len(sys.argv)>1 and 'savefigs' in sys.argv[1:]: tag_new = tag.replace('-noUrban', '') figname = __file__.replace('.py', f'_{tag_new}.png') if 'overwritefig' in sys.argv[1:]: wysavefig(figname, overwritefig=True) else: wysavefig(figname) fig, ax = plt.subplots() daa.isel(icase=1).assign_attrs(long_name='linear trend') \ .plot(vmax=vmax, center=0, levels=levels, cmap='BrBG', extend='both', robust=robust) mapplot() lon0, lon1, lat0, lat1 = pnwbox plt.plot([lon0, lon1, lon1, lon0, lon0], [lat0, lat0, lat1, lat1, lat0], color='gray') plt.title(label.replace(f'{case} - ', '')) #savefig if len(sys.argv)>1 and 'savefigs' in sys.argv[1:]: tag_new = tag.replace(f'{case}-', '') figname = __file__.replace('.py', f'_{tag_new}.png') if 'overwritefig' in sys.argv[1:]: wysavefig(figname, overwritefig=True) else: wysavefig(figname) tt.check(f'**Done**') print() plt.show()