Merge pull request #2263 from dopplershift/moist-lapse

Refactor moist_lapse

Author: dcamron

src/metpy/calc/thermo.py

@@ -295,9 +295,9 @@ def moist_lapse(pressure, temperature, reference_pressure=None):
        Renamed ``ref_pressure`` parameter to ``reference_pressure``
 
     """
-    def dt(t, p):
-        t = units.Quantity(t, temperature.units)
-        p = units.Quantity(p, pressure.units)
+    def dt(p, t):
+        t = units.Quantity(t, 'kelvin')
+        p = units.Quantity(p, 'mbar')
         rs = saturation_mixing_ratio(p, t)
         frac = ((mpconsts.Rd * t + mpconsts.Lv * rs)
                 / (mpconsts.Cp_d + (mpconsts.Lv * mpconsts.Lv * rs * mpconsts.epsilon
@@ -307,39 +307,55 @@ def dt(t, p):
     pressure = np.atleast_1d(pressure)
     if reference_pressure is None:
         reference_pressure = pressure[0]
-
-    if np.isnan(reference_pressure):
-        return units.Quantity(np.full(pressure.shape, np.nan), temperature.units)
-
     pressure = pressure.to('mbar')
     reference_pressure = reference_pressure.to('mbar')
     org_units = temperature.units
-    temperature = np.atleast_1d(temperature).to('kelvin')
+    temperature = np.atleast_1d(temperature).m_as('kelvin')
+
+    if np.isnan(reference_pressure) or np.all(np.isnan(temperature)):
+        return units.Quantity(np.full((temperature.size, pressure.size), np.nan), org_units)
 
     pres_decreasing = (pressure[0] > pressure[-1])
     if pres_decreasing:
         # Everything is easier if pressures are in increasing order
         pressure = pressure[::-1]
 
-    ref_pres_idx = np.searchsorted(pressure.m, reference_pressure.m, side='right')
-    ret_temperatures = np.empty((0, temperature.shape[0]))
+    # It would be preferable to use a regular solver like RK45, but as of scipy 1.8.0
+    # anything other than LSODA goes into an infinite loop when given NaNs for y0.
+    solver_args = {'fun': dt, 'y0': temperature,
+                   'method': 'LSODA', 'atol': 1e-7, 'rtol': 1.5e-8}
 
-    if _greater_or_close(reference_pressure, pressure.min()):
-        # Integrate downward in pressure
-        pres_down = np.append(reference_pressure.m, pressure[(ref_pres_idx - 1)::-1].m)
-        trace_down = si.odeint(dt, temperature.m.squeeze(), pres_down.squeeze())
-        ret_temperatures = np.concatenate((ret_temperatures, trace_down[:0:-1]))
-
-    if reference_pressure < pressure.max():
-        # Integrate upward in pressure
-        pres_up = np.append(reference_pressure.m, pressure[ref_pres_idx:].m)
-        trace_up = si.odeint(dt, temperature.m.squeeze(), pres_up.squeeze())
-        ret_temperatures = np.concatenate((ret_temperatures, trace_up[1:]))
+    # Need to handle close points to avoid an error in the solver
+    close = np.isclose(pressure, reference_pressure)
+    if np.any(close):
+        ret = np.broadcast_to(temperature[:, np.newaxis], (temperature.size, np.sum(close)))
+    else:
+        ret = np.empty((temperature.size, 0), dtype=temperature.dtype)
+
+    # Do we have any points above the reference pressure
+    points_above = (pressure < reference_pressure) & ~close
+    if np.any(points_above):
+        # Integrate upward--need to flip so values are properly ordered from ref to min
+        press_side = pressure[points_above][::-1].m
+
+        # Flip on exit so t values correspond to increasing pressure
+        trace = si.solve_ivp(t_span=(reference_pressure.m, press_side[-1]),
+                             t_eval=press_side, **solver_args).y[..., ::-1]
+        ret = np.concatenate((trace, ret), axis=-1)
+
+    # Do we have any points below the reference pressure
+    points_below = ~points_above & ~close
+    if np.any(points_below):
+        # Integrate downward
+        press_side = pressure[points_below].m
+        trace = si.solve_ivp(t_span=(reference_pressure.m, press_side[-1]),
+                             t_eval=press_side, **solver_args).y
+        ret = np.concatenate((ret, trace), axis=-1)
 
     if pres_decreasing:
-        ret_temperatures = ret_temperatures[::-1]
+        ret = ret[..., ::-1]
 
-    return units.Quantity(ret_temperatures.T.squeeze(), 'kelvin').to(org_units)
+    return units.Quantity(ret.squeeze(), 'kelvin').to(org_units)
 
 
 @exporter.export

src/metpy/plots/skewt.py

@@ -545,8 +545,7 @@ def plot_moist_adiabats(self, t0=None, pressure=None, **kwargs):
             pressure = units.Quantity(np.linspace(*self.ax.get_ylim()), 'mbar')
 
         # Assemble into data for plotting
-        t = moist_lapse(pressure, t0[:, np.newaxis],
-                        units.Quantity(1000., 'mbar')).to(units.degC)
+        t = moist_lapse(pressure, t0, units.Quantity(1000., 'mbar')).to(units.degC)
         linedata = [np.vstack((ti.m, pressure.m)).T for ti in t]
 
         # Add to plot

tests/calc/test_thermo.py

@@ -3,6 +3,8 @@
 # SPDX-License-Identifier: BSD-3-Clause
 """Test the `thermo` module."""
 
+import warnings
+
 import numpy as np
 import pytest
 import xarray as xr
@@ -131,12 +133,31 @@ def test_moist_lapse_ref_pressure():
     assert_array_almost_equal(temp, true_temp, 2)
 
 
+def test_moist_lapse_multiple_temps():
+    """Test moist_lapse with multiple starting temperatures."""
+    temp = moist_lapse(np.array([1050., 800., 600., 500., 400.]) * units.mbar,
+                       np.array([19.85, np.nan, 19.85]) * units.degC, 1000. * units.mbar)
+    true_temp = np.array([[294.76, 284.64, 272.81, 264.42, 252.91],
+                          [np.nan, np.nan, np.nan, np.nan, np.nan],
+                          [294.76, 284.64, 272.81, 264.42, 252.91]]) * units.kelvin
+    assert_array_almost_equal(temp, true_temp, 2)
+
+
 def test_moist_lapse_scalar():
     """Test moist_lapse when given a scalar desired pressure and a reference pressure."""
     temp = moist_lapse(np.array([800.]) * units.mbar, 19.85 * units.degC, 1000. * units.mbar)
     assert_almost_equal(temp, 284.64 * units.kelvin, 2)
 
 
+def test_moist_lapse_close_start():
+    """Test that we behave correctly with a reference pressure close to an actual pressure."""
+    with warnings.catch_warnings(record=True) as record:
+        temp = moist_lapse(units.Quantity(1000, 'hPa'), 0 * units.degC,
+                           units.Quantity(1000., 'mbar'))
+        assert len(record) == 0
+    assert_almost_equal(temp, units.Quantity(0., 'degC'))
+
+
 def test_moist_lapse_uniform():
     """Test moist_lapse when given a uniform array of pressures."""
     temp = moist_lapse(np.array([900., 900., 900.]) * units.hPa, 20. * units.degC)
@@ -158,7 +179,7 @@ def test_moist_lapse_nan_ref_press():
 def test_moist_lapse_downwards():
     """Test moist_lapse when integrating downwards (#2128)."""
     temp = moist_lapse(units.Quantity([600, 700], 'mbar'), units.Quantity(0, 'degC'))
-    assert_almost_equal(temp, units.Quantity([0, 6.47748353], units.degC))
+    assert_almost_equal(temp, units.Quantity([0, 6.47748353], units.degC), 4)
 
 
 @pytest.mark.parametrize('direction', (1, -1))