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):