Merge pull request Unidata#3324 from dopplershift/fix-3315

Improve isentropic interpolation

src/metpy/calc/thermo.py

@@ -2624,7 +2624,7 @@ def isentropic_interpolation(levels, pressure, temperature, *args, vertical_dim=
         One-dimensional array of desired potential temperature surfaces
 
     pressure : array-like
-        One-dimensional array of pressure levels
+        Array of pressure
 
     temperature : array-like
         Array of temperature
@@ -2691,10 +2691,17 @@ def _isen_iter(iter_log_p, isentlevs_nd, ka, a, b, pok):
     pressure = pressure.to('hPa')
     temperature = temperature.to('kelvin')
 
+    # Construct slices for broadcasting with temperature (used for pressure & theta levels)
     slices = [np.newaxis] * temperature.ndim
     slices[vertical_dim] = slice(None)
     slices = tuple(slices)
-    pressure = units.Quantity(np.broadcast_to(pressure[slices].magnitude, temperature.shape),
+
+    # For 1-D pressure, we assume it's the vertical coordinate and know how it should broadcast
+    # to the same shape as temperature. Otherwise, just assume it's ready for broadcast, or
+    # it has the same shape and is a no-op.
+    if pressure.ndim == 1:
+        pressure = pressure[slices]
+    pressure = units.Quantity(np.broadcast_to(pressure.magnitude, temperature.shape),
                               pressure.units)
 
     # Sort input data
@@ -2779,7 +2786,8 @@ def isentropic_interpolation_as_dataset(
     *args,
     max_iters=50,
     eps=1e-6,
-    bottom_up_search=True
+    bottom_up_search=True,
+    pressure=None
 ):
     r"""Interpolate xarray data in isobaric coords to isentropic coords, returning a Dataset.
 
@@ -2799,6 +2807,9 @@ def isentropic_interpolation_as_dataset(
     bottom_up_search : bool, optional
         Controls whether to search for levels bottom-up (starting at lower indices),
         or top-down (starting at higher indices). Defaults to True, which is bottom-up search.
+    pressure : `xarray.DataArray`, optional
+        Array of pressure to use when the vertical coordinate for the passed in data is not
+        pressure (e.g. data using sigma coordinates)
 
     Returns
     -------
@@ -2833,10 +2844,20 @@ def isentropic_interpolation_as_dataset(
                        'The output Dataset includes duplicate levels as a result. '
                        'This may cause xarray to crash when working with this Dataset!')
 
+    if pressure is None:
+        pressure = all_args[0].metpy.vertical
+
+    if (units.get_dimensionality(pressure.metpy.units)
+            != units.get_dimensionality('[pressure]')):
+        raise ValueError('The pressure array/vertical coordinate for the passed in data does '
+                         'not appear to be pressure. In this case you need to pass in '
+                         'pressure values with proper units using the `pressure` keyword '
+                         'argument.')
+
     # Obtain result as list of Quantities
     ret = isentropic_interpolation(
         levels,
-        all_args[0].metpy.vertical,
+        pressure,
         all_args[0].metpy.unit_array,
         *(arg.metpy.unit_array for arg in all_args[1:]),
         vertical_dim=all_args[0].metpy.find_axis_number('vertical'),

tests/calc/test_thermo.py

@@ -1294,9 +1294,12 @@ def test_isentropic_pressure_p_increase_rh_out():
     assert_almost_equal(isentprs[1], truerh, 3)
 
 
-def test_isentropic_pressure_interp():
+@pytest.mark.parametrize('press_3d', [False, True])
+def test_isentropic_pressure_interp(press_3d):
     """Test calculation of isentropic pressure function."""
     lev = [100000., 95000., 90000., 85000.] * units.Pa
+    if press_3d:
+        lev = np.lib.stride_tricks.broadcast_to(lev[:, None, None], (4, 5, 5))
     tmp = np.ones((4, 5, 5))
     tmp[0, :] = 296.
     tmp[1, :] = 292.
@@ -1465,6 +1468,78 @@ def test_isentropic_interpolation_as_dataset_duplicate(xarray_isentropic_data):
                                                 xarray_isentropic_data.rh)
 
 
+@pytest.fixture
+def xarray_sigma_isentropic_data():
+    """Generate test xarray dataset on sigma vertical coords for interpolation functions."""
+    return xr.Dataset(
+        {
+            'temperature': (
+                ('z', 'y', 'x'),
+                [[[296.]], [[292.]], [[290.]], [[288.]]] * units.K
+            ),
+            'rh': (
+                ('z', 'y', 'x'),
+                [[[100.]], [[80.]], [[40.]], [[20.]]] * units.percent
+            ),
+            'pressure': (
+                ('z', 'y', 'x'),
+                [[[1000.]], [[950.]], [[900.]], [[850.]]] * units.hPa
+            )
+        },
+        coords={
+            'z': (('z',), [0.98, 0.928, 0.876, 0.825], {'units': 'dimensionless'}),
+            'time': '2020-01-01T00:00Z'
+        }
+    )
+
+
+def test_isen_interpolation_as_dataset_non_pressure_default(xarray_sigma_isentropic_data):
+    """Test isentropic interpolation with xarray data with non-pressure vertical coord."""
+    isentlev = [296., 297.] * units.kelvin
+    with pytest.raises(ValueError, match='vertical coordinate for the.*does not'):
+        isentropic_interpolation_as_dataset(isentlev,
+                                            xarray_sigma_isentropic_data.temperature,
+                                            xarray_sigma_isentropic_data.rh)
+
+
+def test_isen_interpolation_as_dataset_passing_pressre(xarray_sigma_isentropic_data):
+    """Test isentropic interpolation with xarray when passing a pressure array."""
+    isentlev = [296., 297.] * units.kelvin
+    result = isentropic_interpolation_as_dataset(
+        isentlev, xarray_sigma_isentropic_data.temperature,
+        xarray_sigma_isentropic_data.rh, pressure=xarray_sigma_isentropic_data.pressure)
+    expected = xr.Dataset(
+        {
+            'pressure': (
+                ('isentropic_level', 'y', 'x'),
+                [[[1000.]], [[936.213]]] * units.hPa,
+                {'standard_name': 'air_pressure'}
+            ),
+            'temperature': (
+                ('isentropic_level', 'y', 'x'),
+                [[[296.]], [[291.4579]]] * units.K,
+                {'standard_name': 'air_temperature'}
+            ),
+            'rh': (
+                ('isentropic_level', 'y', 'x'),
+                [[[100.]], [[69.19706]]] * units.percent
+            )
+        },
+        coords={
+            'isentropic_level': (
+                ('isentropic_level',),
+                [296., 297.],
+                {'units': 'kelvin', 'positive': 'up'}
+            ),
+            'time': '2020-01-01T00:00Z'
+        }
+    )
+    xr.testing.assert_allclose(result, expected)
+    assert result['pressure'].attrs == expected['pressure'].attrs
+    assert result['temperature'].attrs == expected['temperature'].attrs
+    assert result['isentropic_level'].attrs == expected['isentropic_level'].attrs
+
+
 @pytest.mark.parametrize('array_class', (units.Quantity, masked_array))
 def test_surface_based_cape_cin(array_class):
     """Test the surface-based CAPE and CIN calculation."""