[Bug]: ds.regridder.vertical() breaks with KeyError: 0 if grid_positions=None and input_grid has multiple Z axis

[tomvothecoder]

What happened?

The code below assumes the input grid has a single set of Z axis coordinates (e.g., just "lev").

xcdat/xcdat/regridder/xgcm.py

Lines 242 to 245 in 870b334

	try:
	coord_z = get_dim_coords(self._input_grid, "Z")
	except KeyError:
	raise RuntimeError("Could not determine 'Z' coordinate in input dataset")

If a dataset has multiple Z axis (e.g., "lev" and "ilev"), then the indexing will break on lines 253-255. This is because coord_z will be an xr.Dataset containing multiple Z axis. xr.Dataset cannot be indexed by position (e.g., coord_z[0]).

xcdat/xcdat/regridder/xgcm.py

Lines 252 to 255 in 870b334

	# handle simple point positions based on point and bounds
	if (coord_z[0] > bounds_z[0][0] and coord_z[0] < bounds_z[0][1]) or (
	coord_z[0] < bounds_z[0][0] and coord_z[0] > bounds_z[0][1]
	):

What did you expect to happen? Are there are possible answers you came across?

The fix:

        try:
            coord_z = get_dim_coords(self._input_grid, data_var, "Z")

coord_z should now be an xr.DataArray of the single Z axis for data_var.

Minimal Complete Verifiable Example (MVCE)

import numpy as np
import xarray as xr
import xcdat as xc

# Define a dataset with multiple Z axis systems.
ds = xr.Dataset(
    coords={
        "lev": xr.DataArray(
            name="lev",
            data=[0, 1, 2],
            dims=["lev"],
            attrs={"units": "Pa", "positive": "down", "axis": "Z"},
        ),
        "ilev": xr.DataArray(
            name="ilev",
            data=[0, 1, 2],
            dims=["ilev"],
            attrs={"units": "Pa", "positive": "down", "axis": "Z"},
        ),
    },
    data_vars={
        "so1": xr.DataArray(name="so1", data=[1, 2, 3], dims=["ilev"]),
        "so2": xr.DataArray(name="so2", data=[1, 2, 3], dims=["ilev"]),
    },
)
ds = ds.bounds.add_bounds("Z")

# Create the output grid to regrid to and add bounds
plevs = [0, 1]
output_grid = xc.create_grid(z=xc.create_axis("lev", plevs))
output_grid = output_grid.bounds.add_bounds("Z")

# Call the vertical regridder -- Raises KeyError 0
ds.regridder.vertical("so1", output_grid)

Relevant log output

KeyError                                  Traceback (most recent call last)
Cell In[102], line 1
----> 1 ds.regridder.vertical("so1", output_grid)  # Raises RunTimeError

File ~/mambaforge/envs/diags_geocat/lib/python3.10/site-packages/xcdat/regridder/accessor.py:389, in RegridderAccessor.vertical(self, data_var, output_grid, tool, **options)
    384     raise ValueError(
    385         f"Tool {e!s} does not exist, valid choices "
    386         f"{list(VERTICAL_REGRID_TOOLS)}"
    387     )
    388 regridder = regrid_tool(self._ds, output_grid, **options)
--> 389 output_ds = regridder.vertical(data_var, self._ds)
    391 return output_ds

File ~/mambaforge/envs/diags_geocat/lib/python3.10/site-packages/xcdat/regridder/xgcm.py:180, in XGCMRegridder.vertical(self, data_var, ds)
    177     raise RuntimeError("Could not determine 'Z' coordinate in output dataset")
    179 if self._grid_positions is None:
--> 180     grid_coords = self._get_grid_positions()
    181 else:
    182     # correctly format argument
    183     grid_coords = {"Z": self._grid_positions}

File ~/mambaforge/envs/diags_geocat/lib/python3.10/site-packages/xcdat/regridder/xgcm.py:253, in XGCMRegridder._get_grid_positions(self)
    250     raise RuntimeError("Could not determine 'Z' bounds in input dataset")
    252 # handle simple point positions based on point and bounds
--> 253 if (coord_z[0] > bounds_z[0][0] and coord_z[0] < bounds_z[0][1]) or (
    254     coord_z[0] < bounds_z[0][0] and coord_z[0] > bounds_z[0][1]
    255 ):
    256     grid_positions = {"center": coord_z.name}
    257 elif coord_z[0] == bounds_z[0][0]:

File ~/mambaforge/envs/diags_geocat/lib/python3.10/site-packages/xarray/core/dataset.py:1484, in Dataset.__getitem__(self, key)
   1482     return self.isel(**key)
   1483 if utils.hashable(key):
-> 1484     return self._construct_dataarray(key)
   1485 if utils.iterable_of_hashable(key):
   1486     return self._copy_listed(key)

File ~/mambaforge/envs/diags_geocat/lib/python3.10/site-packages/xarray/core/dataset.py:1395, in Dataset._construct_dataarray(self, name)
   1393     variable = self._variables[name]
   1394 except KeyError:
-> 1395     _, name, variable = _get_virtual_variable(self._variables, name, self.dims)
   1397 needed_dims = set(variable.dims)
   1399 coords: dict[Hashable, Variable] = {}

File ~/mambaforge/envs/diags_geocat/lib/python3.10/site-packages/xarray/core/dataset.py:192, in _get_virtual_variable(variables, key, dim_sizes)
    189     return key, key, variable
    191 if not isinstance(key, str):
--> 192     raise KeyError(key)
    194 split_key = key.split(".", 1)
    195 if len(split_key) != 2:

KeyError: 0

Anything else we need to know?

No response

Environment

xcdat=0.6.0rc1 (and main branch)