From 0232bd5bf1d0d83dbb2127cd0e6c26eeb44b4020 Mon Sep 17 00:00:00 2001 From: Mathias Hauser Date: Fri, 1 Sep 2023 16:46:02 +0200 Subject: [PATCH] add longitude wrap functions (#270) --- CHANGELOG.rst | 3 + docs/source/api.rst | 2 + mesmer/xarray_utils/grid.py | 70 ++++++++++++++++++++ tests/unit/test_grid_wrap.py | 125 +++++++++++++++++++++++++++++++++++ 4 files changed, 200 insertions(+) create mode 100644 tests/unit/test_grid_wrap.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5f5b8850..c1f09e84 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -61,6 +61,9 @@ New Features - Added functions to calculate the weighted global mean (`#220 `_). By `Mathias Hauser `_. + - Added functions to wrap arrays to [-180, 180) and [0, 360), respectively (`#270 + `_). By `Mathias Hauser + `_. - The aerosol data is now automatically downloaded using `pooch `__. (`#267 `_). By `Mathias Hauser diff --git a/docs/source/api.rst b/docs/source/api.rst index 04ad3132..6fff1c77 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -68,6 +68,8 @@ Data manipulation .. autosummary:: :toctree: generated/ + ~xarray_utils.grid.wrap_to_180 + ~xarray_utils.grid.wrap_to_360 ~xarray_utils.grid.stack_lat_lon ~xarray_utils.grid.unstack_lat_lon_and_align ~xarray_utils.grid.unstack_lat_lon diff --git a/mesmer/xarray_utils/grid.py b/mesmer/xarray_utils/grid.py index ce333397..55e45729 100644 --- a/mesmer/xarray_utils/grid.py +++ b/mesmer/xarray_utils/grid.py @@ -3,6 +3,76 @@ from packaging.version import Version +def _lon_to_180(lon): + + with xr.set_options(keep_attrs=True): + lon = ((lon + 180) % 360) - 180 + + if isinstance(lon, xr.DataArray): + lon = lon.assign_coords({lon.name: lon}) + + return lon + + +def _lon_to_360(lon): + + with xr.set_options(keep_attrs=True): + lon = lon % 360 + + if isinstance(lon, xr.DataArray): + lon = lon.assign_coords({lon.name: lon}) + + return lon + + +def wrap_to_180(obj, lon_name="lon"): + """ + wrap longitude coordinates to [-180..180) + + Parameters + ---------- + obj : xr.Dataset or xr.DataArray + object with longitude coordinates + lon : str, default: "lon" + name of the longitude ('lon', 'longitude', ...) + + Returns + ------- + wrapped : Dataset + Another dataset array wrapped around. + """ + + new_lon = _lon_to_180(obj[lon_name]) + + obj = obj.assign_coords(**{lon_name: new_lon}) + obj = obj.sortby(lon_name) + + return obj + + +def wrap_to_360(obj, lon_name="lon"): + """ + wrap longitude coordinates to [0..360) + Parameters + ---------- + obj : xr.Dataset or xr.DataArray + object with longitude coordinates + lon : str, default: "lon" + name of the longitude ('lon', 'longitude', ...) + Returns + ------- + wrapped : Dataset + Another dataset array wrapped around. + """ + + new_lon = _lon_to_360(obj[lon_name]) + + obj = obj.assign_coords(**{lon_name: new_lon}) + obj = obj.sortby(lon_name) + + return obj + + def stack_lat_lon( data, *, diff --git a/tests/unit/test_grid_wrap.py b/tests/unit/test_grid_wrap.py new file mode 100644 index 00000000..efaa35ac --- /dev/null +++ b/tests/unit/test_grid_wrap.py @@ -0,0 +1,125 @@ +import numpy as np +import pytest +import xarray as xr + +from mesmer import xarray_utils as xru + + +def test_lon_to_180(): + + arr = np.array([-180.1, -180, -1, 0, 179.99, 180, 179 + 2 * 360]) + + expected = np.array([179.9, -180, -1, 0, 179.99, -180, 179]) + + result = xru.grid._lon_to_180(arr) + np.testing.assert_allclose(result, expected) + + # ensure arr is not updated in-place + assert not (arr == result).all() + + attrs = {"name": "test"} + da = xr.DataArray(arr, dims="lon", coords={"lon": arr}, attrs=attrs, name="lon") + expected = xr.DataArray( + expected, dims="lon", coords={"lon": expected}, attrs=attrs, name="lon" + ) + + result = xru.grid._lon_to_180(da) + + xr.testing.assert_allclose(result, expected) + + assert result.attrs == expected.attrs + + +def test_lon_to_360(): + + arr = np.array([-180.1, -180, -1, 0, 179.99, 180, 179 + 2 * 360, 259.9, 360]) + + expected = np.array([179.9, 180, 359, 0, 179.99, 180, 179, 259.9, 0]) + + result = xru.grid._lon_to_360(arr) + np.testing.assert_allclose(result, expected) + + # ensure arr is not updated in-place + assert not (arr == result).all() + + attrs = {"name": "test"} + da = xr.DataArray(arr, dims="lon", coords={"lon": arr}, attrs=attrs, name="lon") + expected = xr.DataArray( + expected, dims="lon", coords={"lon": expected}, attrs=attrs, name="lon" + ) + + result = xru.grid._lon_to_360(da) + + xr.testing.assert_allclose(result, expected) + + assert result.attrs == expected.attrs + + +@pytest.mark.parametrize("as_dataset", (True, False)) +def test_wrap180(as_dataset): + + attrs = {"name": "test"} + obj = xr.DataArray( + [0, 1, 2, 3, 4], + dims="lon", + coords={"lon": [-1, 1, 179, 180, 360]}, + name="data", + attrs=attrs, + ) + obj.lon.attrs = {"coord": "attrs"} + expected = xr.DataArray( + [3, 0, 4, 1, 2], + dims="lon", + coords={"lon": [-180, -1, 0, 1, 179]}, + name="data", + attrs=attrs, + ) + expected.lon.attrs = {"coord": "attrs"} + + if as_dataset: + obj = obj.to_dataset() + expected = expected.to_dataset() + + result = xru.grid.wrap_to_180(obj) + + obj = obj.rename(lon="longitude") + expected = expected.rename(lon="longitude") + + result = xru.grid.wrap_to_180(obj, lon_name="longitude") + + xr.testing.assert_identical(result, expected) + + +@pytest.mark.parametrize("as_dataset", (True, False)) +def test_wrap360(as_dataset): + + attrs = {"name": "test"} + obj = xr.DataArray( + [0, 1, 2, 3, 4], + dims="lon", + coords={"lon": [-5, 1, 180, 359, 360]}, + name="data", + attrs=attrs, + ) + obj.lon.attrs = {"coord": "attrs"} + expected = xr.DataArray( + [4, 1, 2, 0, 3], + dims="lon", + coords={"lon": [0, 1, 180, 355, 359]}, + name="data", + attrs=attrs, + ) + expected.lon.attrs = {"coord": "attrs"} + + if as_dataset: + obj = obj.to_dataset() + expected = expected.to_dataset() + + result = xru.grid.wrap_to_360(obj) + + obj = obj.rename(lon="longitude") + expected = expected.rename(lon="longitude") + + result = xru.grid.wrap_to_360(obj, lon_name="longitude") + + xr.testing.assert_identical(result, expected)