fix(arima): use the configured innovation distribution on every fit and refit

@@ -33,6 +33,8 @@     Distribution,     _VALID_DISTRIBUTIONS,     _quantiles,+    dist_tail_from_params,+    distribution_n_extra_params,     error_params_from_model,     extract_dist_params,     guard_dist_tail,@@ -247,6 +249,7 @@ def arima(     tol=1e-8,     optim_control={"maxiter": 100},     distribution="normal",+    fixed_dist=None, ):     if distribution not in _VALID_DISTRIBUTIONS:         raise ValueError(@@ -786,6 +789,26 @@ def arma_css_op(p, x, coef, mask, arma, ncxreg, xreg, narma):         # objective actually optimized, plus the tail appended to the arma params         # (empty for normal/laplace, which concentrate their scale out)         ml_obj = armafn if distribution == Distribution.NORMAL else armafn_laplace+        fixed_tail = None+        if fixed_dist is not None and distribution_n_extra_params(distribution):+            if not no_optim:+                raise ValueError("fixed_dist requires all coefficients to be fixed")+            fixed_tail = np.array(dist_tail_from_params(distribution, fixed_dist))++        def minimize_tail(obj, init_ext):+            if fixed_tail is None:+                return minimize(+                    obj,+                    init_ext,+                    args=ml_args,+                    method=optim_method,+                    tol=tol,+                    options=optim_control,+                )+            return OptimResult(+                True, 0, fixed_tail, obj(fixed_tail, *ml_args), np.array([])+            )+         dist_tail_fit = np.array([])         nu_t = None         sigma2_t = None@@ -801,14 +824,7 @@ def arma_css_op(p, x, coef, mask, arma, ncxreg, xreg, narma):             init_sn = np.concatenate(                 [init[mask], [log_sigma2_init, alpha_init]]             )-            res_sn = minimize(-                armafn_skewnorm,-                init_sn,-                args=ml_args,-                method=optim_method,-                tol=tol,-                options=optim_control,-            )+            res_sn = minimize_tail(armafn_skewnorm, init_sn)             ml_obj = armafn_skewnorm             dist_tail_fit = res_sn.x[n_arma_free:]             _dp_sn = extract_dist_params("skew-normal", dist_tail_fit)@@ -834,14 +850,7 @@ def arma_css_op(p, x, coef, mask, arma, ncxreg, xreg, narma):             init_t = np.concatenate(                 [init[mask], [log_sigma2_init, log_nu_m2_init]]             )-            res_t = minimize(-                armafn_t,-                init_t,-                args=ml_args,-                method=optim_method,-                tol=tol,-                options=optim_control,-            )+            res_t = minimize_tail(armafn_t, init_t)             ml_obj = armafn_t             dist_tail_fit = res_t.x[n_arma_free:]             _dp_t = extract_dist_params("t", dist_tail_fit)@@ -867,14 +876,7 @@ def arma_css_op(p, x, coef, mask, arma, ncxreg, xreg, narma):             init_ged = np.concatenate(                 [init[mask], [log_sigma_init, log_beta_init]]             )-            res_ged = minimize(-                armafn_ged,-                init_ged,-                args=ml_args,-                method=optim_method,-                tol=tol,-                options=optim_control,-            )+            res_ged = minimize_tail(armafn_ged, init_ged)             ml_obj = armafn_ged             dist_tail_fit = res_ged.x[n_arma_free:]             _dp_ged = extract_dist_params("ged", dist_tail_fit)@@ -1341,6 +1343,10 @@ def arima2(x, model, xreg, method):     sigma2 = model["sigma2"]     # models fitted before SSinit was stored used Gardner1980     SSinit = model.get("SSinit", "Gardner1980")+    distribution = model.get("distribution", Distribution.NORMAL)+    fixed_dist = {+        k: model[k] for k in ("sigma2", "nu", "alpha_dist", "beta_dist") if k in model+    }     if use_drift:         n = len(model["x"])         time = np.arange(0, (n + 1) / m, 1 / m)[:n].reshape(-1, 1)@@ -1375,6 +1381,8 @@ def arima2(x, model, xreg, method):             include_mean=use_intercept,             method=method,             SSinit=SSinit,+            distribution=distribution,+            fixed_dist=fixed_dist,             fixed=coefs,             xreg=xreg if use_xreg else None,         )@@ -1385,6 +1393,8 @@ def arima2(x, model, xreg, method):             include_mean=use_intercept,             method=method,             SSinit=SSinit,+            distribution=distribution,+            fixed_dist=fixed_dist,             fixed=coefs,             xreg=xreg if use_xreg else None,         )@@ -1395,6 +1405,8 @@ def arima2(x, model, xreg, method):             include_mean=False,             method=method,             SSinit=SSinit,+            distribution=distribution,+            fixed_dist=fixed_dist,         )     n_coef = len(refit["coef"])     refit["var_coef"] = np.zeros((n_coef, n_coef), dtype=np.float32)@@ -2137,6 +2149,9 @@ def auto_arima_f( ):     if approximation is None:         approximation = len(x) > 150 or period > 12+    if distribution != Distribution.NORMAL:+        # the CSS approximation is only defined for normal innovations+        approximation = False     if x.ndim > 1:         raise ValueError("auto_arima can only handle univariate time series")     # myarima swallows ValueErrors, so validate before the search@@ -2161,12 +2176,15 @@ def auto_arima_f(     if is_constant(x):         if np.isnan(x).all():             raise ValueError("all data are missing")+        const_method = ArimaMethod.CSS if distribution == Distribution.NORMAL else method         if allowmean:             fit = Arima(                 x,                 order=(0, 0, 0),                 seasonal={"order": (0, 0, 0), "period": m},                 SSinit=SSinit,+                method=const_method,+                distribution=distribution,                 fixed=np.array([np.mean(x)]),             )         else:@@ -2175,6 +2193,8 @@ def auto_arima_f(                 order=(0, 0, 0),                 seasonal={"order": (0, 0, 0), "period": m},                 SSinit=SSinit,+                method=const_method,+                distribution=distribution,                 include_mean=False,             )         fit["x"] = origx@@ -2272,6 +2292,7 @@ def auto_arima_f(                     include_constant=True,                     fixed=np.array([np.mean(dx / m)]),                     method=method,+                    distribution=distribution,                 )             elif D > 0 and d > 0:                 fit = Arima(@@ -2280,6 +2301,7 @@ def auto_arima_f(                     seasonal={"order": (0, D, 0), "period": m},                     SSinit=SSinit,                     method=method,+                    distribution=distribution,                 )             elif d == 2:                 fit = Arima(@@ -2288,6 +2310,7 @@ def auto_arima_f(                     seasonal={"order": (0, 0, 0), "period": m},                     SSinit=SSinit,                     method=method,+                    distribution=distribution,                 )             elif d < 2:                 fit = Arima(@@ -2298,6 +2321,7 @@ def auto_arima_f(                     include_constant=True,                     fixed=np.array([np.mean(dx)]),                     method=method,+                    distribution=distribution,                 )             else:                 raise ValueError(@@ -2312,6 +2336,7 @@ def auto_arima_f(                     SSinit=SSinit,                     xreg=xreg,                     method=method,+                    distribution=distribution,                 )             else:                 fit = Arima(@@ -2321,6 +2346,7 @@ def auto_arima_f(                     SSinit=SSinit,                     xreg=xreg,                     method=method,+                    distribution=distribution,                 )         fit["x"] = origx         fit["lambda"] = blambda@@ -2381,6 +2407,7 @@ def auto_arima_f(             allow_mean=allowmean,             period=m,             SSinit=SSinit,+            distribution=distribution,         )         bestfit["lambda"] = blambda         bestfit["biasadj"] = biasadj

@@ -26,6 +26,7 @@     "dist_tail_bounds",     "guard_dist_tail",     "extract_dist_params",+    "dist_tail_from_params",     "aic_bic_aicc",     "error_params_from_model",     "_quantiles",@@ -155,6 +156,17 @@ def extract_dist_params(distribution: str, fit_par_dist, residuals=None) -> dict     return {}  +def dist_tail_from_params(distribution: str, params: dict) -> list:+    """Inverse of `extract_dist_params` for the distributions with a tail."""+    if distribution == "t":+        return [np.log(params["sigma2"]), np.log(params["nu"] - 2.0)]+    if distribution == "skew-normal":+        return [np.log(params["sigma2"]), params["alpha_dist"]]+    if distribution == "ged":+        return [0.5 * np.log(params["sigma2"]), np.log(params["beta_dist"])]+    return []  # laplace / normal++ def error_params_from_model(model: dict):     """Map a fitted model dict's distribution params to sample_errors() params. 

@@ -314,7 +314,7 @@ class AutoARIMA(_TS):         stepwise (bool, default=True): If True, will do stepwise selection (faster).         nmodels (int, default=94): Number of models considered in stepwise search.         trace (bool, default=False): If True, the searched ARIMA models is reported.-        approximation (Optional[bool], default=False): If True, conditional sums-of-squares estimation, final MLE.+        approximation (Optional[bool], default=False): If True, conditional sums-of-squares estimation, final MLE. Ignored when `distribution` isn't 'normal'.         method (Optional[str], optional): Fitting method between maximum likelihood or sums-of-squares.         truncate (Optional[bool], optional): Observations truncated series used in model selection.         test (str, default="kpss"): Unit root test to use. See `ndiffs` for details.@@ -588,6 +588,7 @@ def forecast(                 biasadj=self.biasadj,                 period=self.season_length,                 SSinit=self.SSinit,+                distribution=self.distribution,             )         fcst = forecast_arima(mod, h, xreg=X_future, level=level)         res = {"mean": fcst["mean"]}@@ -731,6 +732,7 @@ def simulate(                     biasadj=self.biasadj,                     period=self.season_length,                     SSinit=self.SSinit,+                    distribution=self.distribution,                 )         else:             if not hasattr(self, "model_"):@@ -2152,6 +2154,7 @@ def forecast(                 method=self.method,                 fixed=self.fixed,                 SSinit=self.SSinit,+                distribution=self.distribution,             )         fcst = forecast_arima(mod, h, xreg=X_future, level=level)         res = {"mean": fcst["mean"]}

@@ -1009,6 +1009,88 @@ def test_distribution_forecast_intervals(distribution, param_key):     assert np.all(lower_95 < upper_95), "lower must be below upper"  +_DIST_PARAM_KEYS = {"t": "nu", "skew-normal": "alpha_dist", "ged": "beta_dist", "laplace": None}+++@pytest.mark.parametrize("distribution", list(_DIST_PARAM_KEYS))+@pytest.mark.parametrize(+    "order,seasonal",+    [+        ((1, 1, 1), {"order": (0, 0, 0), "period": 1}),+        ((0, 1, 1), {"order": (0, 1, 1), "period": 12}),+        ((0, 1, 0), {"order": (0, 0, 0), "period": 1}),+    ],+)+def test_distribution_forward_arima(distribution, order, seasonal):+    fit = Arima(ap[:120], order=order, seasonal=seasonal, method="ML", distribution=distribution)+    refit = forward_arima(fit, ap)+    assert refit["distribution"] == distribution+    key = _DIST_PARAM_KEYS[distribution]+    if key is not None:+        assert refit[key] == pytest.approx(fit[key])+++@pytest.mark.parametrize("distribution", ["t", "skew-normal", "ged"])+def test_distribution_forward_arima_keeps_params(distribution, monkeypatch):+    fit = Arima(ap[:120], order=(1, 1, 1), method="ML", distribution=distribution)+    coefs = np.array(list(fit["coef"].values()))+    reestimated = Arima(+        ap, order=(1, 1, 1), method="ML", distribution=distribution, fixed=coefs+    )++    def fail(*args, **kwargs):+        raise AssertionError("forward_arima shouldn't re-estimate parameters")++    monkeypatch.setattr("statsforecast.arima.minimize", fail)+    same_data = forward_arima(fit, ap[:120], method="ML")+    assert same_data["loglik"] == pytest.approx(fit["loglik"], rel=1e-10)+    refit = forward_arima(fit, ap, method="ML")+    key = _DIST_PARAM_KEYS[distribution]+    assert refit["sigma2"] == fit["sigma2"]+    assert refit[key] == pytest.approx(fit[key])+    assert refit["loglik"] < reestimated["loglik"]+++@pytest.mark.parametrize("distribution", list(_DIST_PARAM_KEYS))+@pytest.mark.parametrize("approximation", [False, True])+def test_distribution_autoarima_exhaustive(distribution, approximation):+    rng = np.random.default_rng(9)+    y = np.cumsum(rng.standard_normal(100))+    fit = auto_arima_f(+        y,+        stepwise=False,+        max_p=2,+        max_q=2,+        max_order=3,+        approximation=approximation,+        distribution=distribution,+    )+    assert fit["distribution"] == distribution+++@pytest.mark.parametrize("distribution", list(_DIST_PARAM_KEYS))+@pytest.mark.parametrize(+    "y", [np.full(30, 3.0), np.arange(30.0)], ids=["constant", "linear"]+)+def test_distribution_autoarima_constant_series(distribution, y):+    fit = auto_arima_f(y, distribution=distribution)+    assert fit["distribution"] == distribution+++@pytest.mark.parametrize("distribution", list(_DIST_PARAM_KEYS))+@pytest.mark.parametrize(+    "model_cls,kwargs",+    [(ARIMA, {"order": (1, 1, 1)}), (AutoARIMAModel, {"max_p": 2, "max_q": 2})],+    ids=["ARIMA", "AutoARIMA"],+)+def test_distribution_forecast_matches_fit_predict(distribution, model_cls, kwargs):+    model = model_cls(**kwargs, distribution=distribution)+    fcst = model.forecast(ap, h=12, level=[80])+    pred = model.fit(ap).predict(h=12, level=[80])+    for key in ("mean", "lo-80", "hi-80"):+        np.testing.assert_allclose(fcst[key], pred[key])++ def test_issue_649(capsys):     from statsforecast.models import AutoARIMA