新增爆管位置检测模块及相关API接口

2026-03-06 15:27:59 +08:00
parent 63d3458fb4
commit b83b895e2b
11 changed files with 3084 additions and 0 deletions
@@ -0,0 +1,375 @@
+"""管网分区模块。"""
+
+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 pick_center_pipe(node_x, node_y, candidate_pipe, pipe_start_node, pipe_end_node):
+    data_set_t = pd.DataFrame(dtype=object)
+    data_set_t["x"] = (
+        node_x[pipe_start_node[candidate_pipe]].values
+        + node_x[pipe_start_node[candidate_pipe]].values
+    ) / 2
+    data_set_t["y"] = (
+        node_y[pipe_end_node[candidate_pipe]].values
+        + node_y[pipe_end_node[candidate_pipe]].values
+    ) / 2
+    data_set_t.index = list(candidate_pipe)
+    mean_x = data_set_t["x"].mean()
+    mean_y = data_set_t["y"].mean()
+    data_set_t["d"] = abs(data_set_t["x"] - mean_x) + abs(data_set_t["y"] - mean_y)
+    distance_t = data_set_t["d"].sort_values(ascending=True, inplace=False)
+    """if distance_t.index==[]:
+        print(candidate_pipe)
+    else:"""
+    center_t = distance_t.index[0]
+    return center_t
+
+
+def find_new_center_pipe(
+    node_x, node_y, candidate_pipe, pipe_start_node, pipe_end_node, record_center
+):
+    new_candidate_pipe = list(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 = []
+    nodeset = list(nodeset)
+    for i in range(len(nodeset)):
+        temp_node = nodeset[i]
+        pipe = node_pipe_dic[temp_node]
+        pipeset = pipeset + pipe
+    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,
+):
+    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 line in generate_adjlist_with_all_edges(G1, delimiter):
+        temp_node_name = line.split(sep=delimiter)
+        w_temp = []
+        for i in range(len(temp_node_name) - 1):
+            temp_name_1 = temp_node_name[0] + "," + temp_node_name[i + 1]
+            w_temp.append(couple_node_length[temp_name_1])
+        w.append(w_temp)
+        n_t = []
+        for each_node in temp_node_name:
+            n_t.append(node_dict[each_node])
+        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)
+    (edgecuts, parts) = pymetis.part_graph(
+        nparts=group_num, adjncy=final_adjacency_list, xadj=xadj, eweights=w_f
+    )
+    # (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_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 = list(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,
+                )
+
+                # 更新
+                new_center.append(center_t)
+                new_group.append(candidate_pipe)
+                new_all_node.append(nodeset)
+                all_grouped_pipe = all_grouped_pipe + candidate_pipe
+        else:
+            for c in sub_graphs:
+                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 = list(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,
+                    )
+                    # 更新
+                    new_center.append(center_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,
+                    record_center,
+                )
+            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
+
+
+def visualize_metis_partition(
+    G, center_pipes, pipe_groups, node_x, node_y, pipe_start_node_all, pipe_end_node_all
+):
+    """
+    可视化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)
+    """
+    plt.figure(figsize=(9, 10))
+
+    # 生成颜色映射（自动扩展颜色数量）
+    colors = plt.cm.tab20(np.linspace(0, 1, len(pipe_groups)))
+
+    # --- 绘制背景管网（灰色半透明） ---
+    for edge in G.edges():
+        start_node, end_node = edge
+        plt.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 = plt.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]
+            plt.plot(
+                [node_x[start], node_x[end]],
+                [node_y[start], node_y[end]],
+                color="red",
+                linewidth=4,
+                linestyle="--",
+                dash_capstyle="round",
+                zorder=3,  # 确保中心管道在最顶层
+            )
+
+    # --- 添加图例和标注 ---
+    # 分组图例
+    group_labels = [f"Group {i + 1}" for i in range(len(pipe_groups))]
+    plt.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
+            plt.text(
+                x,
+                y,
+                f"C{i + 1}",
+                color="red",
+                fontsize=10,
+                ha="center",
+                va="center",
+                bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
+            )
+
+    # --- 图形美化 ---
+    plt.title("Water Network Partitioning Overview", fontsize=14, pad=20)
+    plt.xlabel("X Coordinate", fontsize=10)
+    plt.ylabel("Y Coordinate", fontsize=10)
+    plt.grid(True, alpha=0.2, linestyle=":")
+    plt.tight_layout()
+
+    # 显示图形
+    plt.show()
+
+
+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
+
+
+class NetworkPartitioner:
+    @staticmethod
+    def metis_grouping_pipe_weight(*args, **kwargs):
+        return metis_grouping_pipe_weight(*args, **kwargs)
+
+    @staticmethod
+    def visualize_metis_partition(*args, **kwargs):
+        return visualize_metis_partition(*args, **kwargs)
+
+    @staticmethod
+    def generate_adjlist_with_all_edges(*args, **kwargs):
+        return generate_adjlist_with_all_edges(*args, **kwargs)
+
+    @staticmethod
+    def pick_center_pipe(*args, **kwargs):
+        return pick_center_pipe(*args, **kwargs)
+
+    @staticmethod
+    def find_new_center_pipe(*args, **kwargs):
+        return find_new_center_pipe(*args, **kwargs)
+
+    @staticmethod
+    def cal_area_node_linked_pipe(*args, **kwargs):
+        return cal_area_node_linked_pipe(*args, **kwargs)
+
+    @staticmethod
+    def cal_group_num(*args, **kwargs):
+        return cal_group_num(*args, **kwargs)