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
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)
else:
# Seasonal mean and daily timeseries
ssn[nm] = var.sel(day=days_ssn).mean(dim='day')
mfcbar = mfc.mean(dim='YDim').mean(dim='XDim')
# Test atm.gradient
a = atm.constants.radius_earth.values
latdim, londim = 1, 2
lat = atm.get_coord(uq_int, 'lat')
latrad = np.radians(lat)
latrad[abs(lat) > 89] = np.nan
coslat = xray.DataArray(np.cos(latrad), coords={'YDim': lat})
lon = atm.get_coord(uq_int, 'lon')
lonrad = np.radians(lon)
mfc_x = atm.gradient(uq_int, lonrad, londim) / (a * coslat)
mfc_y = atm.gradient(vq_int * coslat, latrad, latdim) / (a * coslat)
mfc_test = mfc_x + mfc_y
mfc_test = -atm.precip_convert(mfc_test, 'kg/m2/s', 'mm/day')
mfc_test_bar = mfc_test.mean(dim='YDim').mean(dim='XDim')
diff = mfc_test - mfc
print(diff.max())
print(diff.min())
plt.plot(mfcbar)
plt.plot(mfc_test_bar)
print(mfc_test_bar - mfcbar)
# ----------------------------------------------------------------------
# Vertical gradient du/dp
lon1, lon2 = 40, 120
pmin, pmax = 100, 300
# lat, lon range to extract
lon1, lon2 = 40, 120
lat1, lat2 = -60, 60
lons = atm.latlon_labels([lon1, lon2], 'lon', deg_symbol=False)
lats = atm.latlon_labels([lat1, lat2], 'lat', deg_symbol=False)
latlon = '%s-%s_%s-%s' % (lons[0], lons[1], lats[0], lats[1])
savefile = datadir + ('merra_precip_%s_days%d-%d_%d-%d.nc' %
(latlon, daymin, daymax, years.min(), years.max()))
subset_dict = {'day' : (daymin, daymax), 'lat' : (lat1, lat2),
'lon' : (lon1, lon2)}
for y, year in enumerate(years):
datafile = datadir + 'merra_precip_%d.nc' % year
print('Loading ' + datafile)
with xray.open_dataset(datafile) as ds:
precip1 = atm.subset(ds['PRECTOT'], subset_dict)
precip1 = precip1.load()
precip1.coords['year'] = year
if y == 0:
precip = precip1
else:
precip = xray.concat((precip, precip1), dim='year')
print('Converting to mm/day')
precip.values = atm.precip_convert(precip, precip.units, 'mm/day')
precip.attrs['units'] = 'mm/day'
print('Saving to ' + savefile)
atm.save_nc(savefile, precip)
nroll = {'CMAP' : 3, 'MERRA_MFC' : 7, 'MERRA_PRECIP' : 7}
# Read CMAP pentad precip
cmap = precipdat.read_cmap(cmapfile)
cmapbar = atm.mean_over_geobox(cmap, lat1, lat2, lon1, lon2)
cmapdays = [atm.pentad_to_jday(p, pmin=1) for p in cmap.pentad.values]
# MERRA moisture flux convergence
mfc = atm.combine_daily_years('MFC', mfcfiles, years, yearname='year')
mfcbar = atm.mean_over_geobox(mfc, lat1, lat2, lon1, lon2)
# MERRA precip
subset_dict = {'lon' : (lon1, lon2), 'lat' : (lat1, lat2)}
precip = atm.combine_daily_years('PRECTOT', precipfiles, years, yearname='year',
subset_dict=subset_dict)
precip = atm.precip_convert(precip, precip.attrs['units'], 'mm/day')
precipbar = atm.mean_over_geobox(precip, lat1, lat2, lon1, lon2)
# Compute indices for each dataset
for name in ['CMAP', 'MERRA_MFC', 'MERRA_PRECIP']:
print('****' + name + '******')
if name == 'CMAP':
pcp = cmapbar
days = cmapdays
pentad = True
precip_jan = None # Calculate from pentad data
elif name == 'MERRA_MFC':
pcp = mfcbar
days = mfcbar.day.values
pentad = False
precip_jan = 0.0 # Use zero for now
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
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 = {}
for varnm in varnms:
datafile = savedir + "merra_onset_changepoint_merged_%s.nc" % varnm.upper()
mfcbar = mfc.mean(dim='YDim').mean(dim='XDim')
# Test atm.gradient
a = atm.constants.radius_earth.values
latdim, londim = 1, 2
lat = atm.get_coord(uq_int, 'lat')
latrad = np.radians(lat)
latrad[abs(lat) > 89] = np.nan
coslat = xray.DataArray(np.cos(latrad), coords={'YDim' : lat})
lon = atm.get_coord(uq_int, 'lon')
lonrad = np.radians(lon)
mfc_x = atm.gradient(uq_int, lonrad, londim) / (a*coslat)
mfc_y = atm.gradient(vq_int * coslat, latrad, latdim) / (a*coslat)
mfc_test = mfc_x + mfc_y
mfc_test = - atm.precip_convert(mfc_test, 'kg/m2/s', 'mm/day')
mfc_test_bar = mfc_test.mean(dim='YDim').mean(dim='XDim')
diff = mfc_test - mfc
print(diff.max())
print(diff.min())
plt.plot(mfcbar)
plt.plot(mfc_test_bar)
print(mfc_test_bar - mfcbar)
# ----------------------------------------------------------------------
# Vertical gradient du/dp
lon1, lon2 = 40, 120
pmin, pmax = 100, 300
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)
chpm[varnm] = xray.open_dataset(datafile)
def plotyear(chp, y, xlims, ylims):