Repository: xCDAT/xcdat

Issue:
[FEATURE]: Update supported spatial `axis` arg keys  generic format
### Is your feature request related to a problem?

The valid `axis` arg values of `"lat"` and `"lon"` do not follow the axis naming convention that we adopted with our APIs.

https://github.com/XCDAT/xcdat/blob/79c488ea890febc422fe2e38b85c4e9dc7c72565/xcdat/spatial.py#L47

We implemented a mapping system for axis and coordinates names to their generic axis name.
https://github.com/XCDAT/xcdat/blob/79c488ea890febc422fe2e38b85c4e9dc7c72565/xcdat/axis.py#L14-L25


### Describe the solution you'd like

1. Update valid `axis` arg values from `"lat"` and `"lon"` to `"Y"`/`"y"` and `"X"`/`"x"` 
3. Convert the `axis` arg value to `.upper()` 
4. Map the `axis` arg to the dimension/coordinates in the `xr.Dataset`
5. Update `KeyError` if the CF `axis` attribute is not set for the X and Y axis coordinates
  * Related to https://github.com/XCDAT/xcdat/issues/166#issuecomment-1099382979
    >     We can probably improve the KeyError so that it says something like:
    >     KeyError: "Could not find an X and/or Y axis for spatial averaging. Make sure the data variable 'tas' has X and Y axis  coordinates and the 'axis' attribute is set for both."
  * A reusable convenience function/method might be useful for checking if the `axis` attribute is set for the desired axis in the Dataset. If it isn't, then raise a `KeyError`. `cf_xarray` already throws an error if we try something like `ds.cf["X"]` when the `"X"` axis attr is not set, but the error is ambiguous.

### Describe alternatives you've considered

### Additional context

* `cdutil.averager()`'s `axis` arg accepts lowercase generic axis names (`"x"`, `"y"`, `"t"`, etc.)



Discussion:
tomvothecoder: Thanks for the review!  

> In CDAT, you specify the axes [as a single string](https://github.com/CDAT/genutil/blob/master/Lib/averager.py#L888), e.g., 'xy'. I think it would be worth including that as a method for specifying the axes.

It is more Pythonic (and probably standard coding practice) for a function argument to accept a list of strings (`axis=["x", "y", "z", "t"]`) over a single string (`axis="xyzt"`), and to limit the accepted argument types as much as possible.

In this case, **each string should be treated as an independent instance**, rather than coupled as a single instance that has to be parsed. There is also less logic involved with accepting just a list of strings because type checking isn't necessary, the behavior of the API is more predictable, and it is more explicit and readable (in my opinion).

Posts I found on this subject:

----
https://stackoverflow.com/a/51737249
> **Parameters that can be either a thing or an iterable or things are a code smell**. It’s even worse when the thing is a string, because a string is an iterable, and even a sequence (so your test for isinstance(names, Iterable) would do the wrong thing).

This is what we're currently doing for a single axis (`axis="x"`).

> But if the best design really is a string or a (non-string) iterable of strings, or a string or a tuple of strings, the standard way to do that is type switching. It may look a bit ugly, but it calls attention to the fact that you’re doing exactly what you’re doing, and it does it in the most idiomatic way.
>
> def spam(names):
>     if isinstance(names, str):
>         names = [names]
>     dostuff(names)

I do agree with this person that it is ugly to have to perform this check for every function that has an argument that accepts a string or a list of strings.


https://stackoverflow.com/questions/998938/handle-either-a-list-or-single-integer-as-an-argument/998949#998949
>
> expect that the argument will always be a list - even if its just a list of one element.
> 
> > Remember:
> > 
> >     It is easier to ask for forgiveness than permission.
> > 

---

I suggest that we accept just a list of strings, but allow for upper and lower case axis names.

Valid `axis` arg values:
* `["x", "y", "t"]`
* `["X", "Y", "T"]`
* `["X"]`
* `["x"]`

Invalid `axis` arg values:
* `"xyt"`
* `"XYT"`
* `"X"`
* `"x"`
pochedls: Thanks for the detailed reply.

Patch:
diff --git a/xcdat/spatial.py b/xcdat/spatial.py
index 97d2b55..65c0581 100644
--- a/xcdat/spatial.py
+++ b/xcdat/spatial.py
@@ -17,18 +17,13 @@ import numpy as np
 import xarray as xr
 from dask.array.core import Array
 
-from xcdat.axis import (
-    GENERIC_AXIS_MAP,
-    GenericAxis,
-    _align_lon_bounds_to_360,
-    _get_prime_meridian_index,
-)
+from xcdat.axis import _align_lon_bounds_to_360, _get_prime_meridian_index
 from xcdat.dataset import get_data_var
 
 #: Type alias for a dictionary of axis keys mapped to their bounds.
 AxisWeights = Dict[Hashable, xr.DataArray]
 #: Type alias for supported spatial axis keys.
-SpatialAxis = Literal["lat", "lon"]
+SpatialAxis = Literal["X", "Y"]
 SPATIAL_AXES: Tuple[SpatialAxis, ...] = get_args(SpatialAxis)
 #: Type alias for a tuple of floats/ints for the regional selection bounds.
 RegionAxisBounds = Tuple[float, float]
@@ -44,13 +39,13 @@ class SpatialAccessor:
     def average(
         self,
         data_var: str,
-        axis: Union[List[SpatialAxis], SpatialAxis] = ["lat", "lon"],
+        axis: List[SpatialAxis] = ["X", "Y"],
         weights: Union[Literal["generate"], xr.DataArray] = "generate",
         lat_bounds: Optional[RegionAxisBounds] = None,
         lon_bounds: Optional[RegionAxisBounds] = None,
     ) -> xr.Dataset:
         """
-        Calculate the spatial average for a rectilinear grid over a (optional)
+        Calculates the spatial average for a rectilinear grid over an optionally
         specified regional domain.
 
         Operations include:
@@ -62,14 +57,21 @@ class SpatialAccessor:
         - Adjust weights to conform to the specified regional boundary.
         - Compute spatial weighted average.
 
+        This method requires that the dataset's coordinates have the 'axis'
+        attribute set to the keys in ``axis``. For example, the latitude
+        coordinates should have its 'axis' attribute set to 'Y' (which is also
+        CF-compliant). This 'axis' attribute is used to retrieve the related
+        coordinates via `cf_xarray`. Refer to this method's examples for more
+        information.
+
         Parameters
         ----------
         data_var: str
             The name of the data variable inside the dataset to spatially
             average.
-        axis : Union[List[SpatialAxis], SpatialAxis]
-            List of axis dimensions or single axis dimension to average over.
-            For example, ["lat", "lon"]  or "lat", by default ["lat", "lon"].
+        axis : List[SpatialAxis]
+            List of axis dimensions to average over, by default ["X", "Y"].
+            Valid axis keys include "X" and "Y".
         weights : Union[Literal["generate"], xr.DataArray], optional
             If "generate", then weights are generated. Otherwise, pass a
             DataArray containing the regional weights used for weighted
@@ -104,30 +106,36 @@ class SpatialAccessor:
 
         >>> import xcdat
 
-        Open a dataset and limit to a single variable:
+        Check the 'axis' attribute is set on the required coordinates:
+
+        >>> ds.lat.attrs["axis"]
+        >>> Y
+        >>>
+        >>> ds.lon.attrs["axis"]
+        >>> X
+
+        Set the 'axis' attribute for the required coordinates if it isn't:
 
-        >>> ds = xcdat.open_dataset("path/to/file.nc", var="tas")
+        >>> ds.lat.attrs["axis"] = "Y"
+        >>> ds.lon.attrs["axis"] = "X"
 
         Call spatial averaging method:
 
-        >>> # First option
         >>> ds.spatial.average(...)
-        >>> # Second option
-        >>> ds.xcdat.average(...)
 
         Get global average time series:
 
-        >>> ts_global = ds.spatial.average("tas", axis=["lat", "lon"])["tas"]
+        >>> ts_global = ds.spatial.average("tas", axis=["X", "Y"])["tas"]
 
         Get time series in Nino 3.4 domain:
 
-        >>> ts_n34 = ds.spatial.average("ts", axis=["lat", "lon"],
+        >>> ts_n34 = ds.spatial.average("ts", axis=["X", "Y"],
         >>>     lat_bounds=(-5, 5),
         >>>     lon_bounds=(-170, -120))["ts"]
 
         Get zonal mean time series:
 
-        >>> ts_zonal = ds.spatial.average("tas", axis=['lon'])["tas"]
+        >>> ts_zonal = ds.spatial.average("tas", axis=["X"])["tas"]
 
         Using custom weights for averaging:
 
@@ -138,18 +146,18 @@ class SpatialAccessor:
         >>>     dims=["lat", "lon"],
         >>> )
         >>>
-        >>> ts_global = ds.spatial.average("tas", axis=["lat","lon"],
+        >>> ts_global = ds.spatial.average("tas", axis=["X", "Y"],
         >>>     weights=weights)["tas"]
         """
         dataset = self._dataset.copy()
         dv = get_data_var(dataset, data_var)
-        axis = self._validate_axis(dv, axis)
+        self._validate_axis_arg(axis)
 
         if isinstance(weights, str) and weights == "generate":
             if lat_bounds is not None:
-                self._validate_region_bounds("lat", lat_bounds)
+                self._validate_region_bounds("Y", lat_bounds)
             if lon_bounds is not None:
-                self._validate_region_bounds("lon", lon_bounds)
+                self._validate_region_bounds("X", lon_bounds)
             dv_weights = self._get_weights(axis, lat_bounds, lon_bounds)
         elif isinstance(weights, xr.DataArray):
             dv_weights = weights
@@ -158,51 +166,39 @@ class SpatialAccessor:
         dataset[dv.name] = self._averager(dv, axis, dv_weights)
         return dataset
 
-    def _validate_axis(
-        self, data_var: xr.DataArray, axis: Union[List[SpatialAxis], SpatialAxis]
-    ) -> List[SpatialAxis]:
-        """Validates if ``axis`` arg is supported and exists in the data var.
+    def _validate_axis_arg(self, axis: List[SpatialAxis]):
+        """
+        Validates that the ``axis`` dimension(s) exists in the dataset.
 
         Parameters
         ----------
-        data_var : xr.DataArray
-            The data variable.
-        axis : Union[List[SpatialAxis], SpatialAxis]
-            List of axis dimensions or single axis dimension to average over.
-
-        Returns
-        -------
-        List[SpatialAxis]
-            List of axis dimensions or single axis dimension to average over.
+        axis : List[SpatialAxis]
+            List of axis dimensions to average over.
 
         Raises
         ------
         ValueError
-            If any key in ``axis`` is not supported for spatial averaging.
+            If a key in ``axis`` is not a supported value.
         KeyError
-            If any key in ``axis`` does not exist in the ``data_var``.
+            If the dataset does not have coordinates for the ``axis`` dimension,
+            or the `axis` attribute is not set for those coordinates.
         """
-        if isinstance(axis, str):
-            axis = [axis]
-
         for key in axis:
             if key not in SPATIAL_AXES:
                 raise ValueError(
-                    "Incorrect `axis` argument. Supported axes include: "
+                    "Incorrect `axis` argument value. Supported values include: "
                     f"{', '.join(SPATIAL_AXES)}."
                 )
 
-            generic_axis_key = GENERIC_AXIS_MAP[key]
             try:
-                data_var.cf.axes[generic_axis_key]
+                self._dataset.cf.axes[key]
             except KeyError:
                 raise KeyError(
-                    f"The data variable '{data_var.name}' is missing the '{axis}' "
-                    "dimension, which is required for spatial averaging."
+                    f"A '{key}' axis dimension was not found in the dataset. Make sure "
+                    f"the dataset has '{key}' axis coordinates and the coordinates' "
+                    f"'axis' attribute is set to '{key}'."
                 )
 
-        return axis
-
     def _validate_domain_bounds(self, domain_bounds: xr.DataArray):
         """Validates the ``domain_bounds`` arg based on a set of criteria.
 
@@ -244,7 +240,7 @@ class SpatialAccessor:
         TypeError
             If the ``bounds`` upper bound is not a float or integer.
         ValueError
-            If the ``axis`` is "lat" and the ``bounds`` lower value is larger
+            If the ``axis`` is "Y" and the ``bounds`` lower value is larger
             than the upper value.
         """
         if not isinstance(bounds, tuple):
@@ -269,12 +265,12 @@ class SpatialAccessor:
                 f"The regional {axis} upper bound is not a float or an integer."
             )
 
-        # For latitude, require that the upper bound be larger than the lower
-        # bound. Note that this does not apply to longitude (since it is
-        # a circular axis).
-        if axis == "lat" and lower >= upper:
+        # For the "Y" axis (latitude), require that the upper bound be larger
+        # than the lower bound. Note that this does not apply to the "X" axis
+        # (longitude) since it is circular.
+        if axis == "Y" and lower >= upper:
             raise ValueError(
-                f"The regional {axis} lower bound is greater than the upper. "
+                "The regional latitude lower bound is greater than the upper. "
                 "Pass a tuple with the format (lower, upper)."
             )
 
@@ -299,9 +295,8 @@ class SpatialAccessor:
 
         Parameters
         ----------
-        axis : Union[List[SpatialAxis], SpatialAxis]
-            List of axis dimensions or single axis dimension to average over.
-            For example, ["lat", "lon"]  or "lat".
+        axis : List[SpatialAxis]
+            List of axis dimensions to average over.
         lat_bounds : Optional[RegionAxisBounds]
             Tuple of latitude boundaries for regional selection.
         lon_bounds : Optional[RegionAxisBounds]
@@ -326,17 +321,18 @@ class SpatialAccessor:
             {"domain": xr.DataArray, "region": Optional[RegionAxisBounds]},
         )
         axis_bounds: Dict[SpatialAxis, Bounds] = {
-            "lat": {
-                "domain": self._dataset.bounds.get_bounds("lat").copy(),
-                "region": lat_bounds,
-            },
-            "lon": {
+            "X": {
                 "domain": self._dataset.bounds.get_bounds("lon").copy(),
                 "region": lon_bounds,
             },
+            "Y": {
+                "domain": self._dataset.bounds.get_bounds("lat").copy(),
+                "region": lat_bounds,
+            },
         }
 
         axis_weights: AxisWeights = {}
+
         for key in axis:
             d_bounds = axis_bounds[key]["domain"]
             self._validate_domain_bounds(d_bounds)
@@ -347,9 +343,9 @@ class SpatialAccessor:
             if r_bounds is not None:
                 r_bounds = np.array(r_bounds, dtype="float")
 
-            if key == "lon":
+            if key == "X":
                 weights = self._get_longitude_weights(d_bounds, r_bounds)
-            elif key == "lat":
+            elif key == "Y":
                 weights = self._get_latitude_weights(d_bounds, r_bounds)
 
             weights.attrs = d_bounds.attrs
@@ -357,6 +353,7 @@ class SpatialAccessor:
             axis_weights[key] = weights
 
         weights = self._combine_weights(axis_weights)
+
         return weights
 
     def _get_longitude_weights(
@@ -386,9 +383,9 @@ class SpatialAccessor:
         Parameters
         ----------
         domain_bounds : xr.DataArray
-            The array of bounds for the latitude domain.
+            The array of bounds for the longitude domain.
         region_bounds : Optional[np.ndarray]
-            The array of bounds for latitude regional selection.
+            The array of bounds for longitude regional selection.
 
         Returns
         -------
@@ -655,14 +652,22 @@ class SpatialAccessor:
             If the axis dimension sizes between ``weights`` and ``data_var``
             are misaligned.
         """
-        # Check that the supplied weights include lat and lon dimensions.
-        lat_key = data_var.cf.axes["Y"][0]
-        lon_key = data_var.cf.axes["X"][0]
-
-        if "lat" in axis and lat_key not in weights.dims:
-            raise KeyError(f"Check weights DataArray includes {lat_key} dimension.")
-        if "lon" in axis and lon_key not in weights.dims:
-            raise KeyError(f"Check weights DataArray includes {lon_key} dimension.")
+        # Check that the supplied weights include x and y dimensions.
+        x_key = data_var.cf.axes["X"][0]
+        y_key = data_var.cf.axes["Y"][0]
+
+        if "X" in axis and x_key not in weights.dims:
+            raise KeyError(
+                "The weights DataArray either does not include an X axis, "
+                "or the X axis coordinates does not have the 'axis' attribute "
+                "set to 'X'."
+            )
+        if "Y" in axis and y_key not in weights.dims:
+            raise KeyError(
+                "The weights DataArray either does not include an Y axis, "
+                "or the Y axis coordinates does not have the 'axis' attribute "
+                "set to 'Y'."
+            )
 
         # Check the weight dim sizes equal data var dim sizes.
         dim_sizes = {key: data_var.sizes[key] for key in weights.sizes.keys()}
@@ -692,8 +697,7 @@ class SpatialAccessor:
         data_var : xr.DataArray
             Data variable inside a Dataset.
         axis : List[SpatialAxis]
-            List of axis dimensions or single axis dimension to average over.
-            For example, ["lat", "lon"]  or "lat".
+            List of axis dimensions to average over.
         weights : xr.DataArray
             A DataArray containing the region area weights for averaging.
             ``weights`` must include the same spatial axis dimensions and have
@@ -710,34 +714,8 @@ class SpatialAccessor:
         Missing values are replaced with 0 using ``weights.fillna(0)``.
         """
         weights = weights.fillna(0)
-        with xr.set_options(keep_attrs=True):
-            weighted_mean = data_var.cf.weighted(weights).mean(
-                self._get_generic_axis_keys(axis)
-            )
-            return weighted_mean
 
-    def _get_generic_axis_keys(self, axis: List[SpatialAxis]) -> List[GenericAxis]:
-        """Converts supported axis keys to their generic CF representations.
-
-        Since xCDAT's spatial averaging accepts the CF short version of axes
-        keys, attempting to index a Dataset/DataArray on the short key through
-        cf_xarray might fail for cases where the long key is used instead (e.g.,
-        "latitude" instead of "lat"). This method handles this edge case by
-        converting the list of axis keys to their generic representations (e.g.,
-        "Y" instead of "lat") for indexing operations.
-
-        Parameters
-        ----------
-        axis_keys : List[SpatialAxis]
-            List of axis dimension(s) to average over.
-
-        Returns
-        -------
-        List[GenericAxis]
-            List of axis dimension(s) to average over.
-        """
-        generic_axis_keys = []
-        for key in axis:
-            generic_axis_keys.append(GENERIC_AXIS_MAP[key])
+        with xr.set_options(keep_attrs=True):
+            weighted_mean = data_var.cf.weighted(weights).mean(axis)
 
-        return generic_axis_keys
+        return weighted_mean
diff --git a/xcdat/temporal.py b/xcdat/temporal.py
index 702b15e..ad13d93 100644
--- a/xcdat/temporal.py
+++ b/xcdat/temporal.py
@@ -106,8 +106,9 @@ class TemporalAccessor:
             dataset.cf["T"]
         except KeyError:
             raise KeyError(
-                "This dataset does not have a time dimension, which is required for "
-                "using the methods in the TemporalAccessor class."
+                "A 'T' axis dimension was not found in the dataset. Make sure the "
+                "dataset has time axis coordinates and its 'axis' attribute is set to "
+                "'T'."
             )
 
         self._dataset: xr.Dataset = dataset
@@ -210,6 +211,15 @@ class TemporalAccessor:
 
         >>> import xcdat
 
+        Check the 'axis' attribute is set on the time coordinates:
+
+        >>> ds.time.attrs["axis"]
+        >>> T
+
+        Set the 'axis' attribute for the time coordinates if it isn't:
+
+        >>> ds.time.attrs["axis"] = "T"
+
         Call ``average()`` method:
 
         >>> ds.temporal.average(...)


Test patch:
diff --git a/tests/test_spatial.py b/tests/test_spatial.py
index dd0f9fb..49e01b8 100644
--- a/tests/test_spatial.py
+++ b/tests/test_spatial.py
@@ -41,7 +41,7 @@ class TestAverage:
         with pytest.raises(KeyError):
             self.ds.spatial.average(
                 "not_a_data_var",
-                axis=["lat", "incorrect_axis"],
+                axis=["Y", "incorrect_axis"],
             )
 
     def test_spatial_average_for_lat_and_lon_region_using_custom_weights(self):
@@ -53,7 +53,7 @@ class TestAverage:
             dims=["lat", "lon"],
         )
         result = ds.spatial.average(
-            axis=["lat", "lon"],
+            axis=["X", "Y"],
             lat_bounds=(-5.0, 5),
             lon_bounds=(-170, -120.1),
             weights=weights,
@@ -72,7 +72,7 @@ class TestAverage:
     def test_spatial_average_for_lat_and_lon_region(self):
         ds = self.ds.copy()
         result = ds.spatial.average(
-            "ts", axis=["lat", "lon"], lat_bounds=(-5.0, 5), lon_bounds=(-170, -120.1)
+            "ts", axis=["X", "Y"], lat_bounds=(-5.0, 5), lon_bounds=(-170, -120.1)
         )
 
         expected = self.ds.copy()
@@ -89,7 +89,7 @@ class TestAverage:
 
         # Specifying axis as a str instead of list of str.
         result = ds.spatial.average(
-            "ts", axis="lat", lat_bounds=(-5.0, 5), lon_bounds=(-170, -120.1)
+            "ts", axis=["Y"], lat_bounds=(-5.0, 5), lon_bounds=(-170, -120.1)
         )
 
         expected = self.ds.copy()
@@ -109,7 +109,7 @@ class TestAverage:
 
         # Specifying axis as a str instead of list of str.
         result = ds.spatial.average(
-            "ts", axis="lat", lat_bounds=(-5.0, 5), lon_bounds=(-170, -120.1)
+            "ts", axis=["Y"], lat_bounds=(-5.0, 5), lon_bounds=(-170, -120.1)
         )
 
         expected = self.ds.copy()
@@ -124,32 +124,26 @@ class TestAverage:
         assert result.identical(expected)
 
 
-class TestValidateAxis:
+class TestValidateAxisArg:
     @pytest.fixture(autouse=True)
     def setup(self):
         self.ds = generate_dataset(cf_compliant=True, has_bounds=True)
 
     def test_raises_error_if_axis_list_contains_unsupported_axis(self):
         with pytest.raises(ValueError):
-            self.ds.spatial._validate_axis(self.ds.ts, axis=["lat", "incorrect_axis"])
+            self.ds.spatial._validate_axis_arg(axis=["Y", "incorrect_axis"])
 
     def test_raises_error_if_lat_axis_does_not_exist(self):
         ds = self.ds.copy()
-        ds["ts"] = xr.DataArray(data=None, coords={"lon": ds.lon}, dims=["lon"])
+        ds.lat.attrs["axis"] = None
         with pytest.raises(KeyError):
-            ds.spatial._validate_axis(ds.ts, axis=["lat", "lon"])
+            ds.spatial._validate_axis_arg(axis=["X", "Y"])
 
     def test_raises_error_if_lon_axis_does_not_exist(self):
         ds = self.ds.copy()
-        ds["ts"] = xr.DataArray(data=None, coords={"lat": ds.lat}, dims=["lat"])
+        ds.lon.attrs["axis"] = None
         with pytest.raises(KeyError):
-            ds.spatial._validate_axis(ds.ts, axis=["lat", "lon"])
-
-    def test_returns_list_of_str_if_axis_is_a_single_supported_str_input(self):
-        result = self.ds.spatial._validate_axis(self.ds.ts, axis="lat")
-        expected = ["lat"]
-
-        assert result == expected
+            ds.spatial._validate_axis_arg(axis=["X", "Y"])
 
 
 class TestValidateRegionBounds:
@@ -178,18 +172,18 @@ class TestValidateRegionBounds:
 
     def test_raises_error_if_lower_bound_is_not_a_float_or_int(self):
         with pytest.raises(TypeError):
-            self.ds.spatial._validate_region_bounds("lat", ("invalid", 1))
+            self.ds.spatial._validate_region_bounds("Y", ("invalid", 1))
 
     def test_raises_error_if_upper_bound_is_not_a_float_or_int(self):
         with pytest.raises(TypeError):
-            self.ds.spatial._validate_region_bounds("lon", (1, "invalid"))
+            self.ds.spatial._validate_region_bounds("X", (1, "invalid"))
 
     def test_raises_error_if_lower_lat_bound_is_bigger_than_upper(self):
         with pytest.raises(ValueError):
-            self.ds.spatial._validate_region_bounds("lat", (2, 1))
+            self.ds.spatial._validate_region_bounds("Y", (2, 1))
 
     def test_does_not_raise_error_if_lon_lower_bound_is_larger_than_upper(self):
-        self.ds.spatial._validate_region_bounds("lon", (2, 1))
+        self.ds.spatial._validate_region_bounds("X", (2, 1))
 
 
 class TestValidateWeights:
@@ -209,7 +203,7 @@ class TestValidateWeights:
             coords={"lat": self.ds.lat, "lon": self.ds.lon},
             dims=["lat", "lon"],
         )
-        self.ds.spatial._validate_weights(self.ds["ts"], axis="lat", weights=weights)
+        self.ds.spatial._validate_weights(self.ds["ts"], axis=["Y"], weights=weights)
 
     def test_error_is_raised_when_lat_axis_is_specified_but_lat_is_not_in_weights_dims(
         self,
@@ -219,7 +213,7 @@ class TestValidateWeights:
         )
         with pytest.raises(KeyError):
             self.ds.spatial._validate_weights(
-                self.ds["ts"], axis=["lon", "lat"], weights=weights
+                self.ds["ts"], axis=["X", "Y"], weights=weights
             )
 
     def test_error_is_raised_when_lon_axis_is_specified_but_lon_is_not_in_weights_dims(
@@ -230,7 +224,7 @@ class TestValidateWeights:
         )
         with pytest.raises(KeyError):
             self.ds.spatial._validate_weights(
-                self.ds["ts"], axis=["lon", "lat"], weights=weights
+                self.ds["ts"], axis=["X", "Y"], weights=weights
             )
 
     def test_error_is_raised_when_weights_lat_and_lon_dims_dont_align_with_data_var_dims(
@@ -247,7 +241,7 @@ class TestValidateWeights:
 
         with pytest.raises(ValueError):
             self.ds.spatial._validate_weights(
-                self.ds["ts"], axis=["lat", "lon"], weights=weights
+                self.ds["ts"], axis=["X", "Y"], weights=weights
             )
 
 
@@ -404,7 +398,7 @@ class TestGetWeights:
 
     def test_weights_for_region_in_lat_and_lon_domains(self):
         result = self.ds.spatial._get_weights(
-            axis=["lat", "lon"], lat_bounds=(-5, 5), lon_bounds=(-170, -120)
+            axis=["Y", "X"], lat_bounds=(-5, 5), lon_bounds=(-170, -120)
         )
         expected = xr.DataArray(
             data=np.array(
@@ -423,7 +417,7 @@ class TestGetWeights:
 
     def test_area_weights_for_region_in_lat_domain(self):
         result = self.ds.spatial._get_weights(
-            axis=["lat", "lon"], lat_bounds=(-5, 5), lon_bounds=None
+            axis=["Y", "X"], lat_bounds=(-5, 5), lon_bounds=None
         )
         expected = xr.DataArray(
             data=np.array(
@@ -454,7 +448,7 @@ class TestGetWeights:
             dims=["lat", "lon"],
         )
         result = self.ds.spatial._get_weights(
-            axis=["lat", "lon"], lat_bounds=None, lon_bounds=(-170, -120)
+            axis=["Y", "X"], lat_bounds=None, lon_bounds=(-170, -120)
         )
 
         xr.testing.assert_allclose(result, expected)
@@ -828,7 +822,7 @@ class TestAverager:
             dims=["lat", "lon"],
         )
 
-        result = ds.spatial._averager(ds.ts, axis=["lat", "lon"], weights=weights)
+        result = ds.spatial._averager(ds.ts, axis=["X", "Y"], weights=weights)
         expected = xr.DataArray(
             name="ts", data=np.ones(15), coords={"time": ds.time}, dims=["time"]
         )
@@ -843,7 +837,7 @@ class TestAverager:
             dims=["lat"],
         )
 
-        result = self.ds.spatial._averager(self.ds.ts, axis=["lat"], weights=weights)
+        result = self.ds.spatial._averager(self.ds.ts, axis=["Y"], weights=weights)
         expected = xr.DataArray(
             name="ts",
             data=np.ones((15, 4)),
@@ -861,7 +855,7 @@ class TestAverager:
             dims=["lon"],
         )
 
-        result = self.ds.spatial._averager(self.ds.ts, axis=["lon"], weights=weights)
+        result = self.ds.spatial._averager(self.ds.ts, axis=["X"], weights=weights)
         expected = xr.DataArray(
             name="ts",
             data=np.ones((15, 4)),
@@ -878,22 +872,9 @@ class TestAverager:
             dims=["lat", "lon"],
         )
 
-        result = self.ds.spatial._averager(
-            self.ds.ts, axis=["lat", "lon"], weights=weights
-        )
+        result = self.ds.spatial._averager(self.ds.ts, axis=["X", "Y"], weights=weights)
         expected = xr.DataArray(
             name="ts", data=np.ones(15), coords={"time": self.ds.time}, dims=["time"]
         )
 
         assert result.identical(expected)
-
-
-class TestGetGenericAxisKeys:
-    @pytest.fixture(autouse=True)
-    def setup(self):
-        self.ds = generate_dataset(cf_compliant=True, has_bounds=True)
-
-    def test_generic_keys(self):
-        result = self.ds.spatial._get_generic_axis_keys(["lat", "lon"])
-        expected = ["Y", "X"]
-        assert result == expected