diff --git a/src/metpy/calc/thermo.py b/src/metpy/calc/thermo.py index 4f17a00fd84..fd7dc34b1be 100644 --- a/src/metpy/calc/thermo.py +++ b/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'), diff --git a/tests/calc/test_thermo.py b/tests/calc/test_thermo.py index a5c66a8fc45..ea97646d49d 100644 --- a/tests/calc/test_thermo.py +++ b/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."""