fix: series shorter than max_horizon (#459)

Nixtla · Nov 28, 2024 · 4dabb31 · 4dabb31
1 parent fdb7772
commit 4dabb31
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/mlforecast/grouped_array.py b/mlforecast/grouped_array.py
@@ -146,9 +146,10 @@ def expand_target(self, max_horizon: int) -> np.ndarray:
         )
         for j in range(max_horizon):
             for i in range(self.n_groups):
-                out[self.indptr[i] : self.indptr[i + 1] - j, j] = self.data[
-                    self.indptr[i] + j : self.indptr[i + 1]
-                ]
+                if self.indptr[i + 1] - self.indptr[i] > j:
+                    out[self.indptr[i] : self.indptr[i + 1] - j, j] = self.data[
+                        self.indptr[i] + j : self.indptr[i + 1]
+                    ]
         return out
 
     def take_from_groups(self, idx: Union[int, slice]) -> "GroupedArray":

diff --git a/nbs/grouped_array.ipynb b/nbs/grouped_array.ipynb
@@ -190,11 +190,12 @@
     "        )\n",
     "        for j in range(max_horizon):\n",
     "            for i in range(self.n_groups):\n",
-    "                out[self.indptr[i] : self.indptr[i + 1] - j, j] = self.data[\n",
-    "                    self.indptr[i] + j : self.indptr[i + 1]\n",
-    "                ]\n",
+    "                if self.indptr[i + 1] - self.indptr[i] > j:\n",
+    "                    out[self.indptr[i] : self.indptr[i + 1] - j, j] = self.data[\n",
+    "                        self.indptr[i] + j : self.indptr[i + 1]\n",
+    "                    ]\n",
     "        return out\n",
-    "        \n",
+    "\n",
     "    def take_from_groups(self, idx: Union[int, slice]) -> 'GroupedArray':\n",
     "        \"\"\"Takes `idx` from each group in the array.\"\"\"\n",
     "        ranges = [\n",