def get_data_rel(varid, plev, years, datafiles, data, onset, npre, npost):
"""Return daily data aligned relative to onset/withdrawal day.
"""
years = atm.makelist(years)
onset = atm.makelist(onset)
datafiles = atm.makelist(datafiles)
daymin = min(onset) - npre
daymax = max(onset) + npost
# For a single year, add extra year before/after, if necessary
wrap_single = False
years_in = years
if len(years) == 1 and var_type(varid) == 'basic':
filenm = datafiles[0]
year = years[0]
if daymin < 1:
wrap_single = True
file_pre = filenm.replace(str(year), str(year - 1))
if os.path.isfile(file_pre):
years_in = [year - 1] + years_in
datafiles = [file_pre] + datafiles
if daymax > len(atm.season_days('ANN', year)):
wrap_single = True
file_post = filenm.replace(str(year), str(year + 1))
if os.path.isfile(file_post):
years_in = years_in + [year + 1]
datafiles = datafiles + [file_post]
var = get_daily_data(varid, plev, years_in, datafiles, data, daymin=daymin,
daymax=daymax)
# Get rid of extra years
if wrap_single:
var = atm.subset(var, {'year' : (years[0], years[0])})
# Make sure year dimension is included for single year
if len(years) == 1 and 'year' not in var.dims:
var = atm.expand_dims(var, 'year', years[0], axis=0)
# Align relative to onset day
# (not needed for calc variables since they're already aligned)
if var_type(varid) == 'basic':
print('Aligning data relative to onset day')
var = daily_rel2onset(var, onset, npre, npost)
return var
def get_strength_indices(years, mfc, precip, onset, retreat, yearnm='year',
daynm='day', varnm1='MFC', varnm2='PCP'):
"""Return various indices of the monsoon strength.
Inputs mfc and precip are the unsmoothed daily values averaged over
the monsoon area.
"""
ssn = xray.Dataset()
coords = {yearnm : years}
ssn['onset'] = xray.DataArray(onset, coords=coords)
ssn['retreat'] = xray.DataArray(retreat, coords=coords)
ssn['length'] = ssn['retreat'] - ssn['onset']
data_in = {}
if mfc is not None:
data_in[varnm1] = mfc
if precip is not None:
data_in[varnm2] = precip
for key in data_in:
for key2 in ['_JJAS_AVG', '_JJAS_TOT', '_LRS_AVG', '_LRS_TOT']:
ssn[key + key2] = xray.DataArray(np.nan * np.ones(len(years)),
coords=coords)
for key in data_in:
for y, year in enumerate(years):
d1 = int(onset.values[y])
d2 = int(retreat.values[y] - 1)
days_jjas = atm.season_days('JJAS', atm.isleap(year))
data = atm.subset(data_in[key], {yearnm : (year, None)})
data_jjas = atm.subset(data, {daynm : (days_jjas, None)})
data_lrs = atm.subset(data, {daynm : (d1, d2)})
ssn[key + '_JJAS_AVG'][y] = data_jjas.mean(dim=daynm).values
ssn[key + '_LRS_AVG'][y] = data_lrs.mean(dim=daynm).values
ssn[key + '_JJAS_TOT'][y] = ssn[key + '_JJAS_AVG'][y] * len(days_jjas)
ssn[key + '_LRS_TOT'][y] = ssn[key + '_LRS_AVG'][y] * ssn['length'][y]
ssn = ssn.to_dataframe()
return ssn
def get_data(datafile, year, d0, npre, npost):
daymin, daymax = d0 - npre, d0 + npost
ndays = len(atm.season_days('ANN', year))
file_pre = datafile.replace(str(year), str(year - 1))
file_post = datafile.replace(str(year), str(year + 1))
if daymin <1 and os.path.isfile(file_pre):
print('---Loading prev year ' + file_pre)
with xray.open_dataset(file_pre) as ds_pre:
ds_pre.load()
else:
ds_pre = None
if daymax > ndays and os.path.isfile(file_post):
print('---Loading next year ' + file_post)
with xray.open_dataset(file_post) as ds_post:
ds_post.load()
else:
ds_post = None
print('Loading ' + datafile)
with xray.open_dataset(datafile) as ds:
data = utils.wrapyear(ds, ds_pre, ds_post, daymin, daymax, year=year)
data.attrs = ds.attrs
return data
def season_days(season, year, d_onset, d_retreat):
if season == 'SSN':
days = range(0, d_retreat - d_onset + 1)
else:
days = atm.season_days(season, atm.isleap(year)) - d_onset
return days
else:
files['PCP'] = [datadir + filestr % ('PRECTOT', yr) for yr in years]
files['EVAP'] = [datadir + filestr % ('EVAP', yr) for yr in years]
files['ANA'] = [datadir + filestr % ('DQVDT_ANA', yr) for yr in years]
files['W'] = [datadir + filestr % ('TQV', yr) for yr in years]
# Lat-lon box for MFC budget
lon1, lon2 = 60, 100
lat1, lat2 = 10, 30
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)
return filn
lon1, lon2 = 40, 100
lat1, lat2 = -15, 35
varlist = ['T']
pmin, pmax = 200e2, 600e2
years = range(1979, 2015)
months = [4, 5, 6, 7, 8, 9]
g = atm.constants.g.values
nperday = 8
subset_dict = {'lon' : (lon1, lon2), 'lat' : (lat1, lat2)}
for year in years:
for mon in months:
dayvals = atm.season_days(atm.month_str(mon), atm.isleap(year))
T = merra.read_daily(varlist, year, mon, subset_dict=subset_dict)
# Daily means of 3-hourly data
T = atm.daily_from_subdaily(T, nperday, dayvals=dayvals)
# Vertical integral
Tbar = atm.int_pres(T, pmin=pmin, pmax=pmax)
Tbar = Tbar * g / (pmax - pmin)
Tbar.name='Tbar'
Tbar.attrs['long_name'] = 'Vertical mean atmospheric temperature'
Tbar.attrs['pmin'] = pmin
Tbar.attrs['pmax'] = pmax
# Save to file
atm.save_nc(savefile(year, mon, pmin, pmax), T, Tbar)
index['CHP_MFC'] = indices.onset_changepoint(mfc_acc)
index['CHP_PCP'] = indices.onset_changepoint(precip_acc)
for key in ['CHP_MFC', 'CHP_PCP']:
index[key].attrs['title'] = key
# ----------------------------------------------------------------------
# Monsoon strength indices
def detrend(vals, index):
vals_det = scipy.signal.detrend(vals)
vals_det = vals_det / np.std(vals_det)
output = pd.Series(vals_det, index=index)
return output
# MERRA MFC
mfc_JJAS = atm.subset(mfcbar, {'day' : (atm.season_days('JJAS'), None)})
mfc_JJAS = mfc_JJAS.mean(dim='day')
# ERA-Interim MFC
era = pd.read_csv(eraIfile, index_col=0)
# All India Rainfall Index, convert to mm/day
air = pd.read_csv(airfile, skiprows=4, index_col=0).loc[years]
air /= len(atm.season_days('JJAS'))
strength = mfc_JJAS.to_series().to_frame(name='MERRA')
strength['ERAI'] = era
strength['AIR'] = air
# Detrended indices
for key in strength.columns:
