implement returning of residuals

mesmer/stats/__init__.py

@@ -6,7 +6,6 @@
     draw_auto_regression_uncorrelated,
     fit_auto_regression,
     fit_auto_regression_monthly,
-    predict_auto_regression_monthly,
     select_ar_order,
 )
 from mesmer.stats._gaspari_cohn import gaspari_cohn, gaspari_cohn_correlation_matrices
@@ -26,7 +25,6 @@
     "fit_auto_regression",
     "select_ar_order",
     "fit_auto_regression_monthly",
-    "predict_auto_regression_monthly",
     "draw_auto_regression_monthly",
     # gaspari cohn
     "gaspari_cohn_correlation_matrices",

mesmer/stats/_auto_regression.py

@@ -818,13 +818,10 @@ def _draw_ar_corr_monthly_xr_internal(
     # make sure non-dimension coords are properly caught
     gridpoint_coords = dict(slope[gridpoint_dim].coords)
 
-    if covariance is not None:
-        covariance = covariance.values
-
     out = _draw_auto_regression_monthly_np(
         intercept=intercept.values,
         slope=slope.transpose(..., gridpoint_dim).values,
-        covariance=covariance,
+        covariance=covariance.values,
         n_samples=n_realisations,
         n_ts=n_ts,
         seed=seed,
@@ -874,7 +871,7 @@ def _draw_auto_regression_monthly_np(
         Predicted time series of the specified AR(1) including spatially correllated innovations.
     """
     intercept = np.asarray(intercept)
-    # covariance = np.atleast_3d(covariance)
+    covariance = np.atleast_3d(covariance)
 
     _, n_gridcells = intercept.shape
 

tests/unit/test_auto_regression.py

@@ -597,34 +597,41 @@ def test_fit_auto_regression_np(lags):
 
 @pytest.mark.parametrize("intercept", [1.0, -4.0])
 @pytest.mark.parametrize("slope", [0.2, -0.3])
-def test_fit_autoregression_monthly_np(slope, intercept):
+def test_fit_autoregression_monthly_np_no_noise(slope, intercept):
     # test if autoregrerssion can fit using previous month as independent variable
     # and current month as dependent variable
-    # NOTE: the fit for the slope is bounded between -1 and 1
+    n_ts = 100
     np.random.seed(0)
-    prev_month = np.random.normal(size=100)
+    prev_month = np.random.normal(size=n_ts)
     cur_month = prev_month * slope + intercept
 
-    result = mesmer.stats._auto_regression._fit_auto_regression_monthly_np(
+    slope_fit, intercept_fit, residuals = mesmer.stats._auto_regression._fit_auto_regression_monthly_np(
         cur_month, prev_month
     )
-    expected = (slope, intercept)
-    np.testing.assert_allclose(result, expected)
+    expected_resids = np.zeros_like(cur_month)
 
+    np.testing.assert_allclose(slope, slope_fit)
+    np.testing.assert_allclose(intercept, intercept_fit)
+    np.testing.assert_allclose(residuals, expected_resids, atol=1e-6)
 
-@pytest.mark.parametrize("slope", [2.0, -2])
-def test_fit_autoregression_monthly_np_outside_bounds(slope):
-    # test that given a slope outside the bounds of -1 and 1, the slope is clipped
+@pytest.mark.parametrize("intercept", [1.0, -4.0])
+@pytest.mark.parametrize("slope", [0.2, -0.3])
+@pytest.mark.parametrize("variance", [1, 0.5])
+def test_fit_autoregression_monthly_np_with_noise(slope, intercept, variance):
+    # test if autoregrerssion can fit using previous month as independent variable
+    # and current month as dependent variable
+    n_ts = 5000
     np.random.seed(0)
-    prev_month = np.random.normal(size=100)
-    cur_month = prev_month * slope + 3
+    prev_month = np.random.normal(size=n_ts)
+    cur_month = prev_month * slope + intercept + np.random.normal(0, variance, size=n_ts)
 
-    result = mesmer.stats._auto_regression._fit_auto_regression_monthly_np(
+    slope_fit, intercept_fit, residuals = mesmer.stats._auto_regression._fit_auto_regression_monthly_np(
         cur_month, prev_month
     )
-    expected_slope = 1.0 * np.sign(slope)
-    np.testing.assert_allclose(result[0], expected_slope)
 
+    np.testing.assert_allclose(slope, slope_fit, atol=1e-1)
+    np.testing.assert_allclose(intercept, intercept_fit, atol=1e-1)
+    np.testing.assert_allclose(np.var(residuals), variance, rtol=1)
 
 def test_fit_auto_regression_monthly():
     n_years = 20
@@ -661,103 +668,14 @@ def test_fit_auto_regression_monthly():
         mesmer.stats.fit_auto_regression_monthly(data.values)
 
 
-def test_predict_auto_regression_monthly_intercept():
-    n_gridcells = 2
-    slope = np.zeros((12, n_gridcells))
-    slope = xr.DataArray(
-        slope,
-        dims=("month", "gridcell"),
-        coords={"month": np.arange(1, 13), "gridcell": np.arange(n_gridcells)},
-    )
-    intercept = np.tile(np.arange(1, 13), [n_gridcells, 1]).T
-    intercept = xr.DataArray(
-        intercept,
-        dims=("month", "gridcell"),
-        coords={"month": np.arange(1, 13), "gridcell": np.arange(n_gridcells)},
-    )
-    ar_params = xr.Dataset({"intercept": intercept, "slope": slope})
-
-    n_years = 10
-    time = pd.date_range("2000-01-01", periods=12 * n_years, freq="M")
-    time = xr.DataArray(time, dims="time", coords={"time": time})
-
-    result = mesmer.stats.predict_auto_regression_monthly(ar_params, time, buffer=0)
-
-    expected = np.tile(np.arange(1, 13), [n_gridcells, n_years]).T
-    expected[0] = 0
-    expected = xr.DataArray(
-        expected,
-        dims=("time", "gridcell"),
-        coords={"time": time, "gridcell": np.arange(n_gridcells)},
-    )
-    np.testing.assert_allclose(result, expected)
-
-
-def test_predict_auto_regression_monthly():
-    n_gridcells = 10
-    n_years = 10
-    np.random.seed(0)
-    slopes = xr.DataArray(
-        np.random.normal(-0.99, 0.99, size=(12, n_gridcells)),
-        dims=("month", "gridcell"),
-        coords={"month": np.arange(1, 13), "gridcell": np.arange(n_gridcells)},
-    )
-    intercepts = xr.DataArray(
-        np.random.normal(-10, 10, size=(12, n_gridcells)),
-        dims=("month", "gridcell"),
-        coords={"month": np.arange(1, 13), "gridcell": np.arange(n_gridcells)},
-    )
-    ar_params = xr.Dataset({"intercept": intercepts, "slope": slopes})
-
-    time = pd.date_range("2000-01-01", periods=n_years * 12, freq="M")
-    time = xr.DataArray(time, dims="time", coords={"time": time})
-
-    result = mesmer.stats.predict_auto_regression_monthly(ar_params, time, buffer=10)
-
-    _check_dataarray_form(
-        result,
-        "result",
-        ndim=2,
-        required_dims={"time", "gridcell"},
-        shape=(len(time), n_gridcells),
-    )
-
-
-def test_fit_predict_autoregression_monthly_roundtrip():
-    n_gridcells = 10
-    n_years = 150
-    buffer = 30 * 12
-    np.random.seed(0)
-    slopes = xr.DataArray(
-        np.random.uniform(-0.99, 0.99, size=(12, n_gridcells)),
-        dims=("month", "gridcell"),
-        coords={"month": np.arange(1, 13), "gridcell": np.arange(n_gridcells)},
-    )
-    intercepts = xr.DataArray(
-        np.random.normal(-10, 10, size=(12, n_gridcells)),
-        dims=("month", "gridcell"),
-        coords={"month": np.arange(1, 13), "gridcell": np.arange(n_gridcells)},
-    )
-    ar_params = xr.Dataset({"intercept": intercepts, "slope": slopes})
-
-    time = pd.date_range("2000-01-01", periods=n_years * 12, freq="M")
-    time = xr.DataArray(time, dims="time", coords={"time": time})
-
-    data = mesmer.stats.predict_auto_regression_monthly(ar_params, time, buffer)
-    AR_fit = mesmer.stats.fit_auto_regression_monthly(data)
-    predicted = mesmer.stats.predict_auto_regression_monthly(AR_fit, data.time, buffer)
-
-    np.testing.assert_allclose(predicted, data)
-
-
 @pytest.mark.parametrize("buffer", [1, 10, 20])
 def test_draw_auto_regression_monthly_np_buffer(buffer):
     n_realisations = 1
     n_gridcells = 10
     seed = 0
     slope = np.random.uniform(-1, 1, size=(12, n_gridcells))
     intercept = np.ones((12, n_gridcells))
-    covariance = None
+    covariance = np.zeros((12, n_gridcells, n_gridcells))
     n_ts = 120
 
     res_wo_buffer = mesmer.stats._auto_regression._draw_auto_regression_monthly_np(