From b671267c64b19e1888105c235d2395a322a01cdd Mon Sep 17 00:00:00 2001
From: thefriedbee <diyi93@outlook.com>
Date: Wed, 27 Mar 2024 13:08:05 -0500
Subject: [PATCH] implementing pipeline to prepare dataset

---
 carpoolsim/prepare_input.py | 34 ++++++++++++++++++++++----------
 notebooks/test.ipynb        | 39 +++++++++++++++++++++++++++++++------
 2 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/carpoolsim/prepare_input.py b/carpoolsim/prepare_input.py
index 71e8135..b9aab0b 100644
--- a/carpoolsim/prepare_input.py
+++ b/carpoolsim/prepare_input.py
@@ -27,8 +27,13 @@ def __init__(
         self.gdf_nodes = network_nodes
         self.gdf_links = network_links
         self.tazs = tazs
+        self.tazs_ids = []
         self.network_dict = {}
 
+    def get_taz_id_list(self):
+        self.tazs.taz_id = self.tazs.taz_id.astype(int)
+        self.tazs_ids = sorted(self.tazs.taz_id.tolist())
+
     def convert_abm_links(self):
         gdf_links = net_prep.initialize_abm15_links(
             self.gdf_nodes,
@@ -49,20 +54,22 @@ def build_network(self):
         # print(network_dict['DG']['65666']['1'])
         self.network_dict = network_dict
 
-    def run_batch(self):
+    def prepare_taz_lists(self):
+        # break taz list to chunks for multiprocessing
+
+        pass
+
+    def run_batch(self, source_id=1):
         network_dict = self.network_dict
 
         # only record results to centroids
-        destination_lst = [str(i) for i in range(1, 5874)]
-        t1 = time.perf_counter()
         dists_dict, paths_dict = run_batch_ods(
             network_dict['DG'],
-            source='1',
-            destination_lst=destination_lst,
-            weight='backward'
+            source=f'{source_id}',
+            destination_lst=self.tazs_ids,
+            weight='forward'
         )
-        delta_t = time.perf_counter() - t1
-        print('It takes {} seconds to run'.format(delta_t))
+        return dists_dict, paths_dict
 
 
 # Define a new function to simply run shortest path given origin and destination node id.
@@ -129,10 +136,17 @@ def main():
     taz_centriod_nodes = gpd.read_file(os.environ["taz"])
     df_nodes_raw = gpd.read_file(os.environ['network_nodes'])
     df_links_raw = gpd.read_file(os.environ['network_links'])
-    # load pnr nodes to park and ride trip
-    pnr_nodes = gpd.read_file(os.environ["parking_lots"])
 
     # init object
+    traffic_network = TrafficNetwork(
+        network_links=df_links_raw,
+        network_nodes=df_nodes_raw,
+        tazs=taz_centriod_nodes
+    )
+
+    traffic_network.get_taz_id_list()
+    traffic_network.convert_abm_links()
+    traffic_network.build_network()
 
     pass
 
diff --git a/notebooks/test.ipynb b/notebooks/test.ipynb
index 339ceff..04655e3 100644
--- a/notebooks/test.ipynb
+++ b/notebooks/test.ipynb
@@ -2662,19 +2662,46 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "id": "dbdb4e11-3bf7-4d9d-95de-8753e505f95b",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[0 1 2 3] [4 5 6] [7 8 9]\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(*np.array_split(range(10), 3))"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "id": "2adf7931-efea-449a-bc21-2edb9f449669",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[array(['a', 'b'], dtype='<U21'),\n",
+       " array(['1'], dtype='<U21'),\n",
+       " array(['123123'], dtype='<U21')]"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "np.array_split(\n",
+    "    [\"a\", \"b\", 1, 123123], \n",
+    "    3)"
+   ]
   },
   {
    "cell_type": "code",