From fa1e239f446bbff1634d3e6218432ee6364fd8aa Mon Sep 17 00:00:00 2001 From: Tom Vo Date: Mon, 26 Sep 2022 15:47:49 -0700 Subject: [PATCH] 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 | 2 +- xcdat/axis.py | 30 ++++++++++++++++++++++++++---- xcdat/spatial.py | 10 +++++----- 3 files changed, 32 insertions(+), 10 deletions(-) diff --git a/tests/test_axis.py b/tests/test_axis.py index 885546c7..2cad2117 100644 --- a/tests/test_axis.py +++ b/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) diff --git a/xcdat/axis.py b/xcdat/axis.py index b7126b85..02733ec8 100644 --- a/xcdat/axis.py +++ b/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. diff --git a/xcdat/spatial.py b/xcdat/spatial.py index b1f88294..9b312ec3 100644 --- a/xcdat/spatial.py +++ b/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