amachine.am_optimize.am_optimize_node_permutation_ils

  1import networkx as nx
  2import numpy as np
  3from numba import njit
  4
  5from ..am_random import resolve_rng
  6
  7def _calculate_score(
  8    G: nx.DiGraph,
  9    instance_cohesion: np.ndarray,
 10    perm: list[int]
 11) -> float:
 12
 13    nodes = list(G.nodes())
 14    node_to_idx = {node: i for i, node in enumerate(nodes)}
 15
 16    total_score = 0.0
 17    for u, v in G.edges():
 18
 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
 29@njit
 30def evaluate_swap_delta(a, b, perm, A, C, N):
 31    """Calculates the exact change in objective if nodes A and B swap their instance assignments."""
 32    
 33    p_a = perm[a]
 34    p_b = perm[b]
 35    
 36    delta = 0.0
 37    
 38    for i in range(N):
 39
 40        if i == a or i == b:
 41            continue
 42        
 43        p_i = perm[i]
 44        
 45        # Edges connected to A
 46        if A[i, a] > 0: delta += A[i, a] * (C[p_i, p_b] - C[p_i, p_a])
 47        if A[a, i] > 0: delta += A[a, i] * (C[p_b, p_i] - C[p_a, p_i])
 48        
 49        # Edges connected to B
 50        if A[i, b] > 0: delta += A[i, b] * (C[p_i, p_a] - C[p_i, p_b])
 51        if A[b, i] > 0: delta += A[b, i] * (C[p_a, p_i] - C[p_b, p_i])
 52
 53    # Handle A and B connected directly to each other
 54    if A[a, b] > 0: delta += A[a, b] * (C[p_b, p_a] - C[p_a, p_b])
 55    if A[b, a] > 0: delta += A[b, a] * (C[p_a, p_b] - C[p_b, p_a])
 56    
 57    # Handle self-loops correctly if they exist in the graph
 58    if A[a, a] > 0: delta += A[a, a] * (C[p_b, p_b] - C[p_a, p_a])
 59    if A[b, b] > 0: delta += A[b, b] * (C[p_a, p_a] - C[p_b, p_b])
 60    
 61    return delta
 62
 63@njit
 64def steepest_descent(perm, A, C, N):
 65    """Explores the full neighborhood and applies the steepest ascent swap until stuck."""
 66    improved = True
 67    current_perm = perm.copy()
 68    
 69    while improved:
 70        improved = False
 71        best_delta = 1e-9  # Must strictly improve
 72        best_a = -1
 73        best_b = -1
 74        
 75        for a in range(N):
 76            for b in range(a + 1, N):
 77                delta = evaluate_swap_delta(a, b, current_perm, A, C, N)
 78                if delta > best_delta:
 79                    best_delta = delta
 80                    best_a = a
 81                    best_b = b
 82                    
 83        # Apply the best move found
 84        if best_a != -1:
 85            tmp = current_perm[best_a]
 86            current_perm[best_a] = current_perm[best_b]
 87            current_perm[best_b] = tmp
 88            improved = True
 89            
 90    return current_perm
 91
 92@njit
 93def set_numba_seed(seed):
 94    np.random.seed(seed)
 95
 96@njit
 97def perturb(perm, N, kick_strength):
 98    """Executes a random walk of K valid swaps to escape local optima."""
 99    new_perm = perm.copy()
100    swaps_done = 0
101    attempts = 0
102
103    while swaps_done < kick_strength and attempts < 200:
104        attempts += 1
105        
106        a = np.random.randint(0, N)
107        b = np.random.randint(0, N)
108        
109        # In a pure permutation without graph state restrictions, 
110        # the only invalid swap is swapping an index with itself.
111        if a == b:
112            continue
113            
114        # Swap
115        tmp = new_perm[a]
116        new_perm[a] = new_perm[b]
117        new_perm[b] = tmp
118        swaps_done += 1
119        
120    return new_perm
121
122@njit
123def ils_engine(initial_permutation, A, C, N, max_iterations, kick_strength, random_seed=None):
124    
125    if random_seed is not None:
126        set_numba_seed(random_seed)
127
128    best_permutation = steepest_descent(initial_permutation, A, C, N)
129    
130    # Calculate baseline absolute energy (objective is maximization)
131    best_energy = 0.0
132    for i in range(N):
133        for j in range(N):
134            if A[i, j] > 0:
135                best_energy += A[i, j] * C[best_permutation[i], best_permutation[j]]
136                
137    for i in range(max_iterations):
138        perturbed_perm = perturb(best_permutation, N, kick_strength)
139        candidate_perm = steepest_descent(perturbed_perm, A, C, N)
140        
141        candidate_energy = 0.0
142        for x in range(N):
143            for y in range(N):
144                if A[x, y] > 0:
145                    candidate_energy += A[x, y] * C[candidate_perm[x], candidate_perm[y]]
146                    
147        # Accept if it is a new global best
148        if candidate_energy > best_energy:
149            best_energy = candidate_energy
150            best_permutation = candidate_perm.copy()
151            
152    return best_permutation, best_energy
153
154def optimize_node_permutation_ils(
155    G: nx.DiGraph,
156    cohesion_matrix: np.ndarray,
157    ils_iterations: int = 5000, 
158    kick_strength: int = 4,
159    np_rng: np.random.Generator | None = None
160) -> list[int]:
161    """
162    Optimizes the node assignment utilizing Iterated Local Search.
163    Maximizes the summed cohesion weights across the directed graph.
164    """
165    np_rng = resolve_rng( np_rng )
166    numba_seed = np_rng.integers(0, 2**32 - 1)
167
168    # Convert to 64-bit float for Numba compatibility 
169    # (avoiding the 10_000 integer scaling needed for CP-SAT)
170    C = cohesion_matrix.astype(np.float64)
171
172    N = len(list(G.nodes()))
173    node_to_idx = {node: i for i, node in enumerate(G.nodes())}
174
175    # Build an adjacency weight matrix from the graph
176    A = np.zeros((N, N), dtype=np.float64)
177    for u, v in G.edges():
178        i, j = node_to_idx[u], node_to_idx[v]
179        A[i, j] += 1.0
180
181    # Start from a randomized initial state
182    initial_permutation = np.arange(N, dtype=np.int32)
183    np_rng.shuffle(initial_permutation)
184
185    best_permutation, _ = ils_engine(
186        initial_permutation, 
187        A, 
188        C, 
189        N, 
190        ils_iterations, 
191        kick_strength,
192        numba_seed
193    )
194
195    print( f"ils node permuation score {_calculate_score(G,cohesion_matrix,best_permutation)}" )
196
197    return best_permutation.tolist()