153 lines
5.2 KiB
Python
153 lines
5.2 KiB
Python
import ctypes
|
|
import os
|
|
import numpy as np
|
|
import pymetis
|
|
from .database import *
|
|
from .s0_base import get_nodes
|
|
from .s32_region_util import get_nodes_in_region
|
|
from .s32_region_util import Topology
|
|
|
|
|
|
PARTITION_TYPE_RB = 0
|
|
PARTITION_TYPE_KWAY = 1
|
|
|
|
'''
|
|
adjacency_list = [np.array([4, 2, 1]),
|
|
np.array([0, 2, 3]),
|
|
np.array([4, 3, 1, 0]),
|
|
np.array([1, 2, 5, 6]),
|
|
np.array([0, 2, 5]),
|
|
np.array([4, 3, 6]),
|
|
np.array([5, 3])]
|
|
n_cuts, membership = pymetis.part_graph(2, adjacency=adjacency_list)
|
|
# n_cuts = 3
|
|
# membership = [1, 1, 1, 0, 1, 0, 0]
|
|
|
|
nodes_part_0 = np.argwhere(np.array(membership) == 0).ravel() # [3, 5, 6]
|
|
nodes_part_1 = np.argwhere(np.array(membership) == 1).ravel() # [0, 1, 2, 4]
|
|
|
|
print(nodes_part_0)
|
|
print(nodes_part_1)
|
|
'''
|
|
|
|
|
|
def calculate_district_metering_area_for_nodes(name: str, nodes: list[str], part_count: int = 1, part_type: int = PARTITION_TYPE_RB) -> list[list[str]]:
|
|
topology = Topology(name, nodes)
|
|
t_nodes = topology.nodes()
|
|
t_links = topology.links()
|
|
t_node_list = topology.node_list()
|
|
|
|
adjacency_list = []
|
|
|
|
for node in t_node_list:
|
|
links: list[str] = t_nodes[node]['links']
|
|
a_nodes: list[int] = []
|
|
for link in links:
|
|
if t_links[link]['node1'] == node:
|
|
i = t_node_list.index(t_links[link]['node2'])
|
|
a_nodes.append(i)
|
|
elif t_links[link]['node2'] == node:
|
|
i = t_node_list.index(t_links[link]['node1'])
|
|
a_nodes.append(i)
|
|
adjacency_list.append(np.array(a_nodes))
|
|
|
|
recursive = part_type == PARTITION_TYPE_RB
|
|
n_cuts, membership = pymetis.part_graph(nparts=part_count, adjacency=adjacency_list, recursive=recursive, contiguous=True)
|
|
|
|
result: list[list[str]] = []
|
|
for i in range(0, part_count):
|
|
indices: list[int] = list(np.argwhere(np.array(membership) == i).ravel())
|
|
index_strs: list[str] = []
|
|
for index in indices:
|
|
index_strs.append(t_node_list[index])
|
|
result.append(index_strs)
|
|
|
|
return result
|
|
|
|
|
|
def _calculate_district_metering_area_for_nodes(name: str, nodes: list[str], part_count: int = 1, part_type: int = PARTITION_TYPE_RB) -> list[list[str]]:
|
|
if part_type != PARTITION_TYPE_RB and part_type != PARTITION_TYPE_KWAY:
|
|
return []
|
|
if part_count <= 0:
|
|
return []
|
|
elif part_count == 1:
|
|
return [nodes]
|
|
|
|
lib = ctypes.CDLL(os.path.join(os.getcwd(), 'api', 'CMetis.dll'))
|
|
|
|
METIS_NOPTIONS = 40
|
|
c_options = (ctypes.c_int64 * METIS_NOPTIONS)()
|
|
|
|
METIS_OK = 1
|
|
result = lib.set_default_options(c_options)
|
|
if result != METIS_OK:
|
|
return []
|
|
|
|
METIS_OPTION_PTYPE , METIS_OPTION_CONTIG = 0, 13
|
|
c_options[METIS_OPTION_PTYPE] = part_type
|
|
c_options[METIS_OPTION_CONTIG] = 1
|
|
|
|
topology = Topology(name, nodes)
|
|
t_nodes = topology.nodes()
|
|
t_links = topology.links()
|
|
t_node_list = topology.node_list()
|
|
t_link_list = topology.link_list()
|
|
|
|
nedges = len(t_link_list) * 2
|
|
|
|
c_nvtxs = ctypes.c_int64(len(t_node_list))
|
|
c_ncon = ctypes.c_int64(1)
|
|
c_xadj = (ctypes.c_int64 * (c_nvtxs.value + 1))()
|
|
c_adjncy = (ctypes.c_int64 * nedges)()
|
|
c_vwgt = (ctypes.c_int64 * (c_ncon.value * c_nvtxs.value))()
|
|
c_adjwgt = (ctypes.c_int64 * nedges)()
|
|
c_vsize = (ctypes.c_int64 * c_nvtxs.value)()
|
|
|
|
c_xadj[0] = 0
|
|
|
|
l, n = 0, 0
|
|
c_xadj_i = 1
|
|
for node in t_node_list:
|
|
links = t_nodes[node]['links']
|
|
for link in links:
|
|
node1 = t_links[link]['node1']
|
|
node2 = t_links[link]['node2']
|
|
c_adjncy[l] = t_node_list.index(node2) if node2 != node else t_node_list.index(node1)
|
|
c_adjwgt[l] = 1
|
|
l += 1
|
|
if len(links) > 0:
|
|
c_xadj[c_xadj_i] = l # adjncy.size()
|
|
c_xadj_i += 1
|
|
c_vwgt[n] = 1
|
|
c_vsize[n] = 1
|
|
n += 1
|
|
|
|
part_func = lib.part_graph_recursive if part_type == PARTITION_TYPE_RB else lib.part_graph_kway
|
|
|
|
c_nparts = ctypes.c_int64(part_count)
|
|
c_tpwgts = ctypes.POINTER(ctypes.c_double)()
|
|
c_ubvec = ctypes.POINTER(ctypes.c_double)()
|
|
c_out_edgecut = ctypes.c_int64(0)
|
|
c_out_part = (ctypes.c_int64 * c_nvtxs.value)()
|
|
result = part_func(ctypes.byref(c_nvtxs), ctypes.byref(c_ncon), c_xadj, c_adjncy, c_vwgt, c_vsize, c_adjwgt, ctypes.byref(c_nparts), c_tpwgts, c_ubvec, c_options, ctypes.byref(c_out_edgecut), c_out_part)
|
|
if result != METIS_OK:
|
|
return []
|
|
|
|
dmas : list[list[str]]= []
|
|
for i in range(part_count):
|
|
dmas.append([])
|
|
for i in range(c_nvtxs.value):
|
|
dmas[c_out_part[i]].append(t_node_list[i])
|
|
|
|
return dmas
|
|
|
|
|
|
def calculate_district_metering_area_for_region(name: str, region: str, part_count: int = 1, part_type: int = PARTITION_TYPE_RB) -> list[list[str]]:
|
|
nodes = get_nodes_in_region(name, region)
|
|
return calculate_district_metering_area_for_nodes(name, nodes, part_count, part_type)
|
|
|
|
|
|
def calculate_district_metering_area_for_network(name: str, part_count: int = 1, part_type: int = PARTITION_TYPE_RB) -> list[list[str]]:
|
|
nodes = get_nodes(name)
|
|
return calculate_district_metering_area_for_nodes(name, nodes, part_count, part_type)
|