add_series_dimension

Nixtla · Feb 25, 2025 · 53ffe8a · 53ffe8a
1 parent d4944dc
commit 53ffe8a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/nbs/common.base_model.ipynb b/nbs/common.base_model.ipynb
@@ -671,12 +671,12 @@
     "            available_idx = temporal_cols.get_loc('available_mask')           \n",
     "            available_condition = windows[:, :self.input_size, available_idx]\n",
     "            available_condition = torch.sum(available_condition, axis=(1, -1)) # Sum over time & series dimension\n",
-    "            final_condition = (available_condition > self.data_availability_threshold * self.input_size)\n",
+    "            final_condition = (available_condition > self.data_availability_threshold * self.input_size * self.n_series)\n",
     "            \n",
     "            if self.h > 0:\n",
     "                sample_condition = windows[:, self.input_size:, available_idx]\n",
     "                sample_condition = torch.sum(sample_condition, axis=(1, -1)) # Sum over time & series dimension\n",
-    "                final_condition = (sample_condition > self.data_availability_threshold * self.h) & (available_condition > self.data_availability_threshold * self.input_size)\n",
+    "                final_condition = (sample_condition > self.data_availability_threshold * self.h  * self.n_series) & (available_condition > self.data_availability_threshold * self.input_size * self.n_series)\n",
     "            \n",
     "            windows = windows[final_condition]\n",
     "            \n",

diff --git a/neuralforecast/common/_base_model.py b/neuralforecast/common/_base_model.py
@@ -680,7 +680,8 @@ def _create_windows(self, batch, step, w_idxs=None):
                 available_condition, axis=(1, -1)
             )  # Sum over time & series dimension
             final_condition = (
-                available_condition > self.data_availability_threshold * self.input_size
+                available_condition
+                > self.data_availability_threshold * self.input_size * self.n_series
             )
 
             if self.h > 0:
@@ -689,10 +690,11 @@ def _create_windows(self, batch, step, w_idxs=None):
                     sample_condition, axis=(1, -1)
                 )  # Sum over time & series dimension
                 final_condition = (
-                    sample_condition > self.data_availability_threshold * self.h
+                    sample_condition
+                    > self.data_availability_threshold * self.h * self.n_series
                 ) & (
                     available_condition
-                    > self.data_availability_threshold * self.input_size
+                    > self.data_availability_threshold * self.input_size * self.n_series
                 )
 
             windows = windows[final_condition]