Let fit AR return nobs (#395)

* add nobs to AR results

* Changed std to variance in documentation (see #317)

* changed tests to contain nobs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add entry to changelog

---------

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

CHANGELOG.rst

@@ -6,6 +6,8 @@ v0.11.0 - unreleased
 
 New Features
 ^^^^^^^^^^^^
+- Added number of observations to the output of the AR process. `#395 <https://github.com/MESMER-group/mesmer/pull/395>`_.
+  By `Victoria Bauer`_.
 
 Breaking changes
 ^^^^^^^^^^^^^^^^
@@ -367,3 +369,4 @@ v0.8.0 - 2021-07-13
 .. _`Mathias Hauser`: https://github.com/mathause
 .. _`Yann Quilcaille`: https://github.com/yquilcaille
 .. _`Zeb Nicholls`: https://github.com/znicholls
+.. _`Victoria Bauer`: https://github.com/veni-vidi-vici-dormivi

mesmer/stats/_auto_regression.py

@@ -511,20 +511,20 @@ def fit_auto_regression(data, dim, lags):
     -------
     :obj:`xr.Dataset`
         Dataset containing the estimated parameters of the ``intercept``, the AR
-        ``coeffs`` and the ``variance`` of the residuals.
+        ``coeffs``, the ``variance`` of the residuals and the number of observations ``nobs``.
     """
 
     if not isinstance(data, xr.DataArray):
         raise TypeError(f"Expected a `xr.DataArray`, got {type(data)}")
 
     # NOTE: this is slowish, see https://github.com/MESMER-group/mesmer/pull/290
-    intercept, coeffs, variance = xr.apply_ufunc(
+    intercept, coeffs, variance, nobs = xr.apply_ufunc(
         _fit_auto_regression_np,
         data,
         input_core_dims=[[dim]],
-        output_core_dims=((), ("lags",), ()),
+        output_core_dims=((), ("lags",), (), ()),
         vectorize=True,
-        output_dtypes=[float, float, float],
+        output_dtypes=[float, float, float, int],
         kwargs={"lags": lags},
     )
 
@@ -536,6 +536,7 @@ def fit_auto_regression(data, dim, lags):
         "coeffs": coeffs,
         "variance": variance,
         "lags": lags,
+        "nobs": nobs,
     }
 
     return xr.Dataset(data_vars)
@@ -558,8 +559,10 @@ def _fit_auto_regression_np(data, lags):
         Intercept of the fitted AR model.
     coeffs : :obj:`np.array`
         Coefficients if the AR model. Will have as many entries as ``lags``.
-    std : :obj:`np.array`
-        Standard deviation of the residuals.
+    variance : :obj:`np.array`
+       Variance of the residuals.
+    nobs: :obj:`np.array``
+        Number of observations.
     """
 
     from statsmodels.tsa.ar_model import AutoReg
@@ -573,4 +576,6 @@ def _fit_auto_regression_np(data, lags):
     # variance of the residuals
     variance = AR_result.sigma2
 
-    return intercept, coeffs, variance
+    nobs = AR_result.nobs
+
+    return intercept, coeffs, variance, nobs

tests/unit/test_auto_regression.py

@@ -456,14 +456,16 @@ def test_fit_auto_regression_xr_1D_values():
     # statsmodels.tsa.ar_model.AutoReg
 
     data = trend_data_1D()
-    result = mesmer.stats.fit_auto_regression(data, "time", lags=1)
+    lags = 1
+    result = mesmer.stats.fit_auto_regression(data, "time", lags=lags)
 
     expected = xr.Dataset(
         {
             "intercept": 1.04728995,
             "coeffs": ("lags", [0.99682459]),
             "variance": 1.05381192,
             "lags": [1],
+            "nobs": data["time"].size - lags,
         }
     )
 
@@ -475,14 +477,16 @@ def test_fit_auto_regression_xr_1D_values_lags():
     # statsmodels.tsa.ar_model.AutoReg
 
     data = trend_data_1D()
-    result = mesmer.stats.fit_auto_regression(data, "time", lags=[2])
+    lags = 2
+    result = mesmer.stats.fit_auto_regression(data, "time", lags=[lags])
 
     expected = xr.Dataset(
         {
             "intercept": 2.08295035,
             "coeffs": ("lags", [0.99318256]),
             "variance": 1.18712735,
             "lags": [2],
+            "nobs": data["time"].size - lags,
         }
     )