优化爆管定位算法，增加多进程支持

app/algorithms/burst_location/leak_simulator.py

@@ -1,12 +1,16 @@
 """漏损模拟模块。"""
 
 import math
+import multiprocessing as mp
 import sys
 
 import pandas as pd
 import wntr
 
 _PIPE2LEAKNODE = None
+_SIGNATURE_WN = None
+_SIGNATURE_LEAK_MAG = None
+_SIGNATURE_SENSOR_NAME = None
 
 
 def simple_add_leak(wn, leak_mag, leak_pipe):
@@ -387,7 +391,7 @@ def cal_sum_demand(demand):
 
 
 def cal_signature_pipe_multi_pf(
-    wn, leak_mag, candidate_center, timestep_list, sensor_name
+    wn, leak_mag, candidate_center, timestep_list, sensor_name, n_workers=1
 ):
     candidate_center_num = len(candidate_center)
     pressure_leak = pd.DataFrame(
@@ -398,17 +402,45 @@ def cal_signature_pipe_multi_pf(
     #                          columns=sensor_f_name)
     pressure_leak = pressure_leak.sort_index()
     # flow_leak = flow_leak.sort_index()
-    for i in range(candidate_center_num):
-        wn, pressure_output = leak_simulation_pipe_dd_multi_pf(
-            wn, leak_mag, candidate_center[i], sensor_name
-        )
-        # leak_or_not_list.append(leak_or_not)
-        pressure_leak.loc[(candidate_center[i], slice(None)), :] = pressure_output.to_numpy()
-        # flow_leak.loc[candidate_center[i]].loc[:, :] = flow_output
-        sys.stdout.write("\r" + "已经完成计算" + str(i + 1) + "个特征中心")
+    can_fork = "fork" in mp.get_all_start_methods()
+    if n_workers > 1 and candidate_center_num > 1 and can_fork:
+        _set_signature_worker_context(wn, leak_mag, sensor_name)
+        worker_count = min(n_workers, candidate_center_num)
+        with mp.get_context("fork").Pool(processes=worker_count) as pool:
+            for i, (center_name, pressure_array) in enumerate(
+                pool.imap(_run_signature_for_center, candidate_center)
+            ):
+                pressure_leak.loc[(center_name, slice(None)), :] = pressure_array
+                sys.stdout.write("\r" + "已经完成计算" + str(i + 1) + "个特征中心")
+        _set_signature_worker_context(None, None, None)
+    else:
+        for i in range(candidate_center_num):
+            wn, pressure_output = leak_simulation_pipe_dd_multi_pf(
+                wn, leak_mag, candidate_center[i], sensor_name
+            )
+            # leak_or_not_list.append(leak_or_not)
+            pressure_leak.loc[(candidate_center[i], slice(None)), :] = (
+                pressure_output.to_numpy()
+            )
+            # flow_leak.loc[candidate_center[i]].loc[:, :] = flow_output
+            sys.stdout.write("\r" + "已经完成计算" + str(i + 1) + "个特征中心")
     return pressure_leak, candidate_center
 
 
+def _set_signature_worker_context(wn, leak_mag, sensor_name):
+    global _SIGNATURE_WN, _SIGNATURE_LEAK_MAG, _SIGNATURE_SENSOR_NAME
+    _SIGNATURE_WN = wn
+    _SIGNATURE_LEAK_MAG = leak_mag
+    _SIGNATURE_SENSOR_NAME = sensor_name
+
+
+def _run_signature_for_center(center_name):
+    _, pressure_output = leak_simulation_pipe_dd_multi_pf(
+        _SIGNATURE_WN, _SIGNATURE_LEAK_MAG, center_name, _SIGNATURE_SENSOR_NAME
+    )
+    return center_name, pressure_output.to_numpy()
+
+
 def pick_pipe(all_pipes, pipe_diameter, limited_diameter):
     candidate_pipe = []
     for each_pipe in all_pipes: