"""Pure topology partitioning used by DMA leakage estimation.""" import math from collections import deque from typing import Any, Iterable, Mapping import numpy as np def build_dma_partitions( sensor_nodes: list[str], node_coords: Mapping[str, Mapping[str, Any]], link_entries: Iterable[str], dma_count: int | None, ) -> tuple[dict[str, str], list[dict[str, Any]]]: """Assign every topology node to a sensor-seeded virtual DMA.""" all_nodes = list(node_coords) if not all_nodes: raise ValueError("管网中未获取到可分区节点。") available_sensors = [node for node in sensor_nodes if node in node_coords] if not available_sensors: raise ValueError("无可用压力传感器，无法生成虚拟分区。") area_count = _resolve_dma_count(dma_count, available_sensors, all_nodes) sensor_area_map = _cluster_sensors_to_areas( available_sensors, node_coords, area_count ) adjacency = _build_adjacency(link_entries, all_nodes) distance_by_sensor = { sensor: _bfs_distances(adjacency, sensor) for sensor in available_sensors } assignment_count = {sensor: 0 for sensor in available_sensors} area_map: dict[str, str] = {} for node_id in sorted(all_nodes): sensor = _choose_sensor_for_node( node_id, available_sensors, node_coords, distance_by_sensor, assignment_count, ) assignment_count[sensor] += 1 area_map[node_id] = sensor_area_map[sensor] return area_map, _build_area_meta(area_map, sensor_area_map) def _resolve_dma_count( dma_count: int | None, sensor_nodes: list[str], all_nodes: list[str] ) -> int: if dma_count is None: return min(len(sensor_nodes), len(all_nodes)) if dma_count <= 0: raise ValueError("dma_count 必须大于 0。") if dma_count > len(all_nodes): raise ValueError("dma_count 不能大于可分区节点数量。") if dma_count > len(sensor_nodes): raise ValueError("dma_count 不能大于可用传感器数量。") return dma_count def _cluster_sensors_to_areas( sensor_nodes: list[str], node_coords: Mapping[str, Mapping[str, Any]], area_count: int, ) -> dict[str, str]: if area_count >= len(sensor_nodes): return {sensor: str(index + 1) for index, sensor in enumerate(sensor_nodes)} points = np.array( [ [float(node_coords[sensor]["x"]), float(node_coords[sensor]["y"])] for sensor in sensor_nodes ], dtype=float, ) centers = points[:area_count].copy() labels = np.full(points.shape[0], -1, dtype=int) for _ in range(20): distances_squared = ( (points[:, None, :] - centers[None, :, :]) ** 2 ).sum(axis=2) next_labels = distances_squared.argmin(axis=1) if np.array_equal(labels, next_labels): break labels = next_labels for index in range(area_count): cluster_points = points[labels == index] if cluster_points.size > 0: centers[index] = cluster_points.mean(axis=0) labels = _restore_empty_area_labels(labels, points, centers, area_count) return { sensor: str(int(labels[index]) + 1) for index, sensor in enumerate(sensor_nodes) } def _restore_empty_area_labels( labels: np.ndarray, points: np.ndarray, centers: np.ndarray, area_count: int, ) -> np.ndarray: """Keep every requested area represented when coordinates are degenerate.""" labels = labels.copy() for missing_area in sorted(set(range(area_count)) - set(labels.tolist())): area_sizes = { area: int(np.count_nonzero(labels == area)) for area in range(area_count) } donor_area = max( (area for area, size in area_sizes.items() if size > 1), key=lambda area: (area_sizes[area], -area), ) donor_indices = np.flatnonzero(labels == donor_area) replacement_index = max( (int(index) for index in donor_indices), key=lambda index: ( float(np.sum((points[index] - centers[donor_area]) ** 2)), index, ), ) labels[replacement_index] = missing_area centers[missing_area] = points[replacement_index] return labels def _build_adjacency( link_entries: Iterable[str], all_nodes: list[str] ) -> dict[str, set[str]]: adjacency: dict[str, set[str]] = {node: set() for node in all_nodes} for link in link_entries: parts = str(link).split(":") if len(parts) < 4: continue node1, node2 = parts[-2], parts[-1] if node1 in adjacency and node2 in adjacency: adjacency[node1].add(node2) adjacency[node2].add(node1) return adjacency def _bfs_distances(adjacency: Mapping[str, set[str]], start: str) -> dict[str, int]: distances = {start: 0} queue: deque[str] = deque([start]) while queue: node = queue.popleft() for neighbor in adjacency.get(node, set()): if neighbor in distances: continue distances[neighbor] = distances[node] + 1 queue.append(neighbor) return distances def _choose_sensor_for_node( node_id: str, sensors: list[str], node_coords: Mapping[str, Mapping[str, Any]], distance_by_sensor: Mapping[str, Mapping[str, int]], assignment_count: Mapping[str, int], ) -> str: min_distance: int | None = None candidates: list[str] = [] for sensor in sensors: distance = distance_by_sensor.get(sensor, {}).get(node_id) if distance is None: continue if min_distance is None or distance < min_distance: min_distance = distance candidates = [sensor] elif distance == min_distance: candidates.append(sensor) if not candidates: node_coord = node_coords[node_id] return min( sensors, key=lambda sensor: math.hypot( float(node_coord["x"]) - float(node_coords[sensor]["x"]), float(node_coord["y"]) - float(node_coords[sensor]["y"]), ), ) return min(candidates, key=lambda sensor: (assignment_count[sensor], sensor)) def _build_area_meta( area_map: Mapping[str, str], sensor_area_map: Mapping[str, str] ) -> list[dict[str, Any]]: nodes_by_area: dict[str, list[str]] = {} for node_id, area_id in area_map.items(): nodes_by_area.setdefault(area_id, []).append(node_id) sensors_by_area: dict[str, list[str]] = {} for sensor, area_id in sensor_area_map.items(): sensors_by_area.setdefault(area_id, []).append(sensor) return [ { "area_id": area_id, "sensor_nodes": sorted(sensors_by_area.get(area_id, [])), "node_ids": sorted(nodes_by_area[area_id]), "node_count": len(nodes_by_area[area_id]), } for area_id in sorted(nodes_by_area, key=int) ]