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
