Files
TJWaterServerBinary/app/algorithms/burst_location/network_partitioner.py
T

457 lines
15 KiB
Python

"""管网分区模块。"""
import math
import matplotlib.pyplot as plt
import networkx as nx
import networkx as networkx
import numpy as np
import pandas as pd
import pymetis
from scipy.sparse import coo_matrix, csr_matrix
from scipy.sparse.csgraph import connected_components
def _to_metis_edge_weight(edge_weight):
weight = float(edge_weight)
if not math.isfinite(weight):
raise ValueError(f"Invalid non-finite METIS edge weight: {edge_weight}")
# pymetis expects integer edge weights.
return max(1, int(round(weight)))
def _dedupe_preserve_order(items):
seen = set()
output = []
for item in items:
if item in seen:
continue
seen.add(item)
output.append(item)
return output
def pick_center_pipe(node_x, node_y, candidate_pipe, pipe_start_node, pipe_end_node):
candidate_pipe_list = list(candidate_pipe)
start_nodes = pipe_start_node[candidate_pipe_list]
end_nodes = pipe_end_node[candidate_pipe_list]
x_vals = (node_x[start_nodes].to_numpy() + node_x[end_nodes].to_numpy()) / 2.0
y_vals = (node_y[start_nodes].to_numpy() + node_y[end_nodes].to_numpy()) / 2.0
mean_x = float(np.mean(x_vals))
mean_y = float(np.mean(y_vals))
distance = np.abs(x_vals - mean_x) + np.abs(y_vals - mean_y)
center_idx = int(np.argmin(distance))
return candidate_pipe_list[center_idx]
def pick_max_diameter_pipe(candidate_pipe, pipe_diameter):
candidate_pipe_list = list(candidate_pipe)
diameters = pd.to_numeric(
pipe_diameter.reindex(candidate_pipe_list), errors="coerce"
).dropna()
if len(diameters) != len(candidate_pipe_list):
missing = sorted(set(candidate_pipe_list) - set(diameters.index))
preview = ", ".join(map(str, missing[:10]))
raise ValueError(f"Missing or invalid diameter for pipes: {preview}")
max_diameter = float(diameters.max())
max_diameter_pipes = sorted(
[pipe for pipe, diameter in diameters.items() if float(diameter) == max_diameter],
key=str,
)
return max_diameter_pipes[0]
def pick_dual_center_pipes(
node_x, node_y, candidate_pipe, pipe_start_node, pipe_end_node, pipe_diameter
):
geometric_center = pick_center_pipe(
node_x, node_y, candidate_pipe, pipe_start_node, pipe_end_node
)
diameter_center = pick_max_diameter_pipe(candidate_pipe, pipe_diameter)
return _dedupe_preserve_order([geometric_center, diameter_center])
def find_new_center_pipe(
node_x,
node_y,
candidate_pipe,
pipe_start_node,
pipe_end_node,
pipe_diameter,
record_center,
):
new_candidate_pipe = sorted(set(candidate_pipe) - set(record_center))
if new_candidate_pipe == []:
new_candidate_pipe = candidate_pipe
center_t = pick_center_pipe(
node_x,
node_y,
new_candidate_pipe,
pipe_start_node,
pipe_end_node,
)
return center_t
def cal_area_node_linked_pipe(nodeset, node_pipe_dic):
pipeset = []
for temp_node in nodeset:
pipeset.extend(node_pipe_dic[temp_node])
return pipeset
def metis_grouping_pipe_weight(
G0,
wn,
all_node_iter,
candidate_pipe_input,
group_num,
node_x,
node_y,
pipe_start_node_all,
pipe_end_node_all,
node_pipe_dic,
all_node_series,
couple_node_length,
pipe_diameter,
):
all_node_iter_series_new = all_node_series[all_node_iter]
all_node_iter_series_new = all_node_iter_series_new.sort_values(ascending=True)
all_node_iter_new = list(all_node_iter_series_new.index)
G1 = G0.subgraph(all_node_iter_new)
delimiter = " "
adjacency_list = []
node_dict = {}
c_new = 0
for each_node in all_node_iter_new:
node_dict[each_node] = c_new
c_new = c_new + 1
correspond_dic = {}
count_node = 0
w = []
for node_name in all_node_iter_new:
neighbors = G1[node_name]
w_temp = []
n_t = [node_dict[node_name]]
for neighbor_name in sorted(neighbors.keys()):
edge_data = neighbors[neighbor_name]
edge_key = f"{node_name},{neighbor_name}"
reverse_edge_key = f"{neighbor_name},{node_name}"
if edge_key in couple_node_length:
edge_weight = couple_node_length[edge_key]
elif reverse_edge_key in couple_node_length:
edge_weight = couple_node_length[reverse_edge_key]
elif edge_data.get("weight") is not None:
edge_weight = float(edge_data["weight"])
else:
# Ignore graph edges that are outside candidate pipes and have no usable
# partition weight (e.g. some non-pipe links in mixed network graphs).
continue
w_temp.append(_to_metis_edge_weight(edge_weight))
n_t.append(node_dict[neighbor_name])
w.append(w_temp)
correspond_dic[n_t[0]] = count_node
count_node = count_node + 1
# del n_t[0]
adjacency_list.append(n_t)
adjacency_list_new = [[] * 1 for i in range(len(adjacency_list))]
w_new = [[] * 1 for i in range(len(adjacency_list))]
for i in range(len(adjacency_list)):
adjacency_list_new[int(adjacency_list[i][0])] = adjacency_list[i]
w_new[int(adjacency_list[i][0])] = w[i]
for i in range(len(adjacency_list)):
del adjacency_list_new[i][0]
xadj = [0]
w_f = []
final_adjacency_list = []
for i in range(len(adjacency_list_new)):
final_adjacency_list = final_adjacency_list + adjacency_list_new[i]
xadj.append(len(final_adjacency_list))
w_f = w_f + w_new[i]
# (edgecuts, parts) = pymetis.part_graph(nparts=group_num, adjacency=adjacency_list_new)
metis_options = pymetis.Options()
metis_options.seed = 42
(edgecuts, parts) = pymetis.part_graph(
nparts=group_num,
adjncy=final_adjacency_list,
xadj=xadj,
eweights=w_f,
options=metis_options,
)
# (edgecuts, parts) = pymetis.part_graph(nparts=group_num, adjacency=adjacency_list_new)
candidate_group_list = [[] * 1 for i in range(group_num)]
for i in range(len(all_node_iter_new)):
candidate_group_list[parts[i]].append(all_node_iter_new[i])
"""parts_new = np.zeros(len(candidate_node_input), dtype=int)
for i in range(len(candidate_group_list)):
temp_group = candidate_group_list[i]
for each_node in temp_group:
parts_new[node_dict[each_node]] = i
parts_new = list(parts_new)"""
new_center = []
new_group = []
new_center_candidates = []
new_all_node = []
candidate_pipe_set = set(candidate_pipe_input)
all_grouped_pipe = []
for i in range(group_num):
# 构建子图
G_sub = G0.subgraph(candidate_group_list[i])
# 计算联通子图
sub_graphs = networkx.connected_components(G_sub)
if networkx.number_connected_components(G_sub) == 1:
# 求交集
nodeset = G_sub.nodes()
pipeset_set = set(cal_area_node_linked_pipe(nodeset, node_pipe_dic))
candidate_pipe = sorted(pipeset_set.intersection(candidate_pipe_set))
# 判断集合是否保留
if len(candidate_pipe) > 0:
# 保留 计算中心
center_t = pick_center_pipe(
node_x,
node_y,
candidate_pipe,
pipe_start_node_all,
pipe_end_node_all,
)
center_candidates_t = pick_dual_center_pipes(
node_x,
node_y,
candidate_pipe,
pipe_start_node_all,
pipe_end_node_all,
pipe_diameter,
)
# 更新
new_center.append(center_t)
new_center_candidates.append(center_candidates_t)
new_group.append(candidate_pipe)
new_all_node.append(nodeset)
all_grouped_pipe = all_grouped_pipe + candidate_pipe
else:
for c in sorted(sub_graphs, key=lambda c: min(c)):
G_temp = G0.subgraph(c)
nodeset = G_temp.nodes()
pipeset = cal_area_node_linked_pipe(nodeset, node_pipe_dic)
pipeset_set = set(pipeset)
# 求交集
candidate_pipe = sorted(pipeset_set.intersection(candidate_pipe_set))
# print(len(candidate_node))
# 判断集合是否保留
if len(candidate_pipe) > 0:
# 保留 计算中心
center_t = pick_center_pipe(
node_x,
node_y,
candidate_pipe,
pipe_start_node_all,
pipe_end_node_all,
)
center_candidates_t = pick_dual_center_pipes(
node_x,
node_y,
candidate_pipe,
pipe_start_node_all,
pipe_end_node_all,
pipe_diameter,
)
# 更新
new_center.append(center_t)
new_center_candidates.append(center_candidates_t)
new_group.append(candidate_pipe)
new_all_node.append(nodeset)
all_grouped_pipe = all_grouped_pipe + candidate_pipe
record_center = []
c_g = 0
for each_group in new_group:
if len(each_group) < 3:
record_center.append(new_center[c_g])
c_g += 1
c_g = 0
for each_group in new_group:
if len(each_group) >= 3:
if new_center[c_g] in record_center:
new_center[c_g] = find_new_center_pipe(
node_x,
node_y,
each_group,
pipe_start_node_all,
pipe_end_node_all,
pipe_diameter,
record_center,
)
new_center_candidates[c_g] = _dedupe_preserve_order(
[new_center[c_g]] + list(new_center_candidates[c_g])
)
record_center.append(new_center[c_g])
c_g += 1
# visualize_metis_partition(
# G0, new_center, new_group,
# node_x, node_y,
# pipe_start_node_all, pipe_end_node_all
# )
return new_center, new_group, new_all_node, new_center_candidates
def visualize_metis_partition(
G,
center_pipes,
pipe_groups,
node_x,
node_y,
pipe_start_node_all,
pipe_end_node_all,
title: str | None = None,
block: bool = True,
pause_seconds: float | None = None,
):
"""
可视化METIS分区结果(单图模式)
参数:
G: 原始管网图(nx.Graph)
center_pipes: 中心管道列表(list)
pipe_groups: 分组管道列表(list of lists)
node_x: 节点X坐标字典(dict)
node_y: 节点Y坐标字典(dict)
pipe_start_node_all: 管道起点字典(dict)
pipe_end_node_all: 管道终点字典(dict)
"""
fig = plt.figure("metis_partition_convergence", figsize=(22.51, 12.48))
fig.clf()
ax = fig.add_subplot(111)
if not block:
plt.ion()
# 生成颜色映射(自动扩展颜色数量)
colors = plt.cm.tab20(np.linspace(0, 1, len(pipe_groups)))
# --- 绘制背景管网(灰色半透明) ---
for edge in G.edges():
start_node, end_node = edge
ax.plot(
[node_x[start_node], node_x[end_node]],
[node_y[start_node], node_y[end_node]],
color="lightgray",
linewidth=0.5,
alpha=0.3,
zorder=1, # 确保背景在底层
)
# --- 绘制各分区管道(彩色)---
legend_handles = [] # 用于图例的句柄
for i, (group, center) in enumerate(zip(pipe_groups, center_pipes)):
color = colors[i % len(colors)] # 循环使用颜色
# 绘制分组管道
for pipe in group:
start = pipe_start_node_all[pipe]
end = pipe_end_node_all[pipe]
line = ax.plot(
[node_x[start], node_x[end]],
[node_y[start], node_y[end]],
color=color,
linewidth=2.5,
alpha=0.8,
zorder=2,
)
# 只为每个分组的第一个管道添加图例句柄
if pipe == group[0]:
legend_handles.append(line[0])
# 高亮中心管道(红色虚线)
if center in pipe_start_node_all and center in pipe_end_node_all:
start = pipe_start_node_all[center]
end = pipe_end_node_all[center]
ax.plot(
[node_x[start], node_x[end]],
[node_y[start], node_y[end]],
color="red",
linewidth=4,
linestyle="--",
dash_capstyle="round",
zorder=3, # 确保中心管道在最顶层
)
# --- 添加图例和标注 ---
# 分组图例
if legend_handles:
group_labels = [f"Group {i + 1}" for i in range(len(pipe_groups))]
ax.legend(
legend_handles,
group_labels,
loc="upper right",
title="Partitions",
fontsize=8,
title_fontsize=10,
)
# 中心管道标注(可选)
for i, center in enumerate(center_pipes):
if center in pipe_start_node_all:
x = (
node_x[pipe_start_node_all[center]] + node_x[pipe_end_node_all[center]]
) / 2
y = (
node_y[pipe_start_node_all[center]] + node_y[pipe_end_node_all[center]]
) / 2
ax.text(
x,
y,
f"C{i + 1}",
color="red",
fontsize=10,
ha="center",
va="center",
bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
)
# --- 图形美化 ---
ax.set_title(title or "Water Network Partitioning Overview", fontsize=14, pad=20)
ax.set_xlabel("X Coordinate", fontsize=10)
ax.set_ylabel("Y Coordinate", fontsize=10)
ax.grid(True, alpha=0.2, linestyle=":")
fig.tight_layout()
# 显示图形并强制刷新,避免迭代显示滞后一轮。
plt.show(block=block)
if not block:
fig.canvas.draw_idle()
fig.canvas.flush_events()
pause_value = 0.001 if pause_seconds is None else max(0.0, float(pause_seconds))
plt.pause(max(0.001, pause_value))
elif pause_seconds is not None:
plt.pause(max(0.0, float(pause_seconds)))
return fig
def generate_adjlist_with_all_edges(G, delimiter):
for s, nbrs in G.adjacency():
line = str(s) + delimiter
for t, data in nbrs.items():
line += str(t) + delimiter
yield line[: -len(delimiter)]
def cal_group_num(candidate_node_input, cal_group_num):
candidate_node_num = len(candidate_node_input)
if candidate_node_num > 100:
group_num_input = cal_group_num # 30
else:
group_num_input = 10
return group_num_input