"""漏损模拟模块。""" import math import multiprocessing as mp import os import sys import pandas as pd import wntr from app.algorithms._utils import _cleanup_temp_files _PIPE2LEAKNODE = None _SIGNATURE_WORKER_DATA = {} def _make_temp_prefix(tag): temp_dir = os.path.abspath(os.path.join("temp", "burst_location")) os.makedirs(temp_dir, exist_ok=True) safe_tag = str(tag).replace(os.sep, "_").replace(" ", "_") return os.path.join(temp_dir, f"{safe_tag}_{os.getpid()}") def _snapshot_hydraulic_options(wn): options = wn.options return { "demand_model": options.hydraulic.demand_model, "duration": float(options.time.duration), "hydraulic_timestep": float(options.time.hydraulic_timestep), "pattern_timestep": float(options.time.pattern_timestep), "report_timestep": float(options.time.report_timestep), "required_pressure": float(options.hydraulic.required_pressure), "minimum_pressure": float(options.hydraulic.minimum_pressure), } def _apply_hydraulic_options(wn, option_values): options = wn.options options.hydraulic.demand_model = option_values["demand_model"] options.time.duration = option_values["duration"] options.time.hydraulic_timestep = option_values["hydraulic_timestep"] options.time.pattern_timestep = option_values["pattern_timestep"] options.time.report_timestep = option_values["report_timestep"] options.hydraulic.required_pressure = option_values["required_pressure"] options.hydraulic.minimum_pressure = option_values["minimum_pressure"] def simple_add_leak(wn, leak_mag, leak_pipe): whole_inf = dict() leak_pipe_self = wn.get_link(leak_pipe) pipe_diameter = leak_pipe_self.diameter pipe_length = leak_pipe_self.length pipe_roughness = leak_pipe_self.roughness pipe_minor_loss = leak_pipe_self.minor_loss # pipe_status = leak_pipe_self.status # pipe_check_valve = leak_pipe_self.check_valve pipe_start_node = leak_pipe_self.start_node_name pipe_end_node = leak_pipe_self.end_node_name # close the pipe # leak_pipe_self.status = 'Closed' wn.remove_link(leak_pipe) # add the pipe add_pipe1 = leak_pipe + "A" add_pipe2 = leak_pipe + "B" add_node = leak_pipe + "_" start_n = wn.get_node(pipe_start_node) end_n = wn.get_node(pipe_end_node) if start_n.node_type == "Reservoir": end_n_elevation = end_n.elevation start_n_elevation = end_n_elevation elif end_n.node_type == "Reservoir": start_n_elevation = start_n.elevation end_n_elevation = start_n_elevation else: end_n_elevation = end_n.elevation start_n_elevation = start_n.elevation elevation_self = (start_n_elevation + end_n_elevation) / 2 coordinates_self = ( (start_n.coordinates[0] + end_n.coordinates[0]) / 2, (start_n.coordinates[1] + end_n.coordinates[1]), ) wn.add_junction( add_node, base_demand=0, elevation=elevation_self, coordinates=coordinates_self ) leak_node = wn.get_node(add_node) wn.add_pipe( add_pipe1, start_node_name=pipe_start_node, end_node_name=add_node, length=pipe_length / 2, diameter=pipe_diameter, roughness=pipe_roughness, minor_loss=pipe_minor_loss, ) wn.add_pipe( add_pipe2, start_node_name=pipe_end_node, end_node_name=add_node, length=pipe_length / 2, diameter=pipe_diameter, roughness=pipe_roughness, minor_loss=pipe_minor_loss, ) # simulation leak_node.add_demand(base=leak_mag, pattern_name="add_leak") whole_inf["leak_node_name"] = add_node whole_inf["add_pipe1"] = add_pipe1 whole_inf["add_pipe2"] = add_pipe2 whole_inf["leak_pipe"] = leak_pipe whole_inf["pipe_start_node"] = pipe_start_node whole_inf["pipe_end_node"] = pipe_end_node whole_inf["pipe_length"] = pipe_length whole_inf["pipe_diameter"] = pipe_diameter whole_inf["pipe_roughness"] = pipe_roughness whole_inf["pipe_minor_loss"] = pipe_diameter return wn, whole_inf, add_pipe1 def simple_recover_wn(wn, whole_inf): leak_node = wn.get_node(whole_inf["leak_node_name"]) del leak_node.demand_timeseries_list[-1] # update wn.remove_link(whole_inf["add_pipe1"]) wn.remove_link(whole_inf["add_pipe2"]) wn.remove_node(whole_inf["leak_node_name"]) # open the pipe # leak_pipe_self.status = 'Open' wn.add_pipe( whole_inf["leak_pipe"], start_node_name=whole_inf["pipe_start_node"], end_node_name=whole_inf["pipe_end_node"], length=whole_inf["pipe_length"], diameter=whole_inf["pipe_diameter"], roughness=whole_inf["pipe_roughness"], minor_loss=whole_inf["pipe_minor_loss"], ) return wn def disable_all_controls_temporarily(wn): """返回(控制名, 控制对象)的列表，之后可用 restore_controls 还原。""" removed = [] # WNTR 的控制都在 wn.control_name_list / wn.get_control / wn.remove_control for cname in list(wn.control_name_list): ctrl = wn.get_control(cname) removed.append((cname, ctrl)) wn.remove_control(cname) return removed def restore_controls(wn, removed): """把先前禁用的控制全部加回去。""" for cname, ctrl in removed: wn.add_control(cname, ctrl) def set_pipe2leaknode_mapping(mapping): global _PIPE2LEAKNODE _PIPE2LEAKNODE = mapping def _get_or_create_leak_demand_ts(leak_node): """ 返回：泄漏专用 demand 的下标 idx。 若不存在，以 category='leak' 新建一条 base=0.0 的 demand。 """ # 先尝试找到已有的 'leak' 分类 for i, ts in enumerate(leak_node.demand_timeseries_list): # WNTR 的 Demand object 存在 category 属性 if getattr(ts, "category", None) == "leak": return i # 没有则新建（base=0.0，后续临时改 base_value） leak_node.add_demand(base=0.0, pattern_name=None, category="leak") return len(leak_node.demand_timeseries_list) - 1 def ensure_mid_node(wn, leak_pipe): add_pipe1 = f"{leak_pipe}A" add_pipe2 = f"{leak_pipe}B" add_node = f"{leak_pipe}__mid" if add_node in wn.node_name_list: return add_node if leak_pipe in wn.link_name_list: leak_pipe_self = wn.get_link(leak_pipe) pipe_diameter = leak_pipe_self.diameter pipe_length = leak_pipe_self.length pipe_roughness = leak_pipe_self.roughness pipe_minor_loss = leak_pipe_self.minor_loss pipe_start_node = leak_pipe_self.start_node_name pipe_end_node = leak_pipe_self.end_node_name start_n = wn.get_node(pipe_start_node) end_n = wn.get_node(pipe_end_node) if start_n.node_type == "Reservoir": end_elev = end_n.elevation start_elev = end_elev elif end_n.node_type == "Reservoir": start_elev = start_n.elevation end_elev = start_elev else: end_elev = end_n.elevation start_elev = start_n.elevation elev_mid = (start_elev + end_elev) / 2.0 x_mid = (start_n.coordinates[0] + end_n.coordinates[0]) / 2.0 y_mid = (start_n.coordinates[1] + end_n.coordinates[1]) / 2.0 wn.remove_link(leak_pipe) wn.add_junction( add_node, base_demand=0.0, elevation=elev_mid, coordinates=(x_mid, y_mid) ) wn.add_pipe( add_pipe1, start_node_name=pipe_start_node, end_node_name=add_node, length=pipe_length / 2.0, diameter=pipe_diameter, roughness=pipe_roughness, minor_loss=pipe_minor_loss, ) wn.add_pipe( add_pipe2, start_node_name=add_node, end_node_name=pipe_end_node, length=pipe_length / 2.0, diameter=pipe_diameter, roughness=pipe_roughness, minor_loss=pipe_minor_loss, ) return add_node # 若 A/B 已存在但中点不在，建议确认网络一致性 raise KeyError(f"Cannot ensure mid node for pipe '{leak_pipe}'.") def leak_simulation_pipe_dd_multi_pf( wn, leak_mag, leak_pipe, sensor_name, file_prefix=None ): """ 优化版： - 不再 remove/add link/node - 直接在预插入的中点泄漏节点上设置 base_demand = leak_mag；仿真后设回 0 """ wn.options.hydraulic.demand_model = "DD" # 确保中点节点存在 leak_node_name = ensure_mid_node(wn, leak_pipe) leak_node = wn.get_node(leak_node_name) # 拿到泄漏专用的 demand time-series 下标 leak_idx = _get_or_create_leak_demand_ts(leak_node) ts_obj = leak_node.demand_timeseries_list[leak_idx] # 记录原值（通常是 0.0） orig_base = ts_obj.base_value try: # 打开泄漏：只改 base_value，不碰 base_demand（只读） ts_obj.base_value = float(leak_mag) # 仿真 sim = wntr.sim.EpanetSimulator(wn) if file_prefix is None: results = sim.run_sim() else: results = sim.run_sim(file_prefix=file_prefix) # 输出（保持列顺序） pressure_output = results.node["pressure"].loc[:, sensor_name] # flow_output = results.link['flowrate'].loc[:, sensor_f_name] return wn, pressure_output finally: # 关闭泄漏：还原 base_value ts_obj.base_value = orig_base if file_prefix is not None: _cleanup_temp_files(file_prefix) def prepare_leak_infrastructure(wn, candidate_pipes): """ 把 candidate_pipes 每条管段切成两段，并在中点插入一个泄漏节点（base_demand=0）。 返回一个映射：pipe_id -> leak_node_name 注意：只做一次；后续仿真通过在该节点设置 base_demand 实现“打开泄漏”，结束后恢复为 0。 """ pipe2leaknode = {} for leak_pipe in candidate_pipes: if leak_pipe in pipe2leaknode: continue leak_pipe_self = wn.get_link(leak_pipe) pipe_diameter = leak_pipe_self.diameter pipe_length = leak_pipe_self.length pipe_roughness = leak_pipe_self.roughness pipe_minor_loss = leak_pipe_self.minor_loss pipe_start_node = leak_pipe_self.start_node_name pipe_end_node = leak_pipe_self.end_node_name # 计算中点高程/坐标（与原逻辑一致） start_n = wn.get_node(pipe_start_node) end_n = wn.get_node(pipe_end_node) if start_n.node_type == "Reservoir": end_elev = end_n.elevation start_elev = end_elev elif end_n.node_type == "Reservoir": start_elev = start_n.elevation end_elev = start_elev else: end_elev = end_n.elevation start_elev = start_n.elevation elev_mid = (start_elev + end_elev) / 2.0 x_mid = (start_n.coordinates[0] + end_n.coordinates[0]) / 2.0 y_mid = (start_n.coordinates[1] + end_n.coordinates[1]) / 2.0 # 先删原管，再加中点与两段半长管（只做一次） wn.remove_link(leak_pipe) add_pipe1 = f"{leak_pipe}A" add_pipe2 = f"{leak_pipe}B" add_node = f"{leak_pipe}__mid" # 唯一命名，后面直接用它当泄漏节点 wn.add_junction( add_node, base_demand=0.0, elevation=elev_mid, coordinates=(x_mid, y_mid) ) wn.add_pipe( add_pipe1, start_node_name=pipe_start_node, end_node_name=add_node, length=pipe_length / 2.0, diameter=pipe_diameter, roughness=pipe_roughness, minor_loss=pipe_minor_loss, ) wn.add_pipe( add_pipe2, start_node_name=add_node, end_node_name=pipe_end_node, length=pipe_length / 2.0, diameter=pipe_diameter, roughness=pipe_roughness, minor_loss=pipe_minor_loss, ) pipe2leaknode[leak_pipe] = add_node return pipe2leaknode def normal_simulation_pf( wn, drive_mode, sensor_name, sensor_f_name, inp_time, require_p, minimum_p ): # inp_time = 0 if drive_mode == "PDD": # 需水量根据节点压力动态调整 wn.options.hydraulic.demand_model = "PDD" wn.options.hydraulic.required_pressure = require_p wn.options.hydraulic.minimum_pressure = minimum_p elif drive_mode == "DD": # 需水量固定，与压力无关 wn.options.hydraulic.demand_model = "DD" sim = wntr.sim.EpanetSimulator(wn) results = sim.run_sim() pressure_all = results.node["pressure"][sensor_name] pressure = pressure_all.iloc[inp_time] demand_all = results.node["demand"] demand = demand_all.iloc[inp_time] sum_demand = cal_sum_demand(demand) flow_all = results.link["flowrate"][sensor_f_name] flow = flow_all.iloc[inp_time] top_sensor = pressure.idxmin() basic_p = results.node["pressure"] basic_p = basic_p.iloc[inp_time] return pressure, flow, basic_p, top_sensor, sum_demand def normal_simulation_multi_pf( wn, drive_mode, sensor_name, sensor_f_name, require_p, minimum_p ): # inp_time = 0 if drive_mode == "PDD": wn.options.hydraulic.demand_model = "PDD" wn.options.hydraulic.required_pressure = require_p wn.options.hydraulic.minimum_pressure = minimum_p elif drive_mode == "DD": wn.options.hydraulic.demand_model = "DD" sim = wntr.sim.EpanetSimulator(wn) results = sim.run_sim() pressure_all = results.node["pressure"][sensor_name] pressure = pressure_all demand_all = results.node["demand"] demand = demand_all flow = results.link["flowrate"][sensor_f_name] sum_demand = pd.Series(dtype=object) for i in range(len(demand.index)): sum_demand[str(demand.index[i])] = cal_sum_demand(demand.iloc[i]) if type(pressure) == pd.core.series.Series: top_sensor = pressure.idxmin() else: mean_pressure = pressure.mean() top_sensor = mean_pressure.idxmin() basic_p = results.node["pressure"] return pressure, flow, basic_p, top_sensor, sum_demand def simple_simulation_pf(wn, sensor_name, sensor_f_name, leak_pipe, add_pipe1): sim = wntr.sim.EpanetSimulator(wn) results = sim.run_sim() pressure_all = results.node["pressure"][sensor_name] if len(leak_pipe) > 0 and leak_pipe in sensor_f_name: f_sensor_name = [add_pipe1 if i == leak_pipe else i for i in sensor_f_name] flow_all = results.link["flowrate"][f_sensor_name] flow_all.columns = sensor_f_name else: flow_all = results.link["flowrate"][sensor_f_name] return pressure_all, flow_all def cal_sum_demand(demand): sum_demand = 0 for i in range(len(demand)): if demand.iloc[i] > 0: sum_demand += demand.iloc[i] return sum_demand def cal_signature_pipe_multi_pf( wn, leak_mag, candidate_center, timestep_list, sensor_name, n_workers=1, wn_inp_path=None, ): candidate_center_num = len(candidate_center) pressure_leak = pd.DataFrame( index=pd.MultiIndex.from_product([candidate_center, timestep_list]), columns=sensor_name, ) # flow_leak = pd.DataFrame(index=pd.MultiIndex.from_product([candidate_center, timestep_list]), # columns=sensor_f_name) pressure_leak = pressure_leak.sort_index() # flow_leak = flow_leak.sort_index() can_parallel = ( n_workers > 1 and candidate_center_num > 1 and wn_inp_path is not None and len(str(wn_inp_path)) > 0 ) if can_parallel: option_values = _snapshot_hydraulic_options(wn) worker_count = min(n_workers, candidate_center_num) start_methods = mp.get_all_start_methods() context_name = "spawn" if "spawn" in start_methods else start_methods[0] with mp.get_context(context_name).Pool( processes=worker_count, initializer=_signature_worker_init, initargs=( str(wn_inp_path), float(leak_mag), list(sensor_name), option_values, list(candidate_center), ), ) as pool: for i, (center_name, pressure_array) in enumerate( pool.imap(_signature_worker_run_center, candidate_center) ): pressure_leak.loc[(center_name, slice(None)), :] = pressure_array sys.stdout.write("\r" + "已经完成计算" + str(i + 1) + "个特征中心") else: # Pre-insert all mid-nodes so every simulation sees the same topology for center in candidate_center: ensure_mid_node(wn, center) for i in range(candidate_center_num): temp_prefix = _make_temp_prefix(f"sig_{i}") wn, pressure_output = leak_simulation_pipe_dd_multi_pf( wn, leak_mag, candidate_center[i], sensor_name, file_prefix=temp_prefix, ) # 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 _signature_worker_init( inp_path, leak_mag, sensor_name, option_values, candidate_centers=None ): global _SIGNATURE_WORKER_DATA wn = wntr.network.WaterNetworkModel(inp_path) _apply_hydraulic_options(wn, option_values) # Pre-insert ALL mid-nodes so every simulation runs on the same topology, # regardless of which worker handles which task. if candidate_centers is not None: for center in candidate_centers: ensure_mid_node(wn, center) _SIGNATURE_WORKER_DATA = { "wn": wn, "leak_mag": leak_mag, "sensor_name": sensor_name, } def _signature_worker_run_center(center_name): data = _SIGNATURE_WORKER_DATA temp_prefix = _make_temp_prefix(f"sig_worker_{center_name}") _, pressure_output = leak_simulation_pipe_dd_multi_pf( data["wn"], data["leak_mag"], center_name, data["sensor_name"], file_prefix=temp_prefix, ) return center_name, pressure_output.to_numpy() def pick_pipe(all_pipes, pipe_diameter, limited_diameter): candidate_pipe = [] for each_pipe in all_pipes: if pipe_diameter[each_pipe] >= limited_diameter: candidate_pipe.append(each_pipe) return candidate_pipe def cal_possible_pipe(leak_flow, all_pipe, pipe_diameter): basic_pressure = 10 # 基础压力 discharge_coeff = 0.6 # 经验系数 break_area_ratio = 1 # 爆管面积比 0.5 1.25 break_area = leak_flow / ( discharge_coeff * math.sqrt(2 * basic_pressure * 9.81) ) # 爆管面积 m3/h """break_area_diameter = math.sqrt(4 * break_area / math.pi) min_diameter = (math.ceil(1000 * break_area_diameter / break_area_ratio)) / 1000""" break_area_diameter = math.sqrt( 4 * break_area / math.pi / break_area_ratio ) # 爆管直径 min_diameter = (math.ceil(1000 * break_area_diameter)) / 1000 # 向上取整 new_all_pipe = pick_pipe(all_pipe, pipe_diameter, min_diameter) return new_all_pipe, min_diameter def extract_links(data, link_types, direction): return [ link for res_data in data.values() for link_type in link_types for link in res_data[link_type][direction] ]