Add item_id to ERCOT and ETT datasets

awslabs · Mar 22, 2024 · 7b9574f · 7b9574f
1 parent 9f33c1b
commit 7b9574f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/src/gluonts/dataset/repository/_ercot.py b/src/gluonts/dataset/repository/_ercot.py
@@ -34,6 +34,7 @@ def generate_ercot_dataset(dataset_path: Path, dataset_writer: DatasetWriter):
     test = [
         {
             "start": start,
+            "item_id": region,
             "target": df[region].to_numpy(dtype=np.float64),
         }
         for region in regions
@@ -42,6 +43,7 @@ def generate_ercot_dataset(dataset_path: Path, dataset_writer: DatasetWriter):
     train = [
         {
             "start": start,
+            "item_id": region,
             "target": df[region].to_numpy(dtype=np.float64)[
                 :-prediction_length
             ],

diff --git a/src/gluonts/dataset/repository/_ett_small.py b/src/gluonts/dataset/repository/_ett_small.py
@@ -47,6 +47,7 @@ def generate_ett_small_dataset(
             test.append(
                 {
                     "start": start,
+                    "item_id": col,
                     "target": df[col].values,
                 }
             )
@@ -60,6 +61,7 @@ def generate_ett_small_dataset(
             train.append(
                 {
                     "start": start,
+                    "item_id": col,
                     "target": df[col].values[:-prediction_length],
                 }
             )