Refactor abs_error_stats, fix zero-index bug

pred_steps and ref_cut were always equal at every call site, making
using two parameters redundant. ref_cut was also a misleading name.
Collapse both into a single steps parameter and replace negative-index
slicing with (err[n-steps:], err[:n-steps]), which correctly handles
steps=0 (full ref, empty pred) — fixing the err[-0:] bug. Change comment
wording.

Author: isaacaka

Notebooks/Jumper.ipynb

@@ -600,8 +600,7 @@
     "for spec in TERM_SPECS:\n",
     "    error_stats[spec.key] = abs_error_stats(\n",
     "        errs[spec.key],\n",
-    "        pred_steps=steps,\n",
-    "        ref_cut=steps,  # baseline: all but the forecast window\n",
+    "        steps=steps,\n",
     "        axes=spec.err_axes,\n",
     "    )\n",
     "\n",

src/nemo_spinup_forecast/pipeline_utils.py

@@ -133,8 +133,8 @@ def compute_rmse_for_terms(
     sims : Mapping[str, Simulation]
         Prepared and decomposed simulations keyed by :attr:`TermSpec.key`.
     n_components : int or None, default=None
-        Number of components to use for reconstruction.  When ``None``,
-        all fitted components (``len(s.pca.components_)``) are used.
+        Number of components to use for reconstruction.
+        None uses all fitted components (``len(s.pca.components_)``).
 
     Returns
     -------
@@ -291,8 +291,7 @@ def forecast_all(
 def abs_error_stats(
     err: np.ndarray,
     *,
-    pred_steps: int,
-    ref_cut: int,
+    steps: int,
     axes: tuple[int, ...],
 ) -> dict[str, Any]:
     """Compute absolute-error summary statistics for prediction and reference windows.
@@ -301,11 +300,9 @@ def abs_error_stats(
     ----------
     err : np.ndarray
         Absolute-error array, typically ``abs(reference - prediction)``.
-    pred_steps : int
-        Number of trailing time steps considered as the prediction window.
-    ref_cut : int
-        Number of trailing time steps excluded from the reference window.
-        If set to ``0``, the full ``err`` array is used as the reference.
+    steps : int
+        Number of time steps at the end of ``err`` used as the forecast window;
+        the rest of ``err``is the reference.
     axes : tuple[int, ...]
         Axes reduced with ``nanmean`` and ``nanstd``.
 
@@ -314,13 +311,10 @@ def abs_error_stats(
     dict[str, Any]
         Dictionary with keys ``pred_mean``, ``pred_std``, ``ref_mean``,
         and ``ref_std``.
-
-    Notes
-    -----
-    This function prints the prediction and reference window shapes.
     """
-    pred = err[-pred_steps:]
-    ref = err[:-ref_cut] if ref_cut else err
+    n = len(err)
+    pred = err[n - steps :]
+    ref = err[: n - steps]
     print("pred shape:", pred.shape)
     print("ref shape:", ref.shape)
     return {

tests/test_pipeline_utils.py

@@ -0,0 +1,34 @@
+import numpy as np
+import pytest
+
+from nemo_spinup_forecast.pipeline_utils import abs_error_stats
+
+
+def test_3d_spatial_reduction_shape_and_values():
+    """Reducing over spatial axes (y, x) leaves a time-indexed output."""
+    err = np.zeros((8, 3, 4))  # (time, y, x)
+    err[-3:] = 10.0  # pred window: last 3 time steps
+    err[:-3] = 1.0  # ref window: first 5 time steps
+    result = abs_error_stats(err, steps=3, axes=(1, 2))
+    # pred = err[5:] → shape (3, 3, 4), reduce (y, x) → shape (3,)
+    # ref  = err[:5] → shape (5, 3, 4), reduce (y, x) → shape (5,)
+    assert result["pred_mean"].shape == (3,)
+    assert result["ref_mean"].shape == (5,)
+    np.testing.assert_allclose(result["pred_mean"], 10.0)
+    np.testing.assert_allclose(result["ref_mean"], 1.0)
+
+
+def test_steps_zero_gives_full_ref():
+    """steps=0 → pred is empty (nan), ref is the full array (baseline case)."""
+    err = np.ones((10, 4, 5))
+    err[7:] = 3.0
+    result = abs_error_stats(err, steps=0, axes=(0, 1, 2))
+    assert np.isnan(result["pred_mean"])
+    assert result["ref_mean"] == pytest.approx((7 * 1.0 + 3 * 3.0) / 10)
+
+
+def test_steps_equals_full_length():
+    """steps=len(err) → pred covers the entire time axis."""
+    err = np.arange(6.0).reshape(6, 1, 1)
+    result = abs_error_stats(err, steps=6, axes=(0, 1, 2))
+    assert result["pred_mean"] == pytest.approx(2.5)  # mean of 0..5