[Bug]: Regrid2 breaks in _output_axis_sizes() if non-dimension Z axis coords exists (e.g., height)

[tomvothecoder]

What happened?

The Regrid2 API breaks if there are multiple mappings to the "Z" axis, specifically for coordinate variables such as height which are not considered dimensions.

xcdat/xcdat/regridder/regrid2.py

Line 94 in 104a5d5

axis_variable_name_map = {x: y[0] for x, y in input_data_var.cf.axes.items()}

xcdat/xcdat/regridder/regrid2.py

Lines 128 to 139 in 104a5d5

	output_sizes = {}
	axis_name_map = {y[0]: x for x, y in da.cf.axes.items()}
	for standard_name in da.sizes.keys():
	try:
	axis_name = axis_name_map[standard_name]
	except KeyError:
	raise RuntimeError(
	f"Could not find axis {standard_name!r}, ensure {standard_name!r} "
	"exists and the attributes are correct."
	)

In the code above, we use ds.cf.axes mapping table and get the first matching list element for each axis ({y[0]: x for x, y in da.cf.axes.items()}) . Let's say mapping table for Z looks like {"z": ["height", "lev"]}. Our logic will use element 0 height instead of lev, resulting in RuntimeError: Could not find axis 'lev', ensure 'lev' exists and the attributes are correct.

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

Current workaround: The user needs to drop height coordinate variable or any other non-dimensional Z coordinate beforehand.

Solution: This should no longer be an issue on the latest main commit in PR #533 because we now use xc.get_dim_keys() which maps axes by index (aka dimensions). -- NEED TO TEST

Minimal Complete Verifiable Example (MVCE)

Example

# Notice how height is not a dimension
print(da.dims)
('lat', 'lev', 'lon')

# Notice how height is mapped to "Z" for `ds.cf.axes` and is the first element in the list, which breaks Regrid2
print(da.cf) 
Coordinates:
             CF Axes: * X: ['lon']
                      * Y: ['lat']
                        Z: ['height', 'lev']
                        T: n/a

      CF Coordinates: * longitude: ['lon']
                      * latitude: ['lat']
                        vertical: ['height']
                        time: n/a

       Cell Measures:   area, volume: n/a

      Standard Names:   height: ['height']
                      * latitude: ['lat']
                      * longitude: ['lon']

              Bounds:   n/a

       Grid Mappings:   n/a

Relevant log output

Traceback (most recent call last):
  File "/global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xcdat/regridder/regrid2.py", line 134, in _output_axis_sizes
    axis_name = axis_name_map[standard_name]
KeyError: 'lev'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<string>", line 2, in <module>
  File "/global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xcdat/regridder/accessor.py", line 324, in horizontal
    output_ds = regridder.horizontal(data_var, self._ds)
  File "/global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xcdat/regridder/regrid2.py", line 96, in horizontal
    output_axis_sizes = self._output_axis_sizes(input_data_var)
  File "/global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xcdat/regridder/regrid2.py", line 136, in _output_axis_sizes
    raise RuntimeError(
RuntimeError: Could not find axis 'lev', ensure 'lev' exists and the attributes are correct.

Anything else we need to know?

No response

Environment

v0.6.1 and all other versions that have Regrid2 with this logic

[tomvothecoder]

FYI @jasonb5 and @lee1043.

I think it was a good idea that we replaced .cf.axes with xc.get_dim_keys() in #533. I'll need to test if that fixes this issue on the latest main build.

[tomvothecoder]

This looks like it is fixed in PR #533