add rmse in linear trend (#44)

.cruft.json

@@ -1,6 +1,6 @@
 {
   "template": "https://github.com/ecmwf-projects/cookiecutter-conda-package",
-  "commit": "c6665306749b5dd3b4ec0fdcf1cb31d18fe23511",
+  "commit": "d7af7236036e91f530260569aa5b232ed4ffbdd2",
   "checkout": null,
   "context": {
     "cookiecutter": {

.github/workflows/on-push.yml

@@ -23,7 +23,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: 3.x
     - uses: pre-commit/[email protected]
@@ -33,7 +33,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: 3.x
     - name: Install conda-merge
@@ -44,7 +44,7 @@ jobs:
         for SUFFIX in ci integration; do
         conda-merge ci/environment-$SUFFIX.yml environment.yml > ci/combined-environment-$SUFFIX.yml || exit
         done
-    - uses: actions/upload-artifact@v3
+    - uses: actions/upload-artifact@v4
       with:
         name: combined-environments
         path: ci/combined-environment-*.yml
@@ -59,7 +59,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/download-artifact@v3
+    - uses: actions/download-artifact@v4
       with:
         name: combined-environments
         path: ci
@@ -88,7 +88,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/download-artifact@v3
+    - uses: actions/download-artifact@v4
       with:
         name: combined-environments
         path: ci
@@ -117,7 +117,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/download-artifact@v3
+    - uses: actions/download-artifact@v4
       with:
         name: combined-environments
         path: ci
@@ -154,7 +154,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/download-artifact@v3
+    - uses: actions/download-artifact@v4
       with:
         name: combined-environments
         path: ci
@@ -189,7 +189,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: '3.11'
     - name: Install package
@@ -206,7 +206,7 @@ jobs:
         python -m twine check --strict * || exit
         python -c "import c3s_eqc_automatic_quality_control" || exit
         cd ..
-    - uses: actions/upload-artifact@v3
+    - uses: actions/upload-artifact@v4
       with:
         name: distribution
         path: dist
@@ -226,7 +226,7 @@ jobs:
       id-token: write  # IMPORTANT: this permission is mandatory for trusted publish
 
     steps:
-    - uses: actions/download-artifact@v3
+    - uses: actions/download-artifact@v4
       with:
         name: distribution
         path: dist

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
   - id: blackdoc
     additional_dependencies: [black==23.11.0]
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.1.6
+  rev: v0.1.14
   hooks:
   - id: ruff
     args: [--fix, --show-fixes]
@@ -27,7 +27,7 @@ repos:
   hooks:
   - id: mdformat
 - repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
-  rev: v2.11.0
+  rev: v2.12.0
   hooks:
   - id: pretty-format-yaml
     args: [--autofix, --preserve-quotes]

c3s_eqc_automatic_quality_control/_time_weighted.py

@@ -86,8 +86,8 @@ def polyfit(self, **kwargs: Any) -> xr.DataArray | xr.Dataset:
         return self.obj.polyfit(**kwargs)
 
     def linear_trend(
-        self, p_value: bool = False, **kwargs: Any
-    ) -> tuple[xr.DataArray | xr.Dataset, ...]:
+        self, p_value: bool = False, rmse: bool = False, **kwargs: Any
+    ) -> dict[str, xr.DataArray | xr.Dataset]:
         coeff = self.polyfit(deg=1, **kwargs)
         if isinstance(self.obj, xr.DataArray):
             coeff = coeff["polyfit_coefficients"]
@@ -107,20 +107,23 @@ def linear_trend(
                 }
             )
         obj_trend = coeff.sel(degree=1, drop=True)
-        if not p_value:
-            return (obj_trend,)
+        output = {"linear_trend": obj_trend}
+        if not (p_value or rmse):
+            return output
 
         dim = kwargs.get("dim", self.time.name)
         weights = kwargs.get("w", self.obj_weighted.weights if self.weights else None)
         fit = xr.polyval(self.obj[dim], coeff)
-        obj_p_value = xs.pearson_r_p_value(
-            self.obj,
-            fit,
-            dim=dim,
-            weights=weights,
+        xs_kwargs = {
+            "dim": dim,
+            "weights": weights,
             **{k: v for k, v in kwargs.items() if k == "skipna"},
-        )
-        return (obj_trend, obj_p_value)
+        }
+        if p_value:
+            output["p_value"] = xs.pearson_r_p_value(self.obj, fit, **xs_kwargs)
+        if rmse:
+            output["rmse"] = xs.rmse(self.obj, fit, **xs_kwargs)
+        return output
 
     def coverage(self, **kwargs: Any) -> xr.DataArray | xr.Dataset:
         if "dim" not in kwargs:

c3s_eqc_automatic_quality_control/diagnostics.py

@@ -105,12 +105,9 @@ def time_weighted_linear_trend(
     time_name: Hashable | None = None,
     weights: xr.DataArray | bool = True,
     p_value: bool = False,
+    rmse: bool = False,
     **kwargs: Any,
-) -> (
-    xr.DataArray
-    | xr.Dataset
-    | tuple[xr.DataArray | xr.Dataset, xr.DataArray | xr.Dataset]
-):
+) -> xr.DataArray | xr.Dataset | dict[str, xr.DataArray | xr.Dataset]:
     """
     Calculate time weighted linear trend.
 
@@ -127,25 +124,27 @@ def time_weighted_linear_trend(
         - DataArray: custom weights
     p_value: bool, default: False
         Whether to compute 2-tailed Pearson p-value
+    rmse: bool, default: False
+        Whether to compute rmse
 
     Returns
     -------
-    DataArray or Dataset or tuple
-        Reduced object or tuple (trend, p_value)
+    DataArray or Dataset or dict
+        Reduced object or dict (linear_trend, p_value, rmse)
     """
-    obj_tuple = _time_weighted.TimeWeighted(obj, time_name, weights).linear_trend(
-        p_value=p_value, **kwargs
+    output = _time_weighted.TimeWeighted(obj, time_name, weights).linear_trend(
+        p_value=p_value, rmse=rmse, **kwargs
     )
-    trend = obj_tuple[0] * 1.0e9  # 1/ns to 1/s
+    output["linear_trend"] *= 1.0e9  # 1/ns to 1/s
 
     def attrs_func(attrs: dict[str, Any]) -> dict[str, Any]:
         return {
             "long_name": f"Linear trend of {attrs.get('long_name', '')}",
             "units": f"{attrs.get('units', '')} s-1",
         }
 
-    trend = _apply_attrs_func(trend, obj, attrs_func)
-    return trend if not p_value else (trend, obj_tuple[1])
+    output["linear_trend"] = _apply_attrs_func(output["linear_trend"], obj, attrs_func)
+    return output["linear_trend"] if not (p_value or rmse) else output
 
 
 def time_weighted_mean(

c3s_eqc_automatic_quality_control/utils.py

@@ -81,7 +81,7 @@ def regionalise(
         return obj.where(mask.compute(), drop=True)
 
     # Convert longitude
-    lon_limits = xr.DataArray([lon_slice.start, lon_slice.stop], dims=lon_name)
+    lon_limits = xr.DataArray([lon_slice.start, lon_slice.stop], dims=(lon_name,))
     lon_limits = lon_limits.dropna(lon_name)
     sort = False
     if lon_limits.min() < 0 and obj[lon_name].max() > 180:
@@ -99,7 +99,9 @@ def regionalise(
         bounds = obj[name][[0, -1]]
         ascending_bounds = bool(bounds.diff(name) > 0)
         ascending_slice = bool(
-            xr.DataArray([slice.start, slice.stop], dims=name).fillna(bounds).diff(name)
+            xr.DataArray([slice.start, slice.stop], dims=(name,))
+            .fillna(bounds)
+            .diff(name)
             > 0
         )
         if ascending_bounds is not ascending_slice:

tests/test_21_time_diagnostics.py

@@ -158,12 +158,13 @@ def test_time_weighted_linear_trend(
                 ds_trend.rename(Tair_polyfit_coefficients="Tair"), actual
             )
 
-    def test_time_weighted_p_value(
+    def test_time_weighted_linear_trend_stats(
         self, obj: xr.DataArray | xr.Dataset, weights: bool
     ) -> None:
-        actual_tuple = diagnostics.time_weighted_linear_trend(
-            obj, weights=weights, p_value=True
+        actual_dict = diagnostics.time_weighted_linear_trend(
+            obj, weights=weights, p_value=True, rmse=True
         )
+        assert isinstance(actual_dict, dict)
 
         da_weights = obj["time"].dt.days_in_month if weights else None
         coeff = obj.polyfit(dim="time", deg=1, w=da_weights)
@@ -177,12 +178,16 @@ def test_time_weighted_p_value(
                 }
             )
         fit = xr.polyval(obj["time"], coeff)
-        expected_tuple = (
-            diagnostics.time_weighted_linear_trend(obj, weights=weights),
-            xs.pearson_r_p_value(obj, fit, "time", weights=da_weights),
-        )
-        for actual, expected in zip(actual_tuple, expected_tuple):
-            xr.testing.assert_identical(actual, expected)
+        expected_dict = {
+            "linear_trend": diagnostics.time_weighted_linear_trend(
+                obj, weights=weights
+            ),
+            "p_value": xs.pearson_r_p_value(obj, fit, "time", weights=da_weights),
+            "rmse": xs.rmse(obj, fit, "time", weights=da_weights),
+        }
+        assert set(actual_dict) == set(expected_dict)
+        for key in actual_dict:
+            xr.testing.assert_identical(actual_dict[key], actual_dict[key])
 
     def test_time_weighted_coverage(
         self, obj: xr.DataArray | xr.Dataset, weights: bool