def get_mfc_box(mfcfiles, precipfiles, evapfiles, years, nroll, lat1, lat2,
lon1, lon2):
"""Return daily tseries MFC, precip and evap averaged over lat-lon box.
"""
subset_dict = {'lat' : (lat1, lat2), 'lon' : (lon1, lon2)}
databox = {}
if mfcfiles is not None:
mfc = atm.combine_daily_years('MFC', mfcfiles, years, yearname='year',
subset_dict=subset_dict)
databox['MFC'] = mfc
if precipfiles is not None:
pcp = atm.combine_daily_years('PRECTOT', precipfiles, years, yearname='year',
subset_dict=subset_dict)
databox['PCP'] = pcp
if evapfiles is not None:
evap = atm.combine_daily_years('EVAP', evapfiles, years, yearname='year',
subset_dict=subset_dict)
databox['EVAP'] = evap
nms = databox.keys()
for nm in nms:
var = databox[nm]
var = atm.precip_convert(var, var.attrs['units'], 'mm/day')
var = atm.mean_over_geobox(var, lat1, lat2, lon1, lon2)
databox[nm + '_UNSM'] = var
databox[nm + '_ACC'] = np.cumsum(var, axis=1)
if nroll is None:
databox[nm] = var
else:
databox[nm] = atm.rolling_mean(var, nroll, axis=-1, center=True)
tseries = xray.Dataset(databox)
return tseries
def get_onset_indices(onset_nm, datafiles, years, data=None):
"""Return monsoon onset/retreat/length indices.
"""
# Options for CHP_MFC and CHP_PCP
lat1, lat2 = 10, 30
lon1, lon2 = 60, 100
chp_opts = [None, lat1, lat2, lon1, lon2]
if onset_nm == 'HOWI':
maxbreak = 10
npts = 100
ds = atm.combine_daily_years(['uq_int', 'vq_int'], datafiles, years,
yearname='year')
index, _ = indices.onset_HOWI(ds['uq_int'], ds['vq_int'], npts,
maxbreak=maxbreak)
index.attrs['title'] = 'HOWI (N=%d)' % npts
elif onset_nm == 'CHP_MFC':
if data is None:
tseries = get_mfc_box(datafiles, None, None, years, *chp_opts)
data = tseries['MFC_ACC']
index['ts_daily'] = tseries['MFC']
index = indices.onset_changepoint(data)
elif onset_nm == 'CHP_PCP':
if data is None:
tseries = get_mfc_box(None, datafiles, None, years, *chp_opts)
data = tseries['PCP_ACC']
index = indices.onset_changepoint(data)
index['ts_daily'] = tseries['PCP']
# Monsoon retreat and length indices
if 'retreat' in index:
index['length'] = index['retreat'] - index['onset']
else:
index['retreat'] = np.nan * index['onset']
index['length'] = np.nan * index['onset']
return index
import numpy as np import xray import pandas as pd import matplotlib.pyplot as plt import atmos as atm import merra from indices import onset_SJKE, summarize_indices, plot_index_years # ---------------------------------------------------------------------- # Compute SJ indices (Boos and Emmanuel 2009) datadir = atm.homedir() + 'datastore/merra/daily/' years = np.arange(1979, 2015) filestr = 'merra_uv850_40E-120E_60S-60N_' datafiles = [datadir + filestr + '%d.nc' % y for y in years] # Read daily data from each year ds = atm.combine_daily_years(['U', 'V'], datafiles, years) # Remove extra dimension from data u = atm.squeeze(ds['U']) v = atm.squeeze(ds['V']) # Calculate OCI index sjke = onset_SJKE(u, v) # Summary plot and timeseries in individual years summarize_indices(years, sjke['onset']) plot_index_years(sjke, suptitle='SJ', yearnm='Year', daynm='Day')
ds_rel[nm] = utils.daily_rel2onset(var, d_onset, npre, npost)
ds_rel.attrs['d_onset'] = d_onset
ds_rel.attrs['d_retreat'] = d_retreat
savefile = savefiles[plev][y]
print('Saving to ' + savefile)
ds_rel.to_netcdf(savefile)
# ----------------------------------------------------------------------
# Compute climatologies and save
yearstr = '%d-%d' % (years.min(), years.max())
for plev in plevs:
relfiles = savefiles[plev]
savefile = savestr % plev + '_' + yearstr + '.nc'
ds = atm.combine_daily_years(None, relfiles, years, yearname='year')
ds = ds.mean(dim='year')
ds.attrs['years'] = years
print('Saving to ' + savefile)
ds.to_netcdf(savefile)
# ----------------------------------------------------------------------
# Concatenate plevels in climatology and save
files = [savestr % plev + '_' + yearstr + '.nc' for plev in plevs]
ubudget = xray.Dataset()
pname, pdim = 'Height', 1
subset_dict = {'lat' : (-60, 60), 'lon' : (40, 120)}
for i, plev in enumerate(plevs):
filenm = files[i]
print('Loading ' + filenm)
latlonstr = '%d-%dE, %d-%dN' % (lon1, lon2, lat1, lat2)
# ----------------------------------------------------------------------
# Read data
nroll = 7
days_ssn = atm.season_days('JJAS')
subset_dict = {'lat' : (lat1, lat2), 'lon' : (lon1, lon2)}
ts = xray.Dataset()
ssn = xray.Dataset()
varnms = {'PCP' : 'PRECTOT', 'MFC' : 'MFC', 'EVAP' : 'EVAP', 'W' : 'TQV',
'ANA' : 'DQVDT_ANA'}
for nm in files:
var = atm.combine_daily_years(varnms[nm], files[nm], years, yearname='year',
subset_dict=subset_dict)
var = atm.mean_over_geobox(var, lat1, lat2, lon1, lon2)
units = var.attrs.get('units')
if units in ['kg/m2/s', 'kg m-2 s-1']:
var = atm.precip_convert(var, units, 'mm/day')
var_sm = atm.rolling_mean(var, nroll, axis=-1, center=True)
ts[nm] = var_sm
if nm == 'W':
# Difference in precipitable water from beginning to end of season
var_ssn = var_sm.sel(day=days_ssn)
ssn['d' + nm] = (var_ssn[:,-1] - var_ssn[:, 0]) / len(days_ssn)
# dW/dt
dvar = np.nan * np.ones(var.shape)
for y, year in enumerate(years):
dvar[y] = np.gradient(var[y])
ts['d%s/dt' % nm] = xray.DataArray(dvar, coords=var.coords)
if m == 0:
dsyr = ds
else:
dsyr = xray.concat((dsyr, ds), dim='day')
savefile = datafiles[y]
print('Saving to ' + savefile)
dsyr.to_netcdf(savefile)
# Read daily data from each year
plist = [200, 400, 600]
if plev is not None:
plist = np.union1d(plist, [plev])
T_p = {}
for p in plist:
T1 = atm.combine_daily_years('T', datafiles, years, yearname='year',
subset_dict={'plev' : (p, p)})
T_p[p] = atm.squeeze(T1)
Tbar = atm.combine_daily_years('Tbar', datafiles, years, yearname='year')
if plev is None:
T = Tbar
varname = 'TT200-600'
else:
T = T_p[plev]
varname = 'TT%d' % plev
# Calculate TT index
# The north region should go up to 35 N but there is some weirdness
# with the topography so I'm setting it to 30 N for now
north=(5, 30, 40, 100)
south=(-15, 5, 40, 100)
'CHP_MFC' : 'CHP_M',
'CHP_PCP' : 'CHP_P'}
short_inv = {v : k for k, v in short.items()}
# ----------------------------------------------------------------------
def saveclose(name, isave, exts):
if isave:
for ext in exts:
atm.savefigs(name, ext)
plt.close('all')
# ----------------------------------------------------------------------
# HOWI index (Webster and Fasullo 2003)
maxbreak = 10
ds = atm.combine_daily_years(['uq_int', 'vq_int'], vimtfiles, years,
yearname='year')
ds_howi = {}
for npts in [50, 100]:
howi, ds_howi[npts] = indices.onset_HOWI(ds['uq_int'], ds['vq_int'], npts,
maxbreak=maxbreak)
howi.attrs['title'] = 'HOWI (N=%d)' % npts
index['HOWI_%d' % npts] = howi
# ----------------------------------------------------------------------
# Wang & LinHo method
def get_onset_WLH(years, days, pcp_sm, threshold, titlestr, pentad=True,
pcp_jan=None):
nyears = len(years)
i_onset = np.zeros(nyears)
i_retreat = np.zeros(nyears)
def get_daily_data(varid, plev, years, datafiles, data, daymin=1,
daymax=366, yearnm='year'):
"""Return daily data (basic variable or calculated variable).
Data is read from datafiles if varnm is a basic variable.
If varnm is a calculated variable (e.g. potential temperature),
the base variables for calculation are provided in the dict data.
"""
years = atm.makelist(years)
datafiles = atm.makelist(datafiles)
if isinstance(plev, int) or isinstance(plev, float):
pres = atm.pres_convert(plev, 'hPa', 'Pa')
elif plev == 'LML' and 'PS' in data:
pres = data['PS']
else:
pres = None
def get_var(data, varnm, plev=None):
if plev is None:
plev = ''
elif plev == 'LML' and varnm == 'QV':
varnm = 'Q'
return data[varnm + str(plev)]
if var_type(varid) == 'calc':
print('Computing ' + varid)
if varid == 'THETA':
var = atm.potential_temp(get_var(data, 'T', plev), pres)
elif varid == 'THETA_E':
var = atm.equiv_potential_temp(get_var(data, 'T', plev), pres,
get_var(data, 'QV', plev))
elif varid == 'DSE':
var = atm.dry_static_energy(get_var(data, 'T', plev),
get_var(data, 'H', plev))
elif varid == 'MSE':
var = atm.moist_static_energy(get_var(data, 'T', plev),
get_var(data, 'H', plev),
get_var(data, 'QV', plev))
elif varid == 'VFLXMSE':
Lv = atm.constants.Lv.values
var = data['VFLXCPT'] + data['VFLXPHI'] + data['VFLXQV'] * Lv
var.attrs['units'] = data['VFLXCPT'].attrs['units']
var.attrs['long_name'] = 'Vertically integrated MSE meridional flux'
else:
with xray.open_dataset(datafiles[0]) as ds:
if varid not in ds.data_vars:
varid = varid + str(plev)
var = atm.combine_daily_years(varid, datafiles, years, yearname=yearnm,
subset_dict={'day' : (daymin, daymax)})
var = atm.squeeze(var)
# Make sure year dimension is included for single year
if len(years) == 1 and 'year' not in var.dims:
var = atm.expand_dims(var, yearnm, years[0], axis=0)
# Wrap years for extended day ranges
if daymin < 1 or daymax > 366:
var = wrapyear_all(var, daymin, daymax)
# Convert precip and evap to mm/day
if varid in ['precip', 'PRECTOT', 'EVAP']:
var = atm.precip_convert(var, var.attrs['units'], 'mm/day')
return var
index.load()
index = index.sel(year=years)
d0 = index[ind_nm].values
# Precip data
if pcp_nm == 'cmap':
pcp = precipdat.read_cmap(pcpfiles, yearmin=min(years), yearmax=max(years))
# Interpolate to daily resolution
name, attrs, coords, dimnames = atm.meta(pcp)
days = np.arange(3, 364)
interp_func = scipy.interpolate.interp1d(pcp['day'], pcp, axis=1)
vals = interp_func(days)
coords['day'] = xray.DataArray(days, coords={'day' : days})
pcp = xray.DataArray(vals, dims=dimnames, coords=coords, name=name,
attrs=attrs)
else:
pcp = atm.combine_daily_years(None, pcpfiles, years, yearname='year',
subset_dict=subset_dict)
# Wrap from following year to get extended daily range
daymin = min(d0) - npre
daymax = max(d0) + npost
pcp = utils.wrapyear_all(pcp, daymin=daymin, daymax=daymax)
# Daily relative to onset/withdrawal
pcp_rel = utils.daily_rel2onset(pcp, d0, npre, npost)
print('Saving to ' + savefile)
atm.save_nc(savefile, pcp_rel)
varnms = ["MFC", "precip"]
datadir = atm.homedir() + "datastore/merra/daily/"
savedir = atm.homedir() + "datastore/merra/analysis/"
icalc = False
isavefigs = True
if icalc:
subset_dict = {"lat": (lat1, lat2), "lon": (lon1, lon2)}
for varnm in varnms:
if varnm == "MFC":
varid, filestr = "MFC", "MFC_ps-300mb"
elif varnm == "precip":
varid, filestr = "PRECTOT", "precip"
files = [datadir + "merra_%s_%d.nc" % (filestr, year) for year in years]
data = atm.combine_daily_years(varid, files, years, yearname="year", subset_dict=subset_dict)
data = atm.mean_over_geobox(data, lat1, lat2, lon1, lon2)
data = atm.precip_convert(data, data.attrs["units"], "mm/day")
# Accumulated precip or MFC
data_acc = np.cumsum(data, axis=1)
# Compute onset and retreat days
chp = onset_changepoint_merged(data_acc)
# Save to file
savefile = savedir + "merra_onset_changepoint_merged_%s.nc" % varnm.upper()
print("Saving to " + savefile)
chp.to_netcdf(savefile)
chp = {}
datadir2 = atm.homedir() + 'datastore/merra/analysis/'
isave = True
incl_merged = True
data = {}
data_acc = {}
chp = {}
subset_dict = {'lat' : (lat1, lat2), 'lon' : (lon1, lon2)}
for varnm in varnms:
if varnm == 'MFC':
varid, filestr = 'MFC', 'MFC_ps-300mb'
elif varnm == 'precip':
varid, filestr = 'PRECTOT', 'precip'
files = [datadir + 'merra_%s_%d.nc' % (filestr, year) for year in years]
data[varnm] = atm.combine_daily_years(varid, files, years, yearname='year',
subset_dict=subset_dict)
data[varnm] = atm.mean_over_geobox(data[varnm], lat1, lat2, lon1, lon2)
data[varnm] = atm.precip_convert(data[varnm], data[varnm].attrs['units'],
'mm/day')
# Accumulated precip or MFC
data_acc[varnm] = np.cumsum(data[varnm], axis=1)
# Compute onset and retreat days
chp[varnm] = onset_changepoint(data_acc[varnm])
if incl_merged:
chpm = {}
for varnm in varnms:
datafile = datadir2 + 'merra_onset_changepoint_merged_%s.nc' % varnm.upper()
print('Reading ' + datafile)
