amachine.am_create.am_random_machine

  1from collections import defaultdict
  2import random
  3import warnings
  4import math
  5import numpy as np
  6
  7from ..am_hmm          import HMM
  8from ..am_causal_state import CausalState
  9from ..am_transition   import Transition
 10
 11from ..am_random import exp_uniform_blend, resolve_rng
 12from ..am_vocabulary import Vocabulary
 13
 14def random_machine( 
 15    n_states : int, 
 16    symbols  : list[str],
 17    connectedness : float,
 18    randomness : float,
 19    symbol_proportions : list[float] | None = None,
 20    stratification_categories : dict[str,int] | None = None,
 21    ensure_strongly_connected : bool = False,
 22    ensure_minimal : bool = False,
 23    np_rng : np.random.Generator | None = None ) -> HMM  :
 24
 25    np_rng = resolve_rng( np_rng )
 26
 27    states=[
 28        CausalState( name=f"{i}" )
 29        for i in range( n_states  )
 30    ]
 31
 32    n_symbols = len( symbols )
 33    transitions = []
 34    cycle_transitions = {}
 35
 36    # first create a Hamiltonian cycle
 37    if ensure_strongly_connected :
 38        cycle = np_rng.permutation( n_states )
 39        for i in range( n_states ) : 
 40            i_to = ( i + 1 ) % n_states
 41            cycle_transitions[ cycle[ i ] ] = cycle[ i_to ]
 42
 43    for state_idx, state in enumerate( states ) : 
 44
 45        n_transitions = sum( 
 46            np_rng.random() < connectedness 
 47            for _ in range( n_symbols - 1 ) 
 48        ) + 1
 49
 50        n_transitions = min( n_transitions, min( n_states, n_symbols ) )
 51        
 52        # unifilar requirement means we can't choose the cycle state a second time
 53        available = set( range( n_states ) )
 54        if ensure_strongly_connected :
 55            available = available - { cycle_transitions[ state_idx ] }
 56
 57        available = sorted( list( available ) )
 58
 59        transition_to = np_rng.choice( 
 60            available, 
 61            size=n_transitions - ensure_strongly_connected,
 62            replace=False
 63        ).tolist()
 64
 65        if ensure_strongly_connected :
 66            transition_to.append( cycle_transitions[ state_idx ] )
 67
 68        transition_probabilities = exp_uniform_blend( 
 69            n=n_transitions, 
 70            alpha=randomness,
 71            np_rng=np_rng )
 72
 73        # if symbol proportions are given we will override this
 74        # transition_symbols_indices = py_rng.sample( range( n_symbols ), n_transitions )
 75        transition_symbols_indices = np_rng.choice( 
 76            list( range( n_symbols ) ), 
 77            size=n_transitions,
 78            replace=False
 79        )
 80
 81        for i, p in enumerate( transition_probabilities ) :
 82            transitions.append(
 83                Transition(
 84                    origin_state_idx=state_idx,
 85                    target_state_idx=transition_to[ i ],
 86                    prob=p,
 87                    symbol_idx=transition_symbols_indices[ i ],
 88                    pq=None
 89                )
 90            ) 
 91
 92    if symbol_proportions is not None :
 93        
 94        # ensure proporitions sum to 1 and normalize either way
 95
 96        p_sum = sum( symbol_proportions )
 97
 98        if not math.isclose(p_sum, 1.0): 
 99            warnings.warn("Symbol proportions don't sum to 1. Normalizing.")
100
101        for i in range( len(symbol_proportions) ) :
102            symbol_proportions[ i ] /= p_sum
103
104        total_symbols_needed = len( transitions )
105        
106        if stratification_categories is not None:
107
108            category_indices = defaultdict(list)
109            category_proportions = defaultdict(float)
110
111            for i, s in enumerate(symbols) :
112                cat = stratification_categories.get( s, "uncategorized" ) 
113                category_indices[cat].append( i )
114                category_proportions[ cat ] += symbol_proportions[i]
115
116            pool = []
117            
118            for cat, cat_proporition in category_proportions.items():
119                
120                if cat_proporition <= 0.0:
121                    continue
122                
123                # number of symbols from the category we need, and the symbol indices to draw from
124                cat_quota = int( np.ceil( cat_proporition * total_symbols_needed ) )
125                cat_sym_indices = category_indices[cat]
126                
127                # distribute quota among symbols within the category
128                # based on their relative weights to each other
129                internal_weights = [ symbol_proportions[i] / cat_proporition for i in cat_sym_indices ]
130                
131                probabilities = np.array(internal_weights) / np.sum(internal_weights)
132
133                drawn_symbols = np_rng.choice(
134                    cat_sym_indices, 
135                    size=cat_quota, 
136                    replace=True, 
137                    p=probabilities
138                )
139
140                pool.extend(drawn_symbols)
141        else :
142            target_symbol_counts = [ int( np.ceil( p*total_symbols_needed ) ) for p in symbol_proportions ]
143            pool = []
144            for i, s in enumerate( symbols ) : 
145                pool = pool + target_symbol_counts[ i ] * [ i ] 
146        
147        # randomize it
148        np_rng.shuffle( pool )
149
150        for state_idx, _ in enumerate( states ) : 
151            
152            # This states transitions
153            my_tr_idxs = [ 
154                tr_idx for tr_idx, tr in enumerate( transitions )
155                if tr.origin_state_idx == state_idx 
156            ]
157
158            my_n_tr = len( my_tr_idxs )
159
160            # Track which we used
161            used = set()
162            used_indices = set()
163            
164            for i, s in enumerate( pool ) :
165                if s in used :
166                    continue
167                used.add(s)
168                used_indices.add( i )
169                if len( used ) == my_n_tr :
170                    break
171
172            # If the pool was too restrictive sample the remaining from symbols not in the pool
173            if len( used ) < my_n_tr :
174                needed = my_n_tr - len( used )
175                alternatives = list(set(range( n_symbols )) - used)
176                used.update( 
177                    np_rng.choice( alternatives, size=needed, replace=False)
178                )
179
180            # shuffle them and assign them to their transitions
181            ordered = list(used)
182            np_rng.shuffle( ordered )
183
184            for i, tr_idx in enumerate( my_tr_idxs ) :
185                transitions[ tr_idx ] = transitions[ tr_idx ].modified_deep_copy(
186                    symbol_idx = ordered[ i ]
187                )
188
189            # update the pool removing the symbols taken from it
190            pool = [ s for i, s in enumerate( pool ) if i not in used_indices ]
191
192    m = HMM( 
193        states=states,
194        transitions=transitions,
195        start_state=0,
196        alphabet=symbols.copy()
197    )
198
199    # TODO consider best approach to ensure minimal while ensuring target size
200    if ensure_minimal :
201        m.minimize( retain_names=True )
202
203    if not m.is_row_stochastic() :
204        raise Exception( "Random machine is not row stochastic" )
205
206    if not  m.is_unifilar() :
207        raise Exception( "Random machine is not unifilar" )
208
209    return m