#!/usr/bin/env python # Wenchang Yang (wenchang@princeton.edu) # Tue Feb 2 22:16:22 EST 2021 if __name__ == '__main__': from misc.timer import Timer tt = Timer(f'start {__file__}') import sys, os.path, os, glob, datetime import xarray as xr, numpy as np, pandas as pd import matplotlib.pyplot as plt #more imports import cartopy.crs as ccrs, cartopy.feature as cfeature #import xfilter, xlearn, geoxarray #from shared import ess_in_corr #from mystats import p2r import xfilter, xlinregress, geoxarray # if __name__ == '__main__': tt.check('end import') # #start from here # params # data # tc liz #ifile = '/tigress/gvecchi/ANALYSIS/LMR_Paper1/normTCcounts_v0719.txt' #ifile = '/tigress/gvecchi/ANALYSIS/LMR_Paper1/normTCcounts_v1219.txt' #tc_liz = pd.read_csv(ifile, sep='\s+', index_col=0) # FLOR tc #ifile = '/tigress/wenchang/analysis/TC/CTL1860_newdiag_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.TS17.0001-2000.yearly.nc' #tc = xr.open_dataset(ifile)['NA'].load() # lme 0850cntl TC #ifile = 'data/LME_TC.0850cntl.fullForcingRef.nc' #tc = xr.open_dataset(ifile)['HU'].load() # lme fullForcing TC ifile = 'data/LME_TC.fullForcing.13ens.fullForcingRef.nc' tc = xr.open_dataset(ifile)['HU'].load() # sst #ifile = '/tigress/gvecchi/DATA/LMR_2019/sst_MCruns_ensemble_mean.nc' #ds_sst = xr.open_dataset(ifile) #FLOR ts #ifile = '/tigress/wenchang/MODEL_OUT/CTL1860_newdiag_tigercpu_intelmpi_18_576PE/analysis/sst.0001-2000.yearly.nc' #ds_sst = xr.open_dataset(ifile) #lme 0850cntl sst #ifile = 'data/b.e11.B1850C5CN.f19_g16.0850cntl.001.cam.h0.TS.0850-2005.nc' #ds_sst = xr.open_dataset(ifile) #lme fullForcing sst ifile = 'data/b.e11.BLMTRC5CN.f19_g16.001-013.cam.h0.TS.0850-2005.nc' ds_sst = xr.open_dataset(ifile) year0, year1 = 850, 1850 n_window = 40 padtype = 'even' datafile = os.path.join( os.path.dirname(__file__), 'figdata', os.path.basename(__file__).replace('.py', f'.{year0:04d}-{year1:04d}_lp{n_window}{padtype}.nc') ) try: # load dataset for figure if exists ds = xr.open_dataset(datafile) print('[loaded]:', datafile) except: # calculate in place if not exists alpha = 0.05 da = ds_sst['sst'] \ .filter.lowpass(1.0/n_window, dim='year', padtype=padtype) \ .sel(year=slice(year0, year1)) #da_rsst = da.pipe(lambda x: x-x.sel(lat=slice(-30, 30)).geo.fldmean()) ts = tc.filter.lowpass(1/n_window, dim='year', padtype=padtype) \ .sel(year=slice(year0, year1)) print('calculating...') ds = da.linregress.numba(ts, dim='year', ess_on=True) ds.to_netcdf(datafile) print('[saved]:', datafile) def wyplot(ax=None, rsst_on=False, regress_on=False): land_color = 'k' cmap = ['turbo', 'RdBu_r'][-1] hatches = ['...'] hatch_color = 'gray' hatch_color_default = plt.rcParams['hatch.color'] plt.rcParams['hatch.color'] = hatch_color levels_cc = np.arange(-1, 1.01, 0.1) levels_b = np.arange(-0.12, 0.121, 0.02) if ax is None: proj = ccrs.Robinson(central_longitude=180) fig,ax = plt.subplots(subplot_kw=dict(projection=proj)) dproj = ccrs.PlateCarree() if rsst_on:#use relative sst instead of sst if regress_on:#show regression slope instead of correlation coefficient ds.b_rsst.plot.contourf(ax=ax, transform=dproj, add_colorbar=True, cmap=cmap, levels=levels_b, extend='both', cbar_kwargs=dict(label='slope [K per STD of NA TC') ) title = 'b(rSST, TC)' else:#correlation coefficient for rsst ds.r_rsst.plot.contourf(ax=ax, transform=dproj, add_colorbar=True, cmap=cmap, levels=levels_cc, cbar_kwargs=dict(label='correlation') ) title = 'corr(rSST, TC)' #significant info ds.r_rsst.where(np.abs(ds.r_rsst)>=ds.r_sig_rsst).pipe(lambda x: x*0).plot.contourf(ax=ax, transform=dproj, add_colorbar=False, colors='none', hatches=hatches, ) else:#use sst instead of rsst if regress_on:#show regression slope instead of correlation coefficient ds.b.plot.contourf(ax=ax, transform=dproj, add_colorbar=True, cmap=cmap, levels=levels_b, extend='both', cbar_kwargs=dict(label='slope [K per STD of NA TC') ) title = 'b(SST, TC)' else:#correlation coefficient for sst ds.r.plot.contourf(ax=ax, transform=dproj, add_colorbar=True, cmap=cmap, levels=levels_cc, cbar_kwargs=dict(label='correlation') ) title = 'corr(SST, TC)' #significant info ds.r.where(np.abs(ds.r)>=ds.r_sig).pipe(lambda x: x*0).plot.contourf(ax=ax, transform=dproj, add_colorbar=False, colors='none', hatches=hatches, ) ax.add_feature(cfeature.LAND, color=land_color) ax.set_global() ax.set_title(title) plt.rcParams['hatch.color'] = hatch_color_default if __name__ == '__main__': from wyconfig import * #my plot settings from geoplots import mapplot, xticks2lon """ # fig #proj = ccrs.Robinson(central_longitude=220) proj = ccrs.Robinson(central_longitude=180) dproj = ccrs.PlateCarree() #fig, axes = plt.subplots(2, 2, figsize=(10, 7), subplot_kw=dict(projection=proj)) #fig, axes = plt.subplots(1, 2, figsize=(7, 2.5), subplot_kw=dict(projection=proj)) fig, axes = plt.subplots(2, 2, figsize=(10, 5), subplot_kw=dict(projection=proj)) ax = axes[0,0] wyplot(ax=ax, rsst_on=False, regress_on=False) ax = axes[0,1] wyplot(ax=ax, rsst_on=False, regress_on=True) ax = axes[1,0] wyplot(ax=ax, rsst_on=True, regress_on=False) ax = axes[1,1] wyplot(ax=ax, rsst_on=True, regress_on=True) fig.suptitle('LME fullForcing ensemble mean') g = ds.r.plot.contourf(col='en', col_wrap=3, levels=np.arange(-1,1.01,0.1), figsize=(10,6.5)) for ax,en in zip(g.axes.flat, ds.en.values): mapplot(ax=ax, fill_continents=True) ax.set_title(None) ax.text(0.05, 1-0.05, f'{en =}', transform=ax.transAxes, ha='left', va='top') """ fig, axes = plt.subplots(4,4, figsize=(10,6.5), sharex=True, sharey=True) for ax,en in zip(axes.flat[:13], ds.en.values): im = ds.r.sel(en=en).plot.contourf(levels=np.arange(-1, 1.01,0.1), add_colorbar=False, ax=ax, xticks=np.arange(0,360,60), yticks=np.arange(-90,91,30)) ds.r.sel(en=en).where(ds.pvalue.sel(en=en)<0.05).pipe(lambda x: x*0+1).plot.contourf(add_colorbar=False, ax=ax, colors='none', hatches=['....']) #mapplot(ax=ax) mapplot(ax=ax, fill_continents=True) ax.text(0,0, f'{en = }', transform=ax.transAxes, ha='left', va='bottom', color='lightgray') for ax in axes.flat: ax.set_xlabel('') ax.set_ylabel('') ax.set_title('') for ax in axes.flat[-3:]: ax.axis('off') xticks2lon(ax=axes.flat[12], rotation=30) gs = fig.add_gridspec(4*8,4*2) ax = fig.add_subplot(gs[28,3:-1]) #fig.colorbar(im, ax=axes[:-1, 1:], orientation='horizontal', shrink=0.8) #fig.colorbar(im, ax=axes.flat[:13], orientation='horizontal') fig.colorbar(im, cax=ax, orientation='horizontal', label='correlation coefficient') figname = __file__.replace('.py', f'.png') if len(sys.argv) > 1 and sys.argv[1] == 'savefig': wysavefig(figname, dpi=128) else: print(f'figure not saved. to save figure: python {__file__} savefig') tt.check(f'**Done**') plt.show()