def get_data(nms, pres, year, month, day, concat_dim, subset_dict, verbose): # Lists of raw and extended variables ids = list(nms) ext = [] for var in ['theta', 'theta_e']: if var in ids: ext.append(var) ids.remove(var) # Read raw data and calculate extended variables data = read_daily(ids, year, month, day, concat_dim=concat_dim, subset_dict=subset_dict, verbose=verbose) if 'theta' in ext: print_if('Computing potential temperature', verbose) T = data[get_varname('T')] data['theta'] = atm.potential_temp(T, pres) if 'theta_e' in ext: print_if('Computing equivalent potential temperature', verbose) T = data[get_varname('T')] q = data[get_varname('q')] data['theta_e'] = atm.equiv_potential_temp(T, pres, q) return data
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