214 lines
8.2 KiB
Python
214 lines
8.2 KiB
Python
#!/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()
|