def _wrf2xarray(netcdf_data, query_variables):
"""
Gets data from the netcdf wrfout file and uses wrf-python
to create an xarray Dataset.
Parameters
----------
data: netcdf
Data returned from UNIDATA NCSS query, or from your local forecast.
query_variables: list
The variables requested.
start: Timestamp
The start time
end: Timestamp
The end time
Returns
-------
xarray.Dataset
"""
first = True
for key in query_variables:
if key in ['wspd', 'wdir']:
var = get_wspd_wdir(netcdf_data, key)
else:
var = wrf.getvar(netcdf_data, key, timeidx=wrf.ALL_TIMES)
if first:
data = var
first = False
else:
with xr.set_options(keep_attrs=True):
try:
data = xr.merge([data, var])
except ValueError:
data = data.drop_vars('Time')
data = xr.merge([data, var])
# Get global attributes from the NetCDF Dataset
wrfattrs_names = netcdf_data.ncattrs()
wrfattrs = wrf.extract_global_attrs(netcdf_data, wrfattrs_names)
data = data.assign_attrs(wrfattrs)
# Fix a bug in how wrfout data is read in -- attributes must be strings to be written to NetCDF
for var in data.data_vars:
try:
data[var].attrs['projection'] = str(data[var].attrs['projection'])
except KeyError:
pass
# Fix another bug that creates a conflict in the 'coordinates' attribute
for var in data.data_vars:
try:
del data[var].attrs['coordinates']
except KeyError:
pass
return data
def _wrf2xarray(self, netcdf_data, query_variables, start, end):
"""
Transforms data from netcdf to xarray Dataset.
Parameters
----------
data: netcdf
Data returned from UNIDATA NCSS query, or from your local forecast.
query_variables: list
The variables requested.
start: Timestamp
The start time
end: Timestamp
The end time
Returns
-------
xarray.Dataset
"""
first = True
for key in query_variables:
var = wrf.getvar(netcdf_data, key, timeidx=wrf.ALL_TIMES)
if first:
data = var
first = False
else:
with xr.set_options(keep_attrs=True):
try:
data = xr.merge([data, var])
except ValueError:
data = data.drop_vars('Time')
data = xr.merge([data, var])
# Get global attributes from the NetCDF Dataset
wrfattrs_names = netcdf_data.ncattrs()
wrfattrs = wrf.extract_global_attrs(netcdf_data, wrfattrs_names)
data = data.assign_attrs(wrfattrs)
# Fix a bug in how wrfout data is read in -- attributes must be strings to be written to NetCDF
for var in data.data_vars:
try:
data[var].attrs['projection'] = str(
data[var].attrs['projection'])
except KeyError:
pass
# Fix another bug that creates a conflict in the 'coordinates' attribute
for var in data.data_vars:
try:
del data[var].attrs['coordinates']
except KeyError:
pass
# Slice the dataset to only include specified time interval
data = data.sel(Time=slice(start, end))
return data
def extract_proj_params(wrfnc):
attrs = extract_global_attrs(
wrfnc, ("MAP_PROJ", "TRUELAT1", "TRUELAT2", "STAND_LON", "POLE_LAT",
"POLE_LON", "DX", "DY"))
result = {key.lower(): val for key, val in viewitems(attrs)}
if is_multi_file(wrfnc):
wrfnc = wrfnc[0]
result["known_x"] = 0
result["known_y"] = 0
result["ref_lat"] = wrfnc.variables["XLAT"][0, 0, 0]
result["ref_lon"] = wrfnc.variables["XLONG"][0, 0, 0]
return result
def __init__(self, config):
""" construct input wrf file names """
self.ncfile = Dataset(config['INPUT']['input_fn'])
self.resolution = float(
wrf.extract_global_attrs(self.ncfile, 'DX')['DX'])
n_radius = int(
int(config['INPUT']['vis_radius']) / int(self.resolution))
times_dt64 = wrf.getvar(self.ncfile,
'Times',
timeidx=wrf.ALL_TIMES,
method="cat").values
self.tgt_timeframe = datetime.datetime.strptime(
config['INPUT']['vis_timestamp'], '%Y%m%d%H%M%S')
print(print_prefix + 'Selected Timeframe:%s' %
self.tgt_timeframe.strftime('%Y-%m-%d_%H:%M:%S'))
# get idx of the field will be extract
tgt_times_dt64 = np.datetime64(self.tgt_timeframe)
t_idx = np.where(times_dt64 == tgt_times_dt64)[0][0]
print(print_prefix +
'init from single input file for lat2d, lon2d, and hgt')
self.xlat = wrf.getvar(self.ncfile, 'XLAT')
self.xlon = wrf.getvar(self.ncfile, 'XLONG')
print(print_prefix + 'init from single input file for Z4d')
self.SLP = wrf.getvar(self.ncfile, 'slp', timeidx=t_idx)
self.tc_ix, self.tc_iy = find_tc_center(self.SLP)
print(print_prefix + 'init from single input file for Z4d')
self.Z = wrf.getvar(self.ncfile, 'z', timeidx=t_idx)
print(print_prefix + 'init from single input file for U4d')
self.U = wrf.getvar(self.ncfile, 'ua', timeidx=t_idx)
print(print_prefix + 'init from single input file for V4d')
self.V = wrf.getvar(self.ncfile, 'va', timeidx=t_idx)
print(print_prefix + 'init from single input file for W4d')
self.W = wrf.getvar(self.ncfile, 'wa', timeidx=t_idx)
print(print_prefix + 'init multi files successfully!')
def extract_proj_params(wrfnc, timeidx=0):
attrs = extract_global_attrs(
wrfnc, ("MAP_PROJ", "TRUELAT1", "TRUELAT2", "STAND_LON", "POLE_LAT",
"POLE_LON", "DX", "DY"))
result = {key.lower(): val for key, val in viewitems(attrs)}
_timeidx = timeidx
if is_multi_file(wrfnc):
wrfnc0 = wrfnc[0]
num_times_per_file = len(wrfnc0.dimensions["Time"])
file_idx = timeidx // num_times_per_file
_timeidx = timeidx % num_times_per_file
wrfnc = wrfnc[file_idx]
result["known_x"] = 0
result["known_y"] = 0
result["ref_lat"] = wrfnc.variables["XLAT"][_timeidx, 0, 0]
result["ref_lon"] = wrfnc.variables["XLONG"][_timeidx, 0, 0]
return result
def open_mfdataset(fname,
convert_to_ppb=True,
mech="racm_esrl_vcp",
var_list=["o3"],
surf_only=False,
surf_only_nc=False,
**kwargs):
"""Method to open WRF-chem and RAP-chem netcdf files.
Parameters
----------
fname : string or list
fname is the path to the file or files. It will accept hot keys in
strings as well.
convert_to_ppb : boolean
If true the units of the gas species will be converted to ppbv
mech: str
Mechanism to be used for calculating sums. Supported mechanisms
include: "racm_esrl_vcp" and "redhc"
var_list: list
List of variables to include in output. MELODIES-MONET only reads in
variables need to plot in order to save on memory and simulation cost
especially for vertical data
surf_only: boolean
Whether to save only surface data to save on memory and computational
cost (True) or not (False).
surf_only_nc : bool
Whether input data has been subsampled to only contain surface data and
not extra pressure & temperature info (True) or not (False).
Returns
-------
xarray.DataSet
WRF-Chem or RAP-Chem model dataset in standard format for use
in MELODIES-MONET
"""
from netCDF4 import Dataset
from wrf import ALL_TIMES, extract_global_attrs, getvar
# Get dictionary of summed species for the mechanism of choice.
dict_sum = dict_species_sums(mech=mech)
list_calc_sum = []
for var_sum in [
"noy_gas",
"noy_aer",
"nox",
"pm25_cl",
"pm25_ec",
"pm25_na",
"pm25_nh4",
"pm25_no3",
"pm25_so4",
"pm25_om",
]:
if var_sum in var_list:
var_list.extend(dict_sum[var_sum])
var_list.remove(var_sum)
list_calc_sum.append(var_sum)
wrflist = []
for files in fname:
wrflist.append(Dataset(files))
if not surf_only_nc:
# Add some additional defaults needed for aircraft analysis
# Turn this on also if need to convert aerosols
var_list.append("pres")
var_list.append("height")
var_list.append("tk")
var_list.append("height_agl")
var_list.append("PSFC")
# need to calculate surface pressure and dp and optionally dz here.
var_wrf_list = []
for var in var_list:
if var == "pres": # Insert special versions.
var_wrf = getvar(wrflist,
var,
timeidx=ALL_TIMES,
method="cat",
squeeze=False,
units="Pa")
elif var == "height":
var_wrf = getvar(wrflist,
var,
timeidx=ALL_TIMES,
method="cat",
squeeze=False,
units="m")
elif var == "height_agl":
var_wrf = getvar(wrflist,
var,
timeidx=ALL_TIMES,
method="cat",
squeeze=False,
units="m")
else:
var_wrf = getvar(wrflist,
var,
timeidx=ALL_TIMES,
method="cat",
squeeze=False)
var_wrf_list.append(var_wrf)
dset = xr.merge(var_wrf_list)
# Add global attributes needed
a_truelat1 = extract_global_attrs(wrflist[0], "TRUELAT1")
a_truelat2 = extract_global_attrs(wrflist[0], "TRUELAT2")
a_moad_cen_lat = extract_global_attrs(wrflist[0], "MOAD_CEN_LAT")
a_stand_lon = extract_global_attrs(wrflist[0], "STAND_LON")
a_map_proj = extract_global_attrs(wrflist[0], "MAP_PROJ")
a_cen_lat = extract_global_attrs(wrflist[0], "CEN_LAT")
a_cen_lon = extract_global_attrs(wrflist[0], "CEN_LON")
dset = dset.assign_attrs(a_truelat1)
dset = dset.assign_attrs(a_truelat2)
dset = dset.assign_attrs(a_moad_cen_lat)
dset = dset.assign_attrs(a_stand_lon)
dset = dset.assign_attrs(a_map_proj)
dset = dset.assign_attrs(a_cen_lat)
dset = dset.assign_attrs(a_cen_lon)
for i in dset.variables:
dset[i] = dset[i].assign_attrs(a_truelat1)
dset[i] = dset[i].assign_attrs(a_truelat2)
dset[i] = dset[i].assign_attrs(a_moad_cen_lat)
dset[i] = dset[i].assign_attrs(a_stand_lon)
dset[i] = dset[i].assign_attrs(a_map_proj)
dset[i] = dset[i].assign_attrs(a_cen_lat)
dset[i] = dset[i].assign_attrs(a_cen_lon)
# Convert names to standards used in MONET
dset = dset.rename({
"Time": "time",
"south_north": "y",
"west_east": "x",
"XLONG": "longitude",
"XLAT": "latitude",
})
# These sums and conversions are quite expensive and memory intensive,
# so add option to shrink dataset to just surface when needed
if (not surf_only_nc) and surf_only:
dset = dset.isel(bottom_top=0).expand_dims("bottom_top", axis=1)
# convert all gas species to ppbv
if convert_to_ppb:
for i in dset.variables:
if "units" in dset[i].attrs:
if "ppmv" in dset[i].attrs["units"]:
dset[i] = dset[i] * 1000.0
dset[i].attrs["units"] = "ppbv"
# convert "ug/kg-dryair -> ug/m3"
for i in dset.variables:
if "units" in dset[i].attrs:
if "ug/kg-dryair" in dset[i].attrs["units"]:
# ug/kg -> ug/m3 using dry air density
dset[i] = dset[i] * dset["pressure"] / dset["temp"] / 287.05535
dset[i].attrs["units"] = r"$\mu g m^{-3}$"
# assign mapping table for airnow
dset = _predefined_mapping_tables(dset)
# add lazy diagnostic variables
if "noy_gas" in list_calc_sum:
dset = add_lazy_noy_g(dset, dict_sum)
if "noy_aer" in list_calc_sum:
dset = add_lazy_noy_a(dset, dict_sum)
if "nox" in list_calc_sum:
dset = add_lazy_nox(dset, dict_sum)
if "pm25_cl" in list_calc_sum:
dset = add_lazy_cl_pm25(dset, dict_sum)
if "pm25_ec" in list_calc_sum:
dset = add_lazy_ec_pm25(dset, dict_sum)
if "pm25_na" in list_calc_sum:
dset = add_lazy_na_pm25(dset, dict_sum)
if "pm25_nh4" in list_calc_sum:
dset = add_lazy_nh4_pm25(dset, dict_sum)
if "pm25_no3" in list_calc_sum:
dset = add_lazy_no3_pm25(dset, dict_sum)
if "pm25_so4" in list_calc_sum:
dset = add_lazy_so4_pm25(dset, dict_sum)
if "pm25_om" in list_calc_sum:
dset = add_lazy_om_pm25(dset, dict_sum)
dset = dset.reset_index(["XTIME", "datetime"], drop=True)
if not surf_only_nc:
# Reset more variables
dset = dset.rename({
"bottom_top": "z",
"temp": "temperature_k",
"height": "alt_msl_m_mid",
"height_agl": "alt_agl_m_mid",
"PSFC": "surfpres_pa",
"pressure": "pres_pa_mid",
})
dset2 = dset
else:
# Expand into z coordinate so that format is consistent.
if "bottom_top" in dset.dims:
# Workaround for tutorial version of the WRF-Chem example dataset
# TODO: eliminate need
dset2 = dset.rename({"bottom_top": "z"})
else:
dset2 = dset.expand_dims("z", axis=3).copy()
dset2 = dset2.reset_coords()
dset2 = dset2.set_coords(["latitude", "longitude"])
return dset2