#!/usr/bin/env python # Wenchang Yang (wenchang@princeton.edu) # Fri Mar 3 13:21:20 EST 2023 if __name__ == '__main__': import sys,os from misc.timer import Timer tt = Timer(f'[{os.getcwd()}] 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 misc import get_kws_from_argv # if __name__ == '__main__': tt.check('end import') # #start from here model = 'FLOR' case_default = """ p2p0solar p6p0solar m2p0solar m6p0solar co2xp5 co2xp25 co2x4 co2x2 """.split()[-1] TsOption_default = """ glbmean oceanmean tropoceanmean """.split()[-1] #tspan = slice('0101', '0600') idir = '/tigress/wenchang/analysis/FLOR' ifiles = { 'co2x2': 'precip_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc', 'co2x4': 'precip_FLOR_co2x4_CTL1860_tigercpu_intelmpi_18_576PE_0101-0250_glbmean.nc', 'co2xp5': 'precip_FLOR_co2xp5_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc', 'co2xp25': 'precip_FLOR_co2xp25_CTL1860_tigercpu_intelmpi_18_576PE_0101-0250_glbmean.nc', 'p1p0solar': 'precip_FLOR_p1p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc', 'p2p0solar': 'precip_FLOR_p2p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0400_glbmean.nc', 'p6p0solar': 'precip_FLOR_p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc', 'm2p0solar': 'precip_FLOR_m2p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0400_glbmean.nc', 'm6p0solar': 'precip_FLOR_m6p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0597_glbmean.nc', } def wyget(case=case_default, TsOption=TsOption_default): #get data ofile = __file__.replace('.py', f'__{model}_{case}_{TsOption}Ts.nc') if not os.path.exists(ofile) or 'od' in sys.argv: #nonlinear_ratio = True #AM2.5 values: see: http://tigress-web.princeton.edu/~wenchang/pub/AM2.5/fig_bars_dprecip_dgmst_florSST.png #and http://tigress-web.princeton.edu/~wenchang/pub/AM2.5/fig_bars_dprecip_dF_florSST.png if TsOption == 'glbmean': eta = 3.28 # % per K elif TsOption == 'tropmean': eta = 3.24 # % per K else: eta = 3.52 #% per K A_co2 = -2.13 # % per 2xCO2 forcing A_solar = -0.85 # % per 2% solar from AM2.5 CTL1860_florSST #precip ctl ifile = 'precip_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_0001-2301_glbmean.nc' ifile = os.path.join(idir, ifile) print(ifile) da = xr.open_dataarray(ifile) da_ctl = da da_ctl_clim = da.sel(time=slice('0101', '0200')).groupby('time.month').mean('time') #monthly clim da_ref = da_ctl_clim.mean('month') #precip pert ifile = ifiles[case] ifile = os.path.join(idir, ifile) print(ifile) da = xr.open_dataarray(ifile) if case == 'p6p0solar': da = da.sel(time=slice(None, '0600')) #p6p0solar is still running, select years within 0600 da_precip = da#.sel(time=slice('0101', '0250')) #monthly ts #Ts ctl #ifile = glob.glob('t_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_????-????_glbmean.nc')[0] if TsOption == 'glbmean': ifile = glob.glob( os.path.join(idir, 't_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_????-????_glbmean.nc') )[0] else: ifile = glob.glob( os.path.join(idir, TsOption, f't_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_????-????_{TsOption}.nc') )[0] print(ifile) da = xr.open_dataarray(ifile) da_ctl_ts = da da_ctl_clim_ts = da.sel(time=slice('0101', '0200')).groupby('time.month').mean('time') da_ref_ts = da_ctl_clim_ts.mean('month') #Ts pert ifile = ifiles[case].replace('precip', 't_surf') if TsOption == 'glbmean': ifile = os.path.join(idir, ifile) else: ifile = ifile.replace('glbmean', TsOption) ifile = os.path.join(idir, TsOption, ifile) if case == 'p6p0solar': ifile = glob.glob(ifile.replace('0600', '????'))[0] #p6p0solar is still running print(ifile) da_ts = xr.open_dataarray(ifile)#.sel(time=slice('0101', '0250')) if case == 'p6p0solar': da_ts = da_ts.sel(time=slice(None, '0600')) #p6p0solar is still running, select years within 0600 #save ds = xr.Dataset(dict(eta=eta, A_co2=A_co2, A_solar=A_solar, pr=da_precip, pr_ctl_clim=da_ctl_clim, Ts=da_ts, Ts_ctl_clim=da_ctl_clim_ts)) ds.to_netcdf(ofile) print('[saved]:', ofile) else: ds = xr.open_dataset(ofile) print('[loaded]:', ofile) return ds def wyplot(case=case_default, TsOption=TsOption_default, ax=None, legend_on=True, lowpass_on=True, residual_on=True, slide=None, ncol_legend=1, **kws): if ax is None: fig,ax = plt.subplots() ds = wyget(case=case, TsOption=TsOption) #retrieve data #ds = ds.sel(time=tspan) eta = ds['eta'].item() A_co2 = ds['A_co2'].item() A_solar = ds['A_solar'].item() da_precip = ds['pr'] da_ctl_clim = ds['pr_ctl_clim'] da_ref = da_ctl_clim.mean('month') da_ts = ds['Ts'] da_ctl_clim_ts = ds['Ts_ctl_clim'] da_ref_ts = da_ctl_clim_ts.mean('month') #Adjustment component time series (usually constant) if case == 'co2x2': da_fco2 = da_ts*0+1 da_solar = da_ts*0 elif case == 'co2x4': da_fco2 = da_ts*0+2 da_solar = da_ts*0 elif case == 'co2xp5': da_fco2 = da_ts*0-1 da_solar = da_ts*0 elif case == 'co2xp25': da_fco2 = da_ts*0-2 da_solar = da_ts*0 elif case == 'p2p0solar': da_fco2 = da_ts*0 da_solar = da_ts*0+1 elif case == 'p6p0solar': da_fco2 = da_ts*0 da_solar = da_ts*0+3 elif case == 'p1p0solar': da_fco2 = da_ts*0 da_solar = da_ts*0+0.5 elif case == 'm2p0solar': da_fco2 = da_ts*0 da_solar = da_ts*0-1 elif case == 'm6p0solar': da_fco2 = da_ts*0 da_solar = da_ts*0-3 func_lp = lambda x: x.rolling(year=9, center=True, min_periods=1).mean('year') func_pct = lambda x: (x.groupby('time.year').mean('time')/da_ref*100-100).assign_attrs(units='%') #for precip func_anom = lambda x: (x.groupby('time.year').mean('time') - da_ref_ts).assign_attrs(units='K') #for Ts alpha = 0.2 #precip #label = f'precip ({model} {case})' if slide is None or slide>=1: label = '$\Delta P$' dP = da_precip.pipe(func_pct) if lowpass_on: dP.plot(color='k', alpha=alpha, ax=ax) dP.pipe(func_lp).plot(label=label, color='k', ax=ax) else: dP.plot(label=label, color='k', ax=ax) ax.set_prop_cycle(None) y = dP #save data for later use if slide is None or slide>=2: ii = 0; cc = f'C{ii}' #from gmsta #label = 'gmst contribution' label = '$\eta_{+2K}\cdot\Delta T_s$' dTs = da_ts.pipe(func_anom) da_ = eta * dTs if lowpass_on: da_.plot(color=cc, alpha=alpha, ax=ax) #da_.pipe(func_lp).plot(label='gmst contribution', color=cc, ls='--') #da_.pipe(func_lp).plot(label=label, color=cc, ls='--', ax=ax) da_.pipe(func_lp).plot(label=label, color=cc, lw=1, ax=ax) else: da_.plot(label=label, color=cc, lw=1, ax=ax) if slide is None or slide>=3: #from direct rad forcing label = '$A$' ii += 1; cc = f'C{ii}' if case in ('p2p0solar', 'p6p0solar', 'p1p0solar', 'm2p0solar', 'm6p0solar'): #label = 'solar contribution' da_ = da_solar.groupby('time.year').mean('time')*A_solar #da_.plot(color=cc, alpha=alpha) #da_.plot(label='solar contribution', ls='--', color=cc) else: # label = 'co2 contribution' da_ = da_fco2.groupby('time.year').mean('time')*A_co2 #da_.plot(color=cc, alpha=alpha) #da_.plot(label='co2 contribution', ls='--', color=cc) #da_.plot(label=label, ls='--', color=cc, ax=ax) da_.plot(label=label, ls='--', lw=1, color=cc, ax=ax) if slide is None or slide>=4: #modeled label = '$\eta_{+2K}\cdot\Delta T_s + A$' ii += 1; cc = f'C{ii}' dTs = da_ts.pipe(func_anom) if case in ('p2p0solar', 'p6p0solar', 'p1p0solar', 'm2p0solar', 'm6p0solar'): #label = f'modeled by gmst({eta:.1f}%/K) and solar({A_solar:.1f}%/2%solar)' da_ = eta*dTs + da_solar.groupby('time.year').mean('time')*A_solar else: #label = f'modeled by gmst({eta:.1f}%/K) and co2({A_co2:.1f}%/2xCO2)' da_ = eta*dTs + da_fco2.groupby('time.year').mean('time')*A_co2 #da_.pipe(func_lp).plot(label=f'modeled by gmst({eta:.1f}%/K) and co2({A_co2:.1f}%/2xCO2)', color=cc) if lowpass_on: da_.plot(color=cc, alpha=alpha, ax=ax) da_.pipe(func_lp).plot(label=label, color=cc, ax=ax) else: da_.plot(label=label, color=cc, ax=ax) y_hat = da_ #save #residual if residual_on: #show the residual label = 'residual = $\Delta P - (\eta\cdot T_s + A)$' ii += 1; cc = f'C{ii}' da_res = y - y_hat if lowpass_on: da_res.plot(color=cc, alpha=alpha, ax=ax) #da_res.pipe(func_lp).plot(label=label, color=cc, ls='--', ax=ax) da_res.pipe(func_lp).plot(label=label, color=cc, lw=1, ax=ax) else: #da_res.plot(label=label, color=cc, ls='--', ax=ax) da_res.plot(label=label, color=cc, lw=1, ax=ax) if legend_on: ax.legend(loc='upper left', ncol=ncol_legend) ax.set_ylabel('global mean precip change [%]') #ax.set_xticks(range(1000, 1300, 40)) #ax.axvline(1140, color='gray', ls='--') ax.axhline(0, color='gray', ls='--') #if case in ('4xCO2',): # plt.xscale('log') #ax.set_xlim(100, 250) #ax.set_xticks(np.arange(100,251,10)) #ax.set_xticklabels(np.arange(0,151,10)) #ax.set_xticks(np.arange(100,601,50)) #ax.set_xticklabels(np.arange(0,501,50)) ax.set_xlim(100, dP.year.max()) xticks = ax.get_xticks() ax.set_xticklabels([int(x-100) for x in xticks]) if __name__ == '__main__': from wyconfig import * #my plot settings case = get_kws_from_argv('case', default=case_default) TsOption = get_kws_from_argv('TsOption', default=TsOption_default) slide = get_kws_from_argv('slide', None) if slide is not None: slide = int(slide) #wyplot(case=case, TsOption=TsOption, lowpass_on=False) wyplot(case=case, TsOption=TsOption, lowpass_on=True, residual_on=False, slide=slide) ax = plt.gca() ax.set_ylim(-3,11) #savefig if 'savefig' in sys.argv or 's' in sys.argv: figname = __file__.replace('.py', f'__{model}_{case}_{TsOption}.png') if slide is not None: figname = figname.replace('.png', f'_slide{slide:02d}.png') if 'pdf' in sys.argv: figname = figname.replace('.png', '.pdf') if 'overwritefig' in sys.argv or 'o' in sys.argv: wysavefig(figname, overwritefig=True) else: wysavefig(figname) """ # import seaborn as sns import xlinregress fig,ax = plt.subplots() dprecip = y dgmst = da_ts.groupby('time.year').mean('time') - da_ref_ts years = slice(101, 250) #years = slice(101, 121) #years = slice(101, 151) #years = slice(121, 251) dprecip = dprecip.sel(year=years) dgmst = dgmst.sel(year=years) #dprecip = dprecip.isel(year=dgmst<1.35) #dgmst = dgmst.isel(year=dgmst<1.35) reg = dprecip.linregress.on(dgmst) slope, y0 = reg.slope, reg.intercept sns.scatterplot(x=dgmst, y=dprecip, color='.1') #sns.histplot(x=dgmst, y=dprecip) #sns.kdeplot(x=dgmst, y=dprecip, fill=True) ax.axline((0,y0), slope=slope, color='C1') ax.text(0, 1, f' {model} {case} gmst\n y0 = {y0.item():.1f}\n b = {slope.item():.1f}', transform=ax.transAxes, va='top') ax.set_xlabel('Tsa [K]') ax.set_ylabel('precip anom [%]') #savefig if 'savefig' in sys.argv or 's' in sys.argv: figname = __file__.replace('.py', f'__{model}_{case}_linregress{years.start}-{years.stop}.png') if 'pdf' in sys.argv: figname = figname.replace('.png', '.pdf') if 'overwritefig' in sys.argv or 'o' in sys.argv: wysavefig(figname, overwritefig=True) else: wysavefig(figname) """ tt.check(f'**Done**') print() if 'notshowfig' in sys.argv or 'n' in sys.argv: pass else: if 'plt' in globals(): plt.show()