import wntr import numpy as np import pandas as pd import os import time import argparse from multiprocessing import Pool, cpu_count from typing import Any, List, Dict, Union from pymoo.core.problem import Problem from pymoo.core.callback import Callback from pymoo.algorithms.soo.nonconvex.ga import GA from pymoo.operators.crossover.sbx import SBX from pymoo.operators.mutation.pm import PM from pymoo.optimize import minimize as pymoo_minimize from pymoo.termination.default import DefaultSingleObjectiveTermination from app.algorithms._utils import _cleanup_temp_files _worker_data: dict[str, Any] = {} DEFAULT_N_WORKERS = max(1, min(cpu_count() - 1, 4)) def _worker_init( inp_path: str, sensor_nodes: list[str], area_ids: list[str], nodes_by_area: dict[str, list[str]], obs_matrix: np.ndarray, q_sum: float, duration_sec: float, timestep_sec: float, ) -> None: global _worker_data wn = wntr.network.WaterNetworkModel(inp_path) wn.options.hydraulic.demand_model = "DD" wn.options.time.duration = duration_sec wn.options.time.hydraulic_timestep = timestep_sec wn.options.time.pattern_timestep = timestep_sec wn.options.time.report_timestep = timestep_sec demand_objs_by_area = {} allocatable_counts = {} for area_id in area_ids: demand_objs = [] for node_name in nodes_by_area.get(area_id, []): if node_name not in wn.node_name_list: continue node = wn.get_node(node_name) if ( hasattr(node, "demand_timeseries_list") and len(node.demand_timeseries_list) > 0 ): demand_objs.append(node.demand_timeseries_list[0]) demand_objs_by_area[area_id] = demand_objs allocatable_counts[area_id] = len(demand_objs) _worker_data = { "wn": wn, "sensor_nodes": sensor_nodes, "area_ids": area_ids, "nodes_by_area": nodes_by_area, "demand_objs_by_area": demand_objs_by_area, "allocatable_counts": allocatable_counts, "obs_matrix": obs_matrix, "q_sum": q_sum, } def _worker_evaluate(raw_ratios: np.ndarray) -> float: d = _worker_data effective_ratio_map = LeakageIdentifier._effective_area_ratios( raw_ratios, d["area_ids"], d["nodes_by_area"], allocatable_counts=d["allocatable_counts"], ) modifications = [] for area_id in d["area_ids"]: ratio = effective_ratio_map.get(area_id, 0.0) if ratio <= 0: continue demand_objs = d["demand_objs_by_area"].get(area_id, []) if not demand_objs: continue per_node_leak = d["q_sum"] * ratio / len(demand_objs) for demand_obj in demand_objs: original_val = demand_obj.base_value demand_obj.base_value = original_val + per_node_leak modifications.append((demand_obj, original_val)) temp_dir = os.path.abspath(os.path.join("temp", "leakage")) os.makedirs(temp_dir, exist_ok=True) prefix = os.path.join(temp_dir, f"temp_{os.getpid()}") try: sim = wntr.sim.EpanetSimulator(d["wn"]) results = sim.run_sim(file_prefix=prefix) sim_pressure = results.node["pressure"].loc[:, d["sensor_nodes"]] n_steps = min(sim_pressure.shape[0], d["obs_matrix"].shape[0]) sim_vals = sim_pressure.values[:n_steps, :] obs_vals = d["obs_matrix"][:n_steps, :] diff = sim_vals - obs_vals row_max = np.max(np.abs(diff), axis=1, keepdims=True) row_max[row_max == 0] = 1.0 normalized_diff = diff / row_max return float(np.linalg.norm(normalized_diff)) except Exception: return 1e9 finally: for demand_obj, original_val in modifications: demand_obj.base_value = original_val _cleanup_temp_files(prefix) class LeakageIdentifier: FLOW_UNIT_TO_M3S = { "m3/s": 1.0, "m3/h": 1.0 / 3600.0, "L/s": 1.0 / 1000.0, "L/min": 1.0 / 60000.0, } @classmethod def _flow_to_m3s(cls, value: float, unit: str) -> float: if unit not in cls.FLOW_UNIT_TO_M3S: raise ValueError(f"不支持的流量单位: {unit}") return float(value) * cls.FLOW_UNIT_TO_M3S[unit] @classmethod def _flow_from_m3s(cls, value_m3s: float, unit: str) -> float: if unit not in cls.FLOW_UNIT_TO_M3S: raise ValueError(f"不支持的流量单位: {unit}") return float(value_m3s) / cls.FLOW_UNIT_TO_M3S[unit] @staticmethod def _effective_area_ratios( raw_ratios: Union["np.ndarray", Dict[str, float]], area_ids: List[str], nodes_by_area: Dict[str, List[str]], allocatable_counts: Union[Dict[str, int], None] = None, ) -> Dict[str, float]: """将输入比例转换为有效区域比例，确保有效区域比例和为 1。""" area_count = len(area_ids) if area_count == 0: return {} if isinstance(raw_ratios, dict): ratios = np.array( [float(raw_ratios.get(area_id, 0.0)) for area_id in area_ids], dtype=float, ) else: arr = np.asarray(raw_ratios, dtype=float).reshape(-1) ratios = np.zeros(area_count, dtype=float) fill_len = min(area_count, arr.shape[0]) if fill_len > 0: ratios[:fill_len] = arr[:fill_len] # 仅保留非负比例，负值按 0 处理 ratios = np.clip(ratios, a_min=0.0, a_max=None) # 仅在有效区域（存在可分配节点）内归一化 if allocatable_counts is not None: valid_mask = np.array( [int(allocatable_counts.get(area_id, 0)) > 0 for area_id in area_ids], dtype=bool, ) else: valid_mask = np.array( [len(nodes_by_area.get(area_id, [])) > 0 for area_id in area_ids], dtype=bool, ) if not np.any(valid_mask): raise ValueError("没有可分配漏损的有效分区，无法满足漏损总量约束。") effective = np.zeros(area_count, dtype=float) valid_sum = float(np.sum(ratios[valid_mask])) if valid_sum > 0: effective[valid_mask] = ratios[valid_mask] / valid_sum else: # 若输入全为 0，则在有效区域内均分，保证总和仍为 1 valid_count = int(np.sum(valid_mask)) effective[valid_mask] = 1.0 / valid_count return {area_id: float(effective[idx]) for idx, area_id in enumerate(area_ids)} @staticmethod def _normalize_area_map_df(df: pd.DataFrame) -> pd.DataFrame: """标准化区域映射列名为 ID 和 Area。""" if "ID" in df.columns and "Area" in df.columns: return df if "ID" in df.columns and "now" in df.columns: df = df.rename(columns={"now": "Area"}) return df df = df.copy() df.columns = ["ID", "Area"] + list(df.columns[2:]) return df def __init__( self, inp_path: str, sensor_nodes: List[str], area_map: Union[str, Dict[str, str]], start_time: float = 0, duration: float = 24, timestep: float = 5, q_sum: float = 0.2, ): """ 初始化漏损识别器。 参数: inp_path: EPANET .inp 文件路径。 sensor_nodes: 用作压力传感器的节点 ID 列表。 area_map: 节点到区域的映射。可以是 CSV 文件路径（列：ID, Area），也可以是字典 {NodeID: AreaID}。 start_time: 模拟开始时间（小时）。 duration: 模拟持续时间（小时）。 timestep: 模拟时间步长（分钟）。 q_sum: 假设的总漏损流量 (m3/s)。 """ self.inp_path = inp_path self.sensor_nodes = sensor_nodes self.start_time = start_time self.duration = duration self.timestep = timestep self.q_sum = q_sum # 加载管网模型（仅一次） self.wn = wntr.network.WaterNetworkModel(self.inp_path) # 优化 WNTR 设置以提高速度 self.wn.options.hydraulic.demand_model = "DD" self.wn.options.time.duration = float(self.duration) * 3600 self.wn.options.time.hydraulic_timestep = float(self.timestep) * 60 self.wn.options.time.pattern_timestep = float(self.timestep) * 60 self.wn.options.time.report_timestep = float(self.timestep) * 60 # 加载区域映射 if isinstance(area_map, str): self.area_map_df = self._load_area_map(area_map) elif isinstance(area_map, dict): self.area_map_df = self._normalize_area_map_df( pd.DataFrame(list(area_map.items()), columns=["ID", "Area"]) ) else: raise ValueError("area_map 必须是 CSV 文件路径或字典。") self.area_ids = sorted(self.area_map_df["Area"].unique()) self.num_areas = len(self.area_ids) # 按区域对节点进行预分类，以便更快查找 self.nodes_by_area = { area: self.area_map_df[self.area_map_df["Area"] == area]["ID"].tolist() for area in self.area_ids } def _load_area_map(self, path: str) -> pd.DataFrame: """加载并验证节点-区域映射文件。""" df = pd.read_csv(path, dtype={"ID": str, "Area": str}) return self._normalize_area_map_df(df) def run_identification( self, observed_pressure_data: Union[ str, pd.DataFrame, Dict[str, List[Any]], List[Dict[str, Any]] ], output_dir: str = "Results", pop_size: int = 50, max_gen: int = 100, output_flow_unit: str = "m3/s", save_result: bool = True, ftol: float = 1e-3, ftol_period: int = 15, n_workers: int = DEFAULT_N_WORKERS, ): """ 运行遗传算法以识别漏损分布。 参数: observed_pressure_data: 包含 SCADA 压力数据的 CSV 文件路径或 DataFrame/字典列表数据。 output_dir: 结果保存目录。 pop_size: GA 的种群大小。 max_gen: GA 的最大代数。 output_flow_unit: 输出漏损流量的单位。 save_result: 是否保存识别结果到本地 CSV。 ftol: 目标值收敛容差（连续 ftol_period 代改善 < ftol 则停止）。 ftol_period: 收敛检测的窗口代数。 n_workers: 并行工作进程数（1=串行，>1=并行评估）。 """ if save_result: os.makedirs(output_dir, exist_ok=True) # 加载观测数据 if isinstance(observed_pressure_data, str): obs_df = pd.read_csv(observed_pressure_data) observed_name = os.path.basename(observed_pressure_data) elif isinstance(observed_pressure_data, pd.DataFrame): obs_df = observed_pressure_data.copy() observed_name = "observed_pressure.csv" else: obs_df = pd.DataFrame(observed_pressure_data) observed_name = "observed_pressure.csv" # 准备 pymoo 问题实例 problem = LeakageProblem( self.wn, self.nodes_by_area, self.area_ids, self.sensor_nodes, obs_df, q_sum=self.q_sum, n_workers=n_workers, inp_path=os.path.abspath(self.inp_path), ) # 配置 pymoo GA 算法 n_var = self.num_areas algorithm = GA( pop_size=pop_size, crossover=SBX(prob=0.9, eta=15), mutation=PM(prob=1.0 / max(1, n_var), eta=20), eliminate_duplicates=False, ) # 终止条件：收敛检测 + 最大代数 termination = DefaultSingleObjectiveTermination( ftol=ftol, period=ftol_period, n_max_gen=max_gen, ) # 回调：记录每代信息 callback = _ProgressCallback() t0 = time.time() try: res = pymoo_minimize( problem, algorithm, termination, seed=42, verbose=True, callback=callback, ) finally: problem.close() elapsed = time.time() - t0 # 提取最优解 best_ind = res.X # 最优个体（漏损比例原始值） best_obj = float(res.F[0]) # 输出终止信息 print(f"\n优化完成。耗时: {elapsed:.1f}s") print(f"总代数: {res.algorithm.n_gen}, 总评估次数: {problem._eval_count}") print(f"最佳目标值: {best_obj:.6f}") # 保存到文件 effective_ratio_map = self._effective_area_ratios( best_ind, self.area_ids, self.nodes_by_area, allocatable_counts=problem.allocatable_counts, ) normalized_ratios = [ effective_ratio_map.get(area_id, 0.0) for area_id in self.area_ids ] leakage_flow_m3s = [ratio * self.q_sum for ratio in normalized_ratios] leakage_flow_output = [ self._flow_from_m3s(value_m3s, output_flow_unit) for value_m3s in leakage_flow_m3s ] result_df = pd.DataFrame( { "Area": self.area_ids, "LeakageRatioRaw": best_ind, "LeakageRatio": normalized_ratios, "LeakageFlow_m3_per_s": leakage_flow_m3s, f"LeakageFlow_{output_flow_unit.replace('/', '_per_')}": leakage_flow_output, } ) result_path = None if save_result: result_path = os.path.join( output_dir, f"identified_leakage_{observed_name}" ) result_df.to_csv(result_path, index=False) print(f"结果已保存至 {result_path}") result_df.attrs["result_path"] = result_path return result_df class _ProgressCallback(Callback): """每代回调：记录进度。""" def __init__(self): super().__init__() self.gen_times = [] self._t_last = None def notify(self, algorithm): now = time.time() if self._t_last is not None: self.gen_times.append(now - self._t_last) self._t_last = now class LeakageProblem(Problem): """pymoo 批量评估问题定义。 搜索空间：n 维 [0, 1] 实数 -> 通过 _effective_area_ratios 归一化到单纯形。 目标：模拟压力与观测压力之间的归一化误差范数。 无显式约束（sum=1 由归一化自动保证）。 """ def __init__( self, wn, nodes_by_area, area_ids, sensor_nodes, observed_data, q_sum: float = 0.2, n_workers: int = DEFAULT_N_WORKERS, inp_path: str | None = None, ): n_var = len(area_ids) super().__init__( n_var=n_var, n_obj=1, n_ieq_constr=0, xl=np.zeros(n_var), xu=np.ones(n_var), ) self.wn = wn self.nodes_by_area = nodes_by_area self.area_ids = area_ids self.sensor_nodes = sensor_nodes self.q_sum = q_sum self.n_workers = max(1, int(n_workers)) self.inp_path = inp_path # 预处理观测数据以匹配模拟格式 try: missing_sensors = [ s for s in self.sensor_nodes if s not in observed_data.columns ] if not missing_sensors: self.obs_matrix = observed_data[self.sensor_nodes].values else: self.obs_matrix = observed_data.values[:, : len(self.sensor_nodes)] except Exception: self.obs_matrix = observed_data.values[:, : len(self.sensor_nodes)] duration_sec = float(self.wn.options.time.duration) step_sec = float(self.wn.options.time.hydraulic_timestep) if step_sec > 0: max_steps = int(duration_sec / step_sec) + 1 self.obs_matrix = self.obs_matrix[:max_steps, :] # 预先缓存每个区域的需水对象，减少每次适应度计算的节点查找 self.demand_objs_by_area = {} for area_id in self.area_ids: demand_objs = [] for node_name in self.nodes_by_area.get(area_id, []): if node_name not in self.wn.node_name_list: continue node = self.wn.get_node(node_name) if ( hasattr(node, "demand_timeseries_list") and len(node.demand_timeseries_list) > 0 ): demand_objs.append(node.demand_timeseries_list[0]) self.demand_objs_by_area[area_id] = demand_objs self.allocatable_counts = { area_id: len(self.demand_objs_by_area.get(area_id, [])) for area_id in self.area_ids } if not any(count > 0 for count in self.allocatable_counts.values()): raise ValueError("没有可分配漏损的有效分区，无法满足漏损总量约束。") # 评估计数器（诊断用） self._eval_count = 0 self._pool = None if self.n_workers > 1: if not self.inp_path: raise ValueError("并行评估需要提供 inp_path。") duration_sec = float(self.wn.options.time.duration) timestep_sec = float(self.wn.options.time.hydraulic_timestep) self._pool = Pool( processes=self.n_workers, initializer=_worker_init, initargs=( self.inp_path, list(self.sensor_nodes), list(self.area_ids), {k: list(v) for k, v in self.nodes_by_area.items()}, self.obs_matrix.copy(), self.q_sum, duration_sec, timestep_sec, ), ) def _evaluate(self, X, out, *args, **kwargs): """批量评估种群。 X: 形状 (pop_size, n_var) 的决策变量矩阵。 """ n_pop = X.shape[0] self._eval_count += n_pop if self._pool is not None: results = self._pool.map(_worker_evaluate, [X[i] for i in range(n_pop)]) out["F"] = np.array(results, dtype=float).reshape(-1, 1) return F = np.zeros((n_pop, 1)) for i in range(n_pop): F[i, 0] = self._evaluate_single(X[i]) out["F"] = F def _evaluate_single(self, x): """评估单个个体，返回归一化误差范数。""" leak_ratios = x # 将漏损分布归一化 effective_ratio_map = LeakageIdentifier._effective_area_ratios( leak_ratios, self.area_ids, self.nodes_by_area, allocatable_counts=self.allocatable_counts, ) # 跟踪修改以便稍后恢复 modifications = [] for j, area_id in enumerate(self.area_ids): ratio = effective_ratio_map.get(area_id, 0.0) if ratio <= 0: continue demand_objs = self.demand_objs_by_area.get(area_id, []) if not demand_objs: continue per_node_leak = self.q_sum * ratio / len(demand_objs) for demand_obj in demand_objs: original_val = demand_obj.base_value demand_obj.base_value = original_val + per_node_leak modifications.append((demand_obj, original_val)) # 结果保存在根目录的temp/leakage文件夹中 temp_dir = os.path.abspath(os.path.join("temp", "leakage")) os.makedirs(temp_dir, exist_ok=True) prefix = os.path.join(temp_dir, f"temp_{os.getpid()}") try: sim = wntr.sim.EpanetSimulator(self.wn) results = sim.run_sim(file_prefix=prefix) sim_pressure = results.node["pressure"].loc[:, self.sensor_nodes] n_steps = min(sim_pressure.shape[0], self.obs_matrix.shape[0]) sim_vals = sim_pressure.values[:n_steps, :] obs_vals = self.obs_matrix[:n_steps, :] diff = sim_vals - obs_vals # 按行最大值归一化 row_max = np.max(np.abs(diff), axis=1, keepdims=True) row_max[row_max == 0] = 1.0 # 防止除以零 normalized_diff = diff / row_max # 目标：归一化差值矩阵的 2-范数 return float(np.linalg.norm(normalized_diff)) except Exception: return 1e9 finally: for demand_obj, original_val in modifications: demand_obj.base_value = original_val _cleanup_temp_files(prefix) def close(self) -> None: if self._pool is not None: self._pool.close() self._pool.join() self._pool = None def main() -> int: parser = argparse.ArgumentParser(description="漏损区域识别") parser.add_argument("--inp", required=True, help=".inp 文件路径") parser.add_argument("--map", help="节点-区域映射 CSV 路径") parser.add_argument("--scada", help="SCADA 压力 CSV 路径 (观测数据)") parser.add_argument("--sensors", help="传感器节点 ID 列表 (逗号分隔)") parser.add_argument("--output", default="Results", help="输出目录") parser.add_argument("--pop_size", type=int, default=50, help="种群大小") parser.add_argument("--max_gen", type=int, default=100, help="最大代数") parser.add_argument("--duration", type=float, default=24, help="模拟时长(小时)") parser.add_argument("--q_sum", type=float, default=0.241, help="总漏损流量") parser.add_argument( "--q_sum_unit", default="m3/s", choices=list(LeakageIdentifier.FLOW_UNIT_TO_M3S.keys()), help="q_sum 输入单位（建议与现场习惯一致，内部统一换算为 m3/s）", ) args = parser.parse_args() if not args.map or not args.scada or not args.sensors: parser.error("--map、--scada、--sensors 为必填") q_sum_m3s = LeakageIdentifier._flow_to_m3s(args.q_sum, args.q_sum_unit) sensors = [sensor.strip() for sensor in args.sensors.split(",") if sensor.strip()] identifier = LeakageIdentifier( args.inp, sensors, args.map, duration=args.duration, q_sum=q_sum_m3s ) identifier.run_identification( args.scada, args.output, pop_size=args.pop_size, max_gen=args.max_gen, output_flow_unit=args.q_sum_unit, ) return 0 if __name__ == "__main__": raise SystemExit(main())