fix arima trained models results index

Nixtla · Jun 18, 2024 · fa982ea · fa982ea
1 parent d73e43b
commit fa982ea
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/nbs/src/arima.ipynb b/nbs/src/arima.ipynb
@@ -3057,15 +3057,15 @@
     "    if fit['ic'] < bestfit['ic']:\n",
     "        bestfit = fit\n",
     "        p = q = P = Q = 0\n",
-    "    k = 1\n",
+    "    k = 2\n",
     "    if max_p > 0 or max_P > 0:\n",
     "        p_ = int(max_p > 0)\n",
     "        P_ = int(m > 1 and max_P > 0)\n",
     "        fit = p_myarima(\n",
     "            order=(p_, d, 0),\n",
     "            seasonal={'order': (P_, D, 0), 'period': m},\n",
     "        )\n",
-    "        results[k + 1] = (p_, d, 0, P_, D, 0, constant, fit['ic'])\n",
+    "        results[k] = (p_, d, 0, P_, D, 0, constant, fit['ic'])\n",
     "        if fit['ic'] < bestfit['ic']:\n",
     "            bestfit = fit\n",
     "            p = p_\n",
@@ -3079,7 +3079,7 @@
     "            order=(0, d, q_),\n",
     "            seasonal={'order': (0, D, Q_), 'period': m},\n",
     "        )\n",
-    "        results[k + 1] = (0, d, q_, 0, D, Q_, constant, fit['ic'])\n",
+    "        results[k] = (0, d, q_, 0, D, Q_, constant, fit['ic'])\n",
     "        if fit['ic'] < bestfit['ic']:\n",
     "            bestfit = fit\n",
     "            p = P = 0\n",
@@ -3092,14 +3092,13 @@
     "            seasonal={'order': (0, D, 0), 'period': m},\n",
     "            constant=False,\n",
     "        )\n",
-    "        results[k + 1] = (0, d, 0, 0, D, 0, 0, fit['ic'])\n",
+    "        results[k] = (0, d, 0, 0, D, 0, 0, fit['ic'])\n",
     "        if fit['ic'] < bestfit['ic']:\n",
     "            bestfit = fit\n",
     "            p = q = P = Q = 0\n",
     "        k += 1\n",
     "        \n",
     "    def try_params(p, d, q, P, D, Q, constant, k, bestfit):\n",
-    "        k += 1\n",
     "        improved = False\n",
     "        if k >= results.shape[0]:\n",
     "            return k, bestfit, improved\n",
@@ -3108,6 +3107,7 @@
     "            seasonal={'order': (P, D, Q), 'period': m},\n",
     "        )\n",
     "        results[k] = (p, d, q, P, D, Q, constant, fit['ic'])\n",
+    "        k += 1\n",
     "        if fit['ic'] < bestfit['ic']:\n",
     "            bestfit = fit\n",
     "            improved = True\n",

diff --git a/statsforecast/arima.py b/statsforecast/arima.py
@@ -2101,15 +2101,15 @@ def auto_arima_f(
     if fit["ic"] < bestfit["ic"]:
         bestfit = fit
         p = q = P = Q = 0
-    k = 1
+    k = 2
     if max_p > 0 or max_P > 0:
         p_ = int(max_p > 0)
         P_ = int(m > 1 and max_P > 0)
         fit = p_myarima(
             order=(p_, d, 0),
             seasonal={"order": (P_, D, 0), "period": m},
         )
-        results[k + 1] = (p_, d, 0, P_, D, 0, constant, fit["ic"])
+        results[k] = (p_, d, 0, P_, D, 0, constant, fit["ic"])
         if fit["ic"] < bestfit["ic"]:
             bestfit = fit
             p = p_
@@ -2123,7 +2123,7 @@ def auto_arima_f(
             order=(0, d, q_),
             seasonal={"order": (0, D, Q_), "period": m},
         )
-        results[k + 1] = (0, d, q_, 0, D, Q_, constant, fit["ic"])
+        results[k] = (0, d, q_, 0, D, Q_, constant, fit["ic"])
         if fit["ic"] < bestfit["ic"]:
             bestfit = fit
             p = P = 0
@@ -2136,14 +2136,13 @@ def auto_arima_f(
             seasonal={"order": (0, D, 0), "period": m},
             constant=False,
         )
-        results[k + 1] = (0, d, 0, 0, D, 0, 0, fit["ic"])
+        results[k] = (0, d, 0, 0, D, 0, 0, fit["ic"])
         if fit["ic"] < bestfit["ic"]:
             bestfit = fit
             p = q = P = Q = 0
         k += 1
 
     def try_params(p, d, q, P, D, Q, constant, k, bestfit):
-        k += 1
         improved = False
         if k >= results.shape[0]:
             return k, bestfit, improved
@@ -2152,6 +2151,7 @@ def try_params(p, d, q, P, D, Q, constant, k, bestfit):
             seasonal={"order": (P, D, Q), "period": m},
         )
         results[k] = (p, d, q, P, D, Q, constant, fit["ic"])
+        k += 1
         if fit["ic"] < bestfit["ic"]:
             bestfit = fit
             improved = True