Let fit_harmonic_model return residuals instead of predictions (#574)

* let fit_harmonic_model return residuals instead of predictions

* update calibration files (drop residuals before saving)

* CHANGELOG

* adjust example

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: veni-vidi-vici-dormivi

CHANGELOG.rst

@@ -178,7 +178,7 @@ Harmonic model
    `#424 <https://github.com/MESMER-group/mesmer/pull/424>`_,
    `#433 <https://github.com/MESMER-group/mesmer/pull/433>`_,
    `#512 <https://github.com/MESMER-group/mesmer/pull/512>`_, and
-   `#512 <https://github.com/MESMER-group/mesmer/pull/512>`_).
+   `#574 <https://github.com/MESMER-group/mesmer/pull/574>`_).
 - add tests (
   `#431 <https://github.com/MESMER-group/mesmer/pull/431>`_, and
   `#458 <https://github.com/MESMER-group/mesmer/pull/458>`_)

examples/example_mesmer_m.ipynb

@@ -12,9 +12,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'0.10.0.post1.dev143+gfc10078.d20241008'"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "import importlib\n",
     "\n",
@@ -41,7 +52,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -52,7 +63,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -80,7 +91,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -120,7 +131,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -140,7 +151,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -153,7 +164,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -174,7 +185,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -196,17 +207,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
-    "resids_after_hm = tas_stacked_m - harmonic_model_fit.predictions\n",
-    "\n",
     "pt_coefficients = mesmer.stats.fit_yeo_johnson_transform(\n",
-    "    tas_stacked_y.tas, resids_after_hm.tas\n",
+    "    tas_stacked_y.tas, harmonic_model_fit.residuals\n",
     ")\n",
     "transformed_hm_resids = mesmer.stats.yeo_johnson_transform(\n",
-    "    tas_stacked_y.tas, resids_after_hm.tas, pt_coefficients\n",
+    "    tas_stacked_y.tas, harmonic_model_fit.residuals, pt_coefficients\n",
     ")"
    ]
   },
@@ -221,7 +230,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -241,7 +250,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -282,7 +291,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -309,7 +318,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -329,7 +338,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -355,7 +364,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -372,7 +381,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -383,7 +392,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -401,7 +410,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -413,7 +422,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -430,7 +439,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -460,7 +469,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -477,7 +486,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 23,
    "metadata": {},
    "outputs": [],
    "source": [

mesmer/stats/_harmonic_model.py

@@ -265,9 +265,9 @@ def fit_harmonic_model(yearly_predictor, monthly_target, max_order=6, time_dim="
     Returns
     -------
     data_vars : `xr.Dataset`
-        Dataset containing the selected order of Fourier Series (selected_order),
-        the estimated coefficients of the Fourier Series (coeffs) and the resulting
-        predictions for monthly values (predictions).
+        Dataset containing the selected order of Fourier Series (`selected_order`),
+        the estimated coefficients of the Fourier Series (`coeffs`) and the resulting
+        residuals of the model (`residuals`).
 
     """
 
@@ -300,12 +300,12 @@ def fit_harmonic_model(yearly_predictor, monthly_target, max_order=6, time_dim="
 
     coeffs = coeffs.assign_coords({"coeff": np.arange(coeffs.sizes["coeff"])})
 
-    preds = yearly_predictor + preds
+    resids = monthly_target - (yearly_predictor + preds)
 
     data_vars = {
         "selected_order": selected_order,
         "coeffs": coeffs,
-        "predictions": preds.transpose(time_dim, ...),
+        "residuals": resids.transpose(time_dim, ...),
     }
 
     return xr.Dataset(data_vars)

tests/integration/test_calibrate_mesmer_m.py

@@ -89,13 +89,12 @@ def mask_and_stack(ds, threshold_land):
     )
 
     # train power transformer
-    resids_after_hm = tas_stacked_m - harmonic_model_fit.predictions
     pt_coefficients = mesmer.stats.fit_yeo_johnson_transform(
         tas_stacked_y.tas,
-        resids_after_hm.tas,
+        harmonic_model_fit.residuals,
     )
     transformed_hm_resids = mesmer.stats.yeo_johnson_transform(
-        tas_stacked_y.tas, resids_after_hm.tas, pt_coefficients
+        tas_stacked_y.tas, harmonic_model_fit.residuals, pt_coefficients
     )
 
     # fit cyclo-stationary AR(1) process
@@ -122,6 +121,11 @@ def mask_and_stack(ds, threshold_land):
 
     # save params
     if update_expected_files:
+        # drop unnecessary variables
+        harmonic_model_fit = harmonic_model_fit.drop_vars(["residuals", "time"])
+        AR1_fit = AR1_fit.drop_vars(["residuals", "time"])
+
+        # save
         harmonic_model_fit.to_netcdf(
             TEST_PATH
             / "harmonic_model"

tests/test-data/output/tas/mon/test-params/harmonic_model/params_harmonic_model_tas_IPSL-CM6A-LR_ssp585.nc

@@ -158,7 +158,9 @@ def test_fit_harmonic_model():
     # test if the model can recover the monthly target from perfect fourier series
     result = mesmer.stats.fit_harmonic_model(yearly_predictor, monthly_target)
     np.testing.assert_equal(result.selected_order.values, orders)
-    xr.testing.assert_allclose(result["predictions"], monthly_target)
+    xr.testing.assert_allclose(
+        result.residuals, xr.zeros_like(monthly_target), atol=1e-6
+    )
 
     # test if the model can recover the underlying cycle with noise on top of monthly target
     rng = np.random.default_rng(0)
@@ -167,35 +169,36 @@ def test_fit_harmonic_model():
     )
 
     result = mesmer.stats.fit_harmonic_model(yearly_predictor, noisy_monthly_target)
-    xr.testing.assert_allclose(result["predictions"], monthly_target, atol=0.1)
+    predictions = mesmer.stats.predict_harmonic_model(
+        yearly_predictor, result.coeffs, time=monthly_time
+    )
+    xr.testing.assert_allclose(predictions, monthly_target, atol=0.1)
 
     # compare numerically one cell of one year
     expected = np.array(
         [
-            7.324277,
-            9.966644,
-            9.972146,
-            7.33931,
-            2.7736,
-            -2.501604,
-            -7.072816,
-            -9.715184,
-            -9.720686,
-            -7.087849,
-            -2.52214,
-            2.753065,
+            0.014026,
+            0.131156,
+            -0.232648,
+            0.040157,
+            0.088749,
+            -0.102724,
+            -0.066836,
+            0.133832,
+            0.180308,
+            -0.12783,
+            -0.042045,
+            0.160109,
         ]
     )
 
-    result_comp = result.predictions.isel(cells=0, time=slice(0, 12)).values
+    result_comp = result.residuals.isel(cells=0, time=slice(0, 12)).values
     np.testing.assert_allclose(result_comp, expected, atol=1e-6)
 
-    # ensure coeffs and predictions are consistent
-    expected = mesmer.stats.predict_harmonic_model(
-        yearly_predictor, result.coeffs, result.time
-    )
+    # ensure predictons and residuals are consistent
+    expected = noisy_monthly_target - predictions
 
-    xr.testing.assert_equal(expected, result.predictions)
+    xr.testing.assert_equal(expected, result.residuals)
 
 
 def test_fit_harmonic_model_checks():