#!/usr/bin/env python3 """ 供水服务范围分析与分区 — 可复用脚本 基于实时水力数据，从水库沿水流方向追溯服务范围，自动发现水库并分区。 用法: python3 service_area_partition.py \ --pipe-props pipes.json \ --reservoirs reservoirs.json \ --links realtime_links.json \ --nodes realtime_nodes.json \ --target-time '2026-04-01T08:00:00+08:00' \ --output ./service_area_partition_wrapper.json """ import argparse import json import os import sys from collections import deque, defaultdict COLORS = [ "rgba(31,119,180,0.7)", "rgba(255,127,14,0.7)", "rgba(44,160,44,0.7)", "rgba(148,103,189,0.7)", "rgba(140,86,75,0.7)", "rgba(227,119,194,0.7)", "rgba(127,127,127,0.7)", "rgba(188,189,34,0.7)", "rgba(23,190,207,0.7)", "rgba(174,199,232,0.7)", "rgba(255,152,150,0.7)", "rgba(196,156,148,0.7)", "rgba(219,64,82,0.7)", "rgba(153,204,153,0.7)", "rgba(255,204,102,0.7)", "rgba(102,102,204,0.7)", "rgba(204,102,102,0.7)", "rgba(102,204,204,0.7)", "rgba(204,153,204,0.7)", "rgba(153,153,153,0.7)" ] def norm_time(t): """把 ISO8601 字符串归一化为 UTC 的 ISO 字符串，用于跨时区比较""" from datetime import datetime return datetime.fromisoformat(t.replace("Z", "+00:00")).astimezone( __import__("datetime").timezone.utc).isoformat() def load_json(path, label): print(f"Loading {label}...", file=sys.stderr) with open(path) as f: return json.load(f) def main(): parser = argparse.ArgumentParser(description="供水服务范围分区分析") parser.add_argument("--pipe-props", required=True, help="管道属性 JSON 文件") parser.add_argument("--reservoirs", required=True, help="水库属性 JSON 文件") parser.add_argument("--links", required=True, help="实时管段数据 JSON 文件") parser.add_argument("--nodes", required=True, help="实时节点数据 JSON 文件") parser.add_argument("--target-time", required=True, help="目标时刻 ISO8601") parser.add_argument("--output", required=True, help="分区结果输出 JSON 路径") args = parser.parse_args() # --- Step 1: Load pipe topology --- pdata = load_json(args.pipe_props, "pipe topology")["data"] pipe_topology = {} node_neighbors = defaultdict(set) for p in pdata: pid = p["id"] n1, n2 = p["node1"], p["node2"] pipe_topology[pid] = (n1, n2) node_neighbors[n1].add(n2) node_neighbors[n2].add(n1) print(f" {len(pdata)} pipes, {len(node_neighbors)} unique nodes", file=sys.stderr) # --- Step 2: Discover reservoirs --- rdata = load_json(args.reservoirs, "reservoirs")["data"] reservoirs = [r["id"] for r in rdata] print(f" {len(reservoirs)} reservoirs: {reservoirs}", file=sys.stderr) # --- Step 3: Load link flow at target time --- ldata = load_json(args.links, "link flows")["data"] target_ts = norm_time(args.target_time) target_links = [l for l in ldata if norm_time(l["time"]) == target_ts] flow_direction = {} pipe_flow = {} for l in target_links: lid = l.get("link_id") or l.get("id") flow_val = l["flow"] pipe_flow[lid] = abs(flow_val) if lid in pipe_topology: n1, n2 = pipe_topology[lid] if flow_val > 1e-6: flow_direction[lid] = (n1, n2) elif flow_val < -1e-6: flow_direction[lid] = (n2, n1) nonzero = len(flow_direction) print(f" {len(target_links)} link records, {nonzero} with non-zero flow", file=sys.stderr) # --- Step 4: Load node data at target time --- ndata = load_json(args.nodes, "node data")["data"] target_nodes = [n for n in ndata if norm_time(n["time"]) == target_ts] node_pressure = {} node_demand = {} for n in target_nodes: nid = n.get("node_id") or n.get("id") node_pressure[nid] = n.get("pressure", 0) node_demand[nid] = n.get("actual_demand", 0) print(f" {len(target_nodes)} nodes", file=sys.stderr) # --- Step 5: Build downstream graph --- downstream = defaultdict(set) for _lid, (up, dn) in flow_direction.items(): downstream[up].add(dn) print(f" downstream graph: {len(downstream)} source nodes", file=sys.stderr) # --- Step 6: Multi-source BFS along flow direction --- reservoir_area = {} node_served_by = {} queue = deque() for rid in reservoirs: reservoir_area[rid] = {rid} node_served_by[rid] = rid queue.append((rid, rid, 0)) while queue: node, source, dist = queue.popleft() for neighbor in downstream.get(node, set()): if neighbor not in node_served_by: node_served_by[neighbor] = source reservoir_area[source].add(neighbor) queue.append((neighbor, source, dist + 1)) directed_count = len(node_served_by) unassigned = set(node_pressure.keys()) - set(node_served_by.keys()) print(f" flow-tracing assigned: {directed_count}, unassigned: {len(unassigned)}", file=sys.stderr) # --- Step 7: Undirected proximity fallback --- if unassigned: print(" running proximity fallback...", file=sys.stderr) ua_queue = deque() ua_visited = {} for nid, src in node_served_by.items(): ua_visited[nid] = src ua_queue.append((nid, src, 0)) while ua_queue: node, source, dist = ua_queue.popleft() for neighbor in node_neighbors.get(node, set()): if neighbor not in ua_visited: ua_visited[neighbor] = source reservoir_area[source].add(neighbor) ua_queue.append((neighbor, source, dist + 1)) still = set(node_pressure.keys()) - set(ua_visited.keys()) if still: print(f" WARNING: {len(still)} nodes still unassigned", file=sys.stderr) node_served_by = ua_visited # --- Step 8: Compute statistics --- print(f"\n=== 供水服务范围分区统计 ===\n", file=sys.stderr) area_stats = [] for rid in reservoirs: nodes_in = reservoir_area.get(rid, set()) pressures = [node_pressure[n] for n in nodes_in if n in node_pressure] demands = [node_demand[n] for n in nodes_in if n in node_demand] area_stats.append({ "reservoir": rid, "node_count": len(nodes_in), "total_demand": round(sum(demands), 4), "avg_pressure": round(sum(pressures)/len(pressures), 2) if pressures else 0, "min_pressure": round(min(pressures), 2) if pressures else 0, "max_pressure": round(max(pressures), 2) if pressures else 0, }) area_stats.sort(key=lambda x: x["node_count"], reverse=True) for s in area_stats: print(f" 水源 {s['reservoir']:>8s}: {s['node_count']:>6d} 节点 | " f"总需水={s['total_demand']:.2f} | " f"压力 avg={s['avg_pressure']:.1f}m [{s['min_pressure']:.1f}–{s['max_pressure']:.1f}m]", file=sys.stderr) # --- Step 9: Assign colors and write output --- area_colors = {} for i, rid in enumerate(reservoirs): area_colors[rid] = COLORS[i % len(COLORS)] output = { "node_area_map": node_served_by, "area_ids": reservoirs, "area_colors": area_colors, "metadata": { "analysis_time": args.target_time, "total_nodes": len(node_served_by), "reservoir_count": len(reservoirs), "directed_assigned": directed_count, "proximity_assigned": len(node_served_by) - directed_count, "method": "flow-direction-source-tracing" } } absolute_output = os.path.abspath(args.output) wrapper = { "metadata": { "generated_by": "service_area_partition.py", "schema_version": 1, }, "location": {"file_path": absolute_output}, "data": output, } with open(absolute_output, "w", encoding="utf-8") as f: json.dump(wrapper, f, ensure_ascii=False) summary = { "total_nodes": len(node_served_by), "reservoirs": len(reservoirs), "areas": area_stats, "output_file": absolute_output } print(json.dumps(summary, ensure_ascii=False)) if __name__ == "__main__": main()