def test_temperature_from_potential_temperature():
"""Test temperature_from_potential_temperature calculation."""
theta = np.array([286.12859679, 288.22362587, 290.31865495, 292.41368403]) * units.kelvin
pres = np.array([850] * 4) * units.mbar
real_t = np.array([273.15, 275.15, 277.15, 279.15]) * units.kelvin
assert_array_almost_equal(temperature_from_potential_temperature(pres, theta),
real_t, 2)
def test(type):
"""Run a single benchmark test for computation of meteorological fields
across a forecast ensemble.
Parameters
----------
type : str
The type of array to test; either 'numpy' or 'dask'.
Returns
-------
tuple of floats
The total wall and process times.
"""
refpath = Path("/lustre/scratch/rmanser/wrfref/wrfoutREFd02")
mempath = Path("/lustre/scratch/rmanser/test_ens/2016052812/")
ref = xr.open_dataset(refpath)
# All members are loaded implicitly as dask arrays, regardless of `type`
members = xr.open_mfdataset(
mempath.glob("mem*/wrfout_d02_2016-05-30_12:00:00"),
concat_dim='members',
combine='nested',
)
# For dask arrays, we select the individual variable dask arrays
if type == 'dask':
pressure = (members.P + ref.PB).data * units(ref.PB.units)
theta = (members.T + ref.T00).data * units(ref.T00.units)
mixing_ratio = members.QVAPOR.data * units('dimensionless')
# For numpy arrays, we force variable arrays to be loaded into memory
elif type == 'numpy':
pressure = (members.P + ref.PB).values * units(ref.PB.units)
theta = (members.T + ref.T00).values * units(ref.T00.units)
mixing_ratio = members.QVAPOR.values * units('dimensionless')
start_wall = time.time()
start_proc = time.process_time()
temperature = mpcalc.temperature_from_potential_temperature(
pressure, theta)
relative_humidity = mpcalc.relative_humidity_from_mixing_ratio(
mixing_ratio, temperature, pressure)
# We don't call metpy's function here because it implicitly triggers
# a dask.compute() call
dewpoint = dewpoint_from_relative_humidity(temperature, relative_humidity)
if type == 'dask':
td = dewpoint.compute()
end_wall = time.time()
end_proc = time.process_time()
total_wall = end_wall - start_wall
total_proc = end_proc - start_proc
return total_wall, total_proc
def test_temperature_from_potential_temperature():
"""Test temperature_from_potential_temperature calculation."""
theta = np.array([286.12859679, 288.22362587, 290.31865495, 292.41368403]) * units.kelvin
pres = np.array([850] * 4) * units.mbar
real_t = np.array([273.15, 275.15, 277.15, 279.15]) * units.kelvin
assert_array_almost_equal(temperature_from_potential_temperature(pres, theta),
real_t, 2)
def calc_T_from_theta(theta, p):
"""
Input :
theta : potential temperature (K)
p : pressure (hPa)
Output :
T : Temperature values
Function to estimate temperature from potential temperature and pressure,
in the given dataset. This function uses MetPy's
functions to get T:
(i) mpcalc.temperature_from_potential_temperature()
"""
T = (mpcalc.temperature_from_potential_temperature(
p * units.hPa,
theta * units.kelvin,
).magnitude - 273.15)
return T
def calc_T_from_theta(dataset):
"""
Input :
dataset : Dataset
Output :
T : Temperature values
Function to estimate temperature from potential temperature and pressure,
in the given dataset. This function uses MetPy's
functions to get T:
(i) mpcalc.temperature_from_potential_temperature()
"""
ta = mpcalc.temperature_from_potential_temperature(
dataset.p.values * units.Pa,
dataset.theta.values * units.kelvin,
).magnitude
return ta