lowess: allow missing input_core_dim (#283)

* lowess: allow missing input_core_dim

* Apply suggestions from code review

* CHANGELOG

CHANGELOG.rst

@@ -48,7 +48,8 @@ New Features
 
 - Other refactorings:
    - Extract the LOWESS smoothing for xarray objects: :py:func:`mesmer.stats.smoothing.lowess`.
-     (`#193 <https://github.com/MESMER-group/mesmer/pull/193>`_).
+     (`#193 <https://github.com/MESMER-group/mesmer/pull/193>`_ and `#283
+     <https://github.com/MESMER-group/mesmer/pull/283>`_).
      By `Mathias Hauser <https://github.com/mathause>`_.
 
 - Added helper functions to process xarray-based model data:

mesmer/stats/smoothing.py

@@ -9,17 +9,17 @@ def lowess(data, dim, *, frac, use_coords_as_x=False, it=0):
 
     Parameters
     ----------
-    data : xr.DataArray
+    data : xr.DataArray | xr.Dataset
         Data to smooth (y-values).
     dim : str
         Dimension along which to smooth (x-dimension)
     frac : float
         Between 0 and 1. The fraction of the data used when estimating each y-value.
     use_coords_as_x : boolean, default: False
-        If True uses ``data[dim]`` as x-values else uses ``np.arange(data[dim.size)``
+        If True uses ``data[dim]`` as x-values else uses ``np.arange(data[dim].size)``
         (useful if ``dim`` are time coordinates).
     it : int, default: 0
-        The number of residual-based reweightings to perform.
+        The number of residual-based re-weightings to perform.
     """
 
     from statsmodels.nonparametric.smoothers_lowess import lowess
@@ -36,14 +36,25 @@ def lowess(data, dim, *, frac, use_coords_as_x=False, it=0):
         x = xr.ones_like(coords)
         x.data = np.arange(coords.size)
 
-    out = xr.apply_ufunc(
-        lowess,
-        data,
-        x,
-        input_core_dims=[[dim], [dim]],
-        output_core_dims=[[dim]],
-        vectorize=True,
-        kwargs={"frac": frac, "it": it, "return_sorted": False},
-    )
-
-    return out
+    def _lowess(da: xr.DataArray) -> xr.DataArray:
+
+        # skip data var if input_core_dim is missing
+        if dim not in da.dims:
+            return da
+
+        out = xr.apply_ufunc(
+            lowess,
+            da,
+            x,
+            input_core_dims=[[dim], [dim]],
+            output_core_dims=[[dim]],
+            vectorize=True,
+            kwargs={"frac": frac, "it": it, "return_sorted": False},
+        )
+
+        return out
+
+    if isinstance(data, xr.Dataset):
+        return data.map(_lowess)
+
+    return _lowess(data)

tests/unit/test_smoothing.py

@@ -50,6 +50,27 @@ def test_lowess_dataset():
     xr.testing.assert_allclose(result, expected)
 
 
+def test_lowess_dataset_missing_core_dims():
+
+    data = trend_data_1D()
+    da1 = xr.DataArray(1, name="extra1")
+    da2 = xr.DataArray([3, 2, 1], dims="y", name="perpendicular")
+
+    ds = xr.merge([data, da1, da2])
+
+    result = mesmer.stats.smoothing.lowess(ds, "time", frac=0.3)
+
+    expected = lowess(
+        data.values, data.time.values, frac=0.3, it=0, return_sorted=False
+    )
+    expected = xr.DataArray(
+        expected, dims="time", coords={"time": data.time}, name="data"
+    )
+    expected = xr.merge([expected, da1, da2])
+
+    xr.testing.assert_allclose(result, expected)
+
+
 def test_lowess_2D():
     data = trend_data_2D()