def test_geopotential_to_height_32bit(): """Test conversion from geopotential to height with 32-bit values.""" geopot = np.arange(201590, 201600, dtype=np.float32) * units('J/kg') truth = np.array([20623.000, 20623.102, 20623.203, 20623.307, 20623.408, 20623.512, 20623.615, 20623.717, 20623.820, 20623.924], dtype=np.float32) * units.m assert_almost_equal(geopotential_to_height(geopot), truth, 2)
def add_metpy(option, filename): """ Adds the variables possible through metpy (theta, pv, n2) """ with xr.load_dataset(filename) as xin: if option.theta or option.pv: print("Adding potential temperature...") xin["pt"] = potential_temperature(xin["pressure"], xin["t"]) xin["pt"].data = np.array(xin["pt"].data) xin["pt"].attrs["units"] = "K" xin["pt"].attrs["standard_name"] = VARIABLES["pt"][2] if option.pv: print("Adding potential vorticity...") xin = xin.metpy.assign_crs(grid_mapping_name='latitude_longitude', earth_radius=6.356766e6) xin["pv"] = potential_vorticity_baroclinic(xin["pt"], xin["pressure"], xin["u"], xin["v"]) xin["pv"].data = np.array(xin["pv"].data * 10 ** 6) xin = xin.drop("metpy_crs") xin["pv"].attrs["units"] = "kelvin * meter ** 2 / kilogram / second" xin["pv"].attrs["standard_name"] = VARIABLES["pv"][2] xin["mod_pv"] = xin["pv"] * ((xin["pt"] / 360) ** (-4.5)) xin["mod_pv"].attrs["standard_name"] = VARIABLES["mod_pv"][2] if option.n2: print("Adding N2...") xin["n2"] = brunt_vaisala_frequency_squared(geopotential_to_height(xin["zh"]), xin["pt"]) xin["n2"].data = np.array(xin["n2"].data) xin["n2"].attrs["units"] = VARIABLES["n2"][1] xin["n2"].attrs["standard_name"] = "square_of_brunt_vaisala_frequency_in_air" xin.to_netcdf(filename)
def plot_vars(f_step, projection, load_all=False): # The one employed for the figure name when exported variable_name = 'gph_t_850' # Build the name of the output image run_string, _ = get_run() if load_all: f_steps = list(range(0, 79)) + list(range(81, 121, 3)) else: f_steps = [f_step] filenames = ['/tmp/' + projection + '_' + variable_name + '_%s_%03d.png' % (run_string, f_step) for f_step in f_steps] test_filenames = [os.path.exists(f) for f in filenames] if all(test_filenames): # means the files already exist return filenames # otherwise do the plots dset = get_dset(vars_3d=['t@850', 'fi@500'], f_times=f_steps).squeeze() # Add a fictictious 1-D time dimension just to avoid problems if 'step' not in dset.dims.keys(): dset = dset.expand_dims('step') # dset = subset_arrays(dset, projection) time = pd.to_datetime(dset.valid_time.values) cum_hour = dset.step.values.astype(int) temp_850 = dset['t'] - 273.15 z_500 = dset['z'] gph_500 = mpcalc.geopotential_to_height(z_500) gph_500 = xr.DataArray(gph_500.magnitude, coords=z_500.coords, attrs={'standard_name': 'geopotential height', 'units': gph_500.units}) levels_temp = np.arange(-30., 30., 1.) levels_gph = np.arange(4700., 6000., 70.) lon, lat = get_coordinates(temp_850) lon2d, lat2d = np.meshgrid(lon, lat) cmap = get_colormap('temp') args = dict(filenames=filenames, projection=projection, levels_temp=levels_temp, cmap=cmap, lon2d=lon2d, lat2d=lat2d, lon=lon, lat=lat, temp_850=temp_850.values, gph_500=gph_500.values, levels_gph=levels_gph, time=time, run_string=run_string) if load_all: single_plot_param = partial(single_plot, **args) iterator = range(0, len(f_steps)) pool = Pool(cpu_count()) results = pool.map(single_plot_param, iterator) pool.close() pool.join() else: results = single_plot(0, **args) return results
def test_geopotential_to_height(): """Test conversion from geopotential to height.""" geopotential = units.Quantity( [0, 9817.70342881, 19632.32592389, 29443.86893527], units('m**2 / second**2')) height = geopotential_to_height(geopotential) assert_array_almost_equal(height, units.Quantity([0, 1000, 2000, 3000], units.m), 0)
def test_geopotential_to_height_32bit(): """Test conversion from geopotential to height with 32-bit values.""" geopot = np.arange(201590, 201600, dtype=np.float32) * units('J/kg') truth = np.array([ 20596.957, 20597.059, 20597.162, 20597.266, 20597.365, 20597.469, 20597.570, 20597.674, 20597.777, 20597.881 ], dtype=np.float32) * units.m assert_almost_equal(geopotential_to_height(geopot), truth, 2)
def test_geopotential_to_height(): """Test conversion from geopotential to height.""" geopotential = units.Quantity( [0., 9805.11102602, 19607.14506998, 29406.10358006], units('m**2 / second**2')) height = geopotential_to_height(geopotential) assert_array_almost_equal(height, units.Quantity([0, 1000, 2000, 3000], units.m), 0)
def test_geopotential_to_height(array_type): """Test conversion from geopotential to height.""" mask = [False, True, False, True] geopotential = array_type( [0., 9805.11102602, 19607.14506998, 29406.10358006], 'm**2 / second**2', mask=mask, ) height = geopotential_to_height(geopotential) truth = array_type([0, 1000, 2000, 3000], 'meter', mask=mask) assert_array_almost_equal(height, truth, 0)
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 compute_geopot_height(dset, zvar='z', level=None): if level: zlevel = dset[zvar].sel(plev=level) else: zlevel = dset[zvar] gph = mpcalc.geopotential_to_height(zlevel) gph = xr.DataArray(gph.magnitude, coords=zlevel.coords, attrs={ 'standard_name': 'geopotential height', 'units': gph.units }, name='geop') return xr.merge([dset, gph])
def main(): """In the main function we basically read the files and prepare the variables to be plotted. This is not included in utils.py as it can change from case to case.""" file = glob(input_file) print('Using file '+file[0]) dset = xr.open_dataset(file[0]) dset = dset.metpy.parse_cf() # Select 850 hPa level using metpy temp_850 = dset['t'].metpy.sel(vertical=850 * units.hPa).metpy.unit_array.to('degC') gph_500 = mpcalc.geopotential_to_height(dset['z'].metpy.sel(vertical=500 * units.hPa)) lon, lat = get_coordinates(dset) time = pd.to_datetime(dset.time.values) cum_hour=np.array((time-time[0]) / pd.Timedelta('1 hour')).astype("int") levels_temp = np.arange(-35., 30., 1.) levels_gph = np.arange(4800., 5800., 70.) cmap = get_colormap("temp") for projection in projections:# This works regardless if projections is either single value or array fig = plt.figure(figsize=(figsize_x, figsize_y)) ax, x, y = get_projection_cartopy(plt, lon, lat, projection) # Create a mask to retain only the points inside the globe # to avoid a bug in matplotlib mask = np.invert(np.logical_or(np.isinf(x), np.isinf(y))) x = np.compress(mask, x) y = np.compress(mask, y) # All the arguments that need to be passed to the plotting function args=dict(x=x, y=y, ax=ax, temp_850=np.compress(mask, temp_850, axis=1), gph_500=np.compress(mask, gph_500, axis=1), levels_temp=levels_temp, cmap=cmap, levels_gph=levels_gph, time=time, projection=projection, cum_hour=cum_hour) print('Pre-processing finished, launching plotting scripts') if debug: plot_files(time[0:1], **args) else: # Parallelize the plotting by dividing into chunks and processes dates = chunks(time, chunks_size) plot_files_param=partial(plot_files, **args) p = Pool(processes) p.map(plot_files_param, dates)
def main(): ''' Match ozonesonde data with tslist ouput file Input: tslist file, wrfout*, ozonesonde file Output: dict of matched profile of h, PR, sonde_O3, QV and O3 or saved in txt ''' from metpy.units import units # --------------- input --------------- # # model paras wrf_path = '../XZ_model/data/20190725/' tslist_path = '../XZ_model/data/20190725/tslist/' domain = 'd01' wrf_file = 'wrfout_d01_2019-07-25_00_00_00' model_start = datetime(2019, 7, 25, 0) step = timedelta(seconds=5) # unit: second # sonde paras sonde = './data/ozonesonde/9_201907251434.txt' p_surf = 998.4 # read data sonde_profile, sonde_dates = read_sonde(sonde, p_surf, model_start, step) station_xs, station_ys, headers = get_sonde_indices( tslist_path, wrf_path, wrf_file, domain, sonde_profile) # get paired profiles model_profiles = {} for index, station_x in enumerate(station_xs): tslist_file = headers[(station_x, station_ys[index])] # get vertical data of specific tslist file model_profile = get_vert_data(tslist_file, model_start, sonde_dates[index], step.seconds) model_profile['O3'] *= 1e3 # ppbv model_profile['QV'] *= 1e6 # ppmv model_profile['h'] = mpcalc.geopotential_to_height( model_profile['PH'] * 9.80665 * (units.meter**2) / (units.second**2)) # assign sonde h, PR and O3 first for key in ['h', 'PR', 'O3']: if index == 0: if key == 'O3': model_profiles['sonde_O3'] = sonde_profile[key][index] else: model_profiles[key] = sonde_profile[key][index] else: if key == 'O3': model_profiles['sonde_O3'] = np.append( model_profiles['sonde_O3'], sonde_profile[key][index]) else: model_profiles[key] = np.append(model_profiles[key], sonde_profile[key][index]) # iterate requested model variables in tslists for key in ['QV', 'O3']: f = interpolate.interp1d(model_profile['h'], model_profile[key]) #, fill_value=np.nan) if index == 0: model_profiles[key] = f(sonde_profile['h'][index]) else: model_profiles[key] = np.append(model_profiles[key], f(sonde_profile['h'][index])) # convert to df and save to txt file df = pd.DataFrame.from_dict(model_profiles) ouput_name = 'sonde_tslist.txt' df.to_csv(ouput_name, sep=',', index=False) # add units info f = open(ouput_name, "r") contents = f.readlines() units = 'm,hPa,ppbv,ppmv,ppbv \n' #h,PR,sonde_O3,QV,O3 contents.insert(1, units) f.close() # write again f = open(ouput_name, "w") f.writelines(contents) f.close()
if not sys.argv[1:]: print_message( 'Projection not defined, falling back to default (euratl, it, de)') projection = 'euratl' else: projection = sys.argv[1:] """In the main function we basically read the files and prepare the variables to be plotted. This is not included in utils.py as it can change from case to case.""" dset = get_dset(vars_3d=['t@850', 'fi@500']).squeeze() time = pd.to_datetime(dset.valid_time.values) cum_hour = dset.step.values.astype(int) temp_850 = dset['t'] - 273.15 z_500 = dset['z'] gph_500 = mpcalc.geopotential_to_height(z_500) gph_500 = xr.DataArray(gph_500.magnitude, coords=z_500.coords, attrs={ 'standard_name': 'geopotential height', 'units': gph_500.units }) levels_temp = np.arange(-30., 30., 1.) levels_gph = np.arange(4700., 6000., 70.) cmap = get_colormap('temp') fig = plt.figure(figsize=(figsize_x, figsize_y)) ax = plt.gca()
def test_geopotential_to_height(): """Test conversion from geopotential to height.""" geopotential = units.Quantity([0, 9817.70342881, 19632.32592389, 29443.86893527], units('m**2 / second**2')) height = geopotential_to_height(geopotential) assert_array_almost_equal(height, units.Quantity([0, 1000, 2000, 3000], units.m), 0)
def plot_var(f_step, projection): # NOTE! # If we are inside this function it means that the picture does not exist # The one employed for the figure name when exported variable_name = 'gph_t_850' # Build the name of the output image run_string, _ = get_run() filename = '/tmp/' + projection + '_' + \ variable_name + '_%s_%03d.png' % (run_string, f_step) """In the main function we basically read the files and prepare the variables to be plotted. This is not included in utils.py as it can change from case to case.""" dset = get_dset(vars_3d=['t@850', 'fi@500'], f_times=f_step).squeeze() dset = subset_arrays(dset, projection) time = pd.to_datetime(dset.valid_time.values) cum_hour = dset.step.values.astype(int) temp_850 = dset['t'] - 273.15 z_500 = dset['z'] gph_500 = mpcalc.geopotential_to_height(z_500) gph_500 = xr.DataArray(gph_500.magnitude, coords=z_500.coords, attrs={'standard_name': 'geopotential height', 'units': gph_500.units}) levels_temp = np.arange(-30., 30., 1.) levels_gph = np.arange(4700., 6000., 70.) cmap = get_colormap('temp') fig = plt.figure(figsize=(figsize_x, figsize_y)) ax = plt.gca() lon, lat = get_coordinates(temp_850) lon2d, lat2d = np.meshgrid(lon, lat) ax = get_projection_cartopy(plt, projection, compute_projection=True) if projection == 'euratl': norm = BoundaryNorm(levels_temp, ncolors=cmap.N) cs = ax.pcolormesh(lon2d, lat2d, temp_850, cmap=cmap, norm=norm) else: cs = ax.contourf(lon2d, lat2d, temp_850, extend='both', cmap=cmap, levels=levels_temp) c = ax.contour(lon2d, lat2d, gph_500, levels=levels_gph, colors='white', linewidths=1.) labels = ax.clabel(c, c.levels, inline=True, fmt='%4.0f', fontsize=6) maxlabels = plot_maxmin_points(ax, lon, lat, gph_500, 'max', 80, symbol='H', color='royalblue', random=True) minlabels = plot_maxmin_points(ax, lon, lat, gph_500, 'min', 80, symbol='L', color='coral', random=True) an_fc = annotation_forecast(ax, time) an_var = annotation( ax, 'Geopotential height @500hPa [m] and temperature @850hPa [C]', loc='lower left', fontsize=6) an_run = annotation_run(ax, time) plt.colorbar(cs, orientation='horizontal', label='Temperature', pad=0.03, fraction=0.04) plt.savefig(filename, **options_savefig) plt.clf() return filename