[WIP] On-the-fly data loading of global models

neuralprophet/data/process.py

@@ -582,7 +582,7 @@ def _create_dataset(model, df, predict_mode, prediction_frequency=None):
     ----------
         df : pd.DataFrame
             dataframe containing column ``ds``, ``y``, and optionally``ID`` and
-            normalized columns normalized columns ``ds``, ``y``, ``t``, ``y_scaled``
+            normalized columns ``ds``, ``y``, ``t``, ``y_scaled``
         predict_mode : bool
             specifies predict mode
 

neuralprophet/time_dataset.py

@@ -26,16 +26,32 @@
             **kwargs : dict
                 Identical to :meth:`tabularize_univariate_datetime`
         """
-        # # TODO (future): vectorize
-        timedatasets = [TimeDataset(df_i, df_name, **kwargs) for df_name, df_i in df.groupby("ID")]
-        self.combined_timedataset = [item for timedataset in timedatasets for item in timedataset]
-        self.length = sum(timedataset.length for timedataset in timedatasets)
+        self.df = df
+        self.kwargs = kwargs
+        self.id_groups = df.groupby("ID")
+
+        # Create a mapping from global index to (df_name, local_idx)
+        self.index_mapping = []
+
+        for df_name, df_i in self.id_groups:
+            n_samples = len(df_i) - kwargs['n_lags'] + 1 - kwargs['n_forecasts']
+            for local_idx in range(n_samples):
+                self.index_mapping.append((df_name, local_idx))
+
+        self.total_length = len(self.index_mapping)
+
 
     def __len__(self):
-        return self.length
+        return self.total_length
+
 
     def __getitem__(self, idx):
-        return self.combined_timedataset[idx]
+        # Directly use the precomputed mapping
+        df_name, local_idx = self.index_mapping[idx]
+        df_i = self.id_groups.get_group(df_name)
+
+        timedataset = TimeDataset(df_i, df_name, **self.kwargs)
+        return timedataset[local_idx]
 
 
 class TimeDataset(Dataset):
@@ -250,6 +266,8 @@
         np.array, float
             Targets to be predicted of same length as each of the model inputs, dims: (num_samples, n_forecasts)
         """
+        if index >= len(self.samples):
+            raise IndexError(f"Index {index} is out of range for sample with length {len(self.samples)}")
         sample = self.samples[index]
         targets = self.targets[index]
         meta = self.meta

tests/test_integration.py

@@ -31,6 +31,21 @@
 PLOT = False
 
 
+def test_minimal():
+    log.info("testing: Plotting")
+    df = pd.read_csv(PEYTON_FILE, nrows=NROWS)
+    m = NeuralProphet(
+        n_forecasts=7,
+        n_lags=14,
+        epochs=EPOCHS,
+        batch_size=BATCH_SIZE,
+        learning_rate=LR,
+    )
+    metrics_df = m.fit(df, freq="D", minimal=True)
+    assert metrics_df is None
+    m.predict(df)
+
+
 def test_names():
     log.info("testing: names")
     m = NeuralProphet()
@@ -1282,20 +1297,6 @@ def test_auto_normalization():
     _ = m.fit(df, freq="D")
 
 
-def test_minimal():
-    log.info("testing: Plotting")
-    df = pd.read_csv(PEYTON_FILE, nrows=NROWS)
-    m = NeuralProphet(
-        n_forecasts=7,
-        n_lags=14,
-        epochs=EPOCHS,
-        batch_size=BATCH_SIZE,
-        learning_rate=LR,
-    )
-    metrics_df = m.fit(df, freq="D", minimal=True)
-    assert metrics_df is None
-    m.predict(df)
-
 
 def test_get_latest_forecast():
     log.info("testing: get_latest_forecast")