#!/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 xfilter, xlearn, geoxarray import cartopy.crs as ccrs, cartopy.feature as cfeature from shared import ess_in_corr from mystats import p2r # 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'].mean('en').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).mean('en') 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) if 'b' not in ds.data_vars: 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)) da_b = da.learn.regress(ts, normalize_x=True).coef_da da_b_rsst = da_rsst.learn.regress(ts, normalize_x=True).coef_da #append to data file ds.close() ds['b'] = da_b ds['b_rsst'] = da_b_rsst ds[['b', 'b_rsst']].to_netcdf(datafile, mode='a') print('[appended]:', 'b and b_rsst to', 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)) da_r = da.learn.regress(ts, normalize_xy=True).coef_da da_r_rsst = da_rsst.learn.regress(ts, normalize_xy=True).coef_da da_b = da.learn.regress(ts, normalize_x=True).coef_da da_b_rsst = da_rsst.learn.regress(ts, normalize_x=True).coef_da nlat, nlon = da.lat.size, da.lon.size zz = np.zeros((nlat, nlon)) # corr of TC and SST zz_n = np.zeros_like(zz) # effective sample size zz_rsst = np.zeros_like(zz) # corr of TC and rSST zz_n_rsst = np.zeros_like(zz) # effective sample size if __name__ == '__main__': tt.check('start calculating ...') for ilat in range(nlat): for ilon in range(nlon): n = ess_in_corr(da.isel(lat=ilat, lon=ilon).values, ts.values) #effective sampel size for the correlation zz[ilat, ilon] = p2r(alpha, n-2) zz_n[ilat, ilon] = n #rsst n = ess_in_corr(da_rsst.isel(lat=ilat, lon=ilon).values, ts.values) #effective sampel size for the correlation zz_rsst[ilat, ilon] = p2r(alpha, n-2) zz_n_rsst[ilat, ilon] = n da_rsig = xr.DataArray(zz, dims=('lat', 'lon'), coords=[da.lat.values, da.lon.values], attrs=dict(long_name=f'correlation coefficient threshold of significance at {alpha} level')) da_n = xr.DataArray(zz_n, dims=('lat', 'lon'), coords=[da.lat.values, da.lon.values], attrs=dict(long_name=f'effective sample size')) da_rsig_rsst = xr.DataArray(zz_rsst, dims=('lat', 'lon'), coords=[da.lat.values, da.lon.values], attrs=dict(long_name=f'correlation coefficient threshold of significance at {alpha} level for rSST')) da_n_rsst = xr.DataArray(zz_n_rsst, dims=('lat', 'lon'), coords=[da.lat.values, da.lon.values], attrs=dict(long_name=f'effective sample size for the rSST case')) if __name__ == '__main__': tt.check('end calculating') ds = xr.Dataset(dict(r=da_r, r_sig=da_rsig, n=da_n, r_rsst=da_r_rsst, r_sig_rsst=da_rsig_rsst, n_rsst=da_n_rsst, b=da_b, b_rsst=da_b_rsst)) 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 # 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') 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()