def getWrfData(wrfPath, dx, dy, dz, levs, debug=True):
# Open the file
perDat = Dataset(wrfPath)
refDat = Dataset("/home/iarsenea/trajs/wrfoutREFd01")
#print(data.variables)
# Pull in the values from the base state that we will add to the perturbations
ref_h = getvar(refDat, "height", units="m", timeidx=0)
thght = np.asarray(refDat.variables["HGT"])[0] * units.meter
lats, lons = latlon_coords(ref_h)
# Pull in the values from the perturbation
ph = np.asarray(destagger(perDat.variables["PH"][0], 0)) * units('m^2/s^2')
phb = np.asarray(destagger(refDat.variables["PHB"][0],
0)) * units('m^2/s^2')
ua = np.asarray(destagger(perDat.variables["U"][0], 2)) * units('m/s')
va = np.asarray(destagger(perDat.variables["V"][0], 1)) * units('m/s')
# Calculate geopotential
print("Converting from perturbation height to height AGL...")
geo = ph + phb
# Convert to height
hght = mcalc.geopotential_to_height(geo)
# Convert to height_agl
h = np.zeros_like(hght)
for i in range(hght.shape[0]):
h[i] = hght[i] - thght
print("Done.\n")
# Get the x and y values of the lat and lon coordinates
x, y = ll_to_xy(refDat, lats, lons)
x = np.arange(0, np.max(x.data) + 1) * dx
y = np.arange(0, np.max(y.data) + 1) * dy
# Interpolate the winds speeds and temps to the heights specified in levs
if debug:
print("\nInterpolating wind values to every " + str(dz.m) +
" meters... \n")
ua_m = metpy.interpolate.interpolate_1d(levs, h, ua)
va_m = metpy.interpolate.interpolate_1d(levs, h, va)
perDat.close()
refDat.close()
return h, x, y, ua_m, va_m
def extract_all_vars(TKE_PBL, all_vars, ncfile):
ds = xr.Dataset()
for v in all_vars:
ds[v] = wrf.getvar(ncfile, v,
timeidx=wrf.ALL_TIMES
)
ds[TKE_PBL] = wrf.destagger(ds[TKE_PBL], 1, meta=True)
return ds
def plot_profile(f, fig, skew):
ds = xr.open_dataset(f)
ds = ds.isel(Time=0)
Theta_m = ds.T.mean(areadims) + basetemp
pm = (ds.P + ds.PB).mean(areadims)
Qvm = (ds.QVAPOR).mean(areadims)
U = wrf.destagger(ds.U, 2, meta=True)
V = wrf.destagger(ds.V, 1, meta=True)
um = U.mean(areadims)
vm = V.mean(areadims)
Tm = Theta_m * (pm / 1e5)**(2 / 7)
p = pm.values / 100. * units.millibar
Tm = Tm.values * units.K
Qv = Qvm.values
vp = mpcalc.vapor_pressure(p, Qv)
svp = mpcalc.saturation_vapor_pressure(Tm)
#svp[svp<1e-16*units.millibar] = 1e-16 * units.millibar
Td = mpcalc.dewpoint_from_relative_humidity(Tm, vp / svp)
Td[np.isnan(
Td)] = -99. * units.degree_Celsius # fill nan with very low temp
# print('Dewpoints: ', Td)
# print('Temperatures: ', Tm)
u = um.values
v = vm.values
skew.plot(p, Tm, 'r.-', ms=1, lw=.5, label='mean T')
skew.plot(p, Td, 'g.-', ms=5, lw=2, label='mean Td')
#skew.plot_barbs(p, u, v)
return Tm, Td, p
eastu_idx = westu_idx + 2
# depth of soil layer
zs = ds_wrf['ZS'][0, :].data
# thickness of soil layer
dzs = ds_wrf['DZS'][0, :].data
# landmask - 1 is land, 0 is water
landmask = ds_wrf['LANDMASK'][0, south_idx:north_idx, west_idx:east_idx].data
# TMN - soil temperature at lower boundary
tmn = ds_wrf['TMN'][0, south_idx:north_idx, west_idx:east_idx].data
# Staggered ASL taking directly from WRF
PH = ds_wrf['PH']
PHB = ds_wrf['PHB']
H_stag = (PH + PHB) / 9.81
H = destagger(H_stag, stagger_dim=1)
wrf_time = nc_wrf.variables['XTIME']
time_var = num2date(nc_wrf.variables['XTIME'][:],
nc_wrf.variables['XTIME'].units)
time_var = np.array(time_var).astype('datetime64[s]')
for i in range(0, time_var.shape[0]):
if time_var[i] == dt_start:
start_idx = i
if time_var[i] == dt_end:
end_idx = i
time_idx = np.arange(start_idx, end_idx + 1, interval)
# round up the end time index so that PALM doesn't crash when the final time step is not given
print('===================================')
#import data
wrfpath = wrfdir + 'wrfout_' + tag
print('Extracting NetCDF data from %s ' % wrfpath)
wrfdata = netcdf.netcdf_file(wrfpath, mode='r')
#prep WRF data----------------------------------------------------
ncdict = wrf.extract_vars(
wrfdata, None,
('GRNHFX', 'W', 'QVAPOR', 'T', 'PHB', 'PH', 'U', 'P', 'PB', 'V'))
#get height and destagger vars
zstag = (ncdict['PHB'] + ncdict['PH']) / 9.81
z = wrf.destagger(zstag, 1)
u = wrf.destagger(ncdict['U'], 3)
v = wrf.destagger(ncdict['V'], 2)
p = ncdict['P'] + ncdict['PB']
interppath = wrfdir + 'interp/wrfinterp_' + tag + '.npy'
if os.path.isfile(interppath):
interpdict = np.load(interppath).item() # load here the above pickle
# qinterp, winterp, interpt = interpdict[()]['QVAPOR'],interpdict[()]['W'], interpdict[()]['T']
print('Interpolated data found at: %s' % interppath)
else:
print('WARNING: no interpolated data found - generating: SLOW ROUTINE!')
nT, nZ, nY, nX = np.shape(z)
qinterp = np.empty((nT, len(lvl), nY, nX)) * np.nan
winterp = np.empty((nT, len(lvl), nY, nX)) * np.nan
uinterp = np.empty((nT, len(lvl), nY, nX)) * np.nan
def create_variables(raw_wrfout_filename):
"""
Description:
Creates variables that are not in the raw wrfout files.
Arguments:
raw_wrfout_file(string): path to raw wrfout file.
Returns:
new_wrfout_file(xarray.core.dataset.Dataset): Saved Dataset with all created variables.
"""
# open raw NetCDF file
raw_wrfout_file_nc = Dataset(raw_wrfout_filename, mode="r")
# extract variables
timestamp = getvar(raw_wrfout_file_nc, "Times")
timestamp = timestamp.rename("timestamp")
timestamp = timestamp.assign_coords(
XTIME=getvar(raw_wrfout_file_nc,
"pres").coords["Time"]) # add both XTIME and time coords
timestamp = timestamp.drop("Time") # drop coord
timestamp = timestamp.expand_dims(dim="Time") # add dim
pres = getvar(raw_wrfout_file_nc, "p")
pres = pres.rename("pres")
pres_sea_level = getvar(raw_wrfout_file_nc, "slp")
pres_sea_level = pres_sea_level.rename("p_sl")
rel_humidity = getvar(raw_wrfout_file_nc, "rh")
rel_humidity_2m = getvar(raw_wrfout_file_nc, "rh2")
rel_humidity_2m = rel_humidity_2m.rename("rh_2m")
temp_dew = getvar(raw_wrfout_file_nc, "td")
temp_dew_2m = getvar(raw_wrfout_file_nc, "td2")
temp_dew_2m = temp_dew_2m.rename("td_2m")
temp = getvar(raw_wrfout_file_nc, "tk")
temp = temp.rename("tk")
temp_potential = getvar(raw_wrfout_file_nc, "th")
temp_potential = temp_potential.rename("th")
terrain = getvar(raw_wrfout_file_nc, "ter")
height = getvar(raw_wrfout_file_nc, "z")
height = height.rename("z")
# convert dew temps to kelvin
temp_dew = temp_dew + 273.15
temp_dew.attrs = pres.attrs # copy attributes as removed when broadcasting
temp_dew.attrs["units"] = "K"
temp_dew.attrs["description"] = "dew point temperature"
temp_dew_2m = temp_dew_2m + 273.15
temp_dew_2m.attrs = pres.attrs # copy attributes as removed when broadcasting
temp_dew_2m.attrs["units"] = "K"
temp_dew_2m.attrs["description"] = "2m dew point temperature"
# rotate wind to Earth coordinates
wind_u = getvar(raw_wrfout_file_nc, "uvmet").isel(u_v=0)
wind_u = wind_u.rename("u_ll")
wind_v = getvar(raw_wrfout_file_nc, "uvmet").isel(u_v=1)
wind_v = wind_v.rename("v_ll")
wind_u_10m = getvar(raw_wrfout_file_nc, "uvmet10").isel(u_v=0)
wind_u_10m = wind_u_10m.rename("u_ll_10m")
wind_v_10m = getvar(raw_wrfout_file_nc, "uvmet10").isel(u_v=1)
wind_v_10m = wind_v_10m.rename("v_ll_10m")
wind_w = getvar(raw_wrfout_file_nc, "wa")
wind_w = wind_w.rename("w_ll")
# calculate layer thickness
geopotential_perturbation = getvar(raw_wrfout_file_nc, "PH")
geopotential_base_state = getvar(raw_wrfout_file_nc, "PHB")
altitude = (geopotential_perturbation + geopotential_base_state) / 9.81
altitude.loc[dict(
bottom_top_stag=0)] = altitude.isel(bottom_top_stag=0) - terrain
altitude.attrs = geopotential_perturbation.attrs # copy attributes
altitude.attrs["units"] = "m"
altitude.attrs["description"] = "altitude"
altitude = altitude.rename("altitude")
layer_thickness = altitude.copy()
for layer in range(0, len(altitude.bottom_top_stag)):
if layer == 0:
pass
else:
layer_thickness.loc[dict(bottom_top_stag=layer)] = altitude.isel(
bottom_top_stag=layer) - altitude.isel(bottom_top_stag=layer -
1)
layer_thickness.attrs = altitude.attrs # copy attributes
layer_thickness.attrs["units"] = "m"
layer_thickness = layer_thickness.rename("dz")
# create destaggered version for use in conversions below
layer_thickness_destaggered = xr.DataArray(
destagger(layer_thickness, stagger_dim=0),
dims=("bottom_top", "south_north", "west_east"),
coords=layer_thickness.coords,
attrs=layer_thickness.attrs,
)
# inverse density
inverse_density = getvar(raw_wrfout_file_nc, "ALT") # m3 kg-1
# cloud liquid water path
cloud_liquid_water = getvar(raw_wrfout_file_nc, "QCLOUD") # kg kg-1
cloud_liquid_water_path = (
cloud_liquid_water * layer_thickness_destaggered / inverse_density *
1_000) # the 1000 at the end has units g/kg, to get from kg/m2 to g/m2
cloud_liquid_water_path = cloud_liquid_water_path.sum(dim="bottom_top")
cloud_liquid_water_path.attrs = cloud_liquid_water.attrs
cloud_liquid_water_path.attrs["units"] = "g m-2"
cloud_liquid_water_path.attrs["description"] = "cloud liquid water path"
cloud_liquid_water_path = cloud_liquid_water_path.rename("CLWP")
# cloud ice path
cloud_ice = getvar(raw_wrfout_file_nc, "QICE") # kg kg-1
cloud_ice_path = cloud_ice * layer_thickness_destaggered / inverse_density * 1_000 # the 1000 at the end has units g/kg, to get from kg/m2 to g/m2
cloud_ice_path = cloud_ice_path.sum(dim="bottom_top")
cloud_ice_path.attrs = cloud_ice.attrs
cloud_ice_path.attrs["units"] = "g m-2"
cloud_ice_path.attrs["description"] = "cloud ice path"
cloud_ice_path = cloud_ice_path.rename("CIP")
# cloud water vapor columns
cloud_water_vapor = getvar(
raw_wrfout_file_nc,
"QVAPOR") # rho_water = 1000 kg m-3, to cm --> H2O in cm3 cm-2
cloud_water_vapor_column = (cloud_water_vapor *
layer_thickness_destaggered / inverse_density /
10)
cloud_water_vapor_column = cloud_water_vapor_column.sum(dim="bottom_top")
cloud_water_vapor_column.attrs = cloud_water_vapor.attrs
cloud_water_vapor_column.attrs["units"] = "cm3 cm-2"
cloud_water_vapor_column.attrs["description"] = "cloud water vapor column"
cloud_water_vapor_column = cloud_water_vapor_column.rename("H2O")
# wind speed and direction
wind_speed = np.sqrt((wind_u**2) + (wind_v**2))
wind_speed.attrs = wind_u.attrs
wind_speed.attrs["description"] = "wind speed"
wind_speed = wind_speed.rename("wspeed")
wind_speed_10m = np.sqrt((wind_u_10m**2) + (wind_v_10m**2))
wind_speed_10m.attrs = wind_u_10m.attrs
wind_speed_10m.attrs["description"] = "wind speed at 10m"
wind_speed_10m = wind_speed_10m.rename("wspeed_10m")
wind_direction = np.rad2deg(np.arctan2(wind_u, wind_v)) + 180.0
wind_direction.attrs = wind_u.attrs
wind_direction.attrs["units"] = "degrees"
wind_direction.attrs["description"] = "wind direction"
wind_direction = wind_direction.rename("wdir")
wind_direction_10m = np.rad2deg(np.arctan2(wind_u_10m, wind_v_10m)) + 180.0
wind_direction_10m.attrs = wind_u_10m.attrs
wind_direction_10m.attrs["units"] = "degrees"
wind_direction_10m.attrs["description"] = "wind direction at 10m"
wind_direction_10m = wind_direction_10m.rename("wdir_10m")
# aerosol optical depth, AOD, column and surface
optical_thickness_a01 = getvar(raw_wrfout_file_nc, "TAUAER1") # 300 nm
optical_thickness_a02 = getvar(raw_wrfout_file_nc, "TAUAER2") # 400 nm
optical_thickness_a03 = getvar(raw_wrfout_file_nc, "TAUAER3") # 600 nm
optical_thickness_a04 = getvar(raw_wrfout_file_nc, "TAUAER4") # 1000 nm
# AOD at 470 nm
angstrom_exponent = (
-1 * np.log(optical_thickness_a02 / optical_thickness_a03) /
np.log(400.0 / 600.0))
AOD470 = optical_thickness_a02 * (470.0 / 600.0)**(-1 * angstrom_exponent)
AOD470.attrs = optical_thickness_a01.attrs
AOD470.attrs["units"] = ""
AOD470.attrs["description"] = "aerosol optical depth at 470 nm"
AOD470 = AOD470.rename("AOD470")
AOD470_sfc = AOD470.sum(
dim="bottom_top"
) # sum up all the values for every level in the atmospheric column
AOD470_sfc.attrs = optical_thickness_a01.attrs
AOD470_sfc.attrs["units"] = ""
AOD470_sfc.attrs["description"] = "surface aerosol optical depth at 470 nm"
AOD470_sfc = AOD470_sfc.rename("AOD470_sfc")
# AOD at 550 nm
angstrom_exponent = (
-1 * np.log(optical_thickness_a02 / optical_thickness_a03) /
np.log(400.0 / 600.0))
AOD550 = optical_thickness_a02 * (550.0 / 600.0)**(-1 * angstrom_exponent)
AOD550.attrs = pres.attrs
AOD550.attrs["units"] = ""
AOD550.attrs["description"] = "aerosol optical depth at 550 nm"
AOD550 = AOD550.rename("AOD550")
AOD550_sfc = AOD550.sum(
dim="bottom_top"
) # sum up all the values for every level in the atmospheric column
AOD550_sfc.attrs = pres.attrs
AOD550_sfc.attrs["units"] = ""
AOD550_sfc.attrs["description"] = "surface aerosol optical depth at 550 nm"
AOD550_sfc = AOD550_sfc.rename("AOD550_sfc")
# AOD at 675 nm
angstrom_exponent = (
-1 * np.log(optical_thickness_a03 / optical_thickness_a04) /
np.log(600.0 / 1000.0))
AOD675 = optical_thickness_a03 * (675.0 / 1000.0)**(-1 * angstrom_exponent)
AOD675.attrs = pres.attrs
AOD675.attrs["units"] = ""
AOD675.attrs["description"] = "aerosol optical depth at 675 nm"
AOD675 = AOD675.rename("AOD675")
AOD675_sfc = AOD675.sum(
dim="bottom_top"
) # sum up all the values for every level in the atmospheric column
AOD675_sfc.attrs = pres.attrs
AOD675_sfc.attrs["units"] = ""
AOD675_sfc.attrs["description"] = "surface aerosol optical depth at 675 nm"
AOD675_sfc = AOD675_sfc.rename("AOD675_sfc")
# convert raw netCDF to xarray dataset
raw_wrfout_file_xr = xr.open_dataset(
xr.backends.NetCDF4DataStore(raw_wrfout_file_nc))
# list of variables
variables = [
timestamp,
pres,
pres_sea_level,
rel_humidity,
rel_humidity_2m,
temp_dew,
temp_dew_2m,
temp,
temp_potential,
terrain,
height,
wind_u,
wind_v,
wind_u_10m,
wind_v_10m,
wind_w,
geopotential_perturbation,
geopotential_base_state,
layer_thickness,
cloud_liquid_water_path,
cloud_ice_path,
cloud_water_vapor_column,
wind_speed,
wind_speed_10m,
wind_direction,
wind_direction_10m,
AOD470,
AOD470_sfc,
AOD550,
AOD550_sfc,
AOD675,
AOD675_sfc,
]
# correct coords and dims
variables_with_time = []
for variable in variables:
if variable.name == "timestamp":
variable["XTIME"] = raw_wrfout_file_xr["XTIME"]
else:
if "u_v" in variable.coords:
variable = variable.drop("u_v")
variable = variable.drop("Time") # remove erroneous time coord
variable = variable.expand_dims(dim="Time") # add time as a dim
variable["XTIME"] = raw_wrfout_file_xr[
"XTIME"] # add time dim to time coord
variable["XLAT"] = raw_wrfout_file_xr[
"XLAT"] # add time dim to lat coord
variable["XLONG"] = raw_wrfout_file_xr[
"XLONG"] # add time dim to lon coord
del variable.attrs["projection"]
del variable.attrs["coordinates"]
if "_FillValue" in variable.attrs:
del variable.attrs["_FillValue"]
if "missing_value" in variable.attrs:
del variable.attrs["missing_value"]
variables_with_time.append(variable)
# new xarray dataset with new variables
new_wrfout_file_xr = xr.merge(
variables_with_time, # combine raw with new
compat="override", # skip comparing and pick variable from first dataset
)
# save new xarray dataset
# base format on the WRFChem version
wrfchem_version = re.findall(r"\d+[\.]?[\d+]?[\.]?\d+]?", sys.argv[3])[-1]
if float(wrfchem_version[0:3]) < 4.0: # pre 4.0 use NETCDF3 CLASSIC
new_wrfout_file_xr.to_netcdf(sys.argv[2], format="NETCDF3_CLASSIC")
else: # post 4.0, use NETCDF4 / HDF5 (default)
new_wrfout_file_xr.to_netcdf(sys.argv[2])
# close file handles
raw_wrfout_file_nc.close()
new_wrfout_file_xr.close()
print("Input file: ", ifile)
print("Output file: ", ofile)
incid = Dataset(ifile, "r")
times = wrf.extract_times(incid,
None,
method='cat',
squeeze=True,
cache=None,
meta=True,
do_xtime=True)
ntimes = times.shape[0]
var = incid.variables["U"][:]
u = wrf.destagger(var, 3, meta=False)
var = incid.variables["V"][:]
v = wrf.destagger(var, 2, meta=False)
var = incid.variables["W"][:]
w = wrf.destagger(var, 1, meta=False)
ph = incid.variables["PH"][:]
phb = incid.variables["PHB"][:]
# slp = wrf.getvar(incid, "slp", timeidx=wrf.ALL_TIMES)
hgt = incid.variables["HGT"][0, :, :]
# Since geopotential hight in WRF goes from mean Earth hight
# We need to substruct hgt to make it be from the real surface
# 1. No correction:
# var = (ph+phb)
# z = wrf.destagger(var, 1, meta=False)
# 2. With correction:
wrfdata_tsk = netcdf.netcdf_file(spinup_tsk, mode='r')
ncdict_tsk = wrf.extract_vars(wrfdata_tsk, None,
('W', 'T', 'PHB', 'PH', 'XTIME'))
time2 = np.where(ncdict_tsk['XTIME'] == testTime)[0][0]
w2 = ncdict_tsk['W'][time2, :blLvl, :, :].ravel()
t2 = (ncdict_tsk['T'][time2, :blLvl, :, :] + 300).ravel()
t2BL = (ncdict_tsk['T'][time2, :blLvl, :, :] + 300).ravel()
#get height and destagger
if os.path.isfile(rx.z_path):
print('.....loading destaggered height array from %s' % rx.z_path)
z = np.load(rx.z_path)
else:
print('.....destaggering height data')
zstag = (ncdict['PHB'] + ncdict['PH']) / 9.81
z = wrf.destagger(zstag, 1)
np.save(rx.z_path, z)
z_ave = np.mean(z, (0, 2, 3))
t_ave = np.mean(ncdict_tsk['T'][time2, :, :, :], (1, 2))
nT, nZ, nY, nX = np.shape(ncdict_tsk['W'])
print('Mean potential temperature (Bubble, TSK): %.3f, %.3f' %
(np.mean(t1BL), np.mean(t2BL)))
print('Meadian potential temperature (Bubble, TSK): %.3f, %.3f' %
(np.median(t1BL), np.median(t2BL)))
print('25th percentile potential temperature (Bubble, TSK): %.3f, %.3f' %
(np.percentile(t1, [25]), np.percentile(t2, [25])))
print('75th percentile potential temperature (Bubble, TSK): %.3f, %.3f' %
(np.percentile(t1, [75]), np.percentile(t2, [75])))
print('===================================')
#import data
wrfpath = wrfdir + 'wrfout_' + tag
print('Extracting NetCDF data from %s ' % wrfpath)
wrfdata = netcdf.netcdf_file(wrfpath, mode='r')
#prep WRF data----------------------------------------------------
ncdict = wrf.extract_vars(
wrfdata, None,
('GRNHFX', 'W', 'QVAPOR', 'T', 'PHB', 'PH', 'U', 'P', 'PB', 'V'))
#get height and destagger vars
zstag = (ncdict['PHB'] + ncdict['PH']) / 9.81
z = wrf.destagger(zstag, 1)
u = wrf.destagger(ncdict['U'], 3)
v = wrf.destagger(ncdict['V'], 2)
w = wrf.destagger(ncdict['W'], 1) #test line
p = ncdict['P'] + ncdict['PB']
# interppath = wrfdir + 'interp/wrfinterp_' + tag + '.npy'
interppath = wrfdir + 'interp/wrfinterp_' + tag + '_test.npy'
if os.path.isfile(interppath):
interpdict = np.load(interppath).item()
print('Interpolated data found at: %s' % interppath)
else:
print('WARNING: no interpolated data found - generating: SLOW ROUTINE!')
nT, nZ, nY, nX = np.shape(z)
qinterp = np.empty((nT, len(lvl), nY, nX)) * np.nan
def wrfout_seriesReader(wrf_path, wrf_file_filter, specified_heights=None):
"""
Construct an a2e-mmc standard, xarrays-based, data structure from a
series of 3-dimensional WRF output files
Note: Base state theta= 300.0 K is assumed by convention in WRF,
this function follow this convention.
Usage
====
wrfpath : string
The path to directory containing wrfout files to be processed
wrf_file_filter : string-glob expression
A string-glob expression to filter a set of 4-dimensional WRF
output files.
specified_heights : list-like, optional
If not None, then a list of static heights to which all data
variables should be interpolated. Note that this significantly
increases the data read time.
"""
TH0 = 300.0 #WRF convention base-state theta = 300.0 K
dims_dict = {
'Time': 'datetime',
'bottom_top': 'nz',
'south_north': 'ny',
'west_east': 'nx',
}
ds = xr.open_mfdataset(os.path.join(wrf_path, wrf_file_filter),
chunks={'Time': 10},
combine='nested',
concat_dim='Time')
dim_keys = ["Time", "bottom_top", "south_north", "west_east"]
horiz_dim_keys = ["south_north", "west_east"]
print('Finished opening/concatenating datasets...')
ds_subset = ds[['XTIME']]
print('Establishing coordinate variables, x,y,z, zSurface...')
zcoord = wrfpy.destagger((ds['PHB'] + ds['PH']) / 9.8,
stagger_dim=1,
meta=False)
#ycoord = ds.DY * np.tile(0.5 + np.arange(ds.dims['south_north']),
# (ds.dims['west_east'],1))
#xcoord = ds.DX * np.tile(0.5 + np.arange(ds.dims['west_east']),
# (ds.dims['south_north'],1))
ycoord = ds.DY * (0.5 + np.arange(ds.dims['south_north']))
xcoord = ds.DX * (0.5 + np.arange(ds.dims['west_east']))
ds_subset['z'] = xr.DataArray(zcoord, dims=dim_keys)
#ds_subset['y'] = xr.DataArray(np.transpose(ycoord), dims=horiz_dim_keys)
#ds_subset['x'] = xr.DataArray(xcoord, dims=horiz_dim_keys)
ds_subset['y'] = xr.DataArray(ycoord, dims='south_north')
ds_subset['x'] = xr.DataArray(xcoord, dims='west_east')
# Assume terrain height is static in time even though WRF allows
# for it to be time-varying for moving grids
ds_subset['zsurface'] = xr.DataArray(ds['HGT'].isel(Time=0),
dims=horiz_dim_keys)
print('Destaggering data variables, u,v,w...')
ds_subset['u'] = xr.DataArray(wrfpy.destagger(ds['U'],
stagger_dim=3,
meta=False),
dims=dim_keys)
ds_subset['v'] = xr.DataArray(wrfpy.destagger(ds['V'],
stagger_dim=2,
meta=False),
dims=dim_keys)
ds_subset['w'] = xr.DataArray(wrfpy.destagger(ds['W'],
stagger_dim=1,
meta=False),
dims=dim_keys)
print('Extracting data variables, p,theta...')
ds_subset['p'] = xr.DataArray(ds['P'] + ds['PB'], dims=dim_keys)
ds_subset['theta'] = xr.DataArray(ds['THM'] + TH0, dims=dim_keys)
# optionally, interpolate to static heights
if specified_heights is not None:
zarr = ds_subset['z']
for var in ['u', 'v', 'w', 'p', 'theta']:
print('Interpolating', var)
interpolated = wrfpy.interplevel(ds_subset[var], zarr,
specified_heights)
ds_subset[var] = interpolated #.expand_dims('Time', axis=0)
#print(ds_subset[var])
ds_subset = ds_subset.drop_dims('bottom_top').rename({'level': 'z'})
dim_keys[1] = 'z'
dims_dict.pop('bottom_top')
print(dims_dict)
# calculate derived variables
print('Calculating derived data variables, wspd, wdir...')
ds_subset['wspd'] = xr.DataArray(np.sqrt(ds_subset['u']**2 +
ds_subset['v']**2),
dims=dim_keys)
ds_subset['wdir'] = xr.DataArray(
180. + np.arctan2(ds_subset['u'], ds_subset['v']) * 180. / np.pi,
dims=dim_keys)
# assign rename coord variable for time, and assign ccordinates
ds_subset = ds_subset.rename({
'XTIME': 'datetime'
}) #Rename after defining the component DataArrays in the DataSet
if specified_heights is None:
ds_subset = ds_subset.assign_coords(z=ds_subset['z'])
ds_subset = ds_subset.assign_coords(y=ds_subset['y'])
ds_subset = ds_subset.assign_coords(x=ds_subset['x'])
ds_subset = ds_subset.assign_coords(zsurface=ds_subset['zsurface'])
ds_subset = ds_subset.rename_vars({'XLAT': 'lat', 'XLONG': 'lon'})
#print(ds_subset)
ds_subset = ds_subset.rename_dims(dims_dict)
#print(ds_subset)
return ds_subset
import numpy as np
import xarray as xr
import wrf
SCRATCH = '/global/cscratch1/sd/qnicolas/'
### MODIFY HERE ###
outfile = "WRF/WRFV4_gw/test/em_beta_plane/input_sounding" #coarse
wind = -10.
###################
wrfinput = xr.open_dataset(
"/global/cscratch1/sd/qnicolas/wrfdata/saved/channel.wrf.100x2.mountain.60lev.dry.3km/wrfinput_d01"
)
z_ref = wrf.destagger(
(wrfinput.PHB + wrfinput.PH).isel(Time=0, west_east=0, south_north=0) /
9.81, 0)
theta_ref = 300 + wrfinput.T.isel(Time=0, west_east=0, south_north=0)
q_ref = wrfinput.QVAPOR.isel(Time=0, west_east=0, south_north=0)
z_stratosphere = np.linspace(z_ref[-1], 35000, 30)
theta_stratosphere = float(
theta_ref[-1]) + 27 * (z_stratosphere - float(z_ref[-1])) / 1000
q_stratosphere = np.array(q_ref[-1]) * np.exp(
-(z_stratosphere - float(z_ref[-1])) / 1.5e3)
z_ref_ext = np.concatenate([np.array(z_ref), z_stratosphere])
q_ref_ext = np.concatenate([np.array(q_ref), q_stratosphere])
theta_ref_ext = np.concatenate([np.array(theta_ref), theta_stratosphere])
z = np.arange(0., 35000, 50.)
thetaz = np.interp(z, z_ref_ext, theta_ref_ext)
#get model data
print('Extracting NetCDF data from %s ' % rx.wrfdata)
nc_data = netcdf.netcdf_file(rx.wrfdata, mode='r')
UTMx = nc_data.variables['XLONG'][0, :, :] + rx.ll_utm[0]
UTMy = nc_data.variables['XLAT'][0, :, :] + rx.ll_utm[1]
ncdict = wrf.extract_vars(nc_data, None, ('PHB', 'PH', 'W'))
#get height and destagger
if os.path.isfile(rx.z_path):
print('.....loading destaggered height array from %s' % rx.z_path)
z = np.load(rx.z_path)
else:
print('.....destaggering height data')
zstag = (ncdict['PHB'] + ncdict['PH']) / 9.81
z = wrf.destagger(zstag, 1)
np.save(rx.z_path, z)
print('.....destaggering model W')
w = wrf.destagger(ncdict['W'], 1)
nT, nZ, nY, nX = np.shape(z)
#create timeseries of micromet-tower wind
print('.....creating OBS timeries')
num_pts = int(freq * ave_int_W)
data_samples = int(np.shape(obs_data)[0] / num_pts)
wMET = []
hMET = []
subset_time = time_data[::num_pts]
for nSample in range(data_samples):
#import all common project variables
import plume
imp.reload(plume) #force load each time
#====================INPUT===================
testLvl = 30 #height level to run the analysis on
#=================end of input===============
print('Extracting NetCDF data from %s ' % plume.wrfdir)
wrfdata = netcdf.netcdf_file(plume.wrfdir + 'wrfout_W5F7R5Tspinup', mode='r')
ncdict = wrf.extract_vars(wrfdata,
None, ('U', 'V', 'W', 'T', 'PHB', 'PH'),
meta=False)
u = wrf.destagger(ncdict['U'], 3)
v = wrf.destagger(ncdict['V'], 2)
w = wrf.destagger(ncdict['W'], 1)
t = ncdict['T']
print('.....destaggering height data')
zstag = (ncdict['PHB'] + ncdict['PH']) / 9.81
zstag_ave = np.mean(zstag, (2, 3))
z = wrf.destagger(zstag_ave, 1)
z_ave = np.mean(z, (0))
nT, nZ1, nY, nX = np.shape(zstag)
nZ = nZ1 - 1
#compare spectra every 2 min
#get stats and plot profile
co2Q1 = np.percentile(stableCO2, 25, axis=1)
co2Q3 = np.percentile(stableCO2, 75, axis=1)
print('\033[93m' + '"True" injection based on LES CWI smoke: %.2f' % rxzCL +
'\033[0m')
#-------------------------------------SOLUTION USING LES SIMULATION-----------------
print('.....extracting NetCDF data from %s ' % plume.rxcadredata)
ncdata = netcdf.netcdf_file(plume.rxcadredata, mode='r')
#get height data
zstag = (ncdata.variables['PHB'][0, :, :, :] +
ncdata.variables['PH'][0, :, :, :]) / g
zdestag = wrf.destagger(zstag, 0)
z0 = np.mean(zdestag, (1, 2))
z0stag = np.mean(zstag, (1, 2))
#get heat flux and extract only the necessasry portion of the domain
temperature = ncdata.variables['T'][0, :, :, :] + 300.
ghfx = ncdata.variables['GRNHFX'][:, 20:70, 170:230] #extract fire heat flux
rotated_fire = rotate(ghfx[:, :, :],
125,
axes=(1, 2),
reshape=False,
mode='constant') #near surface wind is at 125deg (atm)
#do fire averaging: use ~6min as interval: estimated as time during which fire remains withing 40m grid given ROS=0.1m/s
masked_flux = ma.masked_less_equal(rotated_fire[60:80, :, :], 500)
ave_masked_flux = np.nanmean(masked_flux, 0)
f0 = f.mean()
g = mpconst.earth_gravity.magnitude
(nt, ) = ds.RDX.shape
streamfunc = np.zeros(ds.T.shape, dtype=np.float32)
for itime in range(nt):
### set initial precision to float64 so that streamfunction can be closed
prc = np.float64
u = ds.U.values[itime, :]
v = ds.V.values[itime, :]
# destagger for now, todo: work with staggered data
u = destagger(u, stagger_dim=-1).astype(prc)
v = destagger(v, stagger_dim=-2).astype(prc)
Z = getvar(wrfout, "z", timeidx=itime) # geopotential height
nz, ny, nx = Z.shape
psi = Z * g / f0
psi = psi_BC(psi, u, v, rdx, rdy, msfm, msfu, msfv)
prc = np.float32
psi = psi.astype(prc)
avo = getvar(wrfout, 'avo', timeidx=itime, meta=False).astype(prc)
rvo = avo * 1e-5 - f[None, :, :].astype(prc)
streamfunc = np.zeros((nz, ny, nx), dtype=prc)
rdx = rdx.astype(np.int16)
### MODIFY HERE ###
wrfinput = xr.open_dataset(
SCRATCH +
"wrfdata/saved/channel.wrf.100x2.mountain.60lev.3km/wrfinput_d01")
wrfout = xr.open_dataset(
SCRATCH +
"wrfdata/saved/channel.wrf.100x2.mountain.60lev.3km/wrfout_d01_1970-08-29_06_00_00"
)
outfile = "WRF/WRFV4_channelbis/test/em_beta_plane/input_sounding"
wind = -10.
###################
z_stag = (wrfinput.PHB[0, :, 0, 0] + wrfout.PH[-40:, :, :, 2000:].mean(
['Time', 'south_north', 'west_east'])) / 9.81
z_destag = np.array(wrf.destagger(z_stag, 0))
theta = np.array(300 +
wrfout.T[-40:, :, :,
2000:].mean(['Time', 'south_north', 'west_east']))
q = np.array(wrfout.QVAPOR[-40:, :, :,
2000:].mean(['Time', 'south_north', 'west_east']))
z = np.arange(0., 35000, 50.)
thetaz = np.interp(z, z_destag, theta)
qz = np.interp(z, z_destag, q)
i = 0
f = open(SCRATCH + outfile, "w")
print('{:>10.2f}{:>10.2f}{:>12.7f}'.format(1000., thetaz[0], 1000 * qz[0]),
file=f)
for i, zz in enumerate(z):
