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

This commit is contained in:
2026-03-07 15:31:04 +08:00
parent 0f8d33291d
commit 05ca940c9f
3 changed files with 93 additions and 10 deletions
@@ -1,5 +1,6 @@
import argparse
import json
from multiprocessing import cpu_count
from pathlib import Path
from typing import Any, Iterable
@@ -17,6 +18,8 @@ from .network_model import (
read_inf_inp_other,
)
DEFAULT_N_WORKERS = max(1, min(cpu_count() - 1, 4))
def _read_id_list_json(path):
if path is None:
@@ -111,6 +114,7 @@ def run_burst_location(
normal_flow: pd.Series | None = None,
min_dpressure: float = 2.0,
basic_pressure: float = 10.0,
n_workers: int = DEFAULT_N_WORKERS,
) -> dict[str, Any]:
if pressure_scada_ids is None or len(pressure_scada_ids) == 0:
raise ValueError("pressure_scada_ids cannot be empty.")
@@ -171,6 +175,7 @@ def run_burst_location(
similarity_mode = "CAD_new_gy"
max_flow = pd.Series(dtype=float)
stage_timing: dict[str, Any] = {}
located_pipe, elapsed_seconds, simulation_times, _, similarity_series = (
DN_search_multi_simple_add_flow_count_new(
wn=wn,
@@ -203,6 +208,8 @@ def run_burst_location(
if_gy=0,
pressure_threshold=float(min_dpressure),
leak_mag=float(burst_leakage),
n_workers=max(1, int(n_workers)),
stage_timing=stage_timing,
)
)
@@ -213,6 +220,7 @@ def run_burst_location(
"simulation_times": int(simulation_times),
"top_candidates": _build_top_candidates(similarity_series),
"similarity_mode": similarity_mode,
"stage_timing_seconds": stage_timing,
}
@@ -252,6 +260,12 @@ def _parse_args():
default=10.0,
help="(可选)基础服务压力,默认 10.0",
)
parser.add_argument(
"--n-workers",
type=int,
default=DEFAULT_N_WORKERS,
help="(可选)特征中心模拟进程数,默认 max(1, min(cpu_count()-1, 4))",
)
return parser.parse_args()
@@ -268,6 +282,7 @@ def main():
normal_flow=_read_series_csv(args.normal_flow_csv),
min_dpressure=args.min_dpressure,
basic_pressure=args.basic_pressure,
n_workers=args.n_workers,
)
print(json.dumps(result, ensure_ascii=False))
+37 -1
View File
@@ -4,6 +4,7 @@ import copy
import math
import sys
from datetime import datetime
from time import perf_counter
import networkx as nx
import numpy as np
@@ -29,6 +30,7 @@ def _ensure_signatures_for_centers(
pressure_monitor,
flow_monitor, # 用来推断传感器列名
leak_mag, # 泄漏量,比如 400/3600
n_workers=1,
):
"""
只为缺失的中心补算 SLF(调用你现有的 cal_signature_pipe_multi_pf),
@@ -63,7 +65,12 @@ def _ensure_signatures_for_centers(
# 3) 若有缺失中心,仅为这些中心补算一次
if len(need) > 0:
p_new, _ = cal_signature_pipe_multi_pf(
wn, leak_mag, need, timestep_list, sensor_name_all
wn,
leak_mag,
need,
timestep_list,
sensor_name_all,
n_workers=n_workers,
)
# 初始化空缓存时,做一次“同构化”
if pressure_leak_all is None or len(pressure_leak_all) == 0:
@@ -206,6 +213,12 @@ def _dedupe_preserve_order(items):
return output
def _accumulate_stage(stage_timing, stage_name, started_at):
stage_timing[stage_name] = stage_timing.get(stage_name, 0.0) + (
perf_counter() - started_at
)
def cal_DtoTop1(
G0, pipe_leak, located_pipe, pipe_start_node_all, pipe_end_node_all, pipe_length
):
@@ -312,7 +325,11 @@ def DN_search_multi_simple_add_flow_count_new(
if_gy,
pressure_threshold,
leak_mag=400 / 3600,
n_workers=1,
stage_timing=None,
):
if stage_timing is None:
stage_timing = {}
iter_count = 0
all_node_iter = copy.deepcopy(all_node)
candidate_pipe_input = copy.deepcopy(candidate_pipe_input_initial) # 可能漏损管段
@@ -345,6 +362,7 @@ def DN_search_multi_simple_add_flow_count_new(
group_num = cal_group_num(candidate_pipe_input, group_basic_num)
# group 分组,得出候选漏损中心
stage_start = perf_counter()
candidate_center_list, candidate_group_list, new_all_node = (
metis_grouping_pipe_weight(
G0,
@@ -361,12 +379,14 @@ def DN_search_multi_simple_add_flow_count_new(
couple_node_length,
)
)
_accumulate_stage(stage_timing, "group_partitioning", stage_start)
simulation_times = simulation_times + len(candidate_center_list)
# pick_pressure_leak
# pressure_leak = pressure_leak_all.loc[candidate_center_list].loc[:, :]
# flow_leak = flow_leak_all.loc[candidate_center_list].loc[:, :]
# —— 新增泄漏量(保持你现在的一致,或从外部传入)——
# —— 只为缺失中心补算,然后取本轮需要的中心子集 ——
stage_start = perf_counter()
pressure_leak, flow_leak, pressure_leak_all, flow_leak_all = (
_ensure_signatures_for_centers(
wn=wn,
@@ -377,8 +397,10 @@ def DN_search_multi_simple_add_flow_count_new(
pressure_monitor=pressure_monitor,
flow_monitor=flow_monitor,
leak_mag=leak_mag,
n_workers=n_workers,
)
)
_accumulate_stage(stage_timing, "signature_for_candidates", stage_start)
# pressure_leak_f= pressure_leak.swaplevel()
@@ -414,6 +436,7 @@ def DN_search_multi_simple_add_flow_count_new(
# --------------------------------------------------------
# 只为 add_center 里还没算过的中心补算,并与本轮中心合并
if len(add_center) > 0:
stage_start = perf_counter()
pressure_add, flow_add, pressure_leak_all, flow_leak_all = (
_ensure_signatures_for_centers(
wn=wn,
@@ -424,8 +447,10 @@ def DN_search_multi_simple_add_flow_count_new(
pressure_monitor=pressure_monitor,
flow_monitor=flow_monitor,
leak_mag=leak_mag, # 与上面一致
n_workers=n_workers,
)
)
_accumulate_stage(stage_timing, "signature_for_extra_centers", stage_start)
pressure_leak = pd.concat([pressure_leak, pressure_add], axis=0)
if (flow_leak is not None) and (flow_add is not None):
flow_leak = pd.concat([flow_leak, flow_add], axis=0)
@@ -437,6 +462,7 @@ def DN_search_multi_simple_add_flow_count_new(
candidate_center_list_sup = _dedupe_preserve_order(
candidate_center_list + add_center
)
stage_start = perf_counter()
similarity, cos_h, dis_h, dis_f_h, break_flag = (
cal_similarity_all_multi_new_sq_improve_double_lzr(
candidate_center_list_sup,
@@ -463,6 +489,7 @@ def DN_search_multi_simple_add_flow_count_new(
max_flow,
)
)
_accumulate_stage(stage_timing, "similarity_ranking", stage_start)
if break_flag == 1:
break
@@ -473,6 +500,7 @@ def DN_search_multi_simple_add_flow_count_new(
cut_ratio, new_similarity = extra_judge(new_similarity)
else:
cut_ratio = 1
stage_start = perf_counter()
final_area_t, final_center_t, all_node_new_1, if_end = (
area_output_num_ki_improve(
candidate_center_list,
@@ -485,6 +513,7 @@ def DN_search_multi_simple_add_flow_count_new(
cut_ratio,
)
)
_accumulate_stage(stage_timing, "candidate_area_selection", stage_start)
final_area = final_area + final_area_t
final_center = final_center + final_center_t
@@ -522,6 +551,7 @@ def DN_search_multi_simple_add_flow_count_new(
final_area_pipe = list(final_area) # 确保是 list
# 只为还没算过的管段补齐 SLF(按需计算)
stage_start = perf_counter()
pressure_leak_sp, flow_leak_sp, pressure_leak_all, flow_leak_all = (
_ensure_signatures_for_centers(
wn=wn,
@@ -532,12 +562,15 @@ def DN_search_multi_simple_add_flow_count_new(
pressure_monitor=pressure_monitor,
flow_monitor=flow_monitor,
leak_mag=leak_mag,
n_workers=n_workers,
)
)
_accumulate_stage(stage_timing, "signature_for_final_area", stage_start)
# 如果你要精确统计模拟次数,这里可以加上“本次新补的数量”,
# 做法:让 _ensure_signatures_for_centers 额外返回 need_cnt,再 simulation_times += need_cnt
stage_start = perf_counter()
similarity_sp, cos_h, dis_h, dis_f_h, break_flag = (
cal_similarity_all_multi_new_sq_improve_double_lzr(
final_area_pipe,
@@ -564,6 +597,7 @@ def DN_search_multi_simple_add_flow_count_new(
max_flow,
)
)
_accumulate_stage(stage_timing, "similarity_final", stage_start)
else:
dpressure = (pressure_predict - pressure_monitor).mean()
@@ -599,6 +633,8 @@ def DN_search_multi_simple_add_flow_count_new(
similarity_sp = similarity_sp.sort_values(ascending=False)
t2 = datetime.now()
dt = (t2 - t1).seconds
stage_timing["iterations"] = iter_count + 1 if len(dpressure) > 0 else 0
stage_timing["total_elapsed_seconds"] = float(dt)
return similarity_sp.index[0], dt, simulation_times, wn, similarity_sp
@@ -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: