replace second max_flow call with rust

ml-energy · Oct 18, 2024 · 3e07341 · 3e07341
1 parent 5a7d108
commit 3e07341
Showing 1 changed file with 29 additions and 6 deletions.
diff --git a/lowtime/solver.py b/lowtime/solver.py
@@ -486,6 +486,7 @@ def reformat_rust_flow_to_dict(flow_vec: list[tuple[tuple[int, int], float]], da
                 "PROFILING PhillipsDessouky::find_min_cut maximum_flow_1 time: %.10fs",
                 profiling_max_flow,
             )
+        # TODO(ohjun): replace to handle Rust-side pyo3 exception
         except nx.NetworkXUnbounded as e:
             raise LowtimeFlowError("ERROR: Infinite flow for unbounded DAG.") from e
 
@@ -559,19 +560,41 @@ def reformat_rust_flow_to_dict(flow_vec: list[tuple[tuple[int, int], float]], da
 
         # Run max flow on the new residual graph.
         try:
+            # profiling_max_flow = time.time()
+            # _, flow_dict = nx.maximum_flow(
+            #     residual_graph,
+            #     source_node,
+            #     sink_node,
+            #     capacity="capacity",
+            #     flow_func=edmonds_karp,
+            # )
+            # profiling_max_flow = time.time() - profiling_max_flow
+            # logger.info(
+            #     "PROFILING PhillipsDessouky::find_min_cut maximum_flow_2 time: %.10fs",
+            #     profiling_max_flow,
+            # )
             profiling_max_flow = time.time()
-            _, flow_dict = nx.maximum_flow(
-                residual_graph,
-                source_node,
-                sink_node,
-                capacity="capacity",
-                flow_func=edmonds_karp,
+            nodes, edges = format_rust_inputs(residual_graph)
+
+            profiling_data_transfer = time.time()
+            rust_dag = lowtime_rust.PhillipsDessouky(
+                nodes, source_node, sink_node, edges
             )
+            profiling_data_transfer = time.time() - profiling_data_transfer
+            logger.info(
+                "PROFILING PhillipsDessouky::find_min_cut data transfer 2 time: %.10fs",
+                profiling_data_transfer,
+            )
+
+            rust_flow_vec = rust_dag.max_flow()
+            flow_dict = reformat_rust_flow_to_dict(rust_flow_vec, unbound_dag)
+
             profiling_max_flow = time.time() - profiling_max_flow
             logger.info(
                 "PROFILING PhillipsDessouky::find_min_cut maximum_flow_2 time: %.10fs",
                 profiling_max_flow,
             )
+        # TODO(ohjun): replace to handle Rust-side pyo3 exception
         except nx.NetworkXUnbounded as e:
             raise LowtimeFlowError(
                 "ERROR: Infinite flow on capacity residual graph."