Merge pull request #18 from Nixtla/fix/labour-data

[FEAT] Add hierarchical data tests

nbs/hierarchical.ipynb

@@ -252,6 +252,16 @@
     "             download_file(path, HierarchicalData.source_url, decompress=True)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| hide\n",
+    "from fastcore.test import test_close"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -267,7 +277,16 @@
     "    S_hiers = np.vstack(S_hiers)\n",
     "    S_hiers = S_hiers.sum(axis=0)\n",
     "    is_strictly_hierarchical = np.array_equal(S_hiers, np.sort(S_hiers))\n",
-    "    print(f'Is {group} strictly hierarchical? {is_strictly_hierarchical}')"
+    "    print(f'Is {group} strictly hierarchical? {is_strictly_hierarchical}')\n",
+    "    \n",
+    "    # test S recovers Y_df\n",
+    "    for key, hiers in tags.items():\n",
+    "        for ts, bottom_ts in S.loc[hiers].iterrows():\n",
+    "            actual_bottom_ts = bottom_ts.loc[lambda x: x == 1].index\n",
+    "            test_close(\n",
+    "                Y_df.query('unique_id == @ts')['y'].sum(), \n",
+    "                Y_df.query('unique_id in @actual_bottom_ts')['y'].sum()\n",
+    "            )"
    ]
   },
   {