add longitude wrap functions (#270)

CHANGELOG.rst

@@ -61,6 +61,9 @@ New Features
    - Added functions to calculate the weighted global mean (`#220
      <https://github.com/MESMER-group/mesmer/pull/220>`_). By `Mathias Hauser
      <https://github.com/mathause>`_.
+   - Added functions to wrap arrays to [-180, 180) and [0, 360), respectively (`#270
+     <https://github.com/MESMER-group/mesmer/pull/270>`_). By `Mathias Hauser
+     <https://github.com/mathause>`_.
 
 - The aerosol data is now automatically downloaded using `pooch <https://www.fatiando.org/pooch/latest/>`__.
   (`#267 <https://github.com/MESMER-group/mesmer/pull/267>`_). By `Mathias Hauser

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

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,
     *,

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)