amachine.am_optimize.am_optimize_symbol_permutation_ils

  1from collections import defaultdict
  2
  3import numpy as np
  4from numba import njit
  5
  6from ..am_hmm import HMM
  7from ..am_structured_symbol_set import StructuredSymbolSet
  8from ..am_random import resolve_rng
  9
 10@njit
 11def evaluate_swap_delta(a, b, symbols, W, C, n_edges):
 12    """Calculates the exact change in objective if edge A and B swap symbols."""
 13    sym_a = symbols[a]
 14    sym_b = symbols[b]
 15    delta = 0.0
 16    
 17    for i in range(n_edges):
 18
 19        if i == a or i == b:
 20            continue
 21        
 22        sym_i = symbols[i]
 23        
 24        # Edges connected to A
 25        if W[i, a] > 0: delta += W[i, a] * (C[sym_i, sym_b] - C[sym_i, sym_a])
 26        if W[a, i] > 0: delta += W[a, i] * (C[sym_b, sym_i] - C[sym_a, sym_i])
 27        
 28        # Edges connected to B
 29        if W[i, b] > 0: delta += W[i, b] * (C[sym_i, sym_a] - C[sym_i, sym_b])
 30        if W[b, i] > 0: delta += W[b, i] * (C[sym_a, sym_i] - C[sym_b, sym_i])
 31
 32    # Handle A and B connected directly to each other (prevents double counting)
 33    if W[a, b] > 0: delta += W[a, b] * (C[sym_b, sym_a] - C[sym_a, sym_b])
 34    if W[b, a] > 0: delta += W[b, a] * (C[sym_a, sym_b] - C[sym_b, sym_a])
 35
 36    if W[a, a] > 0: delta += W[a, a] * (C[sym_b, sym_b] - C[sym_a, sym_a])
 37    if W[b, b] > 0: delta += W[b, b] * (C[sym_a, sym_a] - C[sym_b, sym_b])
 38
 39    return delta
 40
 41@njit
 42def steepest_descent(
 43    symbols, 
 44    W, 
 45    C, 
 46    origins, 
 47    state_out_edges, 
 48    n_edges ) :
 49
 50    """Explores the neighborhood and applies the best possible swap until stuck."""
 51    
 52    improved = True
 53    current_symbols = symbols.copy()
 54    
 55    while improved:
 56        improved = False
 57        best_delta = 1e-9  # Must strictly improve
 58        best_a = -1
 59        best_b = -1
 60        
 61        for a in range(n_edges):
 62            for b in range(a + 1, n_edges):
 63                if current_symbols[a] == current_symbols[b]:
 64                    continue
 65                    
 66                # unifilarity check
 67                orig_a = origins[a]
 68                orig_b = origins[b]
 69
 70                if orig_a != orig_b:
 71                
 72                    conflict = False
 73                
 74                    for e in state_out_edges[orig_a]:
 75                        if e != -1 and e != a and current_symbols[e] == current_symbols[b]:
 76                            conflict = True
 77                            break
 78                
 79                    if not conflict:
 80                        for e in state_out_edges[orig_b]:
 81                            if e != -1 and e != b and current_symbols[e] == current_symbols[a]:
 82                                conflict = True
 83                                break
 84                
 85                    if conflict:
 86                        continue
 87                        
 88                delta = evaluate_swap_delta(a, b, current_symbols, W, C, n_edges)
 89                if delta > best_delta:
 90                    best_delta = delta
 91                    best_a = a
 92                    best_b = b
 93                    
 94        # Apply the best move found
 95        if best_a != -1:
 96            tmp = current_symbols[best_a]
 97            current_symbols[best_a] = current_symbols[best_b]
 98            current_symbols[best_b] = tmp
 99            improved = True
100            
101    return current_symbols
102
103@njit
104def set_numba_seed( seed ):
105    np.random.seed( seed )
106
107@njit
108def perturb(
109    symbols, 
110    origins, 
111    state_out_edges, 
112    n_edges, 
113    kick_strength ) :
114
115    """Executes a random walk of K valid swaps to escape local optima."""
116
117    new_symbols = symbols.copy()
118    swaps_done = 0
119    attempts = 0
120
121    while swaps_done < kick_strength and attempts < 200:
122        
123        attempts += 1
124
125        a = np.random.randint(0, n_edges)
126        b = np.random.randint(0, n_edges)
127        
128        if a == b or new_symbols[a] == new_symbols[b]:
129            continue
130            
131        orig_a = origins[a]
132        orig_b = origins[b]
133
134        if orig_a != orig_b:
135        
136            conflict = False
137        
138            for e in state_out_edges[orig_a]:
139                if e != -1 and e != a and new_symbols[e] == new_symbols[b]:
140                    conflict = True; break
141        
142            if not conflict:
143                for e in state_out_edges[orig_b]:
144                    if e != -1 and e != b and new_symbols[e] == new_symbols[a]:
145                        conflict = True; break
146           
147            if conflict:
148                continue
149                
150        # Swap
151        tmp = new_symbols[a]
152        new_symbols[a] = new_symbols[b]
153        new_symbols[b] = tmp
154        swaps_done += 1
155        
156    return new_symbols
157
158@njit
159def ils_engine(
160    initial_symbols, 
161    W, 
162    C, 
163    origins, 
164    state_out_edges, 
165    n_edges, 
166    max_iterations, 
167    kick_strength,
168    random_seed=None ) :
169
170    if random_seed is not None :
171        set_numba_seed( random_seed )
172
173    best_symbols = steepest_descent(initial_symbols, W, C, origins, state_out_edges, n_edges)
174    
175    # Calculate baseline absolute energy
176    best_energy = 0.0
177    for i in range(n_edges):
178        for j in range(n_edges):
179            if W[i, j] > 0:
180                best_energy += W[i, j] * C[best_symbols[i], best_symbols[j]]
181                
182    for i in range(max_iterations):
183
184        # Kick from the best known solution
185        perturbed_symbols = perturb(
186            best_symbols, 
187            origins, 
188            state_out_edges, 
189            n_edges, 
190            kick_strength
191        )
192        
193        # Plunge back down to a local optimum
194        candidate_symbols = steepest_descent(
195            perturbed_symbols, 
196            W, 
197            C, 
198            origins, 
199            state_out_edges, 
200            n_edges 
201        )
202        
203        # Evaluate total energy of new local optimum
204        candidate_energy = 0.0
205        for x in range(n_edges):
206            for y in range(n_edges):
207                if W[x, y] > 0:
208                    candidate_energy += W[x, y] * C[candidate_symbols[x], candidate_symbols[y]]
209                    
210        # Accept if it's a new global best
211        if candidate_energy > best_energy:
212            best_energy = candidate_energy
213            best_symbols = candidate_symbols.copy()
214            
215    return best_symbols, best_energy
216
217def optimize_symbol_permutation_ils(
218    m, 
219    symbol_set, 
220    frequency_bias, 
221    ils_iterations=5000, 
222    kick_strength=4,
223    np_rng=None
224):
225
226    rng=resolve_rng( np_rng )
227    numba_seed = rng.integers(0, 2**32 - 1)
228
229    transitions = m.transitions
230    n_edges = len(transitions)
231    n_states = len(m.states)
232
233    state_trs_in = defaultdict(list)
234    state_trs_out = defaultdict(list)
235    for i, tr in enumerate(transitions):
236        state_trs_in[tr.target_state_idx].append((i, tr))
237        state_trs_out[tr.origin_state_idx].append((i, tr))
238
239    max_out_degree = max((len(edges) for edges in state_trs_out.values()), default=0)
240    state_out_edges = np.full((n_states, max_out_degree), -1, dtype=np.int32)
241    origins = np.zeros(n_edges, dtype=np.int32)
242    
243    for state_idx, edges in state_trs_out.items():
244        for col_idx, (edge_idx, _) in enumerate(edges):
245            state_out_edges[state_idx, col_idx] = edge_idx
246            origins[edge_idx] = state_idx
247
248    pi = m.get_stationary_distribution()
249    tr_pair_probs = {}
250    
251    for state_idx in range(n_states):
252        for i_in, t_in in state_trs_in[state_idx]:
253            in_prob = pi[t_in.origin_state_idx] * t_in.prob
254            for i_out, t_out in state_trs_out[state_idx]:
255                tr_pair_probs[(i_in, i_out)] = in_prob * t_out.prob
256
257    n_pairs = len(tr_pair_probs)
258    uniform_w = 1.0 / n_pairs if n_pairs > 0 else 0.0
259    
260    W = np.zeros((n_edges, n_edges), dtype=np.float64)
261    for (i_in, i_out), p in tr_pair_probs.items():
262        w = (1.0 - frequency_bias) * uniform_w + frequency_bias * p
263        if w > 0:
264            W[i_in, i_out] = w
265
266    C = symbol_set.symbol_cohesion.astype(np.float64)
267    initial_symbols = np.array([t.symbol_idx for t in transitions], dtype=np.int32)
268
269    best_symbols, best_energy = ils_engine(
270        initial_symbols, 
271        W, 
272        C, 
273        origins, 
274        state_out_edges, 
275        n_edges, 
276        ils_iterations, 
277        kick_strength,
278        numba_seed
279    )
280    
281    return best_symbols.tolist()