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