From 8a386ed860e36a4fe416ddceb3107263e7068234 Mon Sep 17 00:00:00 2001 From: tomvothecoder Date: Wed, 10 Sep 2025 15:22:54 -0700 Subject: [PATCH 1/5] Update `swap_lon_axis()` to drop extra cell to handle prime meridian --- tests/test_axis.py | 74 ++++++++++++++++--------- tests/test_bounds.py | 38 +------------ tests/test_dataset.py | 5 +- xcdat/axis.py | 123 ++++++++++++++++++++++++++++++++++++++---- xcdat/bounds.py | 48 ----------------- xcdat/temporal.py | 4 +- 6 files changed, 167 insertions(+), 125 deletions(-) diff --git a/tests/test_axis.py b/tests/test_axis.py index b85fb344b..1f0916d27 100644 --- a/tests/test_axis.py +++ b/tests/test_axis.py @@ -5,6 +5,7 @@ from tests.fixtures import generate_dataset from xcdat.axis import ( CFAxisKey, + _get_bounds_dim, center_times, get_coords_by_name, get_dim_coords, @@ -691,12 +692,6 @@ def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_ dims=["lon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, ), - "lon2": xr.DataArray( - name="lon2", - data=np.array([-180, -1, 0, 1, 179]), - dims=["lon"], - attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, - ), }, data_vars={ "ts": xr.DataArray( @@ -726,13 +721,7 @@ def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_ coords={ "lon": xr.DataArray( name="lon", - data=np.array([0, 1, 179, 180, 359, 360]), - dims=["lon"], - attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, - ), - "lon2": xr.DataArray( - name="lon2", - data=np.array([0, 1, 179, 180, 359, 360]), + data=np.array([0, 1, 179, 180, 359]), dims=["lon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, ), @@ -740,7 +729,7 @@ def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_ data_vars={ "ts": xr.DataArray( name="ts", - data=np.array([2, 3, 4, 0, 1, 2]), + data=np.array([2, 3, 4, 0, 1]), dims=["lon"], attrs={"test_attr": "test"}, ), @@ -748,12 +737,11 @@ def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_ name="lon_bnds", data=np.array( [ - [0, 0.5], + [-0.5, 0.5], [0.5, 1.5], [1.5, 179.5], [179.5, 358.5], [358.5, 359.5], - [359.5, 360], ] ), dims=["lon", "bnds"], @@ -762,7 +750,7 @@ def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_ }, ) - assert result.identical(expected) + xr.testing.assert_identical(result, expected) def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lon_bnds( self, @@ -830,13 +818,13 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo coords={ "lon": xr.DataArray( name="lon", - data=np.array([0, 1, 179, 180, 359, 360]), + data=np.array([0, 1, 179, 180, 359]), dims=["lon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, ), "zlon": xr.DataArray( name="zlon", - data=np.array([0, 1, 179, 180, 359, 360]), + data=np.array([0, 1, 179, 180, 359]), dims=["zlon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "zlon_bnds"}, ), @@ -844,13 +832,13 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo data_vars={ "ts": xr.DataArray( name="ts", - data=np.array([2, 3, 4, 0, 1, 2]), + data=np.array([2, 3, 4, 0, 1]), dims=["lon"], attrs={"test_attr": "test"}, ), "ts2": xr.DataArray( name="ts2", - data=np.array([2, 3, 4, 0, 1, 2]), + data=np.array([2, 3, 4, 0, 1]), dims=["zlon"], attrs={"test_attr": "test"}, ), @@ -858,12 +846,11 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo name="lon_bnds", data=np.array( [ - [0, 0.5], + [-0.5, 0.5], [0.5, 1.5], [1.5, 179.5], [179.5, 358.5], [358.5, 359.5], - [359.5, 360], ] ), dims=["lon", "bnds"], @@ -873,12 +860,11 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo name="zlon_bnds", data=np.array( [ - [0, 0.5], + [-0.5, 0.5], [0.5, 1.5], [1.5, 179.5], [179.5, 358.5], [358.5, 359.5], - [359.5, 360], ] ), dims=["zlon", "bnds"], @@ -887,4 +873,40 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo }, ) - assert result.identical(expected) + xr.testing.assert_identical(result, expected) + + +class TestGetBoundsDim: + def test_returns_bounds_dim_for_standard_case(self): + coords = xr.DataArray([0, 1, 2], dims=["lat"]) + bounds = xr.DataArray([[0, 1], [1, 2], [2, 3]], dims=["lat", "bnds"]) + + result = _get_bounds_dim(coords, bounds) + + assert result == "bnds" + + def test_returns_bounds_dim_when_bounds_dim_has_custom_name(self): + coords = xr.DataArray([10, 20, 30], dims=["lon"]) + bounds = xr.DataArray([[5, 15], [15, 25], [25, 35]], dims=["lon", "boundaries"]) + + result = _get_bounds_dim(coords, bounds) + + assert result == "boundaries" + + def test_raises_error_when_bounds_has_no_extra_dim(self): + coords = xr.DataArray([0, 1, 2], dims=["lat"]) + bounds = xr.DataArray([0, 1, 2], dims=["lat"]) + + with pytest.raises( + ValueError, match="No extra dimension found in bounds variable" + ): + _get_bounds_dim(coords, bounds) + + def test_raises_error_when_bounds_has_multiple_extra_dims(self): + coords = xr.DataArray([0, 1, 2], dims=["lat"]) + bounds = xr.DataArray(np.zeros((3, 2, 2)), dims=["lat", "bnds", "extra"]) + + with pytest.raises( + ValueError, match="Bounds variable must have exactly one more dimension" + ): + _get_bounds_dim(coords, bounds) diff --git a/tests/test_bounds.py b/tests/test_bounds.py index 0dbf1de45..00838e324 100644 --- a/tests/test_bounds.py +++ b/tests/test_bounds.py @@ -12,7 +12,7 @@ lat_bnds, lon_bnds, ) -from xcdat.bounds import BoundsAccessor, _get_bounds_dim +from xcdat.bounds import BoundsAccessor logger = logging.getLogger(__name__) @@ -1059,39 +1059,3 @@ def test_add_monthly_bounds_for_end_of_month_set_to_true(self): ) assert result.identical(expected) - - -class TestGetBoundsDim: - def test_returns_bounds_dim_for_standard_case(self): - coords = xr.DataArray([0, 1, 2], dims=["lat"]) - bounds = xr.DataArray([[0, 1], [1, 2], [2, 3]], dims=["lat", "bnds"]) - - result = _get_bounds_dim(coords, bounds) - - assert result == "bnds" - - def test_returns_bounds_dim_when_bounds_dim_has_custom_name(self): - coords = xr.DataArray([10, 20, 30], dims=["lon"]) - bounds = xr.DataArray([[5, 15], [15, 25], [25, 35]], dims=["lon", "boundaries"]) - - result = _get_bounds_dim(coords, bounds) - - assert result == "boundaries" - - def test_raises_error_when_bounds_has_no_extra_dim(self): - coords = xr.DataArray([0, 1, 2], dims=["lat"]) - bounds = xr.DataArray([0, 1, 2], dims=["lat"]) - - with pytest.raises( - ValueError, match="No extra dimension found in bounds variable" - ): - _get_bounds_dim(coords, bounds) - - def test_raises_error_when_bounds_has_multiple_extra_dims(self): - coords = xr.DataArray([0, 1, 2], dims=["lat"]) - bounds = xr.DataArray(np.zeros((3, 2, 2)), dims=["lat", "bnds", "extra"]) - - with pytest.raises( - ValueError, match="Bounds variable must have exactly one more dimension" - ): - _get_bounds_dim(coords, bounds) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index b3841dffd..146938945 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1650,7 +1650,7 @@ def test_orients_longitude_bounds_from_180_to_360_and_sorts_with_prime_meridian_ coords={ "lon": xr.DataArray( name="lon", - data=np.array([0.0, 1.0, 179.0, 180.0, 359.0, 360.0]), + data=np.array([0.0, 1.0, 179.0, 180.0, 359.0]), dims=["lon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, ) @@ -1660,12 +1660,11 @@ def test_orients_longitude_bounds_from_180_to_360_and_sorts_with_prime_meridian_ name="lon_bnds", data=np.array( [ - [0, 0.5], + [-0.5, 0.5], [0.5, 1.5], [1.5, 179.5], [179.5, 358.5], [358.5, 359.5], - [359.5, 360], ] ), dims=["lon", "bnds"], diff --git a/xcdat/axis.py b/xcdat/axis.py index 65782fce8..3f5014e1a 100644 --- a/xcdat/axis.py +++ b/xcdat/axis.py @@ -370,18 +370,23 @@ def _swap_lon_bounds(ds: xr.Dataset, key: str, to: tuple[float, float]): new_bounds = _swap_lon_axis(bounds, to) if not ds[key].identical(new_bounds): - ds[key] = new_bounds - - # Handle cases where a prime meridian cell exists, which can occur - # after swapping longitude bounds to (0, 360). This involves extending - # the longitude and bounds by one cell to take into account the prime - # meridian. It also results in extending the data variables by one - # value. + ds_new = ds.copy() + ds_new[key] = new_bounds + + # Handle cases where a prime meridian cell exists after converting + # longitude bounds to the range (0, 360). This involves extending + # longitude bounds and data variables by one cell to account for the + # prime meridian, updating bounds interval for the prime meridian cell, + # and then trimming the extra cell to maintain the original shape. if to == (0, 360): - p_meridian_index = _get_prime_meridian_index(ds[key]) + p_meridian_index = _get_prime_meridian_index(ds_new[key]) if p_meridian_index is not None: - ds = _align_lon_to_360(ds, ds[key], p_meridian_index) + ds_new = _adjust_bounds_for_prime_meridian( + ds_new, key, p_meridian_index + ) + + return ds_new return ds @@ -470,6 +475,58 @@ def _get_prime_meridian_index(lon_bounds: xr.DataArray) -> np.ndarray | None: return p_meridian_index +def _adjust_bounds_for_prime_meridian( + ds: xr.Dataset, key: str, p_meridian_index: np.ndarray +) -> xr.Dataset: + """Adjusts bounds for the prime meridian cell in longitude coordinates. + + This function modifies the longitude bounds to handle cases where a grid + cell spans the prime meridian. It ensures that the longitude bounds are + aligned to the (0, 360) range by splitting the prime meridian cell into + two parts and adjusting the bounds accordingly. The extra cell created + during this process is removed to restore the original shape of the dataset. + + Parameters + ---------- + ds : xr.Dataset + The Dataset containing longitude bounds and coordinates. + key : str + The name of the longitude bounds variable in the dataset. + p_meridian_index : np.ndarray + An array with a single element representing the index of the prime + meridian cell. + + Returns + ------- + xr.Dataset + The Dataset with adjusted longitude bounds for the prime meridian cell. + """ + ds_new = ds.copy() + + ds_new = _align_lon_to_360(ds_new, ds_new[key], p_meridian_index) + dim = get_dim_keys(ds_new[key], "X") + + # NOTE: The logic for calculating new bounds is similar to the + # implementation in xcdat.bounds.create_bounds. + # Calculate the bounds interval for the prime meridian cell using + # the cell as the midpoint and the difference of the bounds. + bounds = ds_new[key].isel({dim: p_meridian_index}) + bounds_dim = _get_bounds_dim(ds_new[dim], ds_new[key]) + bounds_delta = bounds.diff(bounds_dim).item() + + # Create new bounds for the prime meridian cell. + midpoint = ds_new[dim][p_meridian_index].item() + new_lower = midpoint - bounds_delta + new_upper = midpoint + bounds_delta + + ds_new[key][p_meridian_index, :] = [new_lower, new_upper] + + # Remove the extra cell to restore the original shape. + ds_new = ds_new.isel({dim: slice(0, -1)}) + + return ds_new + + def _align_lon_to_360( ds: xr.Dataset, lon_bounds: xr.DataArray, @@ -596,3 +653,51 @@ def _align_lon_bounds_to_360( bounds[p_meridian_index, 0], bounds[-1, 1] = (0.0, 360.0) return bounds + + +def _get_bounds_dim(da_coords: xr.DataArray, da_bounds: xr.DataArray) -> str: + """Identify the bounds dimension in the given bounds DataArray. + + This function determines which dimension in the bounds DataArray + represents the bounds (e.g., lower and upper limits) for each value + in the coordinate DataArray. According to the CF Conventions, a boundary + variable will have one more dimension than its associated coordinate or + auxiliary coordinate variable. Refer to [6] for more details. + + Parameters + ---------- + da_coords : xr.DataArray + The coordinate DataArray. + da_bounds : xr.DataArray + The bounds DataArray. + + Returns + ------- + str + The name of the bounds dimension. + + Raises + ------ + ValueError + If the bounds variable does not have exactly one more dimension than + the coordinate variable, or if no valid bounds dimension is found. + + References + ---------- + .. [6] https://cfconventions.org/Data/cf-conventions/cf-conventions-1.10/cf-conventions.html#cell-boundaries + """ + bounds_dim = [dim for dim in da_bounds.dims if dim not in da_coords.dims] + + if len(bounds_dim) == 1: + return str(bounds_dim[0]) + elif len(bounds_dim) == 0: + raise ValueError( + "No extra dimension found in bounds variable compared to coordinate " + f"variable. Coordinate dims: {da_coords.dims}, bounds dims: " + f"{da_bounds.dims}" + ) + else: + raise ValueError( + f"Bounds variable must have exactly one more dimension than the coordinate " + f"variable. Coordinate dims: {da_coords.dims}, bounds dims: {da_bounds.dims}" + ) diff --git a/xcdat/bounds.py b/xcdat/bounds.py index 0356405b7..fb4bd81cd 100644 --- a/xcdat/bounds.py +++ b/xcdat/bounds.py @@ -1003,51 +1003,3 @@ def create_bounds(axis: CFAxisKey, coord_var: xr.DataArray) -> xr.DataArray: ) return bounds - - -def _get_bounds_dim(da_coords: xr.DataArray, da_bounds: xr.DataArray) -> str: - """Identify the bounds dimension in the given bounds DataArray. - - This function determines which dimension in the bounds DataArray - represents the bounds (e.g., lower and upper limits) for each value - in the coordinate DataArray. According to the CF Conventions, a boundary - variable will have one more dimension than its associated coordinate or - auxiliary coordinate variable. Refer to [6] for more details. - - Parameters - ---------- - da_coords : xr.DataArray - The coordinate DataArray. - da_bounds : xr.DataArray - The bounds DataArray. - - Returns - ------- - str - The name of the bounds dimension. - - Raises - ------ - ValueError - If the bounds variable does not have exactly one more dimension than - the coordinate variable, or if no valid bounds dimension is found. - - References - ---------- - .. [6] https://cfconventions.org/Data/cf-conventions/cf-conventions-1.10/cf-conventions.html#cell-boundaries - """ - bounds_dim = [dim for dim in da_bounds.dims if dim not in da_coords.dims] - - if len(bounds_dim) == 1: - return str(bounds_dim[0]) - elif len(bounds_dim) == 0: - raise ValueError( - "No extra dimension found in bounds variable compared to coordinate " - f"variable. Coordinate dims: {da_coords.dims}, bounds dims: " - f"{da_bounds.dims}" - ) - else: - raise ValueError( - f"Bounds variable must have exactly one more dimension than the coordinate " - f"variable. Coordinate dims: {da_coords.dims}, bounds dims: {da_bounds.dims}" - ) diff --git a/xcdat/temporal.py b/xcdat/temporal.py index 6de3276b7..96fe30870 100644 --- a/xcdat/temporal.py +++ b/xcdat/temporal.py @@ -17,7 +17,7 @@ from xcdat import bounds # noqa: F401 from xcdat._logger import _setup_custom_logger -from xcdat.axis import get_dim_coords +from xcdat.axis import _get_bounds_dim, get_dim_coords from xcdat.dataset import _get_data_var from xcdat.utils import _get_masked_weights, _validate_min_weight @@ -1726,7 +1726,7 @@ def _get_weights(self, ds: xr.Dataset, data_var: str) -> xr.DataArray: time_bounds = ds.bounds.get_bounds("T", var_key=data_var) time_coords = ds[data_var][self.dim] - bounds_dim = bounds._get_bounds_dim(time_coords, time_bounds) + bounds_dim = _get_bounds_dim(time_coords, time_bounds) time_lengths = time_bounds.diff(dim=bounds_dim).squeeze() # Must be cast dtype from "timedelta64[ns]" to "float64", specifically From 10de20923fd3aeeb528fc9e305bd32c2ab6d5c4c Mon Sep 17 00:00:00 2001 From: tomvothecoder Date: Wed, 10 Sep 2025 15:33:21 -0700 Subject: [PATCH 2/5] Fix bounds delta --- xcdat/axis.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xcdat/axis.py b/xcdat/axis.py index 3f5014e1a..d4baaf7c3 100644 --- a/xcdat/axis.py +++ b/xcdat/axis.py @@ -510,14 +510,14 @@ def _adjust_bounds_for_prime_meridian( # implementation in xcdat.bounds.create_bounds. # Calculate the bounds interval for the prime meridian cell using # the cell as the midpoint and the difference of the bounds. - bounds = ds_new[key].isel({dim: p_meridian_index}) + bounds = ds_new[key].isel({dim: 0}) bounds_dim = _get_bounds_dim(ds_new[dim], ds_new[key]) - bounds_delta = bounds.diff(bounds_dim).item() + bounds_delta = abs(bounds.diff(bounds_dim).item()) # Create new bounds for the prime meridian cell. midpoint = ds_new[dim][p_meridian_index].item() new_lower = midpoint - bounds_delta - new_upper = midpoint + bounds_delta + new_upper = new_lower + bounds_delta ds_new[key][p_meridian_index, :] = [new_lower, new_upper] From b7b243b23feef829339abbd487b35a72a1ccbb6a Mon Sep 17 00:00:00 2001 From: tomvothecoder Date: Wed, 10 Sep 2025 15:37:44 -0700 Subject: [PATCH 3/5] Fix bounds delta --- xcdat/axis.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xcdat/axis.py b/xcdat/axis.py index d4baaf7c3..7cf1a0344 100644 --- a/xcdat/axis.py +++ b/xcdat/axis.py @@ -512,12 +512,12 @@ def _adjust_bounds_for_prime_meridian( # the cell as the midpoint and the difference of the bounds. bounds = ds_new[key].isel({dim: 0}) bounds_dim = _get_bounds_dim(ds_new[dim], ds_new[key]) - bounds_delta = abs(bounds.diff(bounds_dim).item()) + bounds_delta = abs(bounds.diff(bounds_dim).item()) / 2 # Create new bounds for the prime meridian cell. midpoint = ds_new[dim][p_meridian_index].item() new_lower = midpoint - bounds_delta - new_upper = new_lower + bounds_delta + new_upper = midpoint + bounds_delta ds_new[key][p_meridian_index, :] = [new_lower, new_upper] From 68b7c1fe741375c1589d88f243ae58ac27c5fd77 Mon Sep 17 00:00:00 2001 From: tomvothecoder Date: Tue, 23 Sep 2025 15:41:28 -0700 Subject: [PATCH 4/5] Update `swap_lon_axis()` to fix bounds for prime meridian cell - Removes prime meridian handling code to add extra cell - Refactors code in `axis.py` for readability - Moves prime meridian functions from `axis.py` to `spatial.py` --- tests/test_axis.py | 268 ++++++++++++++++++++--------- tests/test_dataset.py | 47 ++++-- tests/test_spatial.py | 18 ++ xcdat/axis.py | 382 +++++++++++++++--------------------------- xcdat/spatial.py | 107 ++++++++++-- 5 files changed, 465 insertions(+), 357 deletions(-) diff --git a/tests/test_axis.py b/tests/test_axis.py index 1f0916d27..98c1ca06f 100644 --- a/tests/test_axis.py +++ b/tests/test_axis.py @@ -511,45 +511,6 @@ def test_raises_error_with_incorrect_lon_orientation_for_swapping(self): with pytest.raises(ValueError): swap_lon_axis(ds, to=9000) # type: ignore - def test_raises_error_if_lon_bounds_contains_more_than_one_prime_meridian_cell( - self, - ): - ds_180 = xr.Dataset( - coords={ - "lon": xr.DataArray( - name="lon", - data=np.array([-180, -1, 0, 1, 179]), - dims=["lon"], - attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, - ) - }, - data_vars={ - "lon_bnds": xr.DataArray( - name="lon_bnds", - data=np.array( - [ - [-180.5, -1.5], - [-1.5, -0.5], - [-0.5, 0.5], - [0.5, 1.5], - [-180.5, 1.5], - ] - ), - dims=["lon", "bnds"], - attrs={"xcdat_bounds": "True"}, - ), - "ts": xr.DataArray( - name="ts", - data=np.array([0, 1, 2, 3, 4]), - dims=["lon"], - attrs={"test_attr": "test"}, - ), - }, - ) - - with pytest.raises(ValueError): - swap_lon_axis(ds_180, to=(0, 360)) - def test_does_not_swap_if_desired_orientation_is_the_same_as_the_existing_orientation( self, ): @@ -617,6 +578,81 @@ def test_does_not_swap_bounds_if_bounds_do_not_exist(self): assert result.identical(expected) + def test_indempotency_when_converting_0_to_360_to_0_to_360(self): + ds_360 = xr.Dataset( + coords={ + "lon": xr.DataArray( + name="lon", + data=np.array([0, 90, 180, 270, 360]), + dims=["lon"], + attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, + ), + }, + data_vars={ + "ts": xr.DataArray( + name="ts", + data=np.array([0, 1, 2, 3, 4]), + dims=["lon"], + attrs={"test_attr": "test"}, + ), + "lon_bnds": xr.DataArray( + name="lon_bnds", + data=np.array( + [ + [0, 90], + [90, 180], + [180, 270], + [270, 360], + [360, 360], + ] + ), + dims=["lon", "bnds"], + attrs={"xcdat_bounds": "True"}, + ), + }, + ) + + result = swap_lon_axis(ds_360, to=(0, 360)) + xr.testing.assert_identical(result, ds_360) + + def test_indempotency_when_converting_minus_180_to_180_to_minus_180_to_180(self): + ds_180 = xr.Dataset( + coords={ + "lon": xr.DataArray( + name="lon", + data=np.array([-180, -90, 0, 90, 180]), + dims=["lon"], + attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, + ), + }, + data_vars={ + "ts": xr.DataArray( + name="ts", + data=np.array([0, 1, 2, 3, 4]), + dims=["lon"], + attrs={"test_attr": "test"}, + ), + "lon_bnds": xr.DataArray( + name="lon_bnds", + data=np.array( + [ + [-180, -135], + [-135, -45], + [-45, 45], + [45, 135], + [135, 180], + ] + ), + dims=["lon", "bnds"], + attrs={"xcdat_bounds": "True"}, + ), + }, + ) + + result = swap_lon_axis(ds_180, to=(-180, 180)) + + xr.testing.assert_identical(result, ds_180) + def test_swaps_single_dim_from_360_to_180_and_sorts(self): ds_360 = xr.Dataset( coords={ @@ -681,14 +717,81 @@ def test_swaps_single_dim_from_360_to_180_and_sorts(self): assert result.identical(expected) - def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lon_bnds( + def test_swaps_single_dim_from_180_to_360_without_prime_meridian_cell(self): + ds_180 = xr.Dataset( + coords={ + "lon": xr.DataArray( + name="lon", + data=np.array([-135, -45, 45, 135]), + dims=["lon"], + attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, + ), + }, + data_vars={ + "ts": xr.DataArray( + name="ts", + data=np.array([0, 1, 2, 3]), + dims=["lon"], + attrs={"test_attr": "test"}, + ), + "lon_bnds": xr.DataArray( + name="lon_bnds", + data=np.array( + [ + [-180, -90], + [-90, 0], + [0, 90], + [90, 180], + ] + ), + dims=["lon", "bnds"], + attrs={"xcdat_bounds": "True"}, + ), + }, + ) + result = swap_lon_axis(ds_180, to=(0, 360)) + expected = xr.Dataset( + coords={ + "lon": xr.DataArray( + name="lon", + data=np.array([45, 135, 225, 315]), + dims=["lon"], + attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, + ), + }, + data_vars={ + "ts": xr.DataArray( + name="ts", + data=np.array([2, 3, 0, 1]), + dims=["lon"], + attrs={"test_attr": "test"}, + ), + "lon_bnds": xr.DataArray( + name="lon_bnds", + data=np.array( + [ + [0, 90], + [90, 180], + [180, 270], + [270, 360], + ] + ), + dims=["lon", "bnds"], + attrs={"xcdat_bounds": "True"}, + ), + }, + ) + + xr.testing.assert_identical(result, expected) + + def test_swaps_single_dim_from_180_to_360_and_normalizes_prime_meridian_cell_in_lon_bnds_to_360( self, ): ds_180 = xr.Dataset( coords={ "lon": xr.DataArray( name="lon", - data=np.array([-180, -1, 0, 1, 179]), + data=np.array([-179.5, -1.5, -0.5, 1.5, 179.5]), dims=["lon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, ), @@ -704,11 +807,11 @@ def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_ name="lon_bnds", data=np.array( [ - [-180.5, -1.5], - [-1.5, -0.5], - [-0.5, 0.5], - [0.5, 1.5], - [1.5, 179.5], + [-180.0, -179], + [-2.0, -1.0], + [-1.0, 0.0], # Prime meridian cell. + [1.0, 2.0], + [179.0, 180.0], ] ), dims=["lon", "bnds"], @@ -721,7 +824,7 @@ def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_ coords={ "lon": xr.DataArray( name="lon", - data=np.array([0, 1, 179, 180, 359]), + data=np.array([1.5, 179.5, 180.5, 358.5, 359.5]), dims=["lon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, ), @@ -729,7 +832,7 @@ def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_ data_vars={ "ts": xr.DataArray( name="ts", - data=np.array([2, 3, 4, 0, 1]), + data=np.array([3, 4, 0, 1, 2]), dims=["lon"], attrs={"test_attr": "test"}, ), @@ -737,11 +840,12 @@ def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_ name="lon_bnds", data=np.array( [ - [-0.5, 0.5], - [0.5, 1.5], - [1.5, 179.5], - [179.5, 358.5], - [358.5, 359.5], + [1.0, 2.0], + [179.0, 180.0], + [180.0, 181.0], + [358.0, 359.0], + # Instead of [359, 0], normalize to [359, 360]. + [359.0, 360.0], ] ), dims=["lon", "bnds"], @@ -752,20 +856,20 @@ def test_swaps_single_dim_from_180_to_360_and_sorts_with_prime_meridian_cell_in_ xr.testing.assert_identical(result, expected) - def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lon_bnds( + def test_swaps_multiple_dims_from_180_to_360_and_normalizes_prime_meridian_cell_in_lon_bnds_to_360( self, ): ds_180 = xr.Dataset( coords={ "lon": xr.DataArray( name="lon", - data=np.array([-180, -1, 0, 1, 179]), + data=np.array([-179.5, -1.5, -0.5, 1.5, 179.5]), dims=["lon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, ), "zlon": xr.DataArray( name="zlon", - data=np.array([-180, -1, 0, 1, 179]), + data=np.array([-179.5, -1.5, -0.5, 1.5, 179.5]), dims=["zlon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "zlon_bnds"}, ), @@ -787,11 +891,11 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo name="lon_bnds", data=np.array( [ - [-180.5, -1.5], - [-1.5, -0.5], - [-0.5, 0.5], - [0.5, 1.5], - [1.5, 179.5], + [-180.0, -179], + [-2.0, -1.0], + [-1.0, 0.0], # Prime meridian cell. + [1.0, 2.0], + [179.0, 180.0], ] ), dims=["lon", "bnds"], @@ -801,11 +905,11 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo name="zlon_bnds", data=np.array( [ - [-180.5, -1.5], - [-1.5, -0.5], - [-0.5, 0.5], - [0.5, 1.5], - [1.5, 179.5], + [-180.0, -179], + [-2.0, -1.0], + [-1.0, 0.0], # Prime meridian cell. + [1.0, 2.0], + [179.0, 180.0], ] ), dims=["zlon", "bnds"], @@ -818,13 +922,13 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo coords={ "lon": xr.DataArray( name="lon", - data=np.array([0, 1, 179, 180, 359]), + data=np.array([1.5, 179.5, 180.5, 358.5, 359.5]), dims=["lon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, ), "zlon": xr.DataArray( name="zlon", - data=np.array([0, 1, 179, 180, 359]), + data=np.array([1.5, 179.5, 180.5, 358.5, 359.5]), dims=["zlon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "zlon_bnds"}, ), @@ -832,13 +936,13 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo data_vars={ "ts": xr.DataArray( name="ts", - data=np.array([2, 3, 4, 0, 1]), + data=np.array([3, 4, 0, 1, 2]), dims=["lon"], attrs={"test_attr": "test"}, ), "ts2": xr.DataArray( - name="ts2", - data=np.array([2, 3, 4, 0, 1]), + name="ts", + data=np.array([3, 4, 0, 1, 2]), dims=["zlon"], attrs={"test_attr": "test"}, ), @@ -846,11 +950,12 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo name="lon_bnds", data=np.array( [ - [-0.5, 0.5], - [0.5, 1.5], - [1.5, 179.5], - [179.5, 358.5], - [358.5, 359.5], + [1.0, 2.0], + [179.0, 180.0], + [180.0, 181.0], + [358.0, 359.0], + # Instead of [359, 0], normalize to [359, 360]. + [359.0, 360.0], ] ), dims=["lon", "bnds"], @@ -860,11 +965,12 @@ def test_swaps_all_dims_from_180_to_360_and_sorts_with_prime_meridian_cell_in_lo name="zlon_bnds", data=np.array( [ - [-0.5, 0.5], - [0.5, 1.5], - [1.5, 179.5], - [179.5, 358.5], - [358.5, 359.5], + [1.0, 2.0], + [179.0, 180.0], + [180.0, 181.0], + [358.0, 359.0], + # Instead of [359, 0], normalize to [359, 360]. + [359.0, 360.0], ] ), dims=["zlon", "bnds"], diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 146938945..25ad1be21 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1614,29 +1614,35 @@ def test_adds_missing_lat_and_lon_and_time_bounds(self): assert "lon_bnds" in result_data_vars assert "time_bnds" in result_data_vars - def test_orients_longitude_bounds_from_180_to_360_and_sorts_with_prime_meridian_cell( + def test_swaps_single_dim_from_180_to_360_and_normalizes_prime_meridian_cell_in_lon_bnds_to_360( self, ): # Chunk the input dataset to test method also works with Dask. - ds = xr.Dataset( + ds_180 = xr.Dataset( coords={ "lon": xr.DataArray( name="lon", - data=np.array([-180, -1, 0, 1, 179]), + data=np.array([-179.5, -1.5, -0.5, 1.5, 179.5]), dims=["lon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, - ) + ), }, data_vars={ + "ts": xr.DataArray( + name="ts", + data=np.array([0, 1, 2, 3, 4]), + dims=["lon"], + attrs={"test_attr": "test"}, + ), "lon_bnds": xr.DataArray( name="lon_bnds", data=np.array( [ - [-180.5, -1.5], - [-1.5, -0.5], - [-0.5, 0.5], - [0.5, 1.5], - [1.5, 179.5], + [-180.0, -179], + [-2.0, -1.0], + [-1.0, 0.0], # Prime meridian cell. + [1.0, 2.0], + [179.0, 180.0], ] ), dims=["lon", "bnds"], @@ -1645,26 +1651,33 @@ def test_orients_longitude_bounds_from_180_to_360_and_sorts_with_prime_meridian_ }, ).chunk({"lon": 2}) - result = _postprocess_dataset(ds, lon_orient=(0, 360)) + result = _postprocess_dataset(ds_180, lon_orient=(0, 360)) expected = xr.Dataset( coords={ "lon": xr.DataArray( name="lon", - data=np.array([0.0, 1.0, 179.0, 180.0, 359.0]), + data=np.array([1.5, 179.5, 180.5, 358.5, 359.5]), dims=["lon"], attrs={"units": "degrees_east", "axis": "X", "bounds": "lon_bnds"}, - ) + ), }, data_vars={ + "ts": xr.DataArray( + name="ts", + data=np.array([3, 4, 0, 1, 2]), + dims=["lon"], + attrs={"test_attr": "test"}, + ), "lon_bnds": xr.DataArray( name="lon_bnds", data=np.array( [ - [-0.5, 0.5], - [0.5, 1.5], - [1.5, 179.5], - [179.5, 358.5], - [358.5, 359.5], + [1.0, 2.0], + [179.0, 180.0], + [180.0, 181.0], + [358.0, 359.0], + # Instead of [359, 0], normalize to [359, 360]. + [359.0, 360.0], ] ), dims=["lon", "bnds"], diff --git a/tests/test_spatial.py b/tests/test_spatial.py index 868071556..6259356a8 100644 --- a/tests/test_spatial.py +++ b/tests/test_spatial.py @@ -604,6 +604,24 @@ def test_dataset_weights_for_region_in_lon_domain_with_region_spanning_p_meridia xr.testing.assert_allclose(result, expected) + def test_dataset_weights_raises_error_when_more_than_one_grid_cell_spans_prime_meridian( + self, + ): + ds = self.ds.copy(deep=True) + + # Modify the dataset to have more than one grid cell spanning the prime meridian. + lon_bnds = ds.lon_bnds.copy() + lon_bnds[0, :] = [359, 1] # First grid cell spans the prime meridian. + lon_bnds[1, :] = [358, 0] # Second grid cell also spans the prime meridian. + ds["lon_bnds"] = lon_bnds + + with pytest.raises( + ValueError, match="More than one grid cell spans the prime meridian" + ): + ds.spatial._get_longitude_weights( + domain_bounds=ds.lon_bnds, region_bounds=np.array([359, 1]) + ) + def test_dataset_weights_all_longitudes_for_equal_region_bounds(self): expected = xr.DataArray( data=np.array( diff --git a/xcdat/axis.py b/xcdat/axis.py index 7cf1a0344..1ca45b12c 100644 --- a/xcdat/axis.py +++ b/xcdat/axis.py @@ -8,8 +8,6 @@ import numpy as np import xarray as xr -from xcdat.utils import _if_multidim_dask_array_then_load - # https://cf-xarray.readthedocs.io/en/latest/coord_axes.html#axis-names CFAxisKey = Literal["X", "Y", "T", "Z"] # https://cf-xarray.readthedocs.io/en/latest/coord_axes.html#coordinate-names @@ -286,30 +284,33 @@ def swap_lon_axis( ------- xr.Dataset The Dataset with swapped lon axes orientation. + Raises + ------ + ValueError + If the ``to`` argument is not one of the supported orientations, + including (-180, 180) or (0, 360). """ ds = dataset.copy() coords = get_dim_coords(ds, "X").coords coord_keys = list(coords.keys()) - # Attempt to swap the orientation for longitude coordinates. - for key in coord_keys: - new_coord = _swap_lon_axis(ds.coords[key], to) - - if ds.coords[key].identical(new_coord): - continue + if to not in [(-180, 180), (0, 360)]: + raise ValueError( + "Only (-180, 180) and (0, 360) longitude ranges are supported." + ) - ds.coords[key] = new_coord + # Attempt to swap the orientation for longitude coordinates and bounds (if + # they exist). + for key in coord_keys: + ds[key] = _orient_lon_coords(ds.coords[key], to) - try: - bounds = ds.bounds.get_bounds("X") - except KeyError: - bounds = None + try: + bounds = ds.bounds.get_bounds("X", ds[key].name) + except KeyError: + bounds = None - if isinstance(bounds, xr.DataArray): - ds = _swap_lon_bounds(ds, str(bounds.name), to) - elif isinstance(bounds, xr.Dataset): - for key in bounds.data_vars.keys(): - ds = _swap_lon_bounds(ds, str(key), to) + if isinstance(bounds, xr.DataArray): + ds[bounds.name] = _orient_lon_bounds(ds[key], bounds, to) if sort_ascending: ds = ds.sortby(list(coords.dims), ascending=True) @@ -365,294 +366,181 @@ def _get_all_coord_keys(obj: xr.Dataset | xr.DataArray, axis: CFAxisKey) -> list return list(set(keys)) -def _swap_lon_bounds(ds: xr.Dataset, key: str, to: tuple[float, float]): - bounds = ds[key].copy() - new_bounds = _swap_lon_axis(bounds, to) - - if not ds[key].identical(new_bounds): - ds_new = ds.copy() - ds_new[key] = new_bounds - - # Handle cases where a prime meridian cell exists after converting - # longitude bounds to the range (0, 360). This involves extending - # longitude bounds and data variables by one cell to account for the - # prime meridian, updating bounds interval for the prime meridian cell, - # and then trimming the extra cell to maintain the original shape. - if to == (0, 360): - p_meridian_index = _get_prime_meridian_index(ds_new[key]) - - if p_meridian_index is not None: - ds_new = _adjust_bounds_for_prime_meridian( - ds_new, key, p_meridian_index - ) - - return ds_new - - return ds - +def _orient_lon_coords(coords: xr.DataArray, to: tuple[float, float]) -> xr.DataArray: + """ + Adjust longitude coordinates to match the specified orientation range. -def _swap_lon_axis(coords: xr.DataArray, to: tuple[float, float]) -> xr.DataArray: - """Swaps the axis orientation for longitude coordinates. + This function ensures that longitude centers are mapped correctly to the + target orientation range and avoids introducing a literal 360 center. If + the coordinates are already in the desired orientation, they are returned + as-is (idempotent). Longitude centers are treated as half-open in the range + [0, 360), ensuring no literal 360 center is introduced. Parameters ---------- coords : xr.DataArray - Coordinates on a longitude axis. + The longitude coordinates to be reoriented. to : tuple[float, float] - The new longitude axis orientation. + The orientation to swap the Dataset's longitude axis to. Supported + orientations include: + + * (-180, 180): represents [-180, 180) in math notation + * (0, 360): represents [0, 360) in math notation Returns ------- xr.DataArray - The longitude coordinates the opposite axis orientation If the - coordinates are already on the specified axis orientation, the same - coordinates are returned. + The longitude coordinates reoriented to the target range, with the + original encoding preserved. """ with xr.set_options(keep_attrs=True): - if to == (-180, 180): - # FIXME: Performance warning produced after swapping and then sorting - # based on how datasets are chunked. - new_coords = ((coords + 180) % 360) - 180 - elif to == (0, 360): - # Example with 180 coords: [-180, -0, 179] -> [0, 180, 360] - # Example with 360 coords: [60, 150, 360] -> [60, 150, 0] - # FIXME: Performance warning produced after swapping and then sorting - # based on how datasets are chunked. - 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] - index_with_360 = np.where(coords == 360) - - if index_with_360[0].size > 0: - _if_multidim_dask_array_then_load(new_coords) - - new_coords[index_with_360] = 360 + if _is_in_orientation(coords, to): + out = coords else: - raise ValueError( - "Currently, only (-180, 180) and (0, 360) are supported longitude axis " - "orientations." - ) - - new_coords.encoding = coords.encoding + out = _map_to_orientation(coords, to) - return new_coords + out.encoding = coords.encoding + return out -def _get_prime_meridian_index(lon_bounds: xr.DataArray) -> np.ndarray | None: - """Gets the index of the prime meridian cell in the longitude bounds. +def _orient_lon_bounds( + da_coords: xr.DataArray, da_bounds: xr.DataArray, to: tuple[float, float] +) -> xr.DataArray: + """Map longitude bounds to the target orientation. - A prime meridian cell can exist when converting the axis orientation - from [-180, 180) to [0, 360). + This function adjusts the longitude bounds of a dataset to match the specified + orientation. It ensures idempotency, meaning if the bounds are already in the + target orientation, they are returned as-is. For the (0, 360) orientation, it + optionally normalizes seam wraps (e.g., [359, 0] → [359, 360]). Parameters ---------- - lon_bounds : xr.DataArray - The longitude bounds. + da_coords : xr.DataArray + The longitude coordinate values of the dataset. + da_bounds : xr.DataArray + The longitude bounds of the dataset. + to : tuple[float, float] + The orientation to swap the Dataset's longitude axis to. Supported + orientations include: + + * (-180, 180): represents [-180, 180) in math notation + * (0, 360): represents [0, 360) in math notation Returns ------- - np.ndarray | None - An array with a single element representing the index of the prime - meridian index if it exists. Otherwise, None if the cell does not exist. - - Raises - ------ - ValueError - If more than one grid cell spans the prime meridian. + xr.DataArray + The longitude bounds adjusted to the target orientation. """ - p_meridian_index = np.where(lon_bounds[:, 1] - lon_bounds[:, 0] < 0)[0] + with xr.set_options(keep_attrs=True): + if _is_in_orientation(da_bounds, to): + out = da_bounds + else: + out = _map_to_orientation(da_bounds, to) - if p_meridian_index.size == 0: # pragma:no cover - return None - elif p_meridian_index.size > 1: - raise ValueError("More than one grid cell spans prime meridian.") + # Adjusts wrap-around bounds in a DataArray from [0, 360) longitude + # range. Example: [359, 0] -> [359, 360]. + if to == (0, 360): + bounds_dim = _get_bounds_dim(da_coords, out) + out = _normalize_wrap_bounds_0_to_360(out, bounds_dim) - return p_meridian_index + out.encoding = da_bounds.encoding + return out -def _adjust_bounds_for_prime_meridian( - ds: xr.Dataset, key: str, p_meridian_index: np.ndarray -) -> xr.Dataset: - """Adjusts bounds for the prime meridian cell in longitude coordinates. - This function modifies the longitude bounds to handle cases where a grid - cell spans the prime meridian. It ensures that the longitude bounds are - aligned to the (0, 360) range by splitting the prime meridian cell into - two parts and adjusting the bounds accordingly. The extra cell created - during this process is removed to restore the original shape of the dataset. +def _is_in_orientation(da: xr.DataArray, to: tuple[float, float]) -> bool: + """ + Check if the values in a DataArray conform to a specified orientation range. Parameters ---------- - ds : xr.Dataset - The Dataset containing longitude bounds and coordinates. - key : str - The name of the longitude bounds variable in the dataset. - p_meridian_index : np.ndarray - An array with a single element representing the index of the prime - meridian cell. + da : xr.DataArray + The input data array to check. + to : tuple[float, float] + The orientation to swap the Dataset's longitude axis to. Supported + orientations include: + + * (-180, 180): represents [-180, 180) in math notation + * (0, 360): represents [0, 360) in math notation Returns ------- - xr.Dataset - The Dataset with adjusted longitude bounds for the prime meridian cell. + bool + True if all values in `da` fall within the specified range, False otherwise. """ - ds_new = ds.copy() - - ds_new = _align_lon_to_360(ds_new, ds_new[key], p_meridian_index) - dim = get_dim_keys(ds_new[key], "X") - - # NOTE: The logic for calculating new bounds is similar to the - # implementation in xcdat.bounds.create_bounds. - # Calculate the bounds interval for the prime meridian cell using - # the cell as the midpoint and the difference of the bounds. - bounds = ds_new[key].isel({dim: 0}) - bounds_dim = _get_bounds_dim(ds_new[dim], ds_new[key]) - bounds_delta = abs(bounds.diff(bounds_dim).item()) / 2 - - # Create new bounds for the prime meridian cell. - midpoint = ds_new[dim][p_meridian_index].item() - new_lower = midpoint - bounds_delta - new_upper = midpoint + bounds_delta - - ds_new[key][p_meridian_index, :] = [new_lower, new_upper] + if to == (-180, 180): + return bool(np.all(da >= -180) and np.all(da <= 180)) - # Remove the extra cell to restore the original shape. - ds_new = ds_new.isel({dim: slice(0, -1)}) + return bool(np.all(da >= 0) and np.all(da <= 360)) - return ds_new - -def _align_lon_to_360( - ds: xr.Dataset, - lon_bounds: xr.DataArray, - p_meridian_index: np.ndarray, -) -> xr.Dataset: - """Handles a prime meridian cell to align longitude axis to (0, 360). - - This method ensures the domain bounds are within 0 to 360 by handling - the grid cell that encompasses the prime meridian (e.g., [359, 1]). - - First, it handles the prime meridian cell within the longitude axis bounds - by splitting the cell into two parts (one east and one west of the prime - meridian, refer to `_align_lon_bounds_to_360()` for more information). Then - it concatenates the 360 coordinate point to the longitude coordinates to - handle the addition of the extra grid cell from the previous step. Finally, - for each data variable associated with the longitude axis, the value of the - data variable at the prime meridian cell is concatenated to the data - variable. +def _map_to_orientation(da: xr.DataArray, to: tuple[float, float]) -> xr.DataArray: + """ + Map the values of a DataArray to a specified longitude orientation. Parameters ---------- - dataset : xr.Dataset - The Dataset. - p_meridian_index : np.ndarray - An array with a single element representing the index of the prime - meridian cell. + da : xr.DataArray + The input DataArray containing longitude values. + to : tuple[float, float] + The orientation to swap the Dataset's longitude axis to. Supported + orientations include: + + * (-180, 180): represents [-180, 180) in math notation + * (0, 360): represents [0, 360) in math notation Returns ------- - xr.Dataset - The Dataset. - """ - dim = get_dim_keys(lon_bounds, "X") - - # Create a dataset to store updated longitude variables. - ds_lon = xr.Dataset() - - # Align the the longitude bounds to the 360 orientation using the prime - # meridian index. This function splits the grid cell into two parts (east - # and west), which appends an extra set of bounds for the 360 coordinate. - ds_lon[lon_bounds.name] = _align_lon_bounds_to_360(lon_bounds, p_meridian_index) - - # After appending the extra set of bounds, update the last coordinate from - # 0 to 360. - for key, coord in ds_lon.coords.items(): - coord.values[-1] = 360 - ds_lon[key] = coord - - # Get the data variables related to the longitude axis and concatenate each - # with the value at the prime meridian. - for key, var in ds.cf.data_vars.items(): - if key != lon_bounds.name and dim in var.dims: - # Concatenate the prime meridian cell to the variable - p_meridian_val = var.isel({dim: p_meridian_index}).copy() - new_var = xr.concat((var, p_meridian_val), dim=dim) - - # Update the longitude coordinates for the variable. - new_var[dim] = ds_lon[dim] - ds_lon[var.name] = new_var + xr.DataArray + The DataArray with longitude values mapped to the specified range. - # Create a new dataset of non-longitude vars and updated longitude vars. - ds_no_lon = ds.get([v for v in ds.data_vars if dim not in ds[v].dims]) # type: ignore - ds_final = xr.merge((ds_no_lon, ds_lon)) + Notes + ----- + If a dataset includes both -180 and +180, this will create a duplicate 180 + when converting to [0, 360). We may need to drop or remap one of them in the + future if this becomes an issue. Datasets that follow CF convetion typically + do not run into this issue, as the longitude centers are within the + half-open interval [0, 360). + """ + if to == (-180, 180): + return ((da + 180) % 360) - 180 - return ds_final + return da % 360 -def _align_lon_bounds_to_360( - bounds: xr.DataArray, p_meridian_index: np.ndarray +def _normalize_wrap_bounds_0_to_360( + da_bounds: xr.DataArray, bounds_dim: str ) -> xr.DataArray: - """Handles a prime meridian cell to align longitude bounds axis to (0, 360). - - This method ensures the domain bounds are within 0 to 360 by handling - the grid cell that encompasses the prime meridian (e.g., [359, 1]). In - this case, calculating longitudinal weights is complicated because the - weights are determined by the difference of the bounds. + """ + Adjusts wrap-around bounds in a DataArray from [0, 360) longitude range. - If this situation exists, the method will split this grid cell into - two parts (one east and west of the prime meridian). The original - grid cell will have domain bounds extending east of the prime meridian - and an extra set of bounds will be concatenated to ``bounds`` - corresponding to the domain bounds west of the prime meridian. For - instance, a grid cell spanning -1 to 1, will be broken into a cell - from 0 to 1 and 359 to 360 (or -1 to 0). + This function modifies rows in the input DataArray where the bounds wrap + around the seam (e.g., [359, 0]) to ensure continuity by converting them + to [359, 360]. Rows that do not wrap are left unchanged. Parameters ---------- - bounds : xr.DataArray - The longitude domain bounds with prime meridian cell. - p_meridian_index : np.ndarray - The index of the prime meridian cell. + da_bounds : xr.DataArray + The input DataArray containing bounds to normalize. It is expected + to have a dimension specified by `bounds_dim` with two elements + representing the lower and upper bounds. + bounds_dim : str + The name of the dimension in `da_bounds` that represents the bounds. Returns ------- xr.DataArray - The longitude domain bounds with split prime meridian cell. - - Raises - ------ - ValueError - If longitude bounds are inclusively between 0 and 360. + A copy of the input DataArray with normalized wrap-around bounds. """ - _if_multidim_dask_array_then_load(bounds) - - # Example array: [[359, 1], [1, 90], [90, 180], [180, 359]] - # Reorient bound to span across zero (i.e., [359, 1] -> [-1, 1]). - # Result: [[-1, 1], [1, 90], [90, 180], [180, 359]] - bounds[p_meridian_index, 0] = bounds[p_meridian_index, 0] - 360.0 - - # Extend the array to nlon+1 by concatenating the grid cell that - # spans the prime meridian to the end. - # Result: [[-1, 1], [1, 90], [90, 180], [180, 359], [-1, 1]] - dim = get_dim_keys(bounds, "X") - bounds = xr.concat((bounds, bounds[p_meridian_index, :]), dim=dim) - - # Add an equivalent bound that spans 360 - # (i.e., [-1, 1] -> [359, 361]) to the end of the array. - # Result: [[-1, 1], [1, 90], [90, 180], [180, 359], [359, 361]] - repeat_bound = bounds[p_meridian_index, :][0] + 360.0 - bounds[-1, :] = repeat_bound - - # Update the lower-most min and upper-most max bounds to [0, 360]. - # Result: [[0, 1], [1, 90], [90, 180], [180, 359], [359, 360]] - bounds[p_meridian_index, 0], bounds[-1, 1] = (0.0, 360.0) - - return bounds + lo = da_bounds.sel({bounds_dim: 0}) + hi = da_bounds.sel({bounds_dim: 1}) + + wrap = lo > hi + + out = da_bounds.copy() + out.loc[{bounds_dim: 1}] = xr.where(wrap, 360.0, hi) + + return out def _get_bounds_dim(da_coords: xr.DataArray, da_bounds: xr.DataArray) -> str: diff --git a/xcdat/spatial.py b/xcdat/spatial.py index 50dd4dc2f..12bc8b2eb 100644 --- a/xcdat/spatial.py +++ b/xcdat/spatial.py @@ -9,8 +9,6 @@ import xarray as xr from xcdat.axis import ( - _align_lon_bounds_to_360, - _get_prime_meridian_index, get_dim_coords, get_dim_keys, ) @@ -430,23 +428,15 @@ def _get_longitude_weights( If the there are multiple instances in which the domain_bounds[:, 0] > domain_bounds[:, 1] """ - p_meridian_index: np.ndarray | None = None d_bounds = domain_bounds.copy() - - pm_cells = np.where(domain_bounds[:, 1] - domain_bounds[:, 0] < 0)[0] - if len(pm_cells) > 1: - raise ValueError( - "More than one longitude bound is out of order. Only one bound " - "value spanning the prime meridian is permitted in data on " - "a rectilinear grid." - ) d_bounds: xr.DataArray = self._swap_lon_axis(d_bounds, to=360) # type: ignore + p_meridian_index = _get_prime_meridian_index(d_bounds) if p_meridian_index is not None: d_bounds = _align_lon_bounds_to_360(d_bounds, p_meridian_index) if region_bounds is not None: - r_bounds: np.ndarray = self._swap_lon_axis(region_bounds, to=360) # type:ignore + r_bounds: np.ndarray = self._swap_lon_axis(region_bounds, to=360) # type: ignore is_region_circular = r_bounds[1] - r_bounds[0] == 0 if is_region_circular: @@ -842,3 +832,96 @@ def _mask_var_with_weight_threshold( dv_new.name = dv_mean.name return dv_new + + +def _align_lon_bounds_to_360( + bounds: xr.DataArray, p_meridian_index: np.ndarray +) -> xr.DataArray: + """Handles a prime meridian cell to align longitude bounds axis to (0, 360). + + This method ensures the domain bounds are within 0 to 360 by handling + the grid cell that encompasses the prime meridian (e.g., [359, 1]). In + this case, calculating longitudinal weights is complicated because the + weights are determined by the difference of the bounds. + + If this situation exists, the method will split this grid cell into + two parts (one east and west of the prime meridian). The original + grid cell will have domain bounds extending east of the prime meridian + and an extra set of bounds will be concatenated to ``bounds`` + corresponding to the domain bounds west of the prime meridian. For + instance, a grid cell spanning -1 to 1, will be broken into a cell + from 0 to 1 and 359 to 360 (or -1 to 0). + + Parameters + ---------- + bounds : xr.DataArray + The longitude domain bounds with prime meridian cell. + p_meridian_index : np.ndarray + The index of the prime meridian cell. + + Returns + ------- + xr.DataArray + The longitude domain bounds with split prime meridian cell. + + Raises + ------ + ValueError + If longitude bounds are inclusively between 0 and 360. + """ + _if_multidim_dask_array_then_load(bounds) + + # Example array: [[359, 1], [1, 90], [90, 180], [180, 359]] + # Reorient bound to span across zero (i.e., [359, 1] -> [-1, 1]). + # Result: [[-1, 1], [1, 90], [90, 180], [180, 359]] + bounds[p_meridian_index, 0] = bounds[p_meridian_index, 0] - 360.0 + + # Extend the array to nlon+1 by concatenating the grid cell that + # spans the prime meridian to the end. + # Result: [[-1, 1], [1, 90], [90, 180], [180, 359], [-1, 1]] + dim = get_dim_keys(bounds, "X") + bounds = xr.concat((bounds, bounds[p_meridian_index, :]), dim=dim) + + # Add an equivalent bound that spans 360 + # (i.e., [-1, 1] -> [359, 361]) to the end of the array. + # Result: [[-1, 1], [1, 90], [90, 180], [180, 359], [359, 361]] + repeat_bound = bounds[p_meridian_index, :][0] + 360.0 + bounds[-1, :] = repeat_bound + + # Update the lower-most min and upper-most max bounds to [0, 360]. + # Result: [[0, 1], [1, 90], [90, 180], [180, 359], [359, 360]] + bounds[p_meridian_index, 0], bounds[-1, 1] = (0.0, 360.0) + + return bounds + + +def _get_prime_meridian_index(lon_bounds: xr.DataArray) -> np.ndarray | None: + """Gets the index of the prime meridian cell in the longitude bounds. + + A prime meridian cell can exist when converting the axis orientation + from [-180, 180) to [0, 360). + + Parameters + ---------- + lon_bounds : xr.DataArray + The longitude bounds. + + Returns + ------- + np.ndarray | None + An array with a single element representing the index of the prime + meridian index if it exists. Otherwise, None if the cell does not exist. + + Raises + ------ + ValueError + If more than one grid cell spans the prime meridian. + """ + p_meridian_index = np.where(lon_bounds[:, 1] - lon_bounds[:, 0] < 0)[0] + + if p_meridian_index.size == 0: # pragma: no cover + return None + elif p_meridian_index.size > 1: # pragma: no cover + raise ValueError("More than one grid cell spans the prime meridian.") + + return p_meridian_index From f9a75b7e4c64677193f70cf8de779520f5a941a8 Mon Sep 17 00:00:00 2001 From: tomvothecoder Date: Wed, 24 Sep 2025 09:23:51 -0700 Subject: [PATCH 5/5] Fix variable name in `test_axis.py` --- tests/test_axis.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_axis.py b/tests/test_axis.py index 98c1ca06f..573067eb0 100644 --- a/tests/test_axis.py +++ b/tests/test_axis.py @@ -941,7 +941,7 @@ def test_swaps_multiple_dims_from_180_to_360_and_normalizes_prime_meridian_cell_ attrs={"test_attr": "test"}, ), "ts2": xr.DataArray( - name="ts", + name="ts2", data=np.array([3, 4, 0, 1, 2]), dims=["zlon"], attrs={"test_attr": "test"},