amachine.am_create.am_random_structured

  1from collections import defaultdict, Counter, deque
  2import warnings
  3
  4import numpy as np
  5import math
  6import random
  7from numba import njit
  8
  9from ..am_hmm import HMM
 10from ..am_structured_symbol_set import StructuredSymbolSet
 11from ..am_optimize.am_optimize_symbol_permutation_ils import(
 12    optimize_symbol_permutation_ils
 13)
 14from ..am_optimize.am_optimize_symbol_permutation_cpsat import(
 15    optimize_symbol_permutation_cpsat
 16)
 17
 18from .am_random_machine import random_machine
 19
 20def get_score( 
 21    m, 
 22    symbol_set,
 23    frequency_bias ) :
 24    
 25    n_symbols = len(symbol_set.symbols)
 26
 27    original_symbols = [t.symbol_idx for t in m.transitions]
 28    symbol_counts = Counter(original_symbols)
 29    present_symbols = sorted(list(symbol_counts.keys()))
 30
 31    state_trs_in  = defaultdict(list)
 32    state_trs_out = defaultdict(list)
 33
 34    for i, tr in enumerate(m.transitions):
 35        state_trs_in[tr.target_state_idx].append((i, tr))
 36        state_trs_out[tr.origin_state_idx].append((i, tr))
 37
 38    pi = m.get_stationary_distribution()
 39    tr_pair_probs = {}
 40
 41    for state_idx in range(len(m.states)):
 42        for i_in, t_in in state_trs_in[state_idx]:
 43            in_prob = pi[t_in.origin_state_idx] * t_in.prob
 44            for i_out, t_out in state_trs_out[state_idx]:
 45                tr_pair_probs[(i_in, i_out)] = in_prob * t_out.prob
 46
 47    n_pairs = len(tr_pair_probs)
 48    uniform_w = 1.0 / n_pairs if n_pairs > 0 else 0.0
 49
 50    pair_weights = {
 51        pair: (1.0 - frequency_bias) * uniform_w + frequency_bias * p
 52        for pair, p in tr_pair_probs.items()
 53    }
 54
 55    original_symbols = [ t.symbol_idx for t in m.transitions ]
 56
 57    rw = 0
 58    for ( i_in, i_out ), _ in tr_pair_probs.items() :
 59
 60        orig_s_in  = original_symbols[ i_in  ]
 61        orig_s_out = original_symbols[ i_out ]
 62
 63        w = pair_weights[ ( i_in, i_out ) ]
 64
 65        rw += w*symbol_set.symbol_cohesion[ orig_s_in, orig_s_out ]
 66
 67    return rw
 68
 69def random_structured(
 70    n_states: int,
 71    symbol_set: StructuredSymbolSet,
 72    connectedness: float,
 73    randomness: float,
 74    frequency_bias: float = 0.8,
 75    ensure_strongly_connected: bool = True,
 76    ensure_minimal: bool = True,
 77    np_rng : np.random.Generator | None = None ) -> HMM:
 78
 79    n_symbols = len(symbol_set.symbols)
 80
 81    # passing proportions will make random_machine attempt to 
 82    # assign symbols to edges as close as possible to those proportions
 83    # instead of just random sampling symbols unformly
 84    symbol_proportions = [ 1.0/n_symbols ]*n_symbols
 85
 86    m = random_machine(
 87        n_states=n_states,
 88        symbols=symbol_set.symbols,
 89        connectedness=connectedness,
 90        randomness=randomness,
 91        symbol_proportions=symbol_proportions,
 92        stratification_categories=symbol_set.symbol_categories,
 93        ensure_strongly_connected=ensure_strongly_connected,
 94        ensure_minimal=True,
 95        np_rng=np_rng
 96    )
 97
 98    if not m.is_unifilar() :
 99        raise ValueError( "random_machine returned a non-unifilar machine." )
100
101    initial_symbols = [ tr.symbol_idx for tr in m.transitions ]
102    E = len( initial_symbols )
103    symbol_counts = Counter( initial_symbols ).values()
104    numerator = math.factorial( E )
105    denominator = math.prod( math.factorial( c ) for c in symbol_counts )
106    estimated_search_space_size = numerator // denominator
107
108    # Note: CP-SAT is about on par with efficiently implemented brute force for this problem
109    # and ILS tends to find optimal soltion anyways when the search space is this small
110    # probably either always use ILS, or consider implementing pure brute force for small cases
111    # would be better, however there is a chance that a better cp-sat formulation chould change the 
112    # calculus. 
113
114    if estimated_search_space_size < 479_001_600 :
115        symbol_assignment = optimize_symbol_permutation_cpsat( 
116            m=m, 
117            symbol_set=symbol_set,
118            frequency_bias=frequency_bias,
119            n_search_workers=3,
120            max_search_time=60,
121            np_rng=np_rng
122        )
123    else :
124        symbol_assignment = optimize_symbol_permutation_ils( 
125            m=m, 
126            symbol_set=symbol_set,
127            frequency_bias=frequency_bias,
128            ils_iterations=3_000,
129            kick_strength=4,
130            np_rng=np_rng
131        )
132
133    new_transitions = []
134    for i, s_idx in enumerate( symbol_assignment ):
135        new_transitions.append( 
136            m.transitions[ i ].modified_deep_copy( symbol_idx=s_idx )  
137        )
138
139    m.set_transitions( new_transitions )
140
141    if ensure_minimal :
142        m.minimize()
143
144    if not m.is_row_stochastic() :
145        raise Exception( "Random structured machine is not row stochastic" )
146
147    if not  m.is_unifilar() :
148        raise Exception( "Random structured machine is not unifilar" )
149
150    print( f"calculated score: {get_score( m, symbol_set, frequency_bias )}" )
151
152    return m