Repository: xCDAT/xcdat

Issue:
[Bug]: averaging a lat-plev grid fails
### What happened?

I was looking at a `[time, plev, lat]` data and wanted to create a global average time series (averaging over `lat`). Spatial averaging fails because the dataset is missing the `X` coordinate.

### What did you expect to happen?

The spatial average should be computed – the `X` dimension was not specified (so it should not be needed). 

### Minimal Complete Verifiable Example

```python
import xcdat as xc

fn = '/p/css03/esgf_publish/CMIP6/CMIP/CNRM-CERFACS/CNRM-CM6-1/historical/r1i1p1f2/AERmonZ/ta/grz/v20180917/ta_AERmonZ_CNRM-CM6-1_historical_r1i1p1f2_grz_185001-201412.nc'
ds = xc.open_dataset(fn)
ds.spatial.average('ta', axis=['Y'])
```


### Relevant log output

```python
KeyError                                  Traceback (most recent call last)                                                                                                                        
Input In [4], in <cell line: 1>()                                                                                                                                                                  
----> 1 ds.spatial.average('ta', axis=['Y'])                                                                                                                                                       
                                                                                                                                                                                                   
File ~/code/xcdat/xcdat/spatial.py:190, in SpatialAccessor.average(self, data_var, axis, weights, keep_weights, lat_bounds, lon_bounds)                                                            
    188     if lon_bounds is not None:                                                                                                                                                             
    189         self._validate_region_bounds("X", lon_bounds)                                                                                                                                      
--> 190     self._weights = self.get_weights(axis, lat_bounds, lon_bounds)                                                                                                                         
    191 elif isinstance(weights, xr.DataArray):                                                                                                                                                    
    192     self._weights = weights                                                                                                                                                                
                                                                                                                                                                                                   
File ~/code/xcdat/xcdat/spatial.py:252, in SpatialAccessor.get_weights(self, axis, lat_bounds, lon_bounds)                                                                                         
    208 """                                                                                                                                                                                        
    209 Get area weights for specified axis keys and an optional target domain.                                                                                                                    
    210                                                                                                                                                                                            
   (...)                                                                                                                                                                                           
    244 and pressure).
    245 """
    246 Bounds = TypedDict(
    247     "Bounds",
    248     {"domain": xr.DataArray, "region": Optional[RegionAxisBounds]},
    249 )
    250 axis_bounds: Dict[SpatialAxis, Bounds] = {
    251     "X": {
--> 252         "domain": self._dataset.bounds.get_bounds("X").copy(),
    253         "region": lon_bounds,
    254     },
    255     "Y": {
    256         "domain": self._dataset.bounds.get_bounds("Y").copy(),
    257         "region": lat_bounds,
    258     },
    259 }
    261 axis_weights: AxisWeights = {}
    263 for key in axis:

File ~/code/xcdat/xcdat/bounds.py:195, in BoundsAccessor.get_bounds(self, axis)
    171 def get_bounds(self, axis: CFAxisName) -> xr.DataArray:
    172     """Get bounds for axis coordinates if both exist.
    173 
    174     Parameters
   (...)
    193         If the coordinate bounds were not found for the ``axis``.
    194     """
--> 195     self._validate_axis_arg(axis)
    196     coord_var = get_axis_coord(self._dataset, axis)
    198     try:

File ~/code/xcdat/xcdat/bounds.py:385, in BoundsAccessor._validate_axis_arg(self, axis)
    380     keys = ", ".join(f"'{key}'" for key in cf_axis_names)
    381     raise ValueError(
    382         f"Incorrect 'axis' argument value. Supported values include {keys}."
    383     )
--> 385 get_axis_coord(self._dataset, axis)

File ~/code/xcdat/xcdat/axis.py:81, in get_axis_coord(obj, axis)
     78         pass
     80 if coord_var is None:
---> 81     raise KeyError(
     82         f"A coordinate variable for the {axis} axis was not found. Make sure "
     83         "the coordinate variable exists and either the (1) 'axis' attr or (2) "
     84         "'standard_name' attr is set, or (3) the dimension name follows the "
     85         "short-hand convention (e.g.,'lat')."
     86     )
     87 return coord_var

KeyError: "A coordinate variable for the X axis was not found. Make sure the coordinate variable exists and either the (1) 'axis' attr or (2) 'standard_name' attr is set, or (3) the dimension name follows the short-hand convention (e.g.,'lat')."
```


### Anything else we need to know?

We probably just need to modify the portion of `get_weights` that [gets the bounds](https://github.com/xCDAT/xcdat/blob/main/xcdat/spatial.py#L250) for `X` and `Y` to only get those bounds if they are needed. Or maybe this dictionary isn't needed and the [logic can be put in this loop](https://github.com/xCDAT/xcdat/blob/main/xcdat/spatial.py#L263) over the user-specified axes.

### Environment

main branch

Discussion:
codecov-commenter: # [Codecov](https://codecov.io/gh/xCDAT/xcdat/pull/340?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT) Report
> Merging [#340](https://codecov.io/gh/xCDAT/xcdat/pull/340?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT) (55f38a2) into [main](https://codecov.io/gh/xCDAT/xcdat/commit/f6bc40944bf3766a335e7db3e064cda863b0a81c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT) (f6bc409) will **not change** coverage.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff            @@
##              main      #340   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           14        14           
  Lines         1196      1204    +8     
=========================================
+ Hits          1196      1204    +8     
```


| [Impacted Files](https://codecov.io/gh/xCDAT/xcdat/pull/340?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT) | Coverage Δ | |
|---|---|---|
| [xcdat/spatial.py](https://codecov.io/gh/xCDAT/xcdat/pull/340/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT#diff-eGNkYXQvc3BhdGlhbC5weQ==) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT)

pochedls: I think this is in good shape and can be reviewed. I do plan on re-running it over a bunch of cases before merging. 
pochedls: I re-did comparisons with CDAT - we still agree with CDAT in all cases where both can load the dataset (loop over 115 model files and 21 domains). 

Patch:
diff --git a/xcdat/spatial.py b/xcdat/spatial.py
index cc4c27a..b1f8829 100644
--- a/xcdat/spatial.py
+++ b/xcdat/spatial.py
@@ -1,6 +1,7 @@
 """Module containing geospatial averaging functions."""
 from functools import reduce
 from typing import (
+    Callable,
     Dict,
     Hashable,
     List,
@@ -244,37 +245,36 @@ class SpatialAccessor:
         and pressure).
         """
         Bounds = TypedDict(
-            "Bounds",
-            {"domain": xr.DataArray, "region": Optional[RegionAxisBounds]},
+            "Bounds", {"weights_method": Callable, "region": Optional[np.ndarray]}
         )
+
         axis_bounds: Dict[SpatialAxis, Bounds] = {
             "X": {
-                "domain": self._dataset.bounds.get_bounds("X").copy(),
-                "region": lon_bounds,
+                "weights_method": self._get_longitude_weights,
+                "region": np.array(lon_bounds, dtype="float")
+                if lon_bounds is not None
+                else None,
             },
             "Y": {
-                "domain": self._dataset.bounds.get_bounds("Y").copy(),
-                "region": lat_bounds,
+                "weights_method": self._get_latitude_weights,
+                "region": np.array(lat_bounds, dtype="float")
+                if lat_bounds is not None
+                else None,
             },
         }
 
         axis_weights: AxisWeights = {}
-
         for key in axis:
-            d_bounds = axis_bounds[key]["domain"]
-            self._validate_domain_bounds(d_bounds)
+            d_bounds = self._dataset.bounds.get_bounds(key).copy()
+            # The logic for generating longitude weights depends on the
+            # bounds being ordered such that d_bounds[:, 0] < d_bounds[:, 1].
+            # They are re-ordered (if need be) for the purpose of creating
+            # weights.
+            d_bounds = self._force_domain_order_low_to_high(d_bounds)
 
-            r_bounds: Union[Optional[RegionAxisBounds], np.ndarray] = axis_bounds[key][
-                "region"
-            ]
-            if r_bounds is not None:
-                r_bounds = np.array(r_bounds, dtype="float")
-
-            if key == "X":
-                weights = self._get_longitude_weights(d_bounds, r_bounds)
-            elif key == "Y":
-                weights = self._get_latitude_weights(d_bounds, r_bounds)
+            r_bounds = axis_bounds[key]["region"]
 
+            weights = axis_bounds[key]["weights_method"](d_bounds, r_bounds)
             weights.attrs = d_bounds.attrs
             axis_weights[key] = weights
 
@@ -309,25 +309,32 @@ class SpatialAccessor:
             # Check the axis coordinate variable exists in the Dataset.
             get_axis_coord(self._dataset, key)
 
-    def _validate_domain_bounds(self, domain_bounds: xr.DataArray):
-        """Validates the ``domain_bounds`` arg based on a set of criteria.
+    def _force_domain_order_low_to_high(self, domain_bounds: xr.DataArray):
+        """Reorders the ``domain_bounds`` low-to-high.
+
+        This method ensures all lower bound values are less than the upper bound
+        values (``domain_bounds[:, 1] < domain_bounds[:, 1]``).
 
         Parameters
         ----------
         domain_bounds: xr.DataArray
-            The bounds of an axis
+            The bounds of an axis.
 
-        Raises
+        Returns
         ------
-        TypeError
-            If the ``domain_bounds`` of a grid cell are not ordered low-to-high.
+        xr.DataArray
+            The bounds of an axis (re-ordered if applicable).
         """
         index_bad_cells = np.where(domain_bounds[:, 1] - domain_bounds[:, 0] < 0)[0]
+
         if len(index_bad_cells) > 0:
-            raise ValueError(
-                "The bounds have unexpected ordering. A lower bound has a "
-                "greater value than an upper bound."
-            )
+            new_domain_bounds = domain_bounds.copy()
+            new_domain_bounds[index_bad_cells, 0] = domain_bounds[index_bad_cells, 1]
+            new_domain_bounds[index_bad_cells, 1] = domain_bounds[index_bad_cells, 0]
+
+            return new_domain_bounds
+
+        return domain_bounds
 
     def _validate_region_bounds(self, axis: SpatialAxis, bounds: RegionAxisBounds):
         """Validates the ``bounds`` arg based on a set of criteria.
@@ -359,7 +366,7 @@ class SpatialAccessor:
                 "upper bounds, (lower, upper)."
             )
 
-        if len(bounds) <= 0 or len(bounds) > 2:
+        if len(bounds) != 2:
             raise ValueError(
                 f"The {axis} regional bounds is not a length of 2 (lower, upper)."
             )
@@ -602,7 +609,7 @@ class SpatialAccessor:
             domain_uppers = d_bounds[:, 1]
             region_lower, region_upper = r_bounds
 
-            # Grid cell stradling lower boundary.
+            # Grid cell straddling lower boundary.
             inds = np.where(
                 (domain_lowers < region_lower) & (domain_uppers > region_lower)
             )[0]
@@ -618,7 +625,7 @@ class SpatialAccessor:
             # and their weight will also be zero.
             d_bounds[inds, :] = region_upper
 
-            # Grid cell stradling upper boundary.
+            # Grid cell straddling upper boundary.
             inds = np.where(
                 (domain_lowers < region_upper) & (domain_uppers > region_upper)
             )[0]


Test patch:
diff --git a/tests/test_spatial.py b/tests/test_spatial.py
index 89e9013..6b5df53 100644
--- a/tests/test_spatial.py
+++ b/tests/test_spatial.py
@@ -254,12 +254,26 @@ class TestGetWeights:
     def setup(self):
         self.ds = generate_dataset(cf_compliant=True, has_bounds=True)
 
-    def test_raises_error_if_domain_lower_bound_exceeds_upper_bound(self):
-        ds = self.ds.copy()
-        ds.lon_bnds.data[:] = np.array([[1, 0], [1, 2], [2, 3], [3, 4]])
+    def test_bounds_reordered_when_upper_indexed_first(self):
+        domain_bounds = xr.DataArray(
+            name="lon_bnds",
+            data=np.array(
+                [[-89.375, -90], [0.0, -89.375], [0.0, 89.375], [89.375, 90]]
+            ),
+            coords={"lat": self.ds.lat},
+            dims=["lat", "bnds"],
+        )
+        result = self.ds.spatial._force_domain_order_low_to_high(domain_bounds)
 
-        with pytest.raises(ValueError):
-            ds.spatial.get_weights(axis=["X"])
+        expected_domain_bounds = xr.DataArray(
+            name="lon_bnds",
+            data=np.array(
+                [[-90, -89.375], [-89.375, 0.0], [0.0, 89.375], [89.375, 90]]
+            ),
+            coords={"lat": self.ds.lat},
+            dims=["lat", "bnds"],
+        )
+        assert result.identical(expected_domain_bounds)
 
     def test_weights_for_region_in_lat_and_lon_domains(self):
         result = self.ds.spatial.get_weights(