keep all unique_ids in aggregate

Nixtla · Sep 13, 2023 · 3b13ffc · 3b13ffc
1 parent dac318b
commit 3b13ffc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/hierarchicalforecast/utils.py b/hierarchicalforecast/utils.py
@@ -215,9 +215,10 @@ def _to_summing_dataframe(
 
     # Match index ordering of S_df and collapse df to Y_bottom_df
     Y_bottom_df = df.copy()
-    Y_bottom_df = Y_bottom_df.groupby(['unique_id', 'ds'])['y'].sum().reset_index()
     Y_bottom_df.unique_id = Y_bottom_df.unique_id.astype('category')
     Y_bottom_df.unique_id = Y_bottom_df.unique_id.cat.set_categories(S_df.columns)
+    # observed=False here makes sure we have all the unique_ids (S_df.columns) in our result
+    Y_bottom_df = Y_bottom_df.groupby(['unique_id', 'ds'], observed=False)['y'].sum().reset_index()
     return Y_bottom_df, S_df, tags
 
 

diff --git a/nbs/utils.ipynb b/nbs/utils.ipynb
@@ -320,9 +320,10 @@
     "\n",
     "    # Match index ordering of S_df and collapse df to Y_bottom_df\n",
     "    Y_bottom_df = df.copy()\n",
-    "    Y_bottom_df = Y_bottom_df.groupby(['unique_id', 'ds'])['y'].sum().reset_index()\n",
     "    Y_bottom_df.unique_id = Y_bottom_df.unique_id.astype('category')\n",
     "    Y_bottom_df.unique_id = Y_bottom_df.unique_id.cat.set_categories(S_df.columns)\n",
+    "    # observed=False here makes sure we have all the unique_ids (S_df.columns) in our result\n",
+    "    Y_bottom_df = Y_bottom_df.groupby(['unique_id', 'ds'], observed=False)['y'].sum().reset_index()\n",
     "    return Y_bottom_df, S_df, tags\n",
     "\n"
    ]