amachine.am_optimize.am_optimize_node_permutation_cpsat

View Source

 1import warnings
 2
 3import networkx as nx
 4import numpy as np
 5
 6from ..am_random import resolve_rng
 7
 8def _calculate_score(
 9    G: nx.DiGraph,
10    instance_cohesion: np.ndarray,
11    perm: list[int]
12) -> float:
13
14    nodes = list(G.nodes())
15    node_to_idx = {node: i for i, node in enumerate(nodes)}
16
17    total_score = 0.0
18    for u, v in G.edges():
19        i = node_to_idx[u]
20        j = node_to_idx[v]
21        
22        assigned_instance_u = perm[i]
23        assigned_instance_v = perm[j]
24        
25        total_score += instance_cohesion[assigned_instance_u, assigned_instance_v]
26
27    return total_score
28
29def optimize_node_permutation_cpsat(
30    G: nx.DiGraph,
31    cohesion_matrix: np.ndarray,
32    max_search_time : int,
33    np_rng : np.random.Generator ) -> list[int] :
34
35    from ortools.sat.python import cp_model
36
37    np_rng = resolve_rng( np_rng )
38    random_seed = np_rng.integers(0, 2**31 - 1)
39
40    N = len(list(G.nodes()))
41    node_to_idx = {node: i for i, node in enumerate(G.nodes())}
42
43    SCALE = 1_000_000
44    cohesion_int = np.round(cohesion_matrix * SCALE).astype(int)
45    flat = cohesion_int.flatten().tolist()
46    score_lo, score_hi = int(cohesion_int.min()), int(cohesion_int.max())
47
48    model = cp_model.CpModel()
49    perm = [model.NewIntVar(0, N - 1, f"p{i}") for i in range(N)]
50    model.AddAllDifferent(perm)
51
52    edge_scores = []
53    for u, v in G.edges():
54        i, j = node_to_idx[u], node_to_idx[v]
55
56        flat_idx = model.NewIntVar(0, N * N - 1, f"idx_{i}_{j}")
57        model.Add(flat_idx == N * perm[i] + perm[j])
58
59        score = model.NewIntVar(score_lo, score_hi, f"s_{i}_{j}")
60        model.AddElement(flat_idx, flat, score)
61        edge_scores.append(score)
62
63    model.Maximize(cp_model.LinearExpr.Sum(edge_scores))
64
65    solver = cp_model.CpSolver()
66    solver.parameters.num_search_workers  = 3
67    solver.parameters.max_time_in_seconds = max_search_time
68    solver.parameters.linearization_level = 2
69    solver.parameters.random_seed = int(random_seed)
70
71    status = solver.Solve(model)
72
73    best_permutation = [ i for i in range( N ) ]
74    print( f"initial node permuation score {_calculate_score(G,cohesion_matrix,best_permutation)}" )
75
76    if status in (cp_model.OPTIMAL, cp_model.FEASIBLE) :
77        msg = "OPTIMAL" if status == cp_model.OPTIMAL else "FEASIBLE"
78        print(f"{msg} solution found. Objective = {solver.ObjectiveValue():.0f}")
79        best_permutation = [solver.Value(perm[i]) for i in range(N)]
80        print( f"cp-sat node permuation score {_calculate_score(G,cohesion_matrix,best_permutation)}" )
81    else:
82        warnings.warn("optimization found no solution within time limit.")
83
84    return best_permutation

def optimize_node_permutation_cpsat( G: networkx.classes.digraph.DiGraph, cohesion_matrix: numpy.ndarray, max_search_time: int, np_rng: numpy.random._generator.Generator) -> list[int]: View Source

30def optimize_node_permutation_cpsat(
31    G: nx.DiGraph,
32    cohesion_matrix: np.ndarray,
33    max_search_time : int,
34    np_rng : np.random.Generator ) -> list[int] :
35
36    from ortools.sat.python import cp_model
37
38    np_rng = resolve_rng( np_rng )
39    random_seed = np_rng.integers(0, 2**31 - 1)
40
41    N = len(list(G.nodes()))
42    node_to_idx = {node: i for i, node in enumerate(G.nodes())}
43
44    SCALE = 1_000_000
45    cohesion_int = np.round(cohesion_matrix * SCALE).astype(int)
46    flat = cohesion_int.flatten().tolist()
47    score_lo, score_hi = int(cohesion_int.min()), int(cohesion_int.max())
48
49    model = cp_model.CpModel()
50    perm = [model.NewIntVar(0, N - 1, f"p{i}") for i in range(N)]
51    model.AddAllDifferent(perm)
52
53    edge_scores = []
54    for u, v in G.edges():
55        i, j = node_to_idx[u], node_to_idx[v]
56
57        flat_idx = model.NewIntVar(0, N * N - 1, f"idx_{i}_{j}")
58        model.Add(flat_idx == N * perm[i] + perm[j])
59
60        score = model.NewIntVar(score_lo, score_hi, f"s_{i}_{j}")
61        model.AddElement(flat_idx, flat, score)
62        edge_scores.append(score)
63
64    model.Maximize(cp_model.LinearExpr.Sum(edge_scores))
65
66    solver = cp_model.CpSolver()
67    solver.parameters.num_search_workers  = 3
68    solver.parameters.max_time_in_seconds = max_search_time
69    solver.parameters.linearization_level = 2
70    solver.parameters.random_seed = int(random_seed)
71
72    status = solver.Solve(model)
73
74    best_permutation = [ i for i in range( N ) ]
75    print( f"initial node permuation score {_calculate_score(G,cohesion_matrix,best_permutation)}" )
76
77    if status in (cp_model.OPTIMAL, cp_model.FEASIBLE) :
78        msg = "OPTIMAL" if status == cp_model.OPTIMAL else "FEASIBLE"
79        print(f"{msg} solution found. Objective = {solver.ObjectiveValue():.0f}")
80        best_permutation = [solver.Value(perm[i]) for i in range(N)]
81        print( f"cp-sat node permuation score {_calculate_score(G,cohesion_matrix,best_permutation)}" )
82    else:
83        warnings.warn("optimization found no solution within time limit.")
84
85    return best_permutation