Update swap_lon_axis() to fix bounds for prime meridian cell
#791
Conversation
tomvothecoder
changed the title
swap_lon_axis() to drop extra cell to handle prime meridianswap_lon_axis() to drop extra cell for prime meridian
There was a problem hiding this comment.
Pull Request Overview
Updates the swap_lon_axis() function to properly handle longitude coordinates when converting from the [-180, 180) range to [0, 360) by dropping an extra cell created during prime meridian handling.
- Refactored
_swap_lon_bounds()to use a new approach that prevents creating an extra longitude cell - Added
_adjust_bounds_for_prime_meridian()function to handle prime meridian bounds adjustment and remove the extra cell - Moved
_get_bounds_dim()function frombounds.pytoaxis.pyand updated import references
Reviewed Changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| xcdat/temporal.py | Updated import to use _get_bounds_dim from axis module instead of bounds module |
| xcdat/bounds.py | Removed _get_bounds_dim function as it was moved to axis module |
| xcdat/axis.py | Added _get_bounds_dim function and new _adjust_bounds_for_prime_meridian function, refactored _swap_lon_bounds logic |
| tests/test_dataset.py | Updated test expectations to reflect the corrected longitude coordinate values and bounds after dropping extra cell |
| tests/test_bounds.py | Removed tests for _get_bounds_dim function since it was moved to axis module |
| tests/test_axis.py | Added tests for _get_bounds_dim function and updated test expectations for longitude swapping operations |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #791 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 16 16
Lines 1782 1764 -18
=========================================
- Hits 1782 1764 -18 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Hey @pochedls, this PR is ready for review. It removes the extra longitude point generated when handling prime meridian cells and rectifies the bounds for the prime meridian cell (after swapping from (-180, 180) to (0, 360)). I still need to fix the unit tests. Not ready yet (#791 (comment)).
|
Running into edge cases with the unit tests. This still needs more work. Converting to draft. |
tomvothecoder
changed the title
swap_lon_axis() to drop extra cell for prime meridianswap_lon_axis() to fix bounds for prime meridian cell
github-actions
bot
added
the
type: bug
- 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`
tomvothecoder
force-pushed
the
bug/790-swap-lon
branch
from
f5ebfdc to
68b7c1f
Compare
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
There was a problem hiding this comment.
Hey @pochedls, this PR is ready for review. Refer to the PR description for more detailed info, if needed.
The main change to focus on is in axis.py under def _normalize_wrap_bounds_0_to_360(), which resolves the incorrect bounds for prime meridian cells. All other code to handle prime meridian cell code was removed (e.g,. cell splitting which adds additional cell). Besides that, refactoring was performed to make the code cleaner.
Unit tests were updated to pass and the minimum example from #790 is now fixed.
Output
Notice that last bounds has been fixed to convert [359, 0] -> [359, 360].
Longitude Values (Last 5):
[355.5 356.5 357.5 358.5 359.5]
Longitude Bounds (Last 5):
[[355. 356.]
[356. 357.]
[357. 358.]
[358. 359.]
[359. 360.]]Script
import xcdat as xc
filename = "/global/cfs/cdirs/m4581/obs4MIPs/obs4MIPs_input/MOHC/HadISST1-1/v20231114/HadISST.1.1.sst.nc"
ds = xc.open_dataset(filename)
ds_swapped = xc.open_dataset(filename, lon_orient=(0, 360))
# Print longitude values
print("Longitude Values:")
print(ds_swapped.longitude)
print("\n")
# Print longitude bounds
print("Longitude Bounds:")
print(ds_swapped.longitude_bnds)| def _orient_lon_coords(coords: xr.DataArray, to: tuple[float, float]) -> xr.DataArray: | ||
| """ | ||
| Adjust longitude coordinates to match the specified orientation range. | ||
| if not ds[key].identical(new_bounds): | ||
| ds[key] = new_bounds | ||
| 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. | ||
| # 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. | ||
| if to == (0, 360): | ||
| p_meridian_index = _get_prime_meridian_index(ds[key]) | ||
| Parameters | ||
| ---------- | ||
| coords : xr.DataArray | ||
| The longitude coordinates to be reoriented. | ||
| to : tuple[float, float] | ||
| The orientation to swap the Dataset's longitude axis to. Supported | ||
| orientations include: | ||
| if p_meridian_index is not None: | ||
| ds = _align_lon_to_360(ds, ds[key], p_meridian_index) | ||
| * (-180, 180): represents [-180, 180) in math notation | ||
| * (0, 360): represents [0, 360) in math notation | ||
| return ds | ||
| Returns | ||
| ------- | ||
| xr.DataArray | ||
| The longitude coordinates reoriented to the target range, with the | ||
| original encoding preserved. | ||
| """ | ||
| with xr.set_options(keep_attrs=True): | ||
| if _is_in_orientation(coords, to): | ||
| out = coords | ||
| else: | ||
| out = _map_to_orientation(coords, to) | ||
|
|
||
| out.encoding = coords.encoding | ||
| return out |
There was a problem hiding this comment.
This function focuses on orienting coordinates (just refactored code).
| 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. | ||
| 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 | ||
| ---------- | ||
| coords : xr.DataArray | ||
| Coordinates on a longitude axis. | ||
| 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 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 bounds adjusted to the target orientation. | ||
| """ | ||
| 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(da_bounds, to): | ||
| out = da_bounds | ||
| else: | ||
| raise ValueError( | ||
| "Currently, only (-180, 180) and (0, 360) are supported longitude axis " | ||
| "orientations." | ||
| ) | ||
| out = _map_to_orientation(da_bounds, to) | ||
|
|
||
| new_coords.encoding = coords.encoding | ||
| # 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 new_coords | ||
| out.encoding = da_bounds.encoding | ||
|
|
||
| return out |
There was a problem hiding this comment.
This function focuses on orienting bounds (just refactored code).
| # 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) |
There was a problem hiding this comment.
This is the function call to normalize bounds to wrap 0 to 360. For example, [359, 0] -> [359, 360].
| 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. |
There was a problem hiding this comment.
Extracted helper function for _orient_lon_coords and _orient_lon_bounds (refactored code).
| def _map_to_orientation(da: xr.DataArray, to: tuple[float, float]) -> xr.DataArray: | ||
| """ | ||
| Map the values of a DataArray to a specified longitude orientation. |
There was a problem hiding this comment.
Extracted helper function for _orient_lon_coords and _orient_lon_bounds (refactored code).
| def test_swaps_single_dim_from_180_to_360_and_normalizes_prime_meridian_cell_in_lon_bnds_to_360( | ||
| self, | ||
| ): |
There was a problem hiding this comment.
Test to swap to (0, 360) and normalizes last prime meridian cell bounds [359, 0] -> [359, 360]
| def test_swaps_multiple_dims_from_180_to_360_and_normalizes_prime_meridian_cell_in_lon_bnds_to_360( | ||
| self, |
There was a problem hiding this comment.
Same test but with multiple longitude dimensions.
| 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( |
There was a problem hiding this comment.
Test update for open_dataset for changes to swap_lon_axis.
|
|
||
| xr.testing.assert_allclose(result, expected) | ||
|
|
||
| def test_dataset_weights_raises_error_when_more_than_one_grid_cell_spans_prime_meridian( |
There was a problem hiding this comment.
New test specifically for spatial average prime meridian cell logic.
|
Tom – this is kind of complicated to deal with. Consider the following case with longitude bounds that do not fall exactly on 0 or 360: # imports
import xcdat as xc
import xarray as xr
import numpy as np
# create plausible dataset with longitude bounds that do not align with 0 / 360
lon = np.arange(-179.75, 180, 1)
lat = np.arange(-89.75, 90, 1)
data = np.random.rand(len(lat), len(lon))
lat = xr.DataArray(lat, dims=['lat'], attrs={'units': 'degrees_north', 'long_name': 'latitude'})
lon = xr.DataArray(lon, dims=['lon'], attrs={'units': 'degrees_east', 'long_name': 'longitude'})
data = xr.DataArray(data, dims=['lat', 'lon'], coords={'lat': lat, 'lon': lon})
ds = data.to_dataset(name='tas')
# lets zoom into the prime meridian region
ds = ds.bounds.add_missing_bounds()
ds = ds.sel(lon=slice(-5, 5))
for i, l in enumerate(ds.lon.values):
print('longitude: ' + str(l) + ' + longitude bounds: ', ds.lon_bnds[i, :].values)
print()
# call the swap function
ds_swapped = xc.swap_lon_axis(ds, to=(0, 360))
for i, l in enumerate(ds_swapped.lon.values):
print('longitude: ' + str(l) + ' + longitude bounds: ', ds_swapped.lon_bnds[i, :].values)These are the original longitude values (and their bounds). Note that the prime meridian grid cell is centered on
But after being swapped, the right bound (
I think we want this bound to be represented either as Maybe we could have a quick chat about this issue? |
Description
This pull request refactors and expands the longitude axis swapping tests, improves coverage for longitude bounds normalization (especially around the prime meridian), and moves the
_get_bounds_dimhelper and its tests to more appropriate locations. The changes also improve test clarity and maintainability by renaming and restructuring test cases.Closes #790
Longitude axis swapping and bounds normalization:
Bounds dimension helper refactor:
_get_bounds_dimhelper function import fromxcdat.boundstoxcdat.axis, and relocated its associated tests fromtests/test_bounds.pytotests/test_axis.py, aligning them with their usage and improving code organization. [1] [2] [3] [4]These changes enhance the reliability of longitude axis operations and clarify the codebase by grouping related functionality and tests.
Checklist
If applicable: