From 87eb2f363e9f6d3434eca224a04a2238efa4916c Mon Sep 17 00:00:00 2001 From: Timothy Raupach Date: Thu, 15 Jul 2021 08:02:49 +1000 Subject: [PATCH 1/2] Added xarray parcel functions and (unintegrated) unit tests. --- tests/calc/test_xarray_parcel.py | 1007 ++++++++++++++++++++++++++++++ 1 file changed, 1007 insertions(+) create mode 100644 tests/calc/test_xarray_parcel.py diff --git a/tests/calc/test_xarray_parcel.py b/tests/calc/test_xarray_parcel.py new file mode 100644 index 00000000000..cf48072a610 --- /dev/null +++ b/tests/calc/test_xarray_parcel.py @@ -0,0 +1,1007 @@ +# Copyright (c) 2008,2015,2016,2017,2018,2019 MetPy Developers. +# Distributed under the terms of the BSD 3-Clause License. +# SPDX-License-Identifier: BSD-3-Clause +# +# Adapted from test_thermo for xarray parcel functions. + +import metpy +import numpy as np +import pytest +import xarray as xarray +import modules.parcel_functions as parcel +from numpy.testing import assert_almost_equal, assert_array_almost_equal +from metpy.units import units + +# Load xarray-parcel's lookup tables. +parcel.load_moist_adiabat_lookups() + +def run_all_tests(): + """Run all the tests.""" + test_dry_lapse() + test_dry_lapse_2_levels() + test_moist_lapse() + test_moist_lapse_ref_pres() + test_moist_lapse_scalar() + test_moist_lapse_uniform() + test_parcel_profile() + test_parcel_profile_lcl() + test_parcel_profile_saturated() + test_lcl() + # test_lcl_nans() -- metpy lcl implementation sometimes fails to converge. + test_lfc_basic() + test_lfc_ml() + test_lfc_ml2() + test_lfc_intersection() + test_no_lfc() + test_lfc_inversion() + test_lfc_equals_lcl() + test_sensitive_sounding() + test_lfc_sfc_precision() + test_lfc_pos_area_below_lcl() + test_el() + test_el_ml() + test_no_el() + test_no_el_multi_crossing() + test_lfc_and_el_below_lcl() + test_el_lfc_equals_lcl() + test_el_small_surface_instability() + test_no_el_parcel_colder() + test_el_below_lcl() + test_cape_cin() + test_cape_cin_no_el() + test_cape_cin_no_lfc() + test_most_unstable_parcel() + test_surface_based_cape_cin() + test_profile_with_nans() + test_most_unstable_cape_cin_surface() + test_mixed_parcel() + test_mixed_layer_cape_cin() + test_multiple_lfcs_el_simple() + test_mixed_layer() + test_lfc_not_below_lcl() + test_cape_cin_custom_profile() + test_parcel_profile_below_lcl() + test_lcl_convergence_issue() + test_cape_cin_value_error() + test_lcl_grid_surface_lcls() + test_lifted_index() + print('All tests passed.') + +def run_moist_lapse_tests_looser(): + """Run all the tests, with looser matching requirements.""" + test_moist_lapse() + test_moist_lapse_ref_pres() + test_moist_lapse_scalar() + test_moist_lapse_uniform(dp=2) + print('Moist lapse tests passed.') + +def metpy_moist_lapse(pressure, parcel_temperature, parcel_pressure=None): + """A wrapper for metpy's moist_lapse().""" + + if isinstance(parcel_pressure, xarray.DataArray): + parcel_pressure = parcel_pressure.values + if isinstance(pressure, xarray.DataArray): + pressure = pressure.values + if isinstance(parcel_temperature, xarray.DataArray): + parcel_temperature = parcel_temperature.values + + if not parcel_pressure is None: + parcel_pressure = parcel_pressure * units.hPa + + return metpy.calc.moist_lapse(pressure=pressure*units.hPa, + temperature=parcel_temperature*units.K, + reference_pressure=parcel_pressure).m + +def vert_array(x, units): + """Make an xarray object with one dimension.""" + return xarray.DataArray(x, dims='model_level_number', + coords={'model_level_number': np.arange(1,len(x)+1)}, + attrs={'units': units}) + +def test_dry_lapse(): + """Test dry_lapse calculation.""" + levels = vert_array([1000, 900, 864.89], 'hPa') + temps = parcel.dry_lapse(pressure=levels, parcel_temperature=303.15) + assert_array_almost_equal(temps.values, np.array([303.15, 294.16, 290.83]), decimal=2) + +def test_dry_lapse_2_levels(): + """Test dry_lapse calculation when given only two levels.""" + levels = vert_array([1000., 500.], 'hPa') + temps = parcel.dry_lapse(pressure=levels, parcel_temperature=293.) + assert_array_almost_equal(temps, [293., 240.3583], 4) + +def test_moist_lapse(): + """Test moist_lapse calculation.""" + levels = vert_array([1000., 800., 600., 500., 400.], 'hPa') + temp = parcel.moist_lapse(pressure=levels, parcel_temperature=293.) + assert_array_almost_equal(temp, [293, 284.64, 272.81, 264.42, 252.91], 2) + +def test_moist_lapse_ref_pres(): + """Test moist_lapse with a reference pressure.""" + levels = vert_array([1050., 800., 600., 500., 400.], 'hPa') + temp = parcel.moist_lapse(pressure=levels, parcel_temperature=293, parcel_pressure=1000) + assert_array_almost_equal(temp, [294.76, 284.64, 272.81, 264.42, 252.91], 2) + +def test_moist_lapse_scalar(): + """Test moist_lapse when given a scalar desired pressure and a reference pressure.""" + levels = vert_array([800.], 'hPa') + temp = parcel.moist_lapse(pressure=levels, parcel_temperature=293, parcel_pressure=1000) + assert_array_almost_equal(temp, [284.64], 2) + +def test_moist_lapse_uniform(dp=7): + """Test moist_lapse when given a uniform array of pressures.""" + levels = vert_array([900., 900., 900.], 'hPa') + temp = parcel.moist_lapse(pressure=levels, parcel_temperature=293.15) + assert_almost_equal(temp, np.array([293.15, 293.15, 293.15]), dp) + +def test_parcel_profile(dp=2): + """Test parcel profile calculation.""" + levels = vert_array([1000., 900., 800., 700., 600., 500., 400.], 'hPa') + true_prof = np.array([303.15, 294.16, 288.026, 283.073, 277.058, 269.402, 258.966]) + + parcel_temperature = xarray.DataArray(303.15, attrs={'units': 'K'}) + parcel_pressure = xarray.DataArray(1000, attrs={'units': 'hPa'}) + parcel_dewpoint = xarray.DataArray(293.15, attrs={'units': 'K'}) + + prof = parcel.parcel_profile(pressure=levels, parcel_pressure=parcel_pressure, + parcel_temperature=parcel_temperature, + parcel_dewpoint=parcel_dewpoint) + + assert_array_almost_equal(prof.temperature, true_prof, dp) + +def test_parcel_profile_lcl(dp=3): + """Test parcel profile with lcl calculation.""" + p = vert_array([1004., 1000., 943., 928., 925., 850., 839., 749., 700., 699.], 'hPa') + t = vert_array([24.2, 24., 20.2, 21.6, 21.4, 20.4, 20.2, 14.4, 13.2, 13.], 'K') + 273.15 + + true_t = np.array([24.2, 24., 22.047, 20.2, 21.6, 21.4, 20.4, 20.2, 14.4, 13.2, 13.]) + 273.15 + true_p = np.array([1004., 1000., 970.711, 943., 928., 925., 850., 839., 749., 700., 699.]) + true_prof = np.array([297.35, 297.01, 294.5, 293.48, 292.92, 292.81, 289.79, 289.32, + 285.15, 282.59, 282.53]) + + parcel_temperature = xarray.DataArray(24.2+273.15, attrs={'units': 'K'}) + parcel_pressure = xarray.DataArray(1004, attrs={'units': 'hPa'}) + parcel_dewpoint = xarray.DataArray(21.9+273.15, attrs={'units': 'K'}) + + prof = parcel.parcel_profile(pressure=p, + parcel_pressure=parcel_pressure, + parcel_temperature=parcel_temperature, + parcel_dewpoint=parcel_dewpoint) + prof = parcel.add_lcl_to_profile(profile=prof, temperature=t) + + assert_array_almost_equal(prof.pressure, true_p, dp) + assert_array_almost_equal(prof.environment_temperature, true_t, dp) + assert_array_almost_equal(prof.temperature, true_prof, dp-1) + +def test_parcel_profile_saturated(dp=2): + """Test parcel_profile works when LCL in levels (issue #232).""" + levels = vert_array([1000., 700., 500.], 'hPa') + true_prof = np.array([296.95, 284.381, 271.123]) + + parcel_temperature = xarray.DataArray(23.8+273.15, attrs={'units': 'K'}) + parcel_pressure = xarray.DataArray(1000, attrs={'units': 'hPa'}) + parcel_dewpoint = xarray.DataArray(23.8+273.15, attrs={'units': 'K'}) + + prof = parcel.parcel_profile(pressure=levels, + parcel_pressure=parcel_pressure, + parcel_temperature=parcel_temperature, + parcel_dewpoint=parcel_dewpoint) + assert_array_almost_equal(prof.temperature, true_prof, dp) + +def test_lcl(): + """Test LCL calculation.""" + parcel_temperature = xarray.DataArray(30+273.15, attrs={'units': 'K'}) + parcel_pressure = xarray.DataArray(1000, attrs={'units': 'hPa'}) + parcel_dewpoint = xarray.DataArray(20+273.15, attrs={'units': 'K'}) + lcl = parcel.lcl(parcel_pressure=parcel_pressure, + parcel_temperature=parcel_temperature, + parcel_dewpoint=parcel_dewpoint) + assert_almost_equal(lcl.lcl_pressure, 864.806, 2) + assert_almost_equal(lcl.lcl_temperature, 17.676+273.15, 2) + +def test_lcl_nans(): + """Test LCL calculation on data with nans.""" + p = xarray.DataArray([900., 900., 900., 900.], attrs={'units': 'hPa'}) + t = xarray.DataArray(np.array([np.nan, 25., 25., 25.])+273.15, attrs={'units': 'K'}) + d = xarray.DataArray(np.array([20., 20., np.nan, 20.])+273.15, attrs={'units': 'K'}) + lcl = parcel.lcl(parcel_pressure=p, parcel_temperature=t, parcel_dewpoint=d) + + assert_array_almost_equal(lcl.lcl_pressure, + np.array([np.nan, 836.4098648012595, + np.nan, 836.4098648012595])) + assert_array_almost_equal(lcl.lcl_temperature, + np.array([np.nan, 18.82281982535794, + np.nan, 18.82281982535794])+273.15) + +def test_lfc_basic(dp=2): + """Test LFC calculation.""" + levels = vert_array([959., 779.2, 751.3, 724.3, 700., 269.], 'hPa') + temperatures = vert_array(np.array([22.2, 14.6, 12., 9.4, 7., -49.])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4, -10., -53.2])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + lfc = parcel.lfc_el(profile=profile) + assert_almost_equal(lfc.lfc_pressure, 727.371, dp) + assert_almost_equal(lfc.lfc_temperature, 9.705+273.15, dp) + +def test_lfc_ml(dp=2): + """Test Mixed-Layer LFC calculation.""" + levels = vert_array([959., 779.2, 751.3, 724.3, 700., 269.], 'hPa') + levels.name = 'pressure' + temperatures = vert_array(np.array([22.2, 14.6, 12., 9.4, 7., -49.])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4, -10., -53.2])+273.15, 'K') + mixed = parcel.mixed_parcel(pressure=levels, temperature=temperatures, dewpoint=dewpoints) + mixed_parcel_prof = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=mixed.pressure, + parcel_temperature=mixed.temperature, + parcel_dewpoint=mixed.dewpoint) + lfc = parcel.lfc_el(profile=mixed_parcel_prof) + assert_almost_equal(lfc.lfc_pressure, 601.225, dp) + assert_almost_equal(lfc.lfc_temperature, -1.90688+273.15, dp) + +def test_lfc_ml2(): + """Test a mixed-layer LFC calculation that previously crashed.""" + levels = vert_array([1024.95703125, 1016.61474609, 1005.33056641, 991.08544922, 973.4163208, + 951.3381958, 924.82836914, 898.25482178, 873.46124268, 848.69830322, + 823.92553711, 788.49304199, 743.44580078, 700.50970459, 659.62017822, + 620.70861816, 583.69421387, 548.49719238, 515.03826904, 483.24401855, + 453.0418396, 424.36477661, 397.1505127, 371.33441162, 346.85922241, + 323.66995239, 301.70935059, 280.92651367, 261.27053833, 242.69168091, + 225.14237976, 208.57781982, 192.95333862, 178.22599792, 164.39630127, + 151.54336548, 139.68635559, 128.74923706, 118.6588974, 109.35111237, + 100.76405334, 92.84288025, 85.53556824, 78.79430389, 72.57549286, + 66.83885193, 61.54678726, 56.66480637, 52.16108322], 'hPa') + levels.name = 'pressure' + temperatures = vert_array(np.array([6.00750732, 5.14892578, 4.177948, 3.00268555, 1.55535889, + -0.25527954, -1.93988037, -3.57766724, -4.40600586, -4.19238281, + -3.71185303, -4.47943115, -6.81280518, -8.08685303, -8.41287231, + -10.79302979, -14.13262939, -16.85784912, -19.51675415, + -22.28689575, -24.99938965, -27.79664612, -30.90414429, + -34.49435425, -38.438797, -42.27981567, -45.99230957, + -49.75340271, -53.58230591, -57.30686951, -60.76026917, + -63.92070007, -66.72470093, -68.97846985, -70.4264679, + -71.16407776, -71.53797913, -71.64375305, -71.52735901, + -71.53523254, -71.61097717, -71.92687988, -72.68682861, + -74.129776, -76.02471924, -76.88977051, -76.26008606, + -75.90351868, -76.15809631])+273.15, 'K') + dewpoints = vert_array(np.array([4.50012302, 3.42483997, 2.78102994, 2.24474645, 1.593485, -0.9440815, + -3.8044982, -3.55629468, -9.7376976, -10.2950449, -9.67498302, + -10.30486488, -8.70559597, -8.71669006, -12.66509628, -18.6697197, + -23.00351334, -29.46240425, -36.82178497, -41.68824768, -44.50320816, + -48.54426575, -52.50753403, -51.09564209, -48.92690659, -49.97380829, + -51.57516098, -52.62096405, -54.24332809, -57.09109879, -60.5596199, + -63.93486404, -67.07530212, -70.01263428, -72.9258728, -76.12271881, + -79.49847412, -82.2350769, -83.91127014, -84.95665741, -85.61238861, + -86.16391754, -86.7653656, -87.34436035, -87.87495422, -88.34281921, + -88.74453735, -89.04680634, -89.26436615])+273.15, 'K') + mixed = parcel.mixed_parcel(pressure=levels, temperature=temperatures, dewpoint=dewpoints) + mixed_parcel_prof = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=mixed.pressure, + parcel_temperature=mixed.temperature, + parcel_dewpoint=mixed.dewpoint) + lfc = parcel.lfc_el(profile=mixed_parcel_prof) + assert_almost_equal(lfc.lfc_pressure, 962.34, 2) + assert_almost_equal(lfc.lfc_temperature, 0.767+273.15, 2) + +def test_lfc_intersection(dp=2): + """Test LFC calculation when LFC is below a tricky intersection.""" + levels = vert_array([1024.957, 930., 924.828, 898.255, 873.461, 848.698, 823.926, 788.493], 'hPa') + levels.name = 'pressure' + temperatures = vert_array(np.array([6.008, -10., -6.94, -8.58, -4.41, -4.19, -3.71, -4.48])+273.15, 'K') + dewpoints = vert_array(np.array([5., -10., -7., -9., -4.5, -4.2, -3.8, -4.5])+273.15, 'K') + + mixed = parcel.mixed_parcel(pressure=levels, temperature=temperatures, dewpoint=dewpoints) + + # Calculate parcel profile without LCL, as per metpy unit tests. + mixed_parcel_prof = parcel.parcel_profile(pressure=levels, + parcel_pressure=mixed.pressure, + parcel_temperature=mixed.temperature, + parcel_dewpoint=mixed.dewpoint) + mixed_parcel_prof['environment_temperature'] = temperatures + lfc = parcel.lfc_el(profile=mixed_parcel_prof) + assert_almost_equal(lfc.lfc_pressure, 981.620, dp) + +def test_no_lfc(): + """Test LFC calculation when there is no LFC in the data.""" + levels = vert_array([959., 867.9, 779.2, 647.5, 472.5, 321.9, 251.], 'hPa') + temperatures = vert_array(np.array([22.2, 17.4, 14.6, 1.4, -17.6, -39.4, -52.5])+273.15, 'K') + dewpoints = vert_array(np.array([9., 4.3, -21.2, -26.7, -31., -53.3, -66.7])+273.15, 'K') + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + lfc = parcel.lfc_el(profile=profile) + assert(np.isnan(lfc.lfc_pressure)) + assert(np.isnan(lfc.lfc_temperature)) + +def test_lfc_inversion(dp=2): + """Test LFC when there is an inversion to be sure we don't pick that.""" + levels = vert_array([963., 789., 782.3, 754.8, 728.1, 727., 700., + 571., 450., 300., 248.], 'hPa') + temperatures = vert_array(np.array([25.4, 18.4, 17.8, 15.4, 12.9, 12.8, + 10., -3.9, -16.3, -41.1, -51.5])+273.15, 'K') + dewpoints = vert_array(np.array([20.4, 0.4, -0.5, -4.3, -8., -8.2, -9., + -23.9, -33.3, -54.1, -63.5])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + lfc = parcel.lfc_el(profile=profile) + + assert_almost_equal(lfc.lfc_pressure, 705.8806 , dp) + assert_almost_equal(lfc.lfc_temperature, 10.6232+273.15, dp) + +def test_lfc_equals_lcl(): + """Test LFC when there is no cap and the lfc is equal to the lcl.""" + levels = vert_array([912., 905.3, 874.4, 850., 815.1, 786.6, 759.1, + 748., 732.2, 700., 654.8], 'hPa') + temperatures = vert_array(np.array([29.4, 28.7, 25.2, 22.4, 19.4, 16.8, + 14.0, 13.2, 12.6, 11.4, 7.1])+273.15, 'K') + dewpoints = vert_array(np.array([18.4, 18.1, 16.6, 15.4, 13.2, 11.4, 9.6, + 8.8, 0., -18.6, -22.9])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + lfc = parcel.lfc_el(profile=profile) + assert_almost_equal(lfc.lfc_pressure, 777.0786, 2) + assert_almost_equal(lfc.lfc_temperature, 15.8714+273.15, 2) + +def test_sensitive_sounding(dp=2): + """Test quantities for a sensitive sounding (#902).""" + # This sounding has a very small positive area in the low level. It's only captured + # properly if the parcel profile includes the LCL, otherwise it breaks LFC and CAPE + levels = vert_array([1004., 1000., 943., 928., 925., 850., 839., 749., 700., 699., + 603., 500., 404., 400., 363., 306., 300., 250., 213., 200., + 176., 150.], 'hPa') + temperatures = vert_array(np.array([24.2, 24., 20.2, 21.6, 21.4, 20.4, 20.2, 14.4, + 13.2, 13., 6.8, -3.3, -13.1, -13.7, -17.9, -25.5, + -26.9, -37.9, -46.7, -48.7, -52.1, -58.9])+273.15, 'K') + dewpoints = vert_array(np.array([21.9, 22.1, 19.2, 20.5, 20.4, 18.4, 17.4, 8.4, -2.8, + -3.0, -15.2, -20.3, -29.1, -27.7, -24.9, -39.5, -41.9, + -51.9, -60.7, -62.7, -65.1, -71.9])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + lfc = parcel.lfc_el(profile=profile) + + assert_almost_equal(lfc.lfc_pressure, 947.422, dp) + assert_almost_equal(lfc.lfc_temperature, 20.498+273.15, dp) + + cape_cin = parcel.surface_based_cape_cin(pressure=levels, temperature=temperatures, + dewpoint=dewpoints) + assert_almost_equal(cape_cin.cape, 0.1115, 3) + assert_almost_equal(cape_cin.cin, -6.0866, 3) + +def test_lfc_sfc_precision(): + """Test LFC when there are precision issues with the parcel path.""" + levels = vert_array([839., 819.4, 816., 807., 790.7, 763., 736.2, 722., 710.1, 700.], 'hPa') + temperatures = vert_array(np.array([20.6, 22.3, 22.6, 22.2, 20.9, 18.7, 16.4, + 15.2, 13.9, 12.8])+273.15, 'K') + dewpoints = vert_array(np.array([10.6, 8., 7.6, 6.2, 5.7, 4.7, 3.7, 3.2, 3., 2.8])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + lfc = parcel.lfc_el(profile=profile) + assert(np.isnan(lfc.lfc_pressure)) + assert(np.isnan(lfc.lfc_temperature)) + +def test_lfc_pos_area_below_lcl(): + """Test LFC when there is positive area below the LCL (#1003).""" + levels = vert_array([902.1554, 897.9034, 893.6506, 889.4047, 883.063, 874.6284, 866.2387, 857.887, + 849.5506, 841.2686, 833.0042, 824.7891, 812.5049, 796.2104, 776.0027, 751.9025, + 727.9612, 704.1409, 680.4028, 656.7156, 629.077, 597.4286, 565.6315, 533.5961, + 501.2452, 468.493, 435.2486, 401.4239, 366.9387, 331.7026, 295.6319, 258.6428, + 220.9178, 182.9384, 144.959, 106.9778, 69.00213], 'hPa') + temperatures = vert_array(np.array([-3.039381, -3.703779, -4.15996, -4.562574, -5.131827, -5.856229, -6.568434, + -7.276881, -7.985013, -8.670911, -8.958063, -7.631381, -6.05927, -5.083627, + -5.11576, -5.687552, -5.453021, -4.981445, -5.236665, -6.324916, -8.434324, + -11.58795, -14.99297, -18.45947, -21.92021, -25.40522, -28.914, -32.78637, + -37.7179, -43.56836, -49.61077, -54.24449, -56.16666, -57.03775, -58.28041, + -60.86264, -64.21677])+273.15, 'K') + dewpoints = vert_array(np.array([-22.08774, -22.18181, -22.2508, -22.31323, -22.4024, -22.51582, -22.62526, + -22.72919, -22.82095, -22.86173, -22.49489, -21.66936, -21.67332, -21.94054, + -23.63561, -27.17466, -31.87395, -38.31725, -44.54717, -46.99218, -43.17544, + -37.40019, -34.3351, -36.42896, -42.1396, -46.95909, -49.36232, -48.94634, + -47.90178, -49.97902, -55.02753, -63.06276, -72.53742, -88.81377, -93.54573, + -92.92464, -91.57479])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + lfc = parcel.lfc_el(profile=profile) + assert(np.isnan(lfc.lfc_pressure)) + assert(np.isnan(lfc.lfc_temperature)) + +def test_el(): + """Test equilibrium layer calculation.""" + levels = vert_array([959., 779.2, 751.3, 724.3, 700., 269.], 'hPa') + temperatures = vert_array(np.array([22.2, 14.6, 12., 9.4, 7., -38.])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4, -10., -53.2])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + el = parcel.lfc_el(profile=profile) + assert_almost_equal(el.el_pressure, 471.83286, 3) + assert_almost_equal(el.el_temperature, -11.5603+273.15, 3) + +def test_el_ml(): + """Test equilibrium layer calculation for a mixed parcel.""" + levels = vert_array([959., 779.2, 751.3, 724.3, 700., 400., 269.], 'hPa') + levels.name = 'pressure' + temperatures = vert_array(np.array([22.2, 14.6, 12., 9.4, 7., -25., -35.])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4, -10., -35., -53.2])+273.15, 'K') + + mixed = parcel.mixed_parcel(pressure=levels, temperature=temperatures, dewpoint=dewpoints) + mixed_parcel_prof = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=mixed.pressure, + parcel_temperature=mixed.temperature, + parcel_dewpoint=mixed.dewpoint) + el = parcel.lfc_el(profile=mixed_parcel_prof) + assert_almost_equal(el.el_pressure, 350.0561, 3) + assert_almost_equal(el.el_temperature, -28.36156+273.15, 3) + +def test_no_el(): + """Test equilibrium layer calculation when there is no EL in the data.""" + levels = vert_array([959., 867.9, 779.2, 647.5, 472.5, 321.9, 251.], 'hPa') + temperatures = vert_array(np.array([22.2, 17.4, 14.6, 1.4, -17.6, -39.4, -52.5])+273.15, 'K') + dewpoints = vert_array(np.array([19., 14.3, -11.2, -16.7, -21., -43.3, -56.7])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + el = parcel.lfc_el(profile=profile) + assert(np.isnan(el.el_pressure)) + assert(np.isnan(el.el_temperature)) + +def test_no_el_multi_crossing(): + """Test el calculation with no el and severel parcel path-profile crossings.""" + levels = vert_array([918., 911., 880., 873.9, 850., 848., 843.5, 818., 813.8, 785., + 773., 763., 757.5, 730.5, 700., 679., 654.4, 645., + 643.9], 'hPa') + temperatures = vert_array(np.array([24.2, 22.8, 19.6, 19.1, 17., 16.8, 16.5, 15., 14.9, 14.4, 16.4, + 16.2, 15.7, 13.4, 10.6, 8.4, 5.7, 4.6, 4.5])+273.15, 'K') + dewpoints = vert_array(np.array([19.5, 17.8, 16.7, 16.5, 15.8, 15.7, 15.3, 13.1, 12.9, 11.9, 6.4, + 3.2, 2.6, -0.6, -4.4, -6.6, -9.3, -10.4, -10.5])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + el = parcel.lfc_el(profile=profile) + assert(np.isnan(el.el_pressure)) + assert(np.isnan(el.el_temperature)) + +def test_lfc_and_el_below_lcl(): + """Test that LFC and EL are returned as NaN if both are below LCL.""" + dewpoints = vert_array([264.5351, 261.13443, 259.0122, 252.30063, 248.58017, 242.66582], 'K') + temperatures = vert_array([273.09723, 268.40173, 263.56207, 260.257, 256.63538, 252.91345], 'K') + levels = vert_array([1017.16, 950, 900, 850, 800, 750], 'hPa') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + el = parcel.lfc_el(profile=profile) + assert(np.isnan(el.el_pressure)) + assert(np.isnan(el.el_temperature)) + assert(np.isnan(el.lfc_pressure)) + assert(np.isnan(el.lfc_temperature)) + +def test_el_lfc_equals_lcl(): + """Test equilibrium layer calculation when the lfc equals the lcl.""" + levels = vert_array([912., 905.3, 874.4, 850., 815.1, 786.6, 759.1, 748., + 732.3, 700., 654.8, 606.8, 562.4, 501.8, 500., 482., + 400., 393.3, 317.1, 307., 300., 252.7, 250., 200., + 199.3, 197., 190., 172., 156.6, 150., 122.9, 112., + 106.2, 100.], 'hPa') + temperatures = vert_array(np.array([29.4, 28.7, 25.2, 22.4, 19.4, 16.8, 14.3, + 13.2, 12.6, 11.4, 7.1, 2.2, -2.7, -10.1, + -10.3, -12.4, -23.3, -24.4, -38., -40.1, -41.1, + -49.8, -50.3, -59.1, -59.1, -59.3, -59.7, -56.3, + -56.9, -57.1, -59.1, -60.1, -58.6, -56.9])+273.15, 'K') + dewpoints = vert_array(np.array([18.4, 18.1, 16.6, 15.4, 13.2, 11.4, 9.6, 8.8, 0., + -18.6, -22.9, -27.8, -32.7, -40.1, -40.3, -42.4, -53.3, + -54.4, -68., -70.1, -70., -70., -70., -70., -70., -70., + -70., -70., -70., -70., -70., -70., -70., -70.])+273.15, 'K') + + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + el = parcel.lfc_el(profile=profile) + + assert_almost_equal(el.el_pressure, 175.7663, 3) + assert_almost_equal(el.el_temperature, -57.03994+273.15, 3) + +def test_el_small_surface_instability(): + """Test that no EL is found when there is a small pocket of instability at the sfc.""" + levels = vert_array([959., 931.3, 925., 899.3, 892., 867.9, 850., 814., + 807.9, 790., 779.2, 751.3, 724.3, 700., 655., 647.5, + 599.4, 554.7, 550., 500.], 'hPa') + temperatures = vert_array(np.array([22.2, 20.2, 19.8, 18.4, 18., 17.4, 17., 15.4, 15.4, + 15.6, 14.6, 12., 9.4, 7., 2.2, 1.4, -4.2, -9.7, + -10.3, -14.9])+273.15, 'K') + dewpoints = vert_array(np.array([20., 18.5, 18.1, 17.9, 17.8, 15.3, 13.5, 6.4, 2.2, + -10.4, -10.2, -9.8, -9.4, -9., -15.8, -15.7, -14.8, -14., + -13.9, -17.9])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + el = parcel.lfc_el(profile=profile) + assert(np.isnan(el.el_pressure)) + assert(np.isnan(el.el_temperature)) + +def test_no_el_parcel_colder(): + """Test no EL when parcel stays colder than environment. INL 20170925-12Z.""" + levels = vert_array([974., 946., 925., 877.2, 866., 850., 814.6, 785., + 756.6, 739., 729.1, 700., 686., 671., 641., 613., + 603., 586., 571., 559.3, 539., 533., 500., 491., + 477.9, 413., 390., 378., 345., 336.], 'hPa') + temperatures = vert_array(np.array([10., 8.4, 7.6, 5.9, 7.2, 7.6, 6.8, 7.1, 7.7, + 7.8, 7.7, 5.6, 4.6, 3.4, 0.6, -0.9, -1.1, -3.1, + -4.7, -4.7, -6.9, -7.5, -11.1, -10.9, -12.1, -20.5, -23.5, + -24.7, -30.5, -31.7])+273.15, 'K') + dewpoints = vert_array(np.array([8.9, 8.4, 7.6, 5.9, 7.2, 7., 5., 3.6, 0.3, + -4.2, -12.8, -12.4, -8.4, -8.6, -6.4, -7.9, -11.1, -14.1, + -8.8, -28.1, -18.9, -14.5, -15.2, -15.1, -21.6, -41.5, -45.5, + -29.6, -30.6, -32.1])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + el = parcel.lfc_el(profile=profile) + assert(np.isnan(el.el_pressure)) + assert(np.isnan(el.el_temperature)) + +def test_el_below_lcl(): + """Test LFC when there is positive area below the LCL (#1003).""" + levels = vert_array([902.1554, 897.9034, 893.6506, 889.4047, 883.063, 874.6284, 866.2387, 857.887, + 849.5506, 841.2686, 833.0042, 824.7891, 812.5049, 796.2104, 776.0027, 751.9025, + 727.9612, 704.1409, 680.4028, 656.7156, 629.077, 597.4286, 565.6315, 533.5961, + 501.2452, 468.493, 435.2486, 401.4239, 366.9387, 331.7026, 295.6319, 258.6428, + 220.9178, 182.9384, 144.959, 106.9778, 69.00213], 'hPa') + temperatures = vert_array(np.array([-3.039381, -3.703779, -4.15996, -4.562574, -5.131827, -5.856229, -6.568434, + -7.276881, -7.985013, -8.670911, -8.958063, -7.631381, -6.05927, -5.083627, + -5.11576, -5.687552, -5.453021, -4.981445, -5.236665, -6.324916, -8.434324, + -11.58795, -14.99297, -18.45947, -21.92021, -25.40522, -28.914, -32.78637, + -37.7179, -43.56836, -49.61077, -54.24449, -56.16666, -57.03775, -58.28041, + -60.86264, -64.21677])+273.15, 'K') + dewpoints = vert_array(np.array([-22.08774, -22.18181, -22.2508, -22.31323, -22.4024, -22.51582, -22.62526, + -22.72919, -22.82095, -22.86173, -22.49489, -21.66936, -21.67332, -21.94054, + -23.63561, -27.17466, -31.87395, -38.31725, -44.54717, -46.99218, -43.17544, + -37.40019, -34.3351, -36.42896, -42.1396, -46.95909, -49.36232, -48.94634, + -47.90178, -49.97902, -55.02753, -63.06276, -72.53742, -88.81377, -93.54573, + -92.92464, -91.57479])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + el = parcel.lfc_el(profile=profile) + assert(np.isnan(el.el_pressure)) + assert(np.isnan(el.el_temperature)) + +def test_cape_cin(): + """Test the basic CAPE and CIN calculation.""" + levels = vert_array([959., 779.2, 751.3, 724.3, 700., 269.], 'hPa') + temperatures = vert_array(np.array([22.2, 14.6, 12., 9.4, 7., -38.])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4, -10., -53.2])+273.15, 'K') + + # Calculate parcel profile without LCL, as per metpy unit tests. + profile = parcel.parcel_profile(pressure=levels, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + profile['environment_temperature'] = temperatures + lfc = parcel.lfc_el(profile=profile) + cape_cin = parcel.cape_cin_base(pressure=levels, + temperature=temperatures, + lfc_pressure=lfc.lfc_pressure, + el_pressure=lfc.el_pressure, + parcel_profile=profile) + + assert_almost_equal(cape_cin.cape, 75.05354, 2) + assert_almost_equal(cape_cin.cin, -89.890078, 2) + +def test_cape_cin_no_el(): + """Test that CAPE works with no EL.""" + levels = vert_array([959., 779.2, 751.3, 724.3], 'hPa') + temperatures = vert_array(np.array([22.2, 14.6, 12., 9.4])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4])+273.15, 'K') + + # Calculate parcel profile without LCL, as per metpy unit tests. + profile = parcel.parcel_profile(pressure=levels, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + profile['environment_temperature'] = temperatures + lfc = parcel.lfc_el(profile=profile) + cape_cin = parcel.cape_cin_base(pressure=levels, + temperature=temperatures, + lfc_pressure=lfc.lfc_pressure, + el_pressure=lfc.el_pressure, + parcel_profile=profile) + + assert_almost_equal(cape_cin.cape, 0.08610409, 2) + assert_almost_equal(cape_cin.cin, -89.8900784, 2) + +def test_cape_cin_no_lfc(): + """Test that CAPE is zero with no LFC.""" + levels = vert_array([959., 779.2, 751.3, 724.3, 700., 269.], 'hPa') + temperatures = vert_array(np.array([22.2, 24.6, 22., 20.4, 18., -10.])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4, -10., -53.2])+273.15, 'K') + + # Calculate parcel profile without LCL, as per metpy unit tests. + profile = parcel.parcel_profile(pressure=levels, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + profile['environment_temperature'] = temperatures + lfc = parcel.lfc_el(profile=profile) + cape_cin = parcel.cape_cin_base(pressure=levels, + temperature=temperatures, + lfc_pressure=lfc.lfc_pressure, + el_pressure=lfc.el_pressure, + parcel_profile=profile) + + assert_almost_equal(cape_cin.cape, 0.0, 2) + assert_almost_equal(cape_cin.cin, 0.0, 2) + +def test_most_unstable_parcel(): + """Test calculating the most unstable parcel.""" + levels = vert_array([1000., 959., 867.9], 'hPa') + levels.name = 'pressure' + temperatures = vert_array(np.array([18.2, 22.2, 17.4])+273.15, 'K') + temperatures.name = 'temperature' + dewpoints = vert_array(np.array([19., 19., 14.3])+273.15, 'K') + dewpoints.name = 'dewpoint' + + ret = parcel.most_unstable_parcel(dat=xarray.merge([levels, temperatures, dewpoints]), + depth=100) + + assert_almost_equal(ret.pressure, 959.0, 6) + assert_almost_equal(ret.temperature, 22.2+273.15, 6) + assert_almost_equal(ret.dewpoint, 19.0+273.15, 6) + +def test_surface_based_cape_cin(): + """Test the surface-based CAPE and CIN calculation.""" + levels = vert_array([959., 779.2, 751.3, 724.3, 700., 269.], 'hPa') + temperatures = vert_array(np.array([22.2, 14.6, 12., 9.4, 7., -38.])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4, -10., -53.2])+273.15, 'K') + + cape_cin = parcel.surface_based_cape_cin(pressure=levels, temperature=temperatures, + dewpoint=dewpoints) + + assert_almost_equal(cape_cin.cape, 75.0535446, 2) + assert_almost_equal(cape_cin.cin, -136.685967, 2) + +def test_profile_with_nans(): + """Test a profile with nans to make sure it calculates functions appropriately (#1187).""" + levels = vert_array([1001, 1000, 997, 977.9, 977, 957, 937.8, 925, 906, 899.3, 887, 862.5, + 854, 850, 800, 793.9, 785, 777, 771, 762, 731.8, 726, 703, 700, 655, + 630, 621.2, 602, 570.7, 548, 546.8, 539, 513, 511, 485, 481, 468, + 448, 439, 424, 420, 412], 'hPa') + levels.name = 'pressure' + temperatures = vert_array(np.array([-22.5, -22.7, -23.1, np.nan, -24.5, -25.1, np.nan, -24.5, -23.9, + np.nan, -24.7, np.nan, -21.3, -21.3, -22.7, np.nan, -20.7, -16.3, + -15.5, np.nan, np.nan, -15.3, np.nan, -17.3, -20.9, -22.5, + np.nan, -25.5, np.nan, -31.5, np.nan, -31.5, -34.1, -34.3, + -37.3, -37.7, -39.5, -42.1, -43.1, -45.1, -45.7, -46.7])+273.15, 'K') + temperatures.name = 'temperature' + dewpoints = vert_array(np.array([-25.1, -26.1, -26.8, np.nan, -27.3, -28.2, np.nan, -27.2, -26.6, + np.nan, -27.4, np.nan, -23.5, -23.5, -25.1, np.nan, -22.9, -17.8, + -16.6, np.nan, np.nan, -16.4, np.nan, -18.5, -21, -23.7, np.nan, + -28.3, np.nan, -32.6, np.nan, -33.8, -35, -35.1, -38.1, -40, + -43.3, -44.6, -46.4, -47, -49.2, -50.7])+273.15, 'K') + dewpoints.name = 'dewpoint' + + # Calculate parcel profile without LCL, as per metpy unit tests. + profile = parcel.parcel_profile(pressure=levels, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + profile['environment_temperature'] = temperatures + lfc = parcel.lfc_el(profile=profile) + + cape_cin_base = parcel.cape_cin_base(pressure=levels, + temperature=temperatures, + lfc_pressure=lfc.lfc_pressure, + el_pressure=lfc.el_pressure, + parcel_profile=profile) + cape_cin_surf = parcel.surface_based_cape_cin(pressure=levels, + temperature=temperatures, + dewpoint=dewpoints) + cape_cin_unstable = parcel.most_unstable_cape_cin(pressure=levels, + temperature=temperatures, + dewpoint=dewpoints) + + assert(np.isnan(lfc.lfc_pressure)) + assert_almost_equal(cape_cin_base.cape, 0, 0) + assert_almost_equal(cape_cin_base.cin, 0, 0) + assert_almost_equal(cape_cin_surf.cape, 0, 0) + assert_almost_equal(cape_cin_surf.cin, 0, 0) + assert_almost_equal(cape_cin_unstable.cape, 0, 0) + assert_almost_equal(cape_cin_unstable.cin, 0, 0) + +def test_most_unstable_cape_cin_surface(): + """Test the most unstable CAPE/CIN calculation when surface is most unstable.""" + levels = vert_array([959., 779.2, 751.3, 724.3, 700., 269.], 'hPa') + temperatures = vert_array(np.array([22.2, 14.6, 12., 9.4, 7., -38.])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4, -10., -53.2])+273.15, 'K') + levels.name = 'pressure' + temperatures.name = 'temperature' + dewpoints.name = 'dewpoint' + + cape_cin = parcel.most_unstable_cape_cin(pressure=levels, temperature=temperatures, + dewpoint=dewpoints) + + assert_almost_equal(cape_cin.cape, 75.0535446, 2) + assert_almost_equal(cape_cin.cin, -136.685967, 2) + +def test_most_unstable_cape_cin(): + """Test the most unstable CAPE/CIN calculation.""" + pressure = np.array([1000., 959., 867.9, 850., 825., 800.]) * units.mbar + temperature = np.array([18.2, 22.2, 17.4, 10., 0., 15]) * units.celsius + dewpoint = np.array([19., 19., 14.3, 0., -10., 0.]) * units.celsius + mucape, mucin = most_unstable_cape_cin(pressure, temperature, dewpoint) + assert_almost_equal(mucape, 157.11404 * units('joule / kilogram'), 4) + assert_almost_equal(mucin, -31.8406578 * units('joule / kilogram'), 4) + +def test_mixed_parcel(): + """Test the mixed parcel calculation.""" + levels = vert_array([959., 779.2, 751.3, 724.3, 700., 269.], 'hPa') + levels.name = 'pressure' + temperatures = vert_array(np.array([22.2, 14.6, 12., 9.4, 7., -38.])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4, -10., -53.2])+273.15, 'K') + + mixed = parcel.mixed_parcel(pressure=levels, temperature=temperatures, dewpoint=dewpoints, + depth=250) + assert_almost_equal(mixed.pressure, 959., 6) + assert_almost_equal(mixed.temperature, 28.7401463+273.15, 6) + assert_almost_equal(mixed.dewpoint, 7.1534658+273.15, 6) + +def test_mixed_layer_cape_cin(): + """Test the calculation of mixed layer cape/cin.""" + levels, temperatures, dewpoints = multiple_intersections() + + cape_cin = parcel.mixed_layer_cape_cin(pressure=levels, temperature=temperatures, + dewpoint=dewpoints) + + assert_almost_equal(cape_cin.cape, 987.7323, 2) + assert_almost_equal(cape_cin.cin, -20.6727628, 2) + +def test_mixed_layer(): + """Test the mixed layer calculation.""" + pressure = vert_array([959., 779.2, 751.3, 724.3, 700., 269.], 'hPa') + pressure.name = 'pressure' + temperature = vert_array(np.array([22.2, 14.6, 12., 9.4, 7., -38.])+273.15, 'K') + temperature.name = 'temperature' + mixed = parcel.mixed_layer(xarray.merge([pressure, temperature]), depth=250) + assert_almost_equal(mixed.temperature, 16.4024930+273.15, 6) + +def test_lfc_not_below_lcl(): + """Test sounding where LFC appears to be (but isn't) below LCL.""" + levels = vert_array([1002.5, 1001.7, 1001., 1000.3, 999.7, 999., 998.2, 977.9, + 966.2, 952.3, 940.6, 930.5, 919.8, 909.1, 898.9, 888.4, + 878.3, 868.1, 858., 848., 837.2, 827., 816.7, 805.4], 'hPa') + temperatures = vert_array(np.array([17.9, 17.9, 17.8, 17.7, 17.7, 17.6, 17.5, 16., + 15.2, 14.5, 13.8, 13., 12.5, 11.9, 11.4, 11., + 10.3, 9.7, 9.2, 8.7, 8., 7.4, 6.8, 6.1])+273.15, 'K') + dewpoints = vert_array(np.array([13.6, 13.6, 13.5, 13.5, 13.5, 13.5, 13.4, 12.5, + 12.1, 11.8, 11.4, 11.3, 11., 9.3, 10., 8.7, 8.9, + 8.6, 8.1, 7.6, 7., 6.5, 6., 5.4])+273.15, 'K') + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + lfc_el = parcel.lfc_el(profile=profile) + assert_almost_equal(lfc_el.lfc_pressure, 811.618879, 3) + assert_almost_equal(lfc_el.lfc_temperature, 6.48644650+273.15, 3) + +def multiple_intersections(): + """Create profile with multiple LFCs and ELs for testing.""" + levels = vert_array([966., 937.2, 925., 904.6, 872.6, 853., 850., 836., 821., 811.6, 782.3, + 754.2, 726.9, 700., 648.9, 624.6, 601.1, 595., 587., 576., 555.7, + 534.2, 524., 500., 473.3, 400., 384.5, 358., 343., 308.3, 300., 276., + 273., 268.5, 250., 244.2, 233., 200.], 'hPa') + levels.name = 'pressure' + + temperatures = vert_array(np.array([18.2, 16.8, 16.2, 15.1, 13.3, 12.2, 12.4, 14., 14.4, + 13.7, 11.4, 9.1, 6.8, 4.4, -1.4, -4.4, -7.3, -8.1, + -7.9, -7.7, -8.7, -9.8, -10.3, -13.5, -17.1, -28.1, -30.7, + -35.3, -37.1, -43.5, -45.1, -49.9, -50.4, -51.1, -54.1, -55., + -56.7, -57.5])+273.15, 'K') + temperatures.name = 'temperature' + + dewpoints = vert_array(np.array([16.9, 15.9, 15.5, 14.2, 12.1, 10.8, 8.6, 0., -3.6, -4.4, + -6.9, -9.5, -12., -14.6, -15.8, -16.4, -16.9, -17.1, -27.9, -42.7, + -44.1, -45.6, -46.3, -45.5, -47.1, -52.1, -50.4, -47.3, -57.1, + -57.9, -58.1, -60.9, -61.4, -62.1, -65.1, -65.6, + -66.7, -70.5])+273.15, 'K') + dewpoints.name = 'dewpoint' + + return levels, temperatures, dewpoints + +def test_multiple_lfcs_el_simple(): + """Test sounding with multiple LFCs.""" + levels, temperatures, dewpoints = multiple_intersections() + + profile = parcel.parcel_profile_with_lcl(pressure=levels, + temperature=temperatures, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + lfc_el = parcel.lfc_el(profile=profile) + + assert_almost_equal(lfc_el.lfc_pressure, 884.14790, 3) + assert_almost_equal(lfc_el.lfc_temperature, 13.95707016+273.15, 3) + assert_almost_equal(lfc_el.el_pressure, 228.151466, 3) + assert_almost_equal(lfc_el.el_temperature, -56.81015490+273.15, 3) + +def test_cape_cin_custom_profile(): + """Test the CAPE and CIN calculation with a custom profile passed to LFC and EL.""" + levels = vert_array([959., 779.2, 751.3, 724.3, 700., 269.], 'hPa') + temperatures = vert_array(np.array([22.2, 14.6, 12., 9.4, 7., -38.])+273.15, 'K') + dewpoints = vert_array(np.array([19., -11.2, -10.8, -10.4, -10., -53.2])+273.15, 'K') + + profile = parcel.parcel_profile(pressure=levels, + parcel_pressure=levels[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + profile['temperature'] = profile.temperature + 5 + profile['environment_temperature'] = temperatures + + lfc = parcel.lfc_el(profile=profile) + cape_cin = parcel.cape_cin_base(pressure=levels, + temperature=temperatures, + lfc_pressure=lfc.lfc_pressure, + el_pressure=lfc.el_pressure, + parcel_profile=profile) + + assert_almost_equal(cape_cin.cape, 1440.463208696, 2) + assert_almost_equal(cape_cin.cin, 0.0, 2) + +def test_parcel_profile_below_lcl(): + """Test parcel profile calculation when pressures do not reach LCL (#827).""" + pressure = vert_array([981, 949.2, 925., 913.9, 903, 879.4, 878, 864, 855, + 850, 846.3, 838, 820, 814.5, 799, 794], 'hPa') + truth = np.array([276.35, 273.760341, 271.747753, 270.812026, 269.885225, + 267.850849, 267.728946, 266.502214, 265.706084, 265.261201, + 264.930782, 264.185801, 262.551884, 262.047526, 260.61294, + 260.145932]) + + parcel_temperature = xarray.DataArray(3.2+273.15, attrs={'units': 'K'}) + parcel_dewpoint = xarray.DataArray(-10.8+273.15, attrs={'units': 'K'}) + + profile = parcel.parcel_profile(pressure=pressure, + parcel_pressure=pressure[0], + parcel_temperature=parcel_temperature, + parcel_dewpoint=parcel_dewpoint) + + assert_array_almost_equal(profile.temperature, truth, 6) + +def test_lcl_convergence_issue(): + """Test profile where LCL wouldn't converge (#1187).""" + pressure = vert_array([990, 973, 931, 925, 905], 'hPa') + temperatures = vert_array(np.array([14.4, 14.2, 13, 12.6, 11.4])+273.15, 'K') + dewpoints = vert_array(np.array([14.4, 11.7, 8.2, 7.8, 7.6])+273.15, 'K') + + lcl = parcel.lcl(parcel_pressure=pressure[0], + parcel_temperature=temperatures[0], + parcel_dewpoint=dewpoints[0]) + assert_almost_equal(lcl.lcl_pressure, 990, 0) + +def test_cape_cin_value_error(): + """Test a profile that originally caused a ValueError in #1190.""" + levels = vert_array([1012.0, 1009.0, 1002.0, 1000.0, 925.0, 896.0, 855.0, 850.0, 849.0, + 830.0, 775.0, 769.0, 758.0, 747.0, 741.0, 731.0, 712.0, 700.0, 691.0, + 671.0, 636.0, 620.0, 610.0, 601.0, 594.0, 587.0, 583.0, 580.0, 571.0, + 569.0, 554.0, 530.0, 514.0, 506.0, 502.0, 500.0, 492.0, 484.0, 475.0, + 456.0, 449.0, 442.0, 433.0, 427.0, 400.0, 395.0, 390.0, 351.0, 300.0, + 298.0, 294.0, 274.0, 250.0], 'hPa') + temperatures = vert_array(np.array([27.8, 25.8, 24.2, 24, 18.8, 16, 13, 12.6, 12.6, 11.6, 9.2, 8.6, + 8.4, 9.2, 10, 9.4, 7.4, 6.2, 5.2, 3.2, -0.3, -2.3, -3.3, -4.5, + -5.5, -6.1, -6.1, -6.1, -6.3, -6.3, -7.7, -9.5, -9.9, -10.3, + -10.9, -11.1, -11.9, -12.7, -13.7, -16.1, -16.9, -17.9, -19.1, + -19.9, -23.9, -24.7, -25.3, -29.5, -39.3, -39.7, -40.5, -44.3, + -49.3])+273.15, 'K') + dewpoints = vert_array(np.array([19.8, 16.8, 16.2, 16, 13.8, 12.8, 10.1, 9.7, 9.7, + 8.6, 4.2, 3.9, 0.4, -5.8, -32, -34.6, -35.6, -34.8, + -32.8, -10.8, -9.3, -10.3, -9.3, -10.5, -10.5, -10, -16.1, + -19.1, -23.3, -18.3, -17.7, -20.5, -27.9, -32.3, -33.9, -34.1, + -35.9, -26.7, -37.7, -43.1, -33.9, -40.9, -46.1, -34.9, -33.9, + -33.7, -33.3, -42.5, -50.3, -49.7, -49.5, -58.3, -61.3])+273.15, 'K') + + cape_cin = parcel.surface_based_cape_cin(pressure=levels, temperature=temperatures, + dewpoint=dewpoints) + + assert_almost_equal(cape_cin.cape, 2007.040698, 3) + assert_almost_equal(cape_cin.cin, 0.0, 3) + +def test_lcl_grid_surface_lcls(): + """Test surface grid where some values have LCLs at the surface.""" + pressure = vert_array([1000, 990, 1010], 'hPa').expand_dims({'x': 3}) + temperature = vert_array(np.array([15, 14, 13])+273.15, 'K').expand_dims({'x': 3}) + dewpoint = vert_array(np.array([15, 10, 13])+273.15, 'K').expand_dims({'x': 3}) + + lcl = parcel.lcl(parcel_pressure=pressure[0], + parcel_temperature=temperature[0], + parcel_dewpoint=dewpoint[0]) + + pres_truth = np.array([1000, 932.1719, 1010]) + temp_truth = np.array([15, 9.10424, 13])+273.15 + assert_array_almost_equal(lcl.lcl_pressure, pres_truth, 4) + assert_array_almost_equal(lcl.lcl_temperature, temp_truth, 4) + +def test_lifted_index(): + """Test the Lifted Index calculation.""" + pressure = vert_array([1014., 1000., 997., 981.2, 947.4, 925., 914.9, 911., + 902., 883., 850., 822.3, 816., 807., 793.2, 770., + 765.1, 753., 737.5, 737., 713., 700., 688., 685., + 680., 666., 659.8, 653., 643., 634., 615., 611.8, + 566.2, 516., 500., 487., 484.2, 481., 475., 460., + 400.], 'hPa') + pressure.name = 'pressure' + + temperature = vert_array(np.array([24.2, 24.2, 24., 23.1, 21., 19.6, 18.7, 18.4, + 19.2, 19.4, 17.2, 15.3, 14.8, 14.4, 13.4, 11.6, + 11.1, 10., 8.8, 8.8, 8.2, 7., 5.6, 5.6, + 5.6, 4.4, 3.8, 3.2, 3., 3.2, 1.8, 1.5, + -3.4, -9.3, -11.3, -13.1, -13.1, -13.1, -13.7, -15.1, + -23.5])+273.15, 'K') + temperature.name = 'temperature' + + dewpoint = vert_array(np.array([23.2, 23.1, 22.8, 22., 20.2, 19., 17.6, 17., + 16.8, 15.5, 14., 11.7, 11.2, 8.4, 7., 4.6, + 5., 6., 4.2, 4.1, -1.8, -2., -1.4, -0.4, + -3.4, -5.6, -4.3, -2.8, -7., -25.8, -31.2, -31.4, + -34.1, -37.3, -32.3, -34.1, -37.3, -41.1, -37.7, -58.1, + -57.5])+273.15, 'K') + + # Use profile without lcl as per metpy unit test. + profile = parcel.parcel_profile(pressure=pressure, + parcel_pressure=pressure[0], + parcel_temperature=temperature[0], + parcel_dewpoint=dewpoint[0]) + profile['environment_temperature'] = temperature + + li = parcel.lifted_index(profile=profile) + assert_almost_equal(li.lifted_index, -7.9176350, 2) From a988dbf7b72da12ef26c812c15a738335c8c0974 Mon Sep 17 00:00:00 2001 From: Timothy Raupach Date: Thu, 15 Jul 2021 08:10:06 +1000 Subject: [PATCH 2/2] Added xarray parcel functions. --- src/metpy/calc/xarray_parcel.py | 1645 +++++++++++++++++++++++++++++++ 1 file changed, 1645 insertions(+) create mode 100644 src/metpy/calc/xarray_parcel.py diff --git a/src/metpy/calc/xarray_parcel.py b/src/metpy/calc/xarray_parcel.py new file mode 100644 index 00000000000..2ca76b79685 --- /dev/null +++ b/src/metpy/calc/xarray_parcel.py @@ -0,0 +1,1645 @@ +# Copyright (c) 2008,2015,2016,2017,2018,2019 MetPy Developers. +# Distributed under the terms of the BSD 3-Clause License. +# SPDX-License-Identifier: BSD-3-Clause +# +# xarray-enabled versions of MetPy functions for atmospheric parcel +# calculations. + +import sys +import metpy +import xarray +import numpy as np +from metpy.units import units +import metpy.constants as mpconsts + +# Global variables to contain moist adiabat lookup tables. +this = sys.modules[__name__] +this.moist_adiabat_lookup = None +this.moist_adiabats = None + +def load_moist_adiabat_lookups(**kwargs): + """ + Load cache adiabat lookup tables ready for use in other functions. + If cache doesn't yet exist, use moist_adiabat_tables to generate + them. + + Arguments: + + - **kwargs: Any arguments to moist_adiabat_tables(). + """ + + this.moist_adiabat_lookup, this.moist_adiabats = moist_adiabat_tables( + regenerate=False, cache=True, **kwargs) + this.moist_adiabat_lookup = this.moist_adiabat_lookup.persist() + this.moist_adiabats = this.moist_adiabats.persist() + +def lookup_tables_loaded(): + """ + Ensure that lookup tables are loaded and throw an error if not. + """ + assert this.moist_adiabat_lookup is not None, 'Call load_moist_adiabat_lookups first.' + assert this.moist_adiabats is not None, 'Call load_moist_adiabat_lookups first.' + +def get_layer(dat, depth=100, drop=False, vert_dim='model_level_number', + interpolate=True): + """ + Return an atmospheric layer from the surface with a given depth. + + Arguments: + + - dat: DataArray, must contain pressure. + - depth: Depth above the bottom of the layer to mix [hPa]. + - drop: Drop unselected elements? + - vert_dim: Vertical dimension name. + - interpolate: Interpolate the bottom/top layers? + + Returns: + + - xarray DataArray with pressure and data variables for the layer. + """ + + # Use the surface (lowest level) pressure as the bottom pressure. + bottom_pressure = dat.pressure.max(dim=vert_dim) + + # Calculate top pressure. + if interpolate: + top_pressure = bottom_pressure-depth + interp_level = log_interp(x=dat, at=top_pressure, + coords=dat.pressure, dim=vert_dim) + interp_level['pressure'] = top_pressure + dat = insert_level(d=dat, level=interp_level, coords='pressure', + vert_dim=vert_dim) + else: + top_pressure = bound_pressure(pressure=dat.pressure, + bound=bottom_pressure-depth, + vert_dim=vert_dim) + + # Select the layer. + layer = dat.where(dat.pressure <= bottom_pressure, drop=drop) + layer = dat.where(dat.pressure >= top_pressure, drop=drop) + + return layer + +def most_unstable_parcel(dat, depth=300, drop=False, + vert_dim='model_level_number'): + """ + Return the most unstable parcel with an atmospheric layer from + with the requested bottom and depth. No interpolation is + performed. If there are multiple 'most unstable' parcels, return + the first in the vertical array. + + Arguments: + + - dat: DataArray, must contain pressure, temperature, and dewpoint. + - bottom: Pressure level to start from [hPa]. + - depth: Depth above the bottom of the layer to mix [hPa]. + - drop: Drop unselected elements? + - vert_dim: Vertical dimension name. + + Returns: + + - xarray DataArray with pressure and data variables for the layer. + + """ + + layer = get_layer(dat=dat, depth=depth, drop=drop, vert_dim=vert_dim, + interpolate=False) + eq = metpy.calc.equivalent_potential_temperature( + pressure=layer.pressure, + temperature=layer.temperature, + dewpoint=layer.dewpoint).metpy.dequantify() + max_eq = eq.max(dim=vert_dim) + pres = layer.where(eq == max_eq).pressure.max(dim=vert_dim) + assert np.all(layer.pressure.where(layer.pressure == pres).count( + dim=vert_dim) == 1), 'Vertical pressures are not unique' + most_unstable = layer.where(layer.pressure == pres).max(dim=vert_dim, + keep_attrs=True) + return most_unstable + +def mixed_layer(dat, depth=100, vert_dim='model_level_number'): + """ + Mix variable(s) over a layer, yielding a mass-weighted average. + + Integrate a data variable with respect to pressure and determine the + average value using the mean value theorem. + + Arguments: + + - dat: The DataArray to mix. Must contain pressure and variables. + - bottom: Pressure above the surface pressure to start from [hPa]. + - depth: Depth above the bottom of the layer to mix [hPa]. + - vert_dim: The name of the vertical dimension. + + Returns: + + - xarray with mixed values of each data variable. + """ + + layer = get_layer(dat=dat, depth=depth, drop=True, vert_dim=vert_dim) + + pressure_depth = np.abs(layer.pressure.min(vert_dim) - + layer.pressure.max(vert_dim)) + + ret = (1. / pressure_depth) * trapz(dat=layer, x='pressure', dim=vert_dim) + return ret + +def trapz(dat, x, dim, mask=None): + """ + Perform trapezoidal rule integration along an axis, ala numpy.trapz. + Estimates int y dx. + + Arguments: + + - dat: Data to process. + - x: The variable that contains 'x' values along dimension 'dim'. + - dim: The dimension along which to integrate 'y' values. + - mask: A mask the size of dx/means (ie dim.size-1) for which + areas to include in the integration. + + Returns: + + - Integrated value along the axis. + """ + + dx = np.abs(dat[x].diff(dim)) + dx = dx.reset_coords(drop=True) + means = dat.rolling({dim: 2}, center=True).mean(keep_attrs=True) + means = means.reset_coords(drop=True) + + dx = dx.assign_coords({dim: dx[dim]-1}) + means = means.assign_coords({dim: means[dim]-1}) + + if mask is not None: + dx = dx.where(mask) + means = means.where(mask) + + return (dx * means).sum(dim) + +def bound_pressure(pressure, bound, vert_dim='model_level_number'): + """ + Calculate the bounding pressure in a layer; returns the closest + pressure to the bound. If two pressures are equally distant from + the bound, the larger pressure is returned. + + Arguments: + + - pressure: Atmospheric pressures [hPa]. + - bound: Bound to retrieve, broadcastable to pressure [hPa]. + + Returns: + + - The bound pressures. + """ + + diffs = np.abs(pressure - bound) + bounds = pressure.where(diffs == diffs.min(dim=vert_dim), drop=True) + bounds = bounds.max(dim=vert_dim).squeeze(drop=True) + return bounds + +def mixed_parcel(pressure, temperature, dewpoint, depth=100, + vert_dim='model_level_number'): + """ + Fully mix a layer of given depth above the surface and find the temparature, + pressure and dewpoint of the parcel. + + Arguments: + + - pressure: Pressure by level [hPa]. + - temperature: Temperature at each level [K]. + - dewpoint: Dewpoint at each level [K]. + - depth: Depth above the surface to mix [hPa]. + - vert_dim: The name of the vertical dimension. + + Returns: + + - DataArray with mixed parcel pressure [hPa], temperature [K] + and dewpoint [K]. + """ + + # Use the surface (lowest level) pressure as the start of the layer to mix. + parcel_start_pressure = pressure.isel({vert_dim: 0}) + + # Calculate the potential temperature over the layer. + theta = metpy.calc.potential_temperature(pressure, temperature) + theta = theta.metpy.dequantify() + theta.name = 'theta' + + # Mixing ratio over the layer. + mixing_ratio = metpy.calc.saturation_mixing_ratio(pressure, dewpoint) + mixing_ratio = mixing_ratio.metpy.dequantify() + mixing_ratio.name = 'mixing_ratio' + + # Mix theta and mixing ratio over the layer. + assert pressure.name is not None, 'pressure requires name pressure.' + mp = mixed_layer(xarray.merge([pressure, theta, mixing_ratio]), depth=depth, + vert_dim=vert_dim) + + # Convert potential temperature back to temperature. + mp['temperature'] = (mp.theta * + metpy.calc.exner_function(parcel_start_pressure)) + mp['temperature'] = mp.temperature.metpy.dequantify() + mp.temperature.attrs['long_name'] = 'Mixed parcel temperature' + mp.temperature.attrs['units'] = 'K' + + # Convert mixing ratio back to dewpoint. + mp['vapour_pressure'] = metpy.calc.vapor_pressure(parcel_start_pressure, + mp.mixing_ratio) + mp['vapour_pressure'] = mp.vapour_pressure.metpy.dequantify() + mp.vapour_pressure.attrs['long_name'] = 'Mixed-parcel vapour pressure' + + mp['dewpoint'] = metpy.calc.dewpoint(mp.vapour_pressure) + mp['dewpoint'] = mp.dewpoint.metpy.convert_units('K') + mp['dewpoint'] = mp.dewpoint.metpy.dequantify() + mp.dewpoint.attrs['long_name'] = 'Mixed-parcel dewpoint' + + # For pressure, use the starting pressure for the layer (following MetPy's + # mixed_parcel function). + mp['pressure'] = parcel_start_pressure + + return mp + +def dry_lapse(pressure, parcel_temperature, parcel_pressure=None, + vert_dim='model_level_number'): + """ + Calculate the temperature of a parcel raised dry-adiabatically (conserving + potential temperature). + + Arguments: + + - pressure: Atmospheric pressure level(s) of interest [hPa]. + - parcel_temperature: Parcel temperature before lifting + (constant or broadcast-able DataArray). + - parcel_pressure: Parcel pressure(s) before lifting. Defaults + to vertical maximum. + - vert_dim: The name of the vertical dimension. + + Returns: + + - Parcel temperature at each pressure level. + """ + + if parcel_pressure is None: + parcel_pressure = pressure.max(vert_dim) + out = parcel_temperature * (pressure / parcel_pressure)**mpconsts.kappa + out.attrs['long_name'] = 'Dry lapse rate temperature' + out.attrs['units'] = 'K' + return out + +def moist_adiabat_tables(regenerate=False, cache=True, chunks=None, base_dir='.', + lookup_cache='/adiabat_lookups/moist_adiabat_lookup.nc', + adiabats_cache='/adiabat_lookups/adiabats_cache.nc', + **kwargs): + """ + Calculate moist adiabat lookup tables. + + Arguments: + + - regenerate: Calculate from scratch and save caches? + - cache: Write cache files? + - chunks: Chunks argument for xarray .chunk() function. + - base_dir: The base directory in which to read/write caches. + - lookup_cache: A cache file (nc) for the adiabat lookup table. + - adiabats_cache: A cache file (nc) for the adiabats cache. + - **kwargs: Keyword arguments to moist_adiabat_lookup(). + + Returns: + + - two DataArrays: 1) a lookup table of pressure/temperature + vs. adiabat number, and 2) a lookup table of adiabat number + to temperature by pressure profiles. + """ + + if not regenerate: + adiabat_lookup = xarray.open_dataset(base_dir + lookup_cache, + chunks=chunks).load() + adiabats = xarray.open_dataset(base_dir + adiabats_cache, + chunks=chunks).load() + return adiabat_lookup, adiabats + + # Generate lookup tables. + adiabat_lookup, adiabats = moist_adiabat_lookup(**kwargs) + + if cache: + adiabats.to_netcdf(base_dir + adiabats_cache) + adiabat_lookup.to_netcdf(base_dir + lookup_cache) + + return adiabat_lookup.chunk(chunks), adiabats.chunk(chunks) + +def round_to(x, to, dp=2): + """ + Round x to the nearest 'to' and return rounded to 'dp' decimal points. + """ + return np.round(np.round(x / to) * to, dp) + +def moist_adiabat_lookup(pressure_levels=np.round(np.arange(1100, 0, + step=-0.5), 1), + temperatures=np.round(np.arange(173, 316, + step=0.02), 2), + pres_step=0.5, temp_step=0.02): + """ + Calculate moist adiabat lookup tables. + + Arguments: + - pressure_levels: Pressure levels to keep in adiabat lookup + table [hPa]. + - temperatures: Temperatures to keep in adiabat lookup table [K]. + - pres_step, temp_step: (Positive) step size for + pressure_levels and temperatures, + respectively. + + Returns: + + - two DataArrays: 1) a lookup table of pressure/temperature + vs. adiabat number, and 2) a lookup table of adiabat number + to temperature by pressure profiles. + """ + + curves = [] + adiabat_lookup = xarray.Dataset({'adiabat': np.nan}) + adiabat_lookup = adiabat_lookup.expand_dims({'pressure': pressure_levels, + 'temperature': + temperatures}).copy(deep=True) + + # Find the adiabat for each starting temperature. + i = 1 + for parcel_temperature in temperatures: + for offset in [0, temp_step/2]: + profile_temps = metpy.calc.moist_lapse( + pressure=pressure_levels*units.hPa, + temperature=(parcel_temperature+offset)*units.K).m + + nearest_temps = round_to(profile_temps, temp_step) + idx = np.isin(nearest_temps, temperatures) + temp_idx = xarray.DataArray(nearest_temps[idx], dims=['idx']) + pres_idx = xarray.DataArray(pressure_levels[idx], dims=['idx']) + adiabat_lookup.adiabat.loc[{'pressure':pres_idx, + 'temperature': temp_idx}] = i + + # In profile_temps we have the adiabat temperature for + # every pressure level. But some temperatures in the + # lookup table may be missing. Interpolate the pressures + # for each required temperature. + pres_per_temp = np.interp(x=temperatures, xp=profile_temps[::-1], + fp=pressure_levels[::-1], right=np.nan, + left=np.nan) + + pres_per_temp = round_to(pres_per_temp, pres_step) + idx = np.isin(pres_per_temp, pressure_levels) + pres_idx = xarray.DataArray(pres_per_temp[idx], dims=['idx']) + temp_idx = xarray.DataArray(temperatures[idx], dims=['idx']) + adiabat_lookup.adiabat.loc[{'pressure':pres_idx, + 'temperature': temp_idx}] = i + + # curves contains the adiabats themselves. + curves.append(xarray.Dataset({'temperature': (['pressure'], + profile_temps)}, + coords={'pressure': pressure_levels})) + curves[-1]['adiabat'] = i + i = i + 1 + + # Combine curves into one dataset. + adiabats = xarray.combine_nested(curves, concat_dim='adiabat') + + for x in [adiabat_lookup, adiabats]: + x.pressure.attrs['long_name'] = 'Pressure' + x.pressure.attrs['units'] = 'hPa' + x.temperature.attrs['long_name'] = 'Temperature' + x.temperature.attrs['units'] = 'K' + x.adiabat.attrs['long_name'] = 'Adiabat index' + + return adiabat_lookup, adiabats + +def moist_lapse(pressure, parcel_temperature, parcel_pressure=None, + vert_dim='model_level_number', chunks=300): + """ + Return the temperature of parcels raised moist-adiabatically + (assuming liquid saturation processes). Note: What is returned + are approximate pseudo-adiabatic moist lapse rates, found using a + lookup table. + + Arguments: + + - pressure: Atmospheric pressure(s) to lift the parcel to [hPa]. + - parcel_temperature: Temperature(s) of parcels to lift [K]. + - parcel_pressure: Parcel pressure before lifting. Defaults to lowest + vertical level. + - vert_dim: The name of the vertical dimension. + - chunks: Chunk size to use in non-vertical dimensions if Dask is + being used. Reduce if memory is being filled. + + Returns: + + - Parcel temperature at each pressure level. + """ + + lookup_tables_loaded() + + if parcel_pressure is None: + parcel_pressure = pressure.isel({vert_dim: 0}) + + # For each starting parcel, find the moist adiabat that intersects + # the parcel pressure and temperature. + adiabat_idx = this.moist_adiabat_lookup.sel({'pressure': parcel_pressure, + 'temperature': parcel_temperature}, + method='nearest') + adiabat_idx = adiabat_idx.adiabat.reset_coords(drop=True) + adiabats = this.moist_adiabats.sel(adiabat=adiabat_idx) + + if isinstance(pressure, xarray.DataArray) and pressure.chunks is not None: + adiabats = adiabats.chunk(chunks) + adiabat_idx = adiabat_idx.chunk(chunks) + adiabats = adiabats.chunk({'pressure': adiabats.pressure.size}) + + # Interpolate the adiabat to get the temperature at each requested + # pressure. + out = adiabats.temperature.interp( + {'pressure': pressure}).reset_coords(drop=True) + return out + +def lcl(parcel_pressure, parcel_temperature, parcel_dewpoint): + """ + Return the lifting condensation level for parcels. + + Arguments: + + - parcel_pressure: Pressure of the parcel to lift [hPa]. + - parcel_temperature: Parcel temperature [K]. + - parfel_dewpoint: Parcel dewpoint [K]. + + Returns: + + - A Dataset with lcl_pressure and lcl_temperature. + """ + + press_lcl, temp_lcl = metpy.calc.lcl(pressure=parcel_pressure, + temperature=parcel_temperature, + dewpoint=parcel_dewpoint) + out = xarray.Dataset({'lcl_pressure': (parcel_temperature.dims, + press_lcl.m), + 'lcl_temperature': (parcel_temperature.dims, + temp_lcl.m)}) + + out.lcl_pressure.attrs['long_name'] = ('Lifting condensation ' + + 'level pressure') + out.lcl_pressure.attrs['units'] = 'hPa' + out.lcl_temperature.attrs['long_name'] = ('Lifting condensation ' + + 'level temperature') + out.lcl_temperature.attrs['units'] = 'K' + + return out + +def parcel_profile(pressure, parcel_pressure, parcel_temperature, parcel_dewpoint): + """ + Calculate temperatures of a lifted parcel. + + Arguments: + + - pressure: Pressure levels to calculate on [hPa]. + - parcel_pressure: Pressure of the parcel [hPa]. + - parcel_temperature: Temperature of the parcel [K]. + - parcel_dewpoint: Dewpoint of the parcel [K]. + + Returns: + + - Dataset with the temperature of the parcel lifted from + parcel_pressure to levels in pressures, plus the LCL + pressure and temperature. + """ + + out = xarray.Dataset() + out['pressure'] = pressure + + # Find the LCL for the selected parcel. + out = xarray.merge([out, lcl(parcel_pressure=parcel_pressure, + parcel_temperature=parcel_temperature, + parcel_dewpoint=parcel_dewpoint)]) + + # Parcels are raised along the dry adiabats from the starting + # point to the LCL. + below_lcl = dry_lapse(pressure=pressure, + parcel_temperature=parcel_temperature, + parcel_pressure=parcel_pressure) + + # Above the LCL parcels follow the moist adiabats from the LCL + # temp/pressure. + above_lcl = moist_lapse(pressure=pressure, + parcel_temperature=out.lcl_temperature, + parcel_pressure=out.lcl_pressure) + + out['temperature'] = below_lcl.where(pressure >= out.lcl_pressure, + other=above_lcl) + out.temperature.attrs['long_name'] = 'Lifted parcel temperature' + out.temperature.attrs['units'] = 'K' + + out = out.reset_coords(drop=True) + return out + +def parcel_profile_with_lcl(pressure, temperature, parcel_pressure, + parcel_temperature, parcel_dewpoint, + vert_dim='model_level_number'): + """ + Calculate temperatures of a lifted parcel, including at the lcl. + + Arguments: + + - pressure: Pressure levels to calculate on [hPa]. + - temperature: Temperature at each pressure level [K]. + - parcel_pressure: Pressure of the parcel [hPa]. + - parcel_temperature: Temperature of the parcel [K]. + - parcel_dewpoint: Dewpoint of the parcel [K]. + - vert_dim: The name of the vertical dimension. + + Returns: + + - Dataset with the temperature of the parcel lifted from + parcel_pressure to levels in pressures, including the LCL, + plus the LCL pressure and temperature, and environmental + temperature including at the LCL. + """ + + profile = parcel_profile(pressure=pressure, + parcel_pressure=parcel_pressure, + parcel_temperature=parcel_temperature, + parcel_dewpoint=parcel_dewpoint) + return add_lcl_to_profile(profile=profile, vert_dim=vert_dim, + temperature=temperature) + +def add_lcl_to_profile(profile, vert_dim='model_level_number', + temperature=None): + """ + Add the LCL to a profile. + + Arguments: + + - profile: Profile as returned from parcel_profile(). + - vert_dim: The vertical dimension to add the LCL pressure/temp to. + - temperature: Environmental temperatures. If provided, + interpolate environment temperature for the LCL + and return as 'env_temperature'. + + Returns: + + - A new profile object with LCL pressure and temperatures + added. Note the vertical coordinate in the new profile is + reindexed. + """ + + level = xarray.Dataset({'pressure': profile.lcl_pressure, + 'temperature': profile.lcl_temperature}) + out = insert_level(d=profile, level=level, coords='pressure', + vert_dim=vert_dim) + out['lcl_pressure'] = profile.lcl_pressure + out['lcl_temperature'] = profile.lcl_temperature + out.temperature.attrs['long_name'] = 'Temperature with LCL' + out.pressure.attrs['long_name'] = 'Pressure with LCL' + + if not temperature is None: + environment = xarray.Dataset({'temperature': temperature, + 'pressure': profile.pressure}) + + # Note: we use linear_interp here even though we are working in pressure + # coordinates, in order to match MetPy's implementation. In future it + # may be more accurate to use log_interp. + temp_at_level = xarray.Dataset({'temperature': + linear_interp(x=temperature, + coords=profile.pressure, + at=level.pressure, + dim=vert_dim), + 'pressure': level.pressure}) + + environment = insert_level(d=environment, level=temp_at_level, + coords='pressure', vert_dim=vert_dim) + out['environment_temperature'] = environment.temperature + out.environment_temperature.attrs['long_name'] = ('Environment ' + + 'temperature') + out.environment_temperature.attrs['units'] = 'K' + + return out + +def insert_level(d, level, coords, vert_dim='model_level_number', + fill_value=-999): + """ + Insert a new level into a vertically sorted dataset. + + Arguments: + + - d: The data to work on. + - coords: The coordinate name in d. + - level: The new values to add; a single layer with values for + 'coord' and any other variables to add. + - vert_dim: The vertical dimension to add new level to. + + Returns: + + - A new object with the new level added. + Note the vertical coordinate in the new profile is reindexed. + """ + + # To conserve nans in the original dataset, replace them with + # fill_value in the coordinate array. + assert not np.any(d[coords] == fill_value), 'dataset d contains fill_value.' + d = d.where(np.logical_not(np.isnan(d[coords])), other=fill_value) + + below = d.where(d[coords] > level[coords]) + above = d.where(d[coords] < level[coords]) + + # Above the new coordinate, shift the vertical coordinate indices + # up one. + above = above.assign_coords({vert_dim: d[vert_dim] + 1}) + + # Use broadcasting to fills regions below the new coordinate. + out, _ = xarray.broadcast(below, above) + + # Fill regions above the new coordinate. + above, _ = xarray.broadcast(above, out) + out = above.where(np.isnan(out[coords]), other=out) + + # Any remaining nan values must now be the new level, so fill + # those regions. + new, _ = xarray.broadcast(level, out) + + # Subset to keys from new only. + out = out[list(new.keys())] + out = new.where(np.isnan(out), other=out) + + # Replace fill_value with nans. + out = out.where(out != fill_value, other=np.nan) + + return out + +def find_intersections(x, a, b, dim, log_x=False): + """ + Find intersections of two lines that share x coordinates. + + Arguments: + + - x: The shared x coordinate values. + - a: y values for line 1. + - b: y values for line 2. + - dim: The dimension along which the coordinates are indexed. + - log_x: Use a logarithmic transform on x coordinates + (e.g. for pressure coords)? + + Returns: + + - Dataset containing x, y coordinates for all intersections, + increasing intersections and decreasing intersections. Note + duplicates are not removed. + """ + + if log_x: + x = np.log(x) + + # Find intersections. Non-zero points in diffs indicates an + # intersection. + diffs = np.sign(a - b).diff(dim=dim) + + # Identify the points after each intersection. + after_intersects = diffs.where(diffs == 0, other=1) + + # And the points just before each intersection.s + before_intersects = xarray.concat([xarray.zeros_like(a.isel({dim: 0})), + after_intersects], dim=dim) + before_intersects = before_intersects.shift({dim: -1}, fill_value=0) + + # The sign of the change for the intersect. + sign_change = np.sign(a.where(after_intersects == 1) - + b.where(after_intersects == 1)) + + x0 = x.where(before_intersects == 1).shift({dim: 1}) + x1 = x.where(after_intersects == 1) + a0 = a.where(before_intersects == 1).shift({dim: 1}) + a1 = a.where(after_intersects == 1) + b0 = b.where(before_intersects == 1).shift({dim: 1}) + b1 = b.where(after_intersects == 1) + + # Calculate the x-intersection. This comes from finding the + # equations of the two lines, one through (x0, a0) and (x1, a1) + # and the other through (x0, b0) and (x1, b1), finding their + # intersection, and reducing with a bunch of algebra. + delta_y0 = a0 - b0 + delta_y1 = a1 - b1 + intersect_x = (delta_y1 * x0 - delta_y0 * x1) / (delta_y1 - delta_y0) + + # Calculate the y-intersection of the lines. Just plug the x above + # into the equation for the line through the a points. + intersect_y = ((intersect_x - x0) / (x1 - x0)) * (a1 - a0) + a0 + + if log_x is True: + intersect_x = np.exp(intersect_x) + + out = xarray.Dataset() + out['all_intersect_x'] = intersect_x + out['all_intersect_y'] = intersect_y + out['increasing_x'] = intersect_x.where(sign_change > 0) + out['increasing_y'] = intersect_y.where(sign_change > 0) + out['decreasing_x'] = intersect_x.where(sign_change < 0) + out['decreasing_y'] = intersect_y.where(sign_change < 0) + out = out.rename({dim: 'offset_dim'}) + + return out + +def lfc_el(profile, vert_dim='model_level_number'): + """ + Calculate the level of free convection (LFC) and equilibrium level + (EL). + + Works by finding the first intersection of the ideal parcel path + and the measured parcel temperature. If this intersection occurs + below the LCL, the LFC is determined to be the same as the LCL, + based upon the conditions set forth in [USAF1990]_, pg 4-14, where + a parcel must be lifted dry adiabatically to saturation before it + can freely rise. The LFC returned is the 'bottom' LFC with + highest pressure; the EL returned is the 'top' EL with the lowest + pressure. + + Arguments: + + - profile: The parcel profile, including the LCL, as returned + from parcel_profile_with_lcl(). + - vert_dim: Vertical dimension name in input arrays. + + Returns: + + - DataArray with LFC pressure (lfc_pressure) and temperature + (lfc_temperature). + """ + + # Find all intersections between parcel and environmental + # temperatures by pressure. + intersections = find_intersections(x=profile.pressure, + a=profile.temperature, + b=profile.environment_temperature, + dim=vert_dim, + log_x=True) + + # Find intersections again, ignoring first level. + intersections_above = find_intersections( + x=profile.pressure.isel({vert_dim: slice(1,None)}), + a=profile.temperature.isel({vert_dim: slice(1,None)}), + b=profile.environment_temperature.isel({vert_dim: slice(1,None)}), + dim=vert_dim, log_x=True).reindex_like(intersections) + + # For points for which the atmosphere and parcel temperatures have + # the same lowest-level value, ignore this point and find the real + # LFC above it. + intersections = intersections.where( + (profile.environment_temperature.isel({vert_dim: 0}) != + profile.temperature.isel({vert_dim: 0})), + other=intersections_above) + + out = xarray.Dataset() + + # By default the first values are the lowest (highest pressure) + # crossings for LFC and the highest (lowest pressure) crossings + # for EL. The LFC also has to be above the LCL. + above_lcl = intersections.increasing_x < profile.lcl_pressure + + out['lfc_pressure'] = intersections.increasing_x.where( + above_lcl).max(dim='offset_dim') + out['lfc_temperature'] = intersections.increasing_y.where( + intersections.increasing_x == out.lfc_pressure).max(dim='offset_dim') + + # Determine equilibrium pressure and temperature. The 'top' + # (lowest pressure) EL is returned. + out['el_pressure'] = intersections_above.decreasing_x.min(dim='offset_dim') + out['el_temperature'] = intersections_above.decreasing_y.where( + intersections.decreasing_x == out.el_pressure).max(dim='offset_dim') + + # If at the top of the atmosphere the parcel profile is warmer + # than the environment, no EL exists. Also if EL is lower than or + # equal to LCL, no EL exists. + top_pressure = profile.pressure == profile.pressure.min(dim=vert_dim) + top_prof_temp = profile.temperature.where(top_pressure).max(dim=vert_dim) + top_env_temp = profile.environment_temperature.where( + top_pressure).max(dim=vert_dim) + + top_colder = top_prof_temp <= top_env_temp + el_above_lcl = out.el_pressure < profile.lcl_pressure + el_exists = np.logical_and(top_colder, el_above_lcl) + out['el_pressure'] = out.el_pressure.where(el_exists, other=np.nan) + out['el_temperature'] = out.el_temperature.where(el_exists, other=np.nan) + + # There should only be one LFC and EL per point. + assert not 'offset_dim' in out.keys(), 'Duplicate crossings detected.' + + # Identify points where no LFC intersections were found. + lfc_missing = np.isnan(intersections.increasing_x.max(dim='offset_dim')) + + # If no intersection was found, but a parcel temperature above the + # LCL is greater than the environmental temperature, return the + # LCL. + above_lcl = profile.pressure < profile.lcl_pressure + pos_parcel = (profile.temperature.where(above_lcl) > + profile.environment_temperature.where(above_lcl)) + pos_parcel = pos_parcel.any(dim=vert_dim) + no_lfc_pos_parcel = np.logical_and(pos_parcel, lfc_missing) + + # Also return LCL if an intersection exists but all intersections + # are below the LCL and EL is above the LCL. + exists_but_na = np.logical_and(np.logical_not(lfc_missing), + np.isnan(out.lfc_pressure)) + el_above_lcl = out.el_pressure < profile.lcl_pressure + lfc_below_el_above = np.logical_and(exists_but_na, el_above_lcl) + + # Do the replacements with LCL. + replace_with_lcl = np.logical_or(no_lfc_pos_parcel, lfc_below_el_above) + out['lfc_pressure'] = profile.lcl_pressure.where(replace_with_lcl, + other=out.lfc_pressure) + out['lfc_temperature'] = profile.lcl_temperature.where( + replace_with_lcl, + other=out.lfc_temperature) + + # Assign metadata. + out.el_pressure.attrs['long_name'] = 'Equilibrium level pressure' + out.el_pressure.attrs['units'] = 'hPa' + out.el_temperature.attrs['long_name'] = 'Equilibrium level temperature' + out.el_temperature.attrs['units'] = 'K' + out.lfc_pressure.attrs['long_name'] = 'Level of free convection pressure' + out.lfc_pressure.attrs['units'] = 'hPa' + out.lfc_temperature.attrs['long_name'] = ('Level of free convection ' + + 'temperature') + out.lfc_temperature.attrs['units'] = 'K' + + return out + +def trap_around_zeros(x, y, dim, log_x=True, start=0): + """ + Calculate dx * y for points just before and after zeros in y. + + Arguments: + + - x: arrays of x along dim. + - y: arrays of y along dim. + - dim: Dimension along which to calculate. + - log_x: Log transform x? + - start: Zero-based position along dim to look for zeros. + + Returns: + + - a Dataset containing the areas and x coordinates for each + rectangular area calculated before and after each zero; and + an array of x coordinates that should be replaced by the new + areas if integrating along x and including these areas + afterwards. + """ + + # Estimate zero crossings. + zeros = xarray.zeros_like(y) + zero_intersections = find_intersections( + x=x.isel({dim: slice(start, None)}), + a=y.isel({dim: slice(start, None)}), + b=zeros.isel({dim: slice(start, None)}), + dim=dim, log_x=log_x) + zero_intersections = zero_intersections.rename({'offset_dim': dim}) + zero_y = zero_intersections.all_intersect_y + zero_x = zero_intersections.all_intersect_x + + # Take log of x if required. + if log_x: + x = np.log(x) + zero_x = np.log(zero_x) + + zero_level = xarray.zeros_like(y.isel({dim: 0})) + + after_zeros_mask = np.logical_not(np.isnan(zero_y)) + before_zeros_mask = xarray.concat([zero_level, zero_y], + dim=dim).shift({dim: -1}) + before_zeros_mask = np.logical_not(np.isnan(before_zeros_mask)) + + def calc_areas(x, y, mask, shift_x=0): + areas = xarray.Dataset({'area': xarray.zeros_like(mask), + 'dx': xarray.zeros_like(mask), + 'x': xarray.zeros_like(mask)}) + + # Get coordinates of zeros. + x_near_zero = x.where(mask) + y_near_zero = y.where(mask) + + # Determine the value of y (mean of y and zero) and dx. + mean_y = y_near_zero / 2 + + dx = x_near_zero.shift({dim:shift_x}) - zero_x + dx = xarray.concat([zero_level, dx], dim=dim).shift({dim: -shift_x}) + + areas['area'] = mean_y * np.abs(dx) + areas['x'] = x_near_zero - dx/2 + areas['dx'] = np.abs(dx) + + if x.chunks is not None: + areas = areas.chunk(200) + areas = areas.chunk({dim: areas[dim].size}) + + areas = areas.reset_coords(drop=True) + return(areas) + + areas_before_zeros = calc_areas(x=x, y=y, mask=before_zeros_mask, shift_x=1) + areas_after_zeros = calc_areas(x=x, y=y, mask=after_zeros_mask, shift_x=0) + + # Concatenate areas before zeros and areas after zeros. + areas = xarray.concat([areas_before_zeros, areas_after_zeros], dim=dim) + + # Determine start/end points on x axis for each area. + areas['x_from'] = areas.x - areas.dx/2 + areas['x_to'] = areas.x + areas.dx/2 + + # Mask is a mask that selects elements that were *not* included in + # the differences; to be used by a CAPE calculation where we don't + # want to count the areas around zeros twice. + mask = xarray.full_like(x, True) + mask, bef = xarray.broadcast(mask, areas_before_zeros) + mask = mask.where(np.isnan(bef.area), other=False) + + return areas, mask + +def cape_cin_base(pressure, temperature, lfc_pressure, el_pressure, + parcel_profile, vert_dim='model_level_number', **kwargs): + """ + Calculate CAPE and CIN. + + Calculate the convective available potential energy (CAPE) and + convective inhibition (CIN) of a given upper air profile and + parcel path. CIN is integrated between the surface and LFC, CAPE + is integrated between the LFC and EL (or top of + sounding). Intersection points of the measured temperature profile + and parcel profile are logarithmically interpolated. + + Uses the bottom (highest-pressure) LFC and the top + (lowest-pressure) EL. + + Arguments: + + - pressure: Pressure level(s) of interest [hPa]. + - temperature: Temperature at each pressure level [K]. + - lfc_pressure: Pressure of level of free convection [hPa]. + - el_pressure: Pressure of equilibrium level [hPa]. + - parcel_profile: The parcel profile as returned from + parcel_profile(). + - vert_dim: The vertical dimension. + + Returns: + + - Dataset with convective available potential energy (cape) + and convective inhibition (cin), both in J kg-1. + + """ + + # Where the EL is nan, use the highest (lowest-pressure) value as + # the EL. + el_pressure = pressure.min(dim=vert_dim).where(np.isnan(el_pressure), + other=el_pressure) + + # Difference between the parcel path and measured temperature + # profiles. + temp_diffs = xarray.Dataset({'temp_diff': (parcel_profile.temperature - + temperature), + 'pressure': pressure, + 'log_pressure': np.log(pressure)}) + + # Integration areas around zero differences. Note MetPy's + # implementation in _find_append_zero_crossings() looks for + # intersections from the 2nd index onward (start=1 in this code); + # but in this implemnetation the whole array needs to be checked + # (start=0) for the unit tests to pass. + areas_around_zeros, trapz_mask = trap_around_zeros(x=temp_diffs.pressure, + y=temp_diffs.temp_diff, + dim=vert_dim, log_x=True) + areas_around_zeros['x'] = np.exp(areas_around_zeros.x) + areas_around_zeros['x_from'] = np.exp(areas_around_zeros.x_from) + areas_around_zeros['x_to'] = np.exp(areas_around_zeros.x_to) + + # Integrate between LFC and EL pressure levels to get CAPE. + diffs_lfc_to_el = temp_diffs.where(pressure <= lfc_pressure) + diffs_lfc_to_el = diffs_lfc_to_el.where(pressure >= el_pressure) + areas_lfc_to_el = areas_around_zeros.where(areas_around_zeros.x <= + lfc_pressure) + areas_lfc_to_el = areas_lfc_to_el.where(areas_around_zeros.x >= el_pressure) + + cape = mpconsts.Rd.m * trapz(dat=diffs_lfc_to_el, x='log_pressure', + dim=vert_dim, mask=trapz_mask) + cape = cape.reset_coords().temp_diff + cape = cape + (mpconsts.Rd.m * areas_lfc_to_el.area.sum(dim=vert_dim)) + cape.name = 'cape' + cape.attrs['long_name'] = 'Convective available potential energy' + cape.attrs['units'] = 'J kg-1' + + # Integrate between surface and LFC to get CIN. + temp_diffs_surf_to_lfc = temp_diffs.where(pressure >= lfc_pressure) + areas_surf_to_lfc = areas_around_zeros.where(areas_around_zeros.x >= + lfc_pressure) + cin = mpconsts.Rd.m * trapz(dat=temp_diffs_surf_to_lfc, x='log_pressure', + dim=vert_dim, mask=trapz_mask) + cin = cin.reset_coords().temp_diff + cin = cin + (mpconsts.Rd.m * areas_surf_to_lfc.area.sum(dim=vert_dim)) + cin.name = 'cin' + cin.attrs['long_name'] = 'Convective inhibition' + cin.attrs['units'] = 'J kg-1' + + # Set any positive values for CIN to 0. + cin = cin.where(cin <= 0, other=0) + + return xarray.merge([cape, cin]) + +def cape_cin(pressure, temperature, parcel_temperature, parcel_pressure, + parcel_dewpoint, vert_dim='model_level_number', + return_profile=False, **kwargs): + """ + Calculate CAPE and CIN; wraps finding of LFC and parcel profile + and call to cape_cin_base. Uses the bottom (highest-pressure) LFC + and the top (lowest-pressure) EL. + + Arguments: + + - pressure: Pressure level(s) of interest [hPa]. + - temperature: Temperature at each pressure level [K]. + - parcel_temperature: The temperature of the starting parcel [K]. + - parcel_pressure: The pressure of the starting parcel [K]. + - parcel_dewpoint: The dewpoint of the starting parcel [K]. + - vert_dim: The vertical dimension. + - return_profile: Also return the lifted profile? + - **kwargs: Optional extra arguments to cape_cin_base. + + Returns: + + - Dataset with convective available potential energy (cape) + and convective inhibition (cin), both in J kg-1, plus the + lifted profile if return_profile is True. + """ + + # Calculate parcel profile. + profile = parcel_profile_with_lcl(pressure=pressure, + temperature=temperature, + parcel_temperature=parcel_temperature, + parcel_pressure=parcel_pressure, + parcel_dewpoint=parcel_dewpoint, + vert_dim=vert_dim) + + # Calculate LFC and EL. + parcel_lfc_el = lfc_el(profile=profile, vert_dim=vert_dim) + + # Calculate CAPE and CIN. + cape_cin = cape_cin_base(pressure=profile.pressure, + temperature=profile.environment_temperature, + lfc_pressure=parcel_lfc_el.lfc_pressure, + el_pressure=parcel_lfc_el.el_pressure, + parcel_profile=profile, + vert_dim=vert_dim, **kwargs) + + if return_profile: + return cape_cin, xarray.merge([profile, parcel_lfc_el]) + else: + return cape_cin + +def surface_based_cape_cin(pressure, temperature, dewpoint, + vert_dim='model_level_number', + return_profile=False, **kwargs): + """ + Calculate surface-based CAPE and CIN. + + Arguments: + + - pressure: Pressure level(s) of interest [hPa]. + - temperature: Temperature at each pressure level [K]. + - dewpoint: Dewpoint at each level [K]. + - vert_dim: The vertical dimension. + - return_profile: Also return the lifted profile? + - **kwargs: Optional extra arguments to cape_cin. + + Returns: + + - Dataset with convective available potential energy (cape) and + convective inhibition (cin), both in J kg-1. + """ + + # Profile for surface-based parcel ascent. + return cape_cin(pressure=pressure, + temperature=temperature, + parcel_temperature=temperature.isel({vert_dim: 0}), + parcel_pressure=pressure.isel({vert_dim: 0}), + parcel_dewpoint=dewpoint.isel({vert_dim: 0}), + vert_dim=vert_dim, + return_profile=return_profile, + **kwargs) + +def from_most_unstable_parcel(pressure, temperature, dewpoint, + vert_dim='model_level_number', depth=300): + """ + Select pressure and temperature data at and above the most unstable + parcel within the first x hPa above the surface. + + Arguments: + + - pressure: Pressure level(s) of interest [hPa]. + - temperature: Temperature at each pressure level [K]. + - dewpoint: Dewpoint at each level [K]. + - vert_dim: The vertical dimension. + - depth: The depth above the surface (lowest-level pressure) + in which to look for the most unstable parcel. + - return_profile: Also return the lifted profile? + + Returns: + + - subset pressure, subset temperature, subset dewpoint, + and most unstable layer. + """ + + assert pressure.name == 'pressure', 'Pressure requires name pressure.' + assert temperature.name == 'temperature', ('Temperature requires ' + + 'name temperature.') + assert dewpoint.name == 'dewpoint', 'Dewpoint requires name dewpoint.' + + dat = xarray.merge([pressure, temperature, dewpoint]) + + # Find the most unstable layer in the lowest 'depth' hPa. + unstable_layer = most_unstable_parcel(dat=dat, depth=depth, + vert_dim=vert_dim) + + # Subset to layers at or above the most unstable parcels. + dat = dat.where(pressure <= unstable_layer.pressure, drop=True) + dat = shift_out_nans(x=dat, name='pressure', dim=vert_dim) + + return dat.pressure, dat.temperature, dat.dewpoint, unstable_layer + +def most_unstable_cape_cin(pressure, temperature, dewpoint, + vert_dim='model_level_number', depth=300, + return_profile=False, **kwargs): + """ + Calculate CAPE and CIN for the most unstable parcel within a given + depth above the surface.. + + Arguments: + + - pressure: Pressure level(s) of interest [hPa]. + - temperature: Temperature at each pressure level [K]. + - dewpoint: Dewpoint at each level [K]. + - vert_dim: The vertical dimension. + - depth: The depth above the surface (lowest-level pressure) + in which to look for the most unstable parcel. + - return_profile: Also return the lifted profile? + - **kwargs: Optional extra arguments to cape_cin. + + Returns: + + - Dataset with convective available potential energy (cape) and + convective inhibition (cin), both in J kg-1. + """ + + pressure, temperature, dewpoint, unstable_layer = from_most_unstable_parcel( + pressure=pressure, temperature=temperature, dewpoint=dewpoint, + vert_dim=vert_dim, depth=depth) + + return cape_cin(pressure=pressure, + temperature=temperature, + parcel_temperature=unstable_layer.temperature, + parcel_pressure=unstable_layer.pressure, + parcel_dewpoint=unstable_layer.dewpoint, + vert_dim=vert_dim, + return_profile=return_profile, + **kwargs) + +def mix_layer(pressure, temperature, dewpoint, vert_dim='model_level_number', + depth=100): + """ + Fully mix the lowest x hPa in vertical profiles. + + Arguments: + + - pressure: Pressure level(s) [hPa]. + - temperature: Temperature at each pressure level [K]. + - dewpoint: Dewpoint at each level [K]. + - vert_dim: The vertical dimension. + - depth: The depth above the surface (lowest-level pressure) + to mix [hPa]. + + Returns: + + - pressure, temperature, and dewpoint profiles and + the mixed layer parcel. + """ + + # Mix the lowest x hPa. + mp = mixed_parcel(pressure=pressure, temperature=temperature, + dewpoint=dewpoint, depth=depth, vert_dim=vert_dim) + + # Remove layers that were part of the mixed layer. + assert pressure.name == 'pressure', 'Pressure requires name pressure.' + assert temperature.name == 'temperature', ('Temperature requires ' + + 'name temperature.') + assert dewpoint.name == 'dewpoint', 'Dewpoint requires name dewpoint.' + + dat = xarray.merge([pressure, temperature, dewpoint]) + dat = dat.where(pressure < (pressure.max(dim=vert_dim) - depth), drop=True) + dat = shift_out_nans(x=dat, name='pressure', dim=vert_dim) + + # Add the mixed layer to the bottom of the profiles. + mp[vert_dim] = dat.pressure[vert_dim].min() - 1 + pressure = xarray.concat([mp.pressure, dat.pressure], dim=vert_dim) + temperature = xarray.concat([mp.temperature, dat.temperature], dim=vert_dim) + dewpoint = xarray.concat([mp.dewpoint, dat.dewpoint], dim=vert_dim) + + return pressure, temperature, dewpoint, mp + +def mixed_layer_cape_cin(pressure, temperature, dewpoint, + vert_dim='model_level_number', + depth=100, return_profile=False, + **kwargs): + """ + Calculate CAPE and CIN for a fully-mixed lowest x hPa parcel. + + Arguments: + + - pressure: Pressure level(s) of interest [hPa]. + - temperature: Temperature at each pressure level [K]. + - dewpoint: Dewpoint at each level [K]. + - vert_dim: The vertical dimension. + - depth: The depth above the surface (lowest-level pressure) + to mix [hPa]. + - return_profile: Also return the lifted profile? + - **kwargs: Optional extra arguments to cape_cin. + + Returns: + + - Dataset with convective available potential energy (cape) + and convective inhibition (cin), both in J kg-1. + """ + + pressure, temperature, dewpoint, mp = mix_layer(pressure=pressure, + temperature=temperature, + dewpoint=dewpoint, + vert_dim=vert_dim, + depth=depth) + + return cape_cin(pressure=pressure, + temperature=temperature, + parcel_temperature=mp.temperature, + parcel_pressure=mp.pressure, + parcel_dewpoint=mp.dewpoint, + vert_dim=vert_dim, + return_profile=return_profile, + **kwargs) + +def shift_out_nans(x, name, dim): + """ + Shift data along a dim to remove all leading nans in that + dimension, element-wise. + + Arguments: + + - x: The data to work on. + - name: The name within data in which to look for nans. + - dim: The dimension to shift. + - pt: The index along the dimension to shift 'to'. + + """ + + while np.any(np.isnan(x[name].isel({dim: 0}))): + shifted = x.shift({dim: -1}) + x = shifted.where(np.isnan(x[name].isel({dim: 0})), other=x) + + return x + +def lifted_index(profile, vert_dim='model_level_number'): + """ + Calculate the lifted index. + + Lifted index formula derived from Galway 1956 and referenced by + DoswellSchultz 2006. + + Arguments: + + - profile: Profile as returned by parcel_profile_with_lcl(). + - vert_dim: The vertical dimension name. + + Returns: + + - Lifted index at each point [K]. + """ + + # Interpolate to get 500 hPa values. + dat = log_interp(x=profile, coords=profile.pressure, at=500, + dim=vert_dim) + dat = dat.reset_coords(drop=True) + + # Calculate lifted index. + li = xarray.Dataset({'lifted_index': (dat.environment_temperature - + dat.temperature)}) + li.lifted_index.attrs['long_name'] = 'Lifted index' + li.lifted_index.attrs['units'] = 'K' + + return li + +def linear_interp(x, coords, at, dim='model_level_number'): + """ + Perform simple linear interpolation to get values at specified + points. + + Arguments: + + - x: Data set to interpolate. + - coords: Coordinate value for each point in x. + - at: Points at which to interpolate. + - dim: The dimension along which to interpolate. + + It is assumed that x[coords] is sorted and does not contain + duplicate values along the selected dimension. + """ + + coords_before = coords.where(coords >= at).min(dim=dim) + coords_after = coords.where(coords <= at).max(dim=dim) + assert dim not in coords_before.coords, 'Duplicates detected in coords.' + assert dim not in coords_after.coords, 'Duplicates detected in coords.' + + x_before = x.where(coords == coords_before).max(dim=dim) + x_after = x.where(coords == coords_after).max(dim=dim) + + # The interpolated values. + res = x_before + (x_after - x_before) * ((at - coords_before) / + (coords_after - coords_before)) + + # When the interpolated x exists already, return it. + res = x_before.where(x_before == x_after, other=res) + + return(res) + +def log_interp(x, coords, at, dim='model_level_number'): + """ + Run linear_interp on logged coordinate values. + + Arguments: + - x: Data set to interpolate. + - coords: Coordinate value for each point in x. + - at: Points at which to interpolate. + - dim: The dimension along which to interpolate. + + It is assumed that x[coords] is sorted and does not contain duplicate + values along the selected dimension. + """ + + return linear_interp(x=x, coords=np.log(coords), at=np.log(at), dim=dim) + +def deep_convective_index(pressure, temperature, dewpoint, lifted_index, + vert_dim='model_level_number'): + """ + Calculate the deep convective index (DCI) as defined by Kunz 2009. + + Arguments: + + - pressure: Pressure values [hPa]. + - temperature: Temperature at each pressure [K]. + - dewpoint: Dewpoint temperature at each pressure. + - lifted_index: The lifted index. + - vert_dim: The vertical dimension name. + + Returns: + + - The DCI [C] per point. + """ + + dat = xarray.merge([pressure, temperature, dewpoint]) + + # Interpolate to get 850 hPa values. + dat = log_interp(x=dat, coords=dat.pressure, at=850, dim=vert_dim) + dat = dat.reset_coords(drop=True) + + # Convert temperature and dewpoint from K to C. + dat['temperature'] = dat.temperature - 273.15 + dat['dewpoint'] = dat.dewpoint - 273.15 + + dci = xarray.Dataset({'dci': dat.temperature + dat.dewpoint - lifted_index}) + dci.dci.attrs['long_name'] = 'Deep convective index' + dci.dci.attrs['units'] = 'C' + return dci + +def conv_properties(dat, vert_dim='model_level_number'): + """ + Calculate selected convective properties for a set of points. + + Arguments: + + - dat: An xarray Dataset containing pressure, temperature, and + specific humidity. + - vert_dim: The name of the vertical dimension in the dataset. + + Returns: + + - Dataset containing mixed parcel (100 hPa depth) cape and + cin, most unstable parcel (250 hPa depth) cape and cin, + lifted index and deep convective index for each point. + """ + + # Calculate dewpoints. + print('Calculating dewpoint...') + dat['dewpoint'] = metpy.calc.dewpoint_from_specific_humidity( + pressure=dat.pressure, + temperature=dat.temperature, + specific_humidity=dat.specific_humidity) + dat['dewpoint'] = dat.dewpoint.metpy.convert_units('K') + dat['dewpoint'] = dat.dewpoint.metpy.dequantify() + + # CAPE and CIN for most unstable parcel. + print('Calculating most-unstable CAPE and CIN...') + max_cape_cin = most_unstable_cape_cin( + pressure=dat.pressure, + temperature=dat.temperature, + dewpoint=dat.dewpoint, + vert_dim=vert_dim, + depth=250) + max_cape_cin = max_cape_cin.rename({'cape': 'max_cape', + 'cin': 'max_cin'}) + max_cape_cin.max_cape.attrs['description'] = ('CAPE for most-unstable ' + + 'parcel in lowest 250 hPa.') + max_cape_cin.max_cin.attrs['description'] = ('CIN for most-unstable ' + + 'parcel in lowest 250 hPa.') + + # Mixed-parcel CAPE and CIN. + print('Calculating mixed-parcel CAPE and CIN...') + mixed_cape_cin, mixed_profile = mixed_layer_cape_cin( + pressure=dat.pressure, + temperature=dat.temperature, + dewpoint=dat.dewpoint, + vert_dim=vert_dim, + depth=100, return_profile=True) + mixed_cape_cin = mixed_cape_cin.rename({'cape': 'mixed_cape', + 'cin': 'mixed_cin'}) + mixed_cape_cin.mixed_cape.attrs['description'] = ('CAPE for fully-mixed ' + + 'lowest 100 hPa parcel.') + mixed_cape_cin.mixed_cin.attrs['description'] = ('CIN for fully-mixed ' + + 'lowest 100 hPa parcel') + + # Lifted index using mixed layer profile. + print('Calculating lifted index...') + li = lifted_index(profile=mixed_profile, vert_dim=vert_dim) + + # Deep convective index for mixed layer profile. + print('Calculating deep convective index...') + dci = deep_convective_index(pressure=dat.pressure, + temperature=dat.temperature, + dewpoint=dat.dewpoint, + lifted_index=li.lifted_index, + vert_dim=vert_dim) + + out = xarray.merge([mixed_cape_cin, + max_cape_cin, + li, + dci]) + + return out + +def lapse_rate(pressure, temperature, height, from_pressure=700, to_pressure=500, + vert_dim='model_level_number'): + """ + Calculate the observed environmental lapse rate between two pressure levels. + + Arguments: + + - pressure: Pressure at each level [hPa]. + - temperature: Temperature at each level [K]. + - height: Height of each level [m] + - from_pressure: Pressure level to calculate from [hPa]. + - to_pressure: Pressure level to calculate to [hPa]. + - vert_dim: Name of vertical dimension. + + Returns: + + - Lapse rate between two levels at each point [K/km]. + """ + + from_temperature = log_interp(x=temperature, coords=pressure, + at=from_pressure, dim=vert_dim) + to_temperature = log_interp(x=temperature, coords=pressure, + at=to_pressure, dim=vert_dim) + from_height = log_interp(x=height, coords=pressure, + at=from_pressure, dim=vert_dim)/1000 + to_height = log_interp(x=height, coords=pressure, + at=to_pressure, dim=vert_dim)/1000 + + lapse = (to_temperature - from_temperature) / (to_height - from_height) + lapse.attrs['long_name'] = f'{from_pressure}-{to_pressure} hPa lapse rate' + lapse.attrs['units'] = 'K/km' + + return lapse + +def freezing_level_height(temperature, height, vert_dim='model_level_number'): + """ + Calculate the freezing level height. + + Arguments: + + - temperature: Temperature at each level [K]. + - height: Height of each level [m]. + - vert_dim: Name of vertical dimension. + + Returns: + + - Freezing level height [m]. + """ + + flh = linear_interp(x=height, coords=temperature, at=273.15, dim=vert_dim) + flh.attrs['long_name'] = f'Freezing level height' + flh.attrs['units'] = 'm' + flh.name = 'freezing_level' + return flh + +def isobar_temperature(pressure, temperature, isobar, vert_dim='model_level_number'): + """ + Calculate the temperature at a given pressure. + + Arguments: + + - pressure: Pressure at each level [hPa]. + - temperature: Temperature at each level [K]. + - isobar: Pressure at which to find temperatures [hPa]. + - vert_dim: Name of vertical dimension. + + Returns: + + - Isobar temperature [K]. + """ + + temp = log_interp(x=temperature, coords=pressure, at=isobar, dim=vert_dim) + temp.attrs['long_name'] = f'Temperature at {isobar} hPa.' + temp.attrs['units'] = 'K' + return temp + +def wind_shear(surface_wind_u, surface_wind_v, wind_u, wind_v, height, shear_height=6000): + """ + Calculate wind shear. + + Arguments: + - surface_wind_u, surface_wind_v: U and V components of surface wind. + - wind_u, wind_v: U and V components of above-surface wind. + - height: The height of every coordinate in wind_u and wind_v. + - shear_height: The wind height to subtract from surface wind [m]. + + Returns: + + - Wind shear = wind at shear_height - wind at surface. + """ + + wind_high_u = linear_interp(x=wind_u, coords=height, at=shear_height) + wind_high_v = linear_interp(x=wind_v, coords=height, at=shear_height) + + shear_u = wind_high_u - surface_wind_u + shear_u.name = 'shear_u' + shear_u.attrs['long_name'] = f'Surface to {shear_height} m wind shear, U component.' + + shear_v = wind_high_v - surface_wind_v + shear_v.name = 'shear_v' + shear_v.attrs['long_name'] = f'Surface to {shear_height} m wind shear, V component.' + + shear_magnitude = np.sqrt(shear_u**2 + shear_v**2) + shear_magnitude.name = 'shear_magnitude' + shear_magnitude.attrs['long_name'] = f'Surface to {shear_height} m bulk wind shear.' + + out = xarray.merge([shear_u, shear_v, shear_magnitude]) + for v in ['shear_u', 'shear_v', 'shear_magnitude']: + out[v].attrs['units'] = 'm s-1' + + return out + +def significant_hail_parameter(mucape, mixing_ratio, lapse, temp_500, shear, flh): + """ + Calculate the significant hail parameter, as given at + https://www.spc.noaa.gov/exper/mesoanalysis/help/help_sigh.html + + Arguments: + + - mucape: Most unstable parcel CAPE [J kg-1]. + - mixing_ratio: Mixing ratio of the most unstable parcel [kg kg-1]. + - lapse: 700-500 hPa lapse rate [K/km]. + - temp_500: Temperature at 500 hPa [K]. + - shear: 0-6 km bulk wind shear [m s-1]. + - flh: Freezing level height [m]. + + Returns: + + - Significant hail parameter. + """ + + # Convert from kg kg-1 to g kg-1 + mixing_ratio = mixing_ratio * 1e3 + + # Use positive values of lapse rate. + lapse = -lapse + + # Convert temperatures from K to C. + temp_500 = temp_500 - 273.15 + + # Apply thresholds on inputs to identify where SHIP is valid. + shear = shear.where(shear >= 7).where(shear <= 27) + mixing_ratio = mixing_ratio.where(mixing_ratio >= 11).where(mixing_ratio <= 13.6) + temp_500 = temp_500.where(temp_500 <= -5.5, other=-5.5) + + # Calculate basic SHIP value. + ship = mucape * mixing_ratio * lapse * -temp_500 * shear / 42000000 + + # Three conditions change the value of SHIP. + ship = ship.where(mucape >= 1300, other=ship * (mucape/1300)) + ship = ship.where(lapse >= 5.8, other=ship * (lapse/5.8)) + ship = ship.where(flh >= 2400, other=ship * (flh/2400)) + + # Metadata. + ship.attrs['long_name'] = 'Significant hail parameter' + ship.attrs['units'] = 'J kg-2 g K^2 km-1 m s-1' + + return ship