amachine.am_optimize.am_optimize_symbol_permutation_toulbar2

  1import numpy as np
  2import warnings
  3from collections import defaultdict, Counter
  4
  5from amachine.am_random import resolve_rng
  6
  7def optimize_symbol_permutation_toulbar2(
  8    m,
  9    symbol_set,
 10    frequency_bias,
 11    max_search_time,
 12    np_rng : np.random.Generator | None = None,
 13    precision_scale: int = 10_000_000 ):
 14
 15    import pytoulbar2
 16
 17    np_rng = resolve_rng( np_rng )
 18    random_seed = int(np_rng.integers(0, 2**31 - 1))
 19
 20    transitions = m.transitions
 21    n_edges     = len(transitions)
 22
 23    print(f"Random machine has {n_edges} transitions.")
 24
 25    state_trs_in  = defaultdict(list)
 26    state_trs_out = defaultdict(list)
 27    for i, tr in enumerate(transitions):
 28        state_trs_in[  tr.target_state_idx ].append((i, tr))
 29        state_trs_out[ tr.origin_state_idx ].append((i, tr))
 30
 31    pi = m.get_stationary_distribution()
 32    tr_pair_probs = {}
 33    for state_idx in range(len(m.states)):
 34        for i_in, t_in in state_trs_in[state_idx]:
 35            in_prob = pi[t_in.origin_state_idx] * t_in.prob
 36            for i_out, t_out in state_trs_out[state_idx]:
 37                tr_pair_probs[(i_in, i_out)] = in_prob * t_out.prob
 38
 39    n_pairs   = len(tr_pair_probs)
 40    uniform_w = 1.0 / n_pairs if n_pairs > 0 else 0.0
 41
 42    # Pre-filter zero-weight pairs, same as the CP-SAT version.
 43    weighted_pairs = {}
 44    for pair, p in tr_pair_probs.items():
 45        coeff = int(((1.0 - frequency_bias) * uniform_w + frequency_bias * p) * precision_scale)
 46        if coeff > 0:
 47            weighted_pairs[pair] = coeff
 48
 49    original_symbols = [t.symbol_idx for t in transitions]
 50    symbol_counts    = Counter(original_symbols)
 51    present_symbols  = sorted(symbol_counts.keys())
 52    n_present        = len(present_symbols)
 53
 54    k_to_sym   = present_symbols
 55    sym_to_k   = {s: k for k, s in enumerate(present_symbols)}
 56    original_k = [sym_to_k[s] for s in original_symbols]
 57
 58    coh_int  = np.round(symbol_set.symbol_cohesion * precision_scale).astype(int)
 59    flat_coh = [int(coh_int[k_to_sym[k_in], k_to_sym[k_out]])
 60                for k_in in range(n_present)
 61                for k_out in range(n_present)]
 62
 63    # toulbar2 minimizes non-negative costs natively, and the docs flag that
 64    # global constraints (alldiff/gcc) have "restricted usage" in maximization
 65    # mode. Rather than fight that, flip the objective: cost(k_in,k_out) =
 66    # coeff * (M - reward), so minimizing total cost is exactly equivalent to
 67    # maximizing the original weighted cohesion sum (M is just a constant
 68    # per-pair offset, so it doesn't change the argmax).
 69    M = max(flat_coh) if flat_coh else 0
 70
 71    # Generous initial upper bound: worst case is every weighted pair paying
 72    # its maximum possible cost (reward = min(flat_coh) instead of M).
 73    worst_per_pair = M - min(flat_coh) if flat_coh else 0
 74    ub_init = sum(weighted_pairs.values()) * worst_per_pair + 1
 75
 76    Problem = pytoulbar2.CFN(ubinit=ub_init, seed=random_seed)
 77
 78    var_idx = [
 79        Problem.AddVariable(f"sym_{i}", list(range(n_present)))
 80        for i in range(n_edges)
 81    ]
 82    # Note: there's no direct equivalent of CP-SAT's AddHint here — pytoulbar2's
 83    # Solve() doesn't take a warm-start assignment in this API. If solve time on
 84    # larger instances becomes an issue, look at Problem.SolveFirst /
 85    # MultipleAssign, or VNS mode (the vns= argument to CFN()), which can use a
 86    # starting solution to drive local search before branch-and-bound.
 87
 88    # Unifilarity: native AllDifferent per state, instead of CP-SAT's
 89    # AddAllDifferent over IntVars (same constraint, no behavioral change).
 90    for out_edges in state_trs_out.values():
 91        if len(out_edges) > 1:
 92            Problem.AddAllDifferent([var_idx[i] for i, _ in out_edges])
 93
 94    # Symbol-count constraint: this is what used to require the whole
 95    # boolean channeling matrix (b[i][k] + OnlyEnforceIf in both directions +
 96    # column-sum constraints) in the CP-SAT version. Here it's one call.
 97    bounds = [
 98        (k, symbol_counts[k_to_sym[k]], symbol_counts[k_to_sym[k]])
 99        for k in range(n_present)
100    ]
101    Problem.AddGlobalCardinalityConstraint(var_idx, bounds)
102
103    # Cohesion reward per transition pair: a single dense binary table cost
104    # function per pair, replacing CP-SAT's idx-IntVar + linear-equality +
105    # AddElement chain. flat_coh is already in the right lexicographic
106    # (k_in, k_out) order for AddFunction's dense-table convention.
107    #
108    # IMPORTANT: a state can have a self-loop transition that is simultaneously
109    # its own "incoming" and "outgoing" edge for that state, producing a pair
110    # with i_in == i_out. CP-SAT's idx == sym[i_in]*n_present + sym[i_out]
111    # formulation tolerates the same IntVar appearing twice; toulbar2's
112    # AddFunction does NOT allow a duplicate variable in scope. Such pairs
113    # must be expressed as a unary cost function over the diagonal of
114    # flat_coh instead of a binary one.
115    for (i_in, i_out), coeff in weighted_pairs.items():
116        if i_in == i_out:
117            diag_costs = [coeff * (M - flat_coh[k * n_present + k]) for k in range(n_present)]
118            Problem.AddFunction([var_idx[i_in]], diag_costs)
119        else:
120            costs = [coeff * (M - r) for r in flat_coh]
121            Problem.AddFunction([var_idx[i_in], var_idx[i_out]], costs)
122
123    res = [tr.symbol_idx for tr in m.transitions]
124
125    solution = Problem.Solve(showSolutions=0, timeLimit=int(max_search_time))
126    if solution is not None:
127        assignment, cost, _ = solution
128
129        # This is the key piece for telling whether the result is provably
130        # optimal or just the best found so far under the time cap.
131        # GetDDualBound() is the proven lower bound (minimization) from
132        # toulbar2's HBFS search. If it equals the incumbent cost, the gap
133        # is zero and optimality is proven — exactly the "upper bound on the
134        # gap" the paper describes. If Solve() returns before timeLimit is
135        # reached at all, that also implies the search tree was exhausted
136        # (same conclusion, sometimes available before computing the bound).
137        dual_bound = Problem.GetDDualBound()
138        gap = cost - dual_bound
139        proven_optimal = gap <= 1e-6
140
141        print(f"Solution found. Cost = {cost:.0f} "
142              f"(equivalent maximized reward = {sum(weighted_pairs.values())*M - cost:.0f}) | "
143              f"dual_bound = {dual_bound:.0f} | gap = {gap:.0f} | "
144              f"{'PROVEN OPTIMAL' if proven_optimal else 'NOT proven optimal — best found under time cap'}")
145
146        for i in range(n_edges):
147            res[i] = k_to_sym[assignment[i]]
148    else:
149        warnings.warn("optimization found no solution within time limit.")
150
151    return res