#!/usr/bin/env python # Wenchang Yang (wenchang@princeton.edu) # Wed Jun 9 13:58:05 EDT 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, matplotlib.pyplot as plt #more imports import geoxarray import xfilter from misc.cim import cim # if __name__ == '__main__': tt.check('end import') # #start from here datafile = __file__.replace('.py', '.nc') if os.path.exists(datafile): ds = xr.open_dataset(datafile) print('[loaded]:', datafile) else: volcs = ['Pinatubo', 'Agung', 'StMaria'] #F daname = 'netrad_toa' print(daname) ifile_ctl = f'data/nudgeclimo_all_model_CTL1860_tigercpu_intelmpi_18_576PE.atmos_month.{daname}.nc' ctl = xr.open_dataarray(ifile_ctl) das_volcs = [] for volc in volcs: ifile = f'data/{volc}_ens_noleap_nudgeclimo_all_model1860.atmos_month.{daname}.nc' da = xr.open_dataarray(ifile) #volc run units = da.attrs['units'] #daa = da - ctl.assign_coords(time=da.time.values) #anomaly from the control run daa = da.assign_coords(time=ctl.time.values) - ctl #anomaly from the control run daa = daa.assign_coords(time=xr.cftime_range('0001-01', '0005-12', freq='MS')) #zonal mean, 3-year time mean if volc == 'Pinatubo': erupt_month = 6 elif volc == 'Agung': erupt_month = 3 elif volc == 'StMaria': erupt_month = 10 da_ = daa.isel(time=slice(erupt_month-1, erupt_month-1+3*12)).mean(['time', 'lon']) das_volcs.append(da_) #concat along volcs da = xr.concat(das_volcs, dim=pd.Index(volcs, name='volc')).assign_attrs(units=units) F = da #N daname = 'netrad_toa' print(daname) ifile_ctl = f'data/CTL1860_noleap_tigercpu_intelmpi_18_576PE.atmos_month.{daname}.nc' ctl = xr.open_dataarray(ifile_ctl) das_volcs = [] for volc in volcs: ifile = f'data/{volc}_PI_ens_noleap.atmos_month.{daname}.nc' da = xr.open_dataarray(ifile) #volc run units = da.attrs['units'] #daa = da - ctl.assign_coords(time=da.time.values) #anomaly from the control run daa = da.assign_coords(time=ctl.time.values) - ctl #anomaly from the control run daa = daa.assign_coords(time=xr.cftime_range('0001-01', '0005-12', freq='MS')) #zonal mean, 3-year time mean if volc == 'Pinatubo': erupt_month = 6 elif volc == 'Agung': erupt_month = 3 elif volc == 'StMaria': erupt_month = 10 da_ = daa.isel(time=slice(erupt_month-1, erupt_month-1+3*12)).mean(['time', 'lon']) das_volcs.append(da_) #concat along volcs da = xr.concat(das_volcs, dim=pd.Index(volcs, name='volc')).assign_attrs(units=units) N = da #Ts daname = 't_surf' print(daname) ifile_ctl = f'data/CTL1860_noleap_tigercpu_intelmpi_18_576PE.atmos_month.{daname}.nc' ctl = xr.open_dataarray(ifile_ctl) das_volcs = [] for volc in volcs: ifile = f'data/{volc}_PI_ens_noleap.atmos_month.{daname}.nc' da = xr.open_dataarray(ifile) #volc run units = da.attrs['units'] #daa = da - ctl.assign_coords(time=da.time.values) #anomaly from the control run daa = da.assign_coords(time=ctl.time.values) - ctl #anomaly from the control run daa = daa.assign_coords(time=xr.cftime_range('0001-01', '0005-12', freq='MS')) #zonal mean, 3-year time mean if volc == 'Pinatubo': erupt_month = 6 elif volc == 'Agung': erupt_month = 3 elif volc == 'StMaria': erupt_month = 10 da_ = daa.isel(time=slice(erupt_month-1, erupt_month-1+3*12)).mean(['time', 'lon']) das_volcs.append(da_) #concat along volcs da = xr.concat(das_volcs, dim=pd.Index(volcs, name='volc')).assign_attrs(units=units) tsa = da #precip daname = 'precip' print(daname) ifile_ctl = f'data/CTL1860_noleap_tigercpu_intelmpi_18_576PE.atmos_month.{daname}.nc' ctl = xr.open_dataarray(ifile_ctl) das_volcs = [] for volc in volcs: ifile = f'data/{volc}_PI_ens_noleap.atmos_month.{daname}.nc' da = xr.open_dataarray(ifile) #volc run #units = da.attrs['units'] #daa = da - ctl.assign_coords(time=da.time.values) #anomaly from the control run daa = da.assign_coords(time=ctl.time.values) - ctl #anomaly from the control run daa = daa.assign_coords(time=xr.cftime_range('0001-01', '0005-12', freq='MS')) #zonal mean, 3-year time mean if volc == 'Pinatubo': erupt_month = 6 elif volc == 'Agung': erupt_month = 3 elif volc == 'StMaria': erupt_month = 10 da_ = daa.isel(time=slice(erupt_month-1, erupt_month-1+3*12)).mean(['time', 'lon']) das_volcs.append(da_) #concat along volcs da = xr.concat(das_volcs, dim=pd.Index(volcs, name='volc')).pipe(lambda x: x*24*3600).assign_attrs(units='mm/day')#units from kg/m^2/s -> mm/day pra = da ds = xr.Dataset(dict(F=F, N=N, tsa=tsa, pra=pra)) ds.to_netcdf(datafile) print(datafile) def wyplot(ax=None, vname='tsa', volc='Pinatubo', color='k', lowpass=None, legend_on=True): if ax is None: fig, ax = plt.subplots(figsize=(4,5)) if lowpass is None: #lowpass = lambda x: x.filter.lowpass(1/9, dim='lat', padtype='constant') lowpass = lambda x: x.rolling(lat=9, center=True, min_periods=1).mean() alpha = 0.2 ylim = (-1, 1) yticks = [-1, -np.sin(np.pi/3), -np.sin(np.pi/4), -.5, -np.sin(np.pi/12), 0, np.sin(np.pi/12), .5, np.sin(np.pi/4), np.sin(np.pi/3), 1] yticklabels = ['90S', '60S', '45S', '30S', '15S', '0', '15N', '30N', '45N', '60N', '90N'] if vname == 'F': units = 'W m$^{-2}$' long_name = 'F' elif vname == 'tsa': units = 'K' long_name = 'T$_s$' elif vname == 'pra': units = 'mm day$^{-1}$' long_name = 'Prcp' da = ds[vname].sel(volc=volc).pipe(lowpass) if vname == 'F': da = da.sel(en=range(1,11)) damean = da.mean('en') dd = da.pipe(cim, dim='en') ci = damean - dd, damean + dd ax.fill_betweenx(np.sin(da.lat/180*np.pi), ci[0], ci[1], color=color, alpha=alpha) ax.plot(damean, np.sin(da.lat/180*np.pi), color=color, label=volc) ax.axhline(0, color='gray', ls='--') ax.axvline(0, color='gray', ls='--') ax.set_ylim(ylim) ax.set_yticks(yticks) ax.set_yticklabels(yticklabels) ax.set_xlabel(f'{long_name} [{units}]') if legend_on: ax.legend(frameon=False) if __name__ == '__main__': from wyconfig import * #my plot settings constrained_layout_off() fig, axes = plt.subplots(1, 3, figsize=(8, 5), sharey=True) legend_on = False ax = axes[0] wyplot(ax=ax, vname='F', volc='Pinatubo', color='k', legend_on=False) wyplot(ax=ax, vname='F', volc='Agung', color='C0', legend_on=False) wyplot(ax=ax, vname='F', volc='StMaria', color='C1', legend_on=False) ax = axes[1] wyplot(ax=ax, vname='tsa', volc='Pinatubo', color='k', legend_on=False) wyplot(ax=ax, vname='tsa', volc='Agung', color='C0', legend_on=False) wyplot(ax=ax, vname='tsa', volc='StMaria', color='C1', legend_on=False) ax = axes[2] wyplot(ax=ax, vname='pra', volc='Pinatubo', color='k', legend_on=False) wyplot(ax=ax, vname='pra', volc='Agung', color='C0', legend_on=False) wyplot(ax=ax, vname='pra', volc='StMaria', color='C1', legend_on=False) ax.legend(loc='upper right', bbox_to_anchor=(1,1.1), ncol=3, frameon=False) for ax,label in zip(axes, list('ABC')): ax.text(0.02, 1-0.01, label, transform=ax.transAxes, ha='left', va='top', fontweight='bold') #savefig if 'savefig' in sys.argv: figname = __file__.replace('.py', f'.png') wysavefig(figname) tt.check(f'**Done**') plt.show() constrained_layout_on()