Fix _pressure_to_plevs

- Fixed incorrect data variable retrieval (need Z coordinates and not ps)
- Remove target_data arg with ds.regridder.vertical
- Add unit tests
- Update comments and docstrings

e3sm_diags/driver/utils/regrid.py

@@ -185,12 +185,12 @@ def _hybrid_to_plevs(
     var_key: str,
     plevs: List[int] | List[float],
 ) -> xr.DataArray:
-    """Regrid the variable's hybrid-sigma levels to the desired pressure levels.
+    """Regrid a variable's hybrid-sigma levels to the desired pressure levels.
 
     Steps:
-        1. Convert hybrid-sigma levels to pressure data
-        2. Regrid the pressure data to the pressure levels (plevs)
-        3. Return the dataset with the Z axis on pressure levels
+        1. Create the output pressure grid using ``plevs``.
+        2. Convert hybrid-sigma levels to pressure coordinates.
+        3. Regrid the pressure coordinates to the output pressure grid (plevs).
 
     Parameters
     ----------
@@ -200,40 +200,39 @@ def _hybrid_to_plevs(
         The variable key.
     plevs : List[int] | List[float]
         A 1-D array of floats or integers representing output pressure levels
-        in mb units. This parameter is usually set by ``CoreParameter.plevs``
-        attribute. For example, ``plevs=[850.0, 200.0]``.
+        in mb units. For example, ``plevs=[850.0, 200.0]``. This parameter is
+        usually set by the ``CoreParameter.plevs`` attribute.
 
     Returns
     -------
     xr.DataArray
         The variable with a Z axis regridded to pressure levels (mb units).
     """
+    # TODO: Do we need to convert the Z axis to mb units if it is in PA?
     ds = dataset.copy()
 
-    # Convert hybrid-sigma levels to pressure data.
-    # Formula: p(k) = hyam(k) * p0 + hybm(k) * ps
-    pressure_data = _hybrid_to_pressure(ds, var_key)
-
-    # Create the output pressure grid to regrid to using the `plevs` array.
+    # 1. Create the output pressure grid using ``plevs``.
     pressure_grid = xc.create_grid(z=xc.create_axis("lev", plevs))
 
+    # 2. Convert hybrid-sigma levels to pressure coordinates.
+    pressure_coords = _hybrid_to_pressure(ds, var_key)
+
     # TODO: Do we need to make sure z is positive down and why?
 
-    # 4. Regrid the variable using the pressure grid, pressure data,
-    # and the log linear method.
+    # 3. Regrid the pressure coordinates to the pressure levels (plevs).
     result = ds.regridder.vertical(
         var_key,
         output_grid=pressure_grid,
         tool="xgcm",
         method="log",
-        target_data=pressure_data,
+        target_data=pressure_coords,
     )
 
     return result
 
 
 def _hybrid_to_pressure(dataset: xr.Dataset, var_key: str) -> xr.DataArray:
-    """Convert hybrid-sigma levels to pressure data (mb).
+    """Regrid hybrid-sigma levels to pressure coordinates (mb).
 
     Formula: p(k) = hyam(k) * p0 + hybm(k) * ps
       * p: pressure data (mb).
@@ -253,37 +252,29 @@ def _hybrid_to_pressure(dataset: xr.Dataset, var_key: str) -> xr.DataArray:
     Returns
     -------
     xr.DataArray
-        The variable's pressure data.
+        The variable with a Z axis on pressure coordinates.
 
     Raises
     ------
     KeyError
         If the dataset does not contain pressure data (ps) or any of the
         hybrid levels (hyam, hymb).
-
-    Notes
-    -----
-
-    Unit conversion formulas:
-        - mb = Pa / 100
-        - Pa = (mb * 100)
     """
     ds = dataset.copy()
 
     # p0 is statically set to mb (1000) instead of retrieved from the dataset.
-    # Usually p0 is either in mb (1000) or Pa (100000).
-    hyam = _get_hybrid_sigma_level(ds, "hyam")
     p0 = 1000
-    hybm = _get_hybrid_sigma_level(ds, "hybm")
     ps = _get_hybrid_sigma_level(ds, "ps")
+    hyam = _get_hybrid_sigma_level(ds, "hyam")
+    hybm = _get_hybrid_sigma_level(ds, "hybm")
 
     if ps is None or hyam is None or hybm is None:
         raise KeyError(
             f"The dataset for '{var_key}' does not contain hybrid-sigma level 'ps', "
             "'hyam' and/or 'hybm' to use for reconstructing to pressure data."
         )
 
-    ps = _convert_units_from_pa_to_mb(ps)
+    ps = _convert_units_to_mb(ps)
 
     pressure_data = hyam * p0 + hybm * ps
     pressure_data.attrs["units"] = "mb"
@@ -328,7 +319,7 @@ def _pressure_to_plevs(
     var_key: str,
     plevs: List[int] | List[float],
 ) -> xr.DataArray:
-    """Regrid pressure data to desired pressure level(s) in mb units.
+    """Regrids pressure coordinates to the desired pressure level(s).
 
     Parameters
     ----------
@@ -344,76 +335,72 @@ def _pressure_to_plevs(
     Returns
     -------
     xr.DataArray
-        The variable with a Z axis regridded to pressure levels (mb units).
+        The variable with a Z axis on pressure levels (mb).
     """
     ds = dataset.copy()
 
     # Create the output pressure grid to regrid to using the `plevs` array.
     pressure_grid = xc.create_grid(z=xc.create_axis("lev", plevs))
 
-    # Get the pressure data (ps) and convert to mb units if not already in mb.
-    ps = ds.get("ps")
-    if ps is None:
-        raise KeyError(
-            f"The dataset for '{var_key}' does not contain 'ps' to convert to the "
-            "desired pressure level(s)."
-        )
+    # Convert pressure coordinates to mb if it is not already in mb.
+    lev_key = xc.get_dim_keys(ds[var_key], axis="Z")
+    ds[lev_key] = _convert_units_to_mb(ds[lev_key])
 
-    ps = _convert_units_from_pa_to_mb(ps)
-
-    # Perform the vertical regridding using log linear method.
     result = ds.regridder.vertical(
         var_key,
         output_grid=pressure_grid,
         tool="xgcm",
         method="log",
-        target_data=ps,
     )
 
     return result
 
 
-def _convert_units_from_pa_to_mb(ps: xr.DataArray) -> xr.DataArray:
-    """Convert ps from Pa (Pascal pressure) to mb (millibars) if not in mb.
+def _convert_units_to_mb(da: xr.DataArray) -> xr.DataArray:
+    """Convert DataArray to mb (millibars) if not in mb.
+
+    Unit conversion formulas:
+      * mb = Pa / 100
+      * Pa = (mb * 100)
 
     The more common unit on weather maps is mb.
 
     Parameters
     ----------
-    ps : xr.DataArray
-        2-D array equal to surface pressure data.
+    da : xr.DataArray
+        An xr.DataArray, usually for "lev" or "ps".
 
     Returns
     -------
     xr.DataArray
-        ps in 'mb' units (if not already).
+        The DataArray in mb units.
 
     Raises
     ------
     ValueError
-        If ``ps`` DataArray has no 'units' attribute.
+        If ``da`` DataArray has no 'units' attribute.
     ValueError
-        If ``ps`` DataArray has units that is not 'Pa' or 'mb'.
+        If ``da`` DataArray has units not in mb or Pa.
     """
-    ps_units = ps.attrs.get("units")
+    units = da.attrs.get("units")
 
-    if ps_units is None:
+    if units is None:
         raise ValueError(
             "'{ps.name}' has no 'units' attribute to determine if data is in 'mb' or "
             "'Pa' units."
         )
 
-    if ps_units == "Pa":
+    if units == "mb":
+        pass
+    elif units == "Pa":
         with xr.set_options(keep_attrs=True):
-            ps = ps / 100.0
+            da = da / 100.0
 
-        ps.attrs["units"] = "mb"
-    elif ps_units == "mb":
-        pass
+        da.attrs["units"] = "mb"
     else:
         raise ValueError(
-            f"'{ps.name}' should be in 'mb' or 'Pa' (which gets converted to 'mb'), "
-            f"not {ps_units}."
+            f"'{da.name}' should be in 'mb' or 'Pa' (which gets converted to 'mb'), "
+            f"not '{units}'."
         )
 
-    return ps
+    return da

tests/e3sm_diags/driver/utils/test_regrid.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 import xarray as xr
+from xarray.testing import assert_identical
 
 from e3sm_diags.driver.utils.regrid import (
     get_z_axis,
@@ -130,7 +131,7 @@ class TestRegridZAxisToPlevs:
     @pytest.fixture(autouse=True)
     def setup(self):
         self.ds = generate_lev_dataset()
-        self.plevs = [8000, 2000]
+        self.plevs = [800, 200]
 
     def test_raises_error_if_long_name_attr_is_None(self):
         ds = generate_lev_dataset("hybrid")
@@ -166,30 +167,107 @@ def test_raises_error_if_ps_variable_units_attr_is_not_mb_or_pa(self):
         with pytest.raises(ValueError):
             regrid_z_axis_to_plevs(ds, "so", self.plevs)
 
-    def test_converts_pressure_coordinates_to_pressure_levels(self):
-        ds_pres = generate_lev_dataset("pressure")
-
-        expected = xr.DataArray()
-        result = regrid_z_axis_to_plevs(ds_pres, "so", self.plevs)
-
-        assert expected.identical(result)
-        # assert result["lev"].attrs == "mb"
-        assert 0
+    def test_regrids_hybrid_levels_to_pressure_levels(self):
+        ds = generate_lev_dataset("hybrid")
 
-        # ds_iso = generate_lev_dataset("isobaric")
+        # Create the expected dataset using the original dataset. This involves
+        # updating the arrays and attributes of data variables and coordinates.
+        expected = ds.sel(lev=[800, 200]).drop_vars(["ps", "hyam", "hybm"])
+        expected["so"].data[:] = np.nan
+        expected["so"].attrs["units"] = "mb"
+        expected["lev"].attrs = {
+            "axis": "Z",
+            "coordinate": "vertical",
+            "bounds": "lev_bnds",
+        }
+        expected["lev_bnds"] = xr.DataArray(
+            name="lev_bnds",
+            data=np.array([[1100.0, 500.0], [500.0, -100.0]]),
+            dims=["lev", "bnds"],
+            attrs={"xcdat_bounds": "True"},
+        )
 
-        # expected = xr.DataArray()
-        # result = convert_z_axis_to_pressure_levels(ds, ds["lev"], self.plevs)
+        result = regrid_z_axis_to_plevs(ds, "so", self.plevs)
 
-        # assert expected.identical(result)
-        # assert result["lev"].attrs == "mb"
-        assert 0
+        assert_identical(expected, result)
 
-    def test_converts_hybrid_levels_to_pressure_levels(self):
+    def test_regrids_hybrid_levels_to_pressure_levels_with_Pa_units(self):
         ds = generate_lev_dataset("hybrid")
 
-        expected = xr.DataArray()
+        # Create the expected dataset using the original dataset. This involves
+        # updating the arrays and attributes of data variables and coordinates.
+        expected = ds.sel(lev=[800, 200]).drop_vars(["ps", "hyam", "hybm"])
+        expected["so"].data[:] = np.nan
+        expected["so"].attrs["units"] = "mb"
+        expected["lev"].attrs = {
+            "axis": "Z",
+            "coordinate": "vertical",
+            "bounds": "lev_bnds",
+        }
+        expected["lev_bnds"] = xr.DataArray(
+            name="lev_bnds",
+            data=np.array([[1100.0, 500.0], [500.0, -100.0]]),
+            dims=["lev", "bnds"],
+            attrs={"xcdat_bounds": "True"},
+        )
+
+        # Update from Pa to mb.
+        ds_pa = ds.copy()
+        with xr.set_options(keep_attrs=True):
+            ds_pa["ps"] = ds_pa.ps * 100
+        ds_pa.ps.attrs["units"] = "Pa"
+
+        result = regrid_z_axis_to_plevs(ds_pa, "so", self.plevs)
+
+        assert_identical(expected, result)
+
+    @pytest.mark.parametrize("long_name", ("pressure", "isobaric"))
+    def test_regrids_pressure_coordinates_to_pressure_levels(self, long_name):
+        ds = generate_lev_dataset(long_name)
+
+        # Create the expected dataset using the original dataset. This involves
+        # updating the arrays and attributes of data variables and coordinates.
+        expected = ds.sel(lev=[800, 200]).drop_vars("ps")
+        expected["lev"].attrs = {
+            "axis": "Z",
+            "coordinate": "vertical",
+            "bounds": "lev_bnds",
+        }
+        expected["lev_bnds"] = xr.DataArray(
+            name="lev_bnds",
+            data=np.array([[1100.0, 500.0], [500.0, -100.0]]),
+            dims=["lev", "bnds"],
+            attrs={"xcdat_bounds": "True"},
+        )
         result = regrid_z_axis_to_plevs(ds, "so", self.plevs)
 
-        assert expected.identical(result)
-        assert result["lev"].attrs == "mb"
+        assert_identical(expected, result)
+
+    @pytest.mark.parametrize("long_name", ("pressure", "isobaric"))
+    def test_regrids_pressure_coordinates_to_pressure_levels_with_Pa_units(
+        self, long_name
+    ):
+        ds = generate_lev_dataset(long_name)
+
+        expected = ds.sel(lev=[800, 200]).drop_vars("ps")
+        expected["lev"].attrs = {
+            "axis": "Z",
+            "coordinate": "vertical",
+            "bounds": "lev_bnds",
+        }
+        expected["lev_bnds"] = xr.DataArray(
+            name="lev_bnds",
+            data=np.array([[1100.0, 500.0], [500.0, -100.0]]),
+            dims=["lev", "bnds"],
+            attrs={"xcdat_bounds": "True"},
+        )
+
+        # Update from Pa to mb.
+        ds_pa = ds.copy()
+        with xr.set_options(keep_attrs=True):
+            ds_pa["lev"] = ds_pa.lev * 100
+        ds_pa.lev.attrs["units"] = "Pa"
+
+        result = regrid_z_axis_to_plevs(ds_pa, "so", self.plevs)
+
+        assert_identical(expected, result)

tests/e3sm_diags/fixtures.py

@@ -35,9 +35,9 @@
 )
 
 lev = xr.DataArray(
-    data=np.flip(np.arange(2000, 10000, 2000, dtype="float64")),
+    data=[800, 600, 400, 200],
     dims=["lev"],
-    attrs={"units": "Pa", "positive": "down", "axis": "Z"},
+    attrs={"units": "mb", "positive": "down", "axis": "Z"},
 )