Update swap_lon_axis() to handle 360 axis swapping

- Update `_swap_lon_axis()` comment on Dask arrays and check for multi-dimensional

tests/test_axis.py

@@ -370,7 +370,7 @@ def test_does_not_swap_if_desired_orientation_is_the_same_as_the_existing_orient
             },
         )
 
-        result = swap_lon_axis(ds_360, to=(0, 360))
+        result = swap_lon_axis(ds_360, to=(0, 360), sort_ascending=False)
 
         assert result.identical(ds_360)
 

xcdat/axis.py

@@ -310,12 +310,12 @@ def _get_axis_coord_keys(
     return list(set(keys))
 
 
-def _swap_lon_axis(coord: xr.DataArray, to: Tuple[float, float]) -> xr.DataArray:
+def _swap_lon_axis(coords: xr.DataArray, to: Tuple[float, float]) -> xr.DataArray:
     """Swaps coordinates on a longitude axis.
 
     Parameters
     ----------
-    lon : xr.DataArray
+    coords : xr.DataArray
         Coordinates on a longitude axis.
     to : Tuple[float, float]
         The new longitude axis orientation.
@@ -326,15 +326,37 @@ def _swap_lon_axis(coord: xr.DataArray, to: Tuple[float, float]) -> xr.DataArray
         The coordinates on a new longitude axis orientation.
     """
     if to == (-180, 180):
-        return ((coord + 180) % 360) - 180
+        new_coords = ((coords + 180) % 360) - 180
     elif to == (0, 360):
-        return coord % 360
+        # Swap the coordinates.
+        # Example with 180 coords: [-180, -0, 179] -> [0, 180, 360]
+        # Example with 360 coords: [60, 150, 360] -> [60, 150, 0]
+        new_coords = coords % 360
+
+        # Check if the original coordinates contain an element with a value of
+        # 360. If this element exists, use its index to revert its swapped
+        # value of 0 (360 % 360 is 0) back to 360. This case usually happens
+        # if the coordinate are already on the (0, 360) axis orientation.
+        # Example with 360 coords: [60, 150, 0] -> [60, 150, 360]
+        circular_index = np.where(coords == 360)
+
+        if len(circular_index) > 0:
+            # Must load lazy, multidimensional Dask arrays into into eager,
+            # in-memory NumPy arrays before manipulating its values. Otherwise,
+            # it raises `NotImplementedError xarray can't set arrays with
+            # multiple array indices to dask yet`.
+            if isinstance(new_coords.data, Array) and new_coords.ndim > 1:
+                new_coords.load()
+
+            new_coords[circular_index] = 360
     else:
         raise ValueError(
             "Currently, only (-180, 180) and (0, 360) are supported longitude axis "
             "orientations."
         )
 
+    return new_coords
+
 
 def _get_prime_meridian_index(lon_bounds: xr.DataArray) -> Optional[np.ndarray]:
     """Gets the index of the prime meridian cell in the longitude bounds.

xcdat/spatial.py

@@ -524,11 +524,11 @@ def _swap_lon_axis(
         """
         lon_swap = lon.copy()
 
-        # If chunking, must convert convert the xarray data structure from lazy
-        # Dask arrays into eager, in-memory NumPy arrays before performing
-        # manipulations on the data. Otherwise, it raises `NotImplementedError
-        # xarray can't set arrays with multiple array indices to dask yet`.
-        if type(lon_swap.data) == Array:
+        # Must load lazy, multidimensional Dask arrays into into eager,
+        # in-memory NumPy arrays before manipulating its values. Otherwise,
+        # it raises `NotImplementedError xarray can't set arrays with multiple
+        # array indices to dask yet`.
+        if type(lon_swap.data) == Array and lon_swap.ndim > 1:
             lon_swap.load()
 
         # Must set keep_attrs=True or the xarray DataArray attrs will get