GitLab Repo

amachine.am_create.am_structured_composition

  1from collections import defaultdict
  2import random
  3import warnings
  4
  5import numpy as np
  6import networkx as nx
  7
  8from ..am_hmm          import HMM
  9from ..am_causal_state import CausalState
 10from ..am_transition   import Transition
 11from ..am_random       import exp_uniform_blend, resolve_rng
 12from ..am_optimize.am_optimize_node_permutation_cpsat import(
 13    optimize_node_permutation_cpsat
 14)
 15from ..am_optimize.am_optimize_node_permutation_ils import(
 16    optimize_node_permutation_ils
 17)
 18
 19from ..am_id_generator import next_composition_id
 20
 21def structured_composition(
 22    core_alphabet : list[str],
 23    composition_rigidity : float,
 24    component_groups : dict[ str, list[HMM] ],
 25    instances_per_component : int | dict[ str, list[int] ],
 26    component_residency_factor : float,
 27    component_repetition_penalty : float,
 28    composition_connectivity_factor : float,
 29    instance_cohesion_randomness : float,
 30    composition_id : str | None = None,
 31    connector_state_pools : dict[ str, list[list[tuple[int,int]]]] | None = None,
 32    end_of_component_symbols : set[str] | None = None,
 33    symbol_cohesion : np.ndarray | None = None,
 34    np_rng : np.random.Generator | None = None ) -> tuple[ HMM, list[tuple[int,int]] ] :
 35
 36    np_rng = resolve_rng( np_rng )
 37
 38    if composition_id is None :
 39        composition_id = str(next_composition_id())
 40
 41    if end_of_component_symbols is not None and len( end_of_component_symbols ) == 0 :
 42        raise ValueError( "end_of_component_symbols is empty" )
 43
 44    if ( end_of_component_symbols is not None 
 45       and len( end_of_component_symbols ) == 1 
 46       and composition_connectivity_factor != 0.0 ) :
 47        warnings.warn( "Only 1 EOC symbol, composition_connectivity_factor will be unused. " )
 48
 49    if end_of_component_symbols is None and symbol_cohesion is None :
 50        raise ValueError( "If using standard symbols instead of EOC symbols, must pass cohesion matrix." )
 51
 52    if isinstance( instances_per_component, int ):
 53        ic_dict = {}
 54        for group_id, components in component_groups.items():
 55            ic_dict[ group_id ] = [ instances_per_component ] * len(components)
 56        instances_per_component = ic_dict
 57
 58    composition_depth = max( (  
 59        m.composition_depth for gr in component_groups.values() for m in gr
 60    ) ) + 1
 61
 62    component_pool = []
 63    for group_id, components in component_groups.items() : 
 64        for c_i, c in enumerate( components ) :
 65            if not c.is_strongly_connected() :
 66                raise ValueError( "Components are not all strongly connected." )
 67            component_pool.append( ( c, group_id, c_i ) )
 68
 69    n_components = len( component_pool )
 70    
 71    # This is the cohesion score for component i flowing into component j
 72    component_cohesion = np.zeros( ( n_components, n_components ) )
 73    for i in range( n_components ) :
 74        component_cohesion[ i, : ] = exp_uniform_blend( 
 75            n=n_components, 
 76            alpha=(1.0-composition_rigidity),
 77            np_rng=np_rng )
 78
 79        component_cohesion[ i, i ] *= ( 1.0 - component_repetition_penalty )
 80
 81    total_subgraphs = sum( [
 82        sum( counts ) for counts in instances_per_component.values()
 83    ] )
 84
 85    component_instance_counts = [
 86        instances_per_component[ group_id ][ c_i ]
 87        for _, group_id, c_i in component_pool
 88    ]
 89
 90    alphabet = list( core_alphabet )
 91    
 92    symbol_idx_map = { s : i for i, s in enumerate( alphabet ) }
 93    symbol_set = set( alphabet )
 94
 95    if end_of_component_symbols is not None :
 96        # it is possible EOC symbols is already in the alphabet
 97        # if not add them
 98        for s in end_of_component_symbols :
 99            if s in symbol_set :
100                continue
101            symbol_idx_map[ s ] = len( alphabet )
102            alphabet.append( s )
103
104    ##############################################################
105    # add the states and transitions
106
107    cstate_name_to_global_idxs = {
108        group_id : defaultdict(list)
109        for group_id in component_groups.keys()
110    }
111    
112    cstate_instance_to_global_idx    = [ {} for _ in range( total_subgraphs ) ]
113    instance_idx_to_component_idx = []
114    global_idx_to_instance_idx = {}
115    global_idx_to_component_state = {}
116    global_idx_to_group = {}
117    g_sidx = 0
118    instance_idx = 0
119
120    #---------------------------------------------------------------------------
121
122    states = []
123    transitions = []
124
125    for m_idx, ( m, group_id, group_offset ) in enumerate( component_pool ) :
126
127        n_instances = instances_per_component[ group_id ][ group_offset ]
128        
129        for _ in range( n_instances ) :
130
131            instance_idx_to_component_idx.append( m_idx )
132
133            for s_idx, state in enumerate( m.states ) : 
134
135                m_classes = {
136                    composition_depth : { 
137                        f"cm_{composition_id}_{m_idx}", 
138                        f"cmi_{composition_id}_{instance_idx}", 
139                        f"c_{composition_id}", 
140                        f"g_{group_id}"
141                    }
142                }
143
144                # Note, later we will have to adjust/add the isomorphs since the names have changed.
145                states.append( CausalState(
146                    name=f"{g_sidx}/{state.name}", 
147                    classes=( state.classes | m_classes )
148                ) )
149
150                global_idx_to_instance_idx[ g_sidx ] = instance_idx
151                global_idx_to_group[ g_sidx ] = group_id
152                global_idx_to_component_state[ g_sidx ] = state 
153                cstate_name_to_global_idxs[ group_id ][ state.name ].append( g_sidx )
154                cstate_instance_to_global_idx[ instance_idx ][ s_idx ] = g_sidx
155
156                g_sidx += 1
157
158            for tr in m.transitions :
159                origin = cstate_instance_to_global_idx[ instance_idx ][ tr.origin_state_idx ]
160                target = cstate_instance_to_global_idx[ instance_idx ][ tr.target_state_idx ]
161                transitions.append(
162                    tr.modified_deep_copy(  
163                        origin_state_idx=origin,
164                        target_state_idx=target,
165                    )
166                ) 
167
168            instance_idx += 1
169    #---------------------------------------------------------------------------
170
171    # Remap and add the isomorphs
172    for g_idx, s in enumerate( states ) :
173        
174        group_id = global_idx_to_group[ g_idx ]
175        c_state  = global_idx_to_component_state[ g_idx ]
176        
177        isomorphs = set()
178        pseudo_isomorphs = set()
179
180        for c_iso in c_state.isomorphs :
181            # Components are duplicated, which means isomorphs are too 
182            g_isos = cstate_name_to_global_idxs[ group_id ][ c_iso ]
183            for g_iso in g_isos :
184                isomorphs.add( states[ g_iso ].name )
185
186        for c_iso in c_state.pseudo_isomorphs :
187            g_isos = cstate_name_to_global_idxs[ group_id ][ c_iso ]
188            for g_iso in g_isos :
189                pseudo_isomorphs.add( states[ g_iso ].name )
190
191        states[ g_idx ] = states[ g_idx ].modified_deep_copy(
192            isomorphs=isomorphs,
193            pseudo_isomorphs=pseudo_isomorphs
194        )
195
196    ###############################################################################################
197    # Construct a cohesive component graph skeleton before adding the actual compononent connections
198
199    G = nx.DiGraph()
200    G.add_nodes_from( [ i for i in range( total_subgraphs ) ] )
201
202    # minus self edges
203    all_edges = [
204        (u, v) 
205        for u in range(total_subgraphs) 
206        for v in range(total_subgraphs) 
207        if u != v
208    ]
209
210    np_rng.shuffle( all_edges )
211
212    n_edges_from = defaultdict(int)
213    n_edges_to = defaultdict(int)
214    edges_added = 0
215
216    cycle_edges = []
217    extra_edges = []
218
219    max_edges = ( 1.0 + composition_connectivity_factor ) * total_subgraphs
220
221    # start with hamiltonian cycle
222    for u, v in all_edges:
223
224        if n_edges_from[u] >= 1 :
225            continue
226
227        if n_edges_to[v] >= 1:
228            continue
229
230        if nx.has_path( G, v, u ) and edges_added < total_subgraphs - 1: 
231            continue
232
233        G.add_edge(u, v)
234        edges_added += 1
235        
236        n_edges_from[ u ] += 1
237        n_edges_to[v] += 1
238
239        cycle_edges.append( ( u,v ) )
240
241        if edges_added >= total_subgraphs :
242            if nx.is_strongly_connected(G):
243                break
244
245    # Get some candidate extra edges
246    for u, v in all_edges :
247
248        if G.has_edge( u , v ) :
249            continue
250
251        if end_of_component_symbols is not None and n_edges_from[u] >= len( end_of_component_symbols ) :
252            continue
253
254        edges_added += 1
255
256        n_edges_from[ u ] += 1
257        n_edges_to[ v ] += 1
258
259        extra_edges.append( (u,v) )
260
261        if edges_added >= max_edges :
262            break
263
264    ##############################################################
265    # calculate best entry and exit nodes for each subgraph
266
267    # lower is better
268    def score_as_from( m, s_idx ) :
269        n_outgoing = 0
270        has_self_loop = False
271        for tr in m.transitions :
272            if tr.origin_state_idx == s_idx :
273                n_outgoing += 1
274            if tr.origin_state_idx == s_idx and tr.target_state_idx == s_idx :
275                has_self_loop = True
276        return ( n_outgoing, not has_self_loop, s_idx )
277
278    # higher is better
279    def score_as_to( m, s_idx, path_length ) :
280        n_incoming = 0
281        has_self_loop = False
282        for tr in m.transitions :
283            if tr.target_state_idx == s_idx :
284                n_incoming += 1
285            if tr.origin_state_idx == s_idx and tr.target_state_idx == s_idx :
286                has_self_loop = True
287
288        return ( path_length, 1.0/n_incoming, has_self_loop, s_idx )
289
290    best_entry_and_exits = []
291
292    if connector_state_pools is None :
293        
294        for cm, _, _ in component_pool:
295
296            n_states = len(cm.states)
297            cm_dg = cm.as_digraph()
298
299            best_from = min(range(n_states), key=lambda s: score_as_from(cm, s))
300
301            lengths_to_from = nx.single_source_shortest_path_length(
302                cm_dg.reverse(), best_from)
303
304            best_to = max(
305                range(n_states),
306                key=lambda s: score_as_to(cm, s, lengths_to_from.get(s, 0)))
307
308            best_entry_and_exits.append( ( best_to, best_from ) )
309    else :
310
311        for cm, g_id, cmi in component_pool:
312          
313            if g_id not in connector_state_pools :
314                raise ValueError( "connector_states missing group id" )
315          
316            if cmi < 0 or cmi >= len( connector_state_pools[ g_id ] ) :
317                raise ValueError( f"Invalid connector state list for group {g_id}." )
318
319            n_states = len(cm.states)
320            pool = connector_state_pools[ g_id ][ cmi ]
321
322            if len( pool ) == 0:
323                raise ValueError( f"Empty connector pool for group {g_id}, component {cmi}." )
324
325            invalid  = [ 
326                ( entry_state, exit_state ) for entry_state, exit_state in pool 
327                if ( entry_state >= n_states 
328                  or exit_state  >= n_states 
329                  or entry_state < 0 
330                  or exit_state  < 0 )
331            ]
332
333            if invalid:
334                raise ValueError( f"connector_state_pools['{g_id}'][{cmi}] has out-of-range indices" )
335
336            entry_options = [ p[ 0 ] for p in pool ]
337            exit_options  = [ p[ 1 ] for p in pool ]
338
339            cm_dg = cm.as_digraph()
340
341            best_from = min( exit_options, key=lambda s: score_as_from(cm, s) )
342
343            lengths_to_from = nx.single_source_shortest_path_length(
344                cm_dg.reverse(), best_from)
345
346            best_to = max(
347                entry_options,
348                key=lambda s: score_as_to(cm, s, lengths_to_from.get(s, 0)))
349
350            best_entry_and_exits.append( ( best_to, best_from ) )
351
352    # The states which are entry or exit states between components are no longer true 
353    # ismorphs of the states they used to be isomorphic to
354    
355    instance_idx = 0
356    for m_idx, (cm, g_idx, g_offset ) in enumerate(component_pool):
357        
358        best_to_local, best_from_local = best_entry_and_exits[m_idx]
359        n_instances = instances_per_component[ g_idx ][ g_offset ]
360        
361        for _ in range(n_instances):
362
363            g_best_to   = cstate_instance_to_global_idx[instance_idx][best_to_local]
364            g_best_from = cstate_instance_to_global_idx[instance_idx][best_from_local]
365            
366            states[g_best_to ] = states[g_best_to].modified_deep_copy(
367                isomorphs = set(),
368                pseudo_isomorphs = states[g_best_to].pseudo_isomorphs | states[g_best_to].isomorphs
369            )
370
371            states[g_best_from ] = states[g_best_from].modified_deep_copy(
372                isomorphs = set(),
373                pseudo_isomorphs = states[g_best_from].pseudo_isomorphs | states[g_best_from].isomorphs
374            )
375
376            instance_idx += 1
377
378    ##############################################################
379
380    # Expand component to instances and apply noise
381
382    nodes = list(G.nodes())
383    N = len(nodes)
384    node_to_idx = {node: i for i, node in enumerate(nodes)}
385
386    assert sum(component_instance_counts) == N
387
388    pos_to_comp = [c for c, n in enumerate(component_instance_counts) for _ in range(n)]
389    instance_cohesion = component_cohesion[np.ix_(pos_to_comp, pos_to_comp)]
390    
391    lo, hi = component_cohesion.min(), component_cohesion.max()
392    noise = np_rng.uniform( lo, hi, ( N, N ) )
393    instance_cohesion = (1 - instance_cohesion_randomness) * instance_cohesion + instance_cohesion_randomness * noise
394
395    ##########################################################################
396
397    # Note: CP-SAT is about on par with efficiently implemented brute force for this problem
398    # and ILS tends to find optimal soltion anyways when the search space is this small
399    # probably either always use ILS, or consider implementing pure brute force for small cases
400    # would be better, however there is a chance that a better cp-sat formulation chould change the 
401    # calculus. 
402
403    # Search space is N!
404    if N < 12 : 
405        node_permutation = optimize_node_permutation_cpsat( 
406            G=G, 
407            cohesion_matrix=instance_cohesion,
408            max_search_time=60,
409            np_rng=np_rng
410        )
411    else :
412        node_permutation = optimize_node_permutation_ils( 
413            G=G, 
414            cohesion_matrix=instance_cohesion,
415            np_rng=np_rng
416        )
417
418    nodes = list( G.nodes() )
419    mapping = { i : node_permutation[i] for i in range(total_subgraphs) }
420
421    cycle_edges  = [ ( mapping[u], mapping[v] ) for u, v in cycle_edges ]
422    extra_edges  = [ ( mapping[u], mapping[v] ) for u, v in extra_edges ]
423
424    # connectors is a list over the component instances of the global enter / exit state of the instances  
425    connectors : list[tuple[int,int]] = []
426
427    for i in range( total_subgraphs ) :
428        m_idx = instance_idx_to_component_idx[i]
429        best_to_local, best_from_local = best_entry_and_exits[ m_idx ]
430        connectors.append((
431            cstate_instance_to_global_idx[i][ best_to_local   ], # entry
432            cstate_instance_to_global_idx[i][ best_from_local ]  # exit
433        ) )
434
435    ##############################################################
436    # Connect the component instances
437
438    # In theory could score to and from nodes also based on what kind of cohesion scores are possible
439    # Given what free symbols they would have, but this is a bit complicated, considering adding it later.
440
441    extra_edges_added = 0
442    for edge_idx, ( u, v ) in enumerate( cycle_edges + extra_edges ) :
443
444        is_cycle_edge = edge_idx < len( cycle_edges )
445
446        m_u_idx = instance_idx_to_component_idx[ u ]
447        m_v_idx = instance_idx_to_component_idx[ v ]
448
449        best_from = best_entry_and_exits[ m_u_idx ][ 1 ]
450        best_to   = best_entry_and_exits[ m_v_idx ][ 0 ]
451
452        origin = cstate_instance_to_global_idx[ u ][ best_from ]
453        target = cstate_instance_to_global_idx[ v ][ best_to   ]
454
455        # if end_of_component_symbols is None we use the whole alphabet, otherwise eoc only
456        free_eoc_symbols    = end_of_component_symbols.copy() if end_of_component_symbols is not None else set()
457        free_normal_symbols = set( alphabet ) if end_of_component_symbols is None else set()
458
459        # reweight existing transitions and figure out which eoc symbols are available
460        for tr_idx, tr in enumerate( transitions ) :
461            if tr.origin_state_idx == origin :
462
463                transitions[ tr_idx  ] = transitions[ tr_idx  ].modified_deep_copy(
464                    prob = transitions[ tr_idx  ].prob * component_residency_factor
465                )
466
467                symbol = alphabet[ tr.symbol_idx ]
468                
469                if symbol in free_eoc_symbols :
470                    free_eoc_symbols.remove( symbol )
471                
472                if symbol in free_normal_symbols :
473                    free_normal_symbols.remove( symbol )
474
475        # Shouldn't happen on a cycle edge
476        if len( free_eoc_symbols ) == 0 and end_of_component_symbols is not None :
477            if is_cycle_edge :
478                raise ValueError( "Ran out of eoc symbols attempting to connect components." )
479            else :
480                continue
481
482        # This could happen on a cycle edge
483        if len( free_normal_symbols ) == 0 and end_of_component_symbols is None :
484            if is_cycle_edge :
485                raise ValueError( "Ran out of symbols attempting to connect components." )
486            else :
487                continue
488
489        if not is_cycle_edge :
490            extra_edges_added += 1
491
492        if len( free_eoc_symbols ) > 0 :
493            symbol = np_rng.choice( sorted( free_eoc_symbols ) )
494
495        elif len( free_normal_symbols ) > 0 :
496            avail = list( free_normal_symbols )
497            cohesion_scores = [ 0 for _ in avail ]
498            for i, s in enumerate( avail ) : 
499                for tr in transitions :
500                    if tr.target_state_idx == origin :
501                        cohesion_scores[ i ] += symbol_cohesion[ tr.symbol_idx, symbol_idx_map[ s ] ]
502                    elif tr.origin_state_idx == target :
503                        cohesion_scores[ i ] += symbol_cohesion[ symbol_idx_map[ s ], tr.symbol_idx ]
504            best_i = np.argmax( np.array( cohesion_scores ) )
505            symbol = avail[ best_i ]
506        else :
507            continue
508
509        transitions.append(
510            Transition(
511                origin_state_idx=origin,
512                target_state_idx=target,
513                prob=( 1.0 - component_residency_factor ),
514                symbol_idx=symbol_idx_map[ symbol ],
515                pq=None,
516                composition_depth=composition_depth
517        ) )
518
519    print( f"Was able to add {extra_edges_added}/{len(extra_edges)} extra cross-component transitions beyond a Hamiltonian cycle" )
520
521    res = HMM( 
522        states=states,
523        transitions=transitions,
524        alphabet=alphabet,
525        start_state=0,
526        composition_depth=composition_depth
527    )
528
529    if not res.is_row_stochastic() :
530        raise Exception( "Structured composition is not row stochastic" )
531
532    if not  res.is_unifilar() :
533        raise Exception( "Structured composition is not unifilar" )
534
535    return res, connectors
def structured_composition( core_alphabet: list[str], composition_rigidity: float, component_groups: dict[str, list[amachine.am_hmm.HMM]], instances_per_component: int | dict[str, list[int]], component_residency_factor: float, component_repetition_penalty: float, composition_connectivity_factor: float, instance_cohesion_randomness: float, composition_id: str | None = None, connector_state_pools: dict[str, list[list[tuple[int, int]]]] | None = None, end_of_component_symbols: set[str] | None = None, symbol_cohesion: numpy.ndarray | None = None, np_rng: numpy.random._generator.Generator | None = None) -> tuple[amachine.am_hmm.HMM, list[tuple[int, int]]]:
 22def structured_composition(
 23    core_alphabet : list[str],
 24    composition_rigidity : float,
 25    component_groups : dict[ str, list[HMM] ],
 26    instances_per_component : int | dict[ str, list[int] ],
 27    component_residency_factor : float,
 28    component_repetition_penalty : float,
 29    composition_connectivity_factor : float,
 30    instance_cohesion_randomness : float,
 31    composition_id : str | None = None,
 32    connector_state_pools : dict[ str, list[list[tuple[int,int]]]] | None = None,
 33    end_of_component_symbols : set[str] | None = None,
 34    symbol_cohesion : np.ndarray | None = None,
 35    np_rng : np.random.Generator | None = None ) -> tuple[ HMM, list[tuple[int,int]] ] :
 36
 37    np_rng = resolve_rng( np_rng )
 38
 39    if composition_id is None :
 40        composition_id = str(next_composition_id())
 41
 42    if end_of_component_symbols is not None and len( end_of_component_symbols ) == 0 :
 43        raise ValueError( "end_of_component_symbols is empty" )
 44
 45    if ( end_of_component_symbols is not None 
 46       and len( end_of_component_symbols ) == 1 
 47       and composition_connectivity_factor != 0.0 ) :
 48        warnings.warn( "Only 1 EOC symbol, composition_connectivity_factor will be unused. " )
 49
 50    if end_of_component_symbols is None and symbol_cohesion is None :
 51        raise ValueError( "If using standard symbols instead of EOC symbols, must pass cohesion matrix." )
 52
 53    if isinstance( instances_per_component, int ):
 54        ic_dict = {}
 55        for group_id, components in component_groups.items():
 56            ic_dict[ group_id ] = [ instances_per_component ] * len(components)
 57        instances_per_component = ic_dict
 58
 59    composition_depth = max( (  
 60        m.composition_depth for gr in component_groups.values() for m in gr
 61    ) ) + 1
 62
 63    component_pool = []
 64    for group_id, components in component_groups.items() : 
 65        for c_i, c in enumerate( components ) :
 66            if not c.is_strongly_connected() :
 67                raise ValueError( "Components are not all strongly connected." )
 68            component_pool.append( ( c, group_id, c_i ) )
 69
 70    n_components = len( component_pool )
 71    
 72    # This is the cohesion score for component i flowing into component j
 73    component_cohesion = np.zeros( ( n_components, n_components ) )
 74    for i in range( n_components ) :
 75        component_cohesion[ i, : ] = exp_uniform_blend( 
 76            n=n_components, 
 77            alpha=(1.0-composition_rigidity),
 78            np_rng=np_rng )
 79
 80        component_cohesion[ i, i ] *= ( 1.0 - component_repetition_penalty )
 81
 82    total_subgraphs = sum( [
 83        sum( counts ) for counts in instances_per_component.values()
 84    ] )
 85
 86    component_instance_counts = [
 87        instances_per_component[ group_id ][ c_i ]
 88        for _, group_id, c_i in component_pool
 89    ]
 90
 91    alphabet = list( core_alphabet )
 92    
 93    symbol_idx_map = { s : i for i, s in enumerate( alphabet ) }
 94    symbol_set = set( alphabet )
 95
 96    if end_of_component_symbols is not None :
 97        # it is possible EOC symbols is already in the alphabet
 98        # if not add them
 99        for s in end_of_component_symbols :
100            if s in symbol_set :
101                continue
102            symbol_idx_map[ s ] = len( alphabet )
103            alphabet.append( s )
104
105    ##############################################################
106    # add the states and transitions
107
108    cstate_name_to_global_idxs = {
109        group_id : defaultdict(list)
110        for group_id in component_groups.keys()
111    }
112    
113    cstate_instance_to_global_idx    = [ {} for _ in range( total_subgraphs ) ]
114    instance_idx_to_component_idx = []
115    global_idx_to_instance_idx = {}
116    global_idx_to_component_state = {}
117    global_idx_to_group = {}
118    g_sidx = 0
119    instance_idx = 0
120
121    #---------------------------------------------------------------------------
122
123    states = []
124    transitions = []
125
126    for m_idx, ( m, group_id, group_offset ) in enumerate( component_pool ) :
127
128        n_instances = instances_per_component[ group_id ][ group_offset ]
129        
130        for _ in range( n_instances ) :
131
132            instance_idx_to_component_idx.append( m_idx )
133
134            for s_idx, state in enumerate( m.states ) : 
135
136                m_classes = {
137                    composition_depth : { 
138                        f"cm_{composition_id}_{m_idx}", 
139                        f"cmi_{composition_id}_{instance_idx}", 
140                        f"c_{composition_id}", 
141                        f"g_{group_id}"
142                    }
143                }
144
145                # Note, later we will have to adjust/add the isomorphs since the names have changed.
146                states.append( CausalState(
147                    name=f"{g_sidx}/{state.name}", 
148                    classes=( state.classes | m_classes )
149                ) )
150
151                global_idx_to_instance_idx[ g_sidx ] = instance_idx
152                global_idx_to_group[ g_sidx ] = group_id
153                global_idx_to_component_state[ g_sidx ] = state 
154                cstate_name_to_global_idxs[ group_id ][ state.name ].append( g_sidx )
155                cstate_instance_to_global_idx[ instance_idx ][ s_idx ] = g_sidx
156
157                g_sidx += 1
158
159            for tr in m.transitions :
160                origin = cstate_instance_to_global_idx[ instance_idx ][ tr.origin_state_idx ]
161                target = cstate_instance_to_global_idx[ instance_idx ][ tr.target_state_idx ]
162                transitions.append(
163                    tr.modified_deep_copy(  
164                        origin_state_idx=origin,
165                        target_state_idx=target,
166                    )
167                ) 
168
169            instance_idx += 1
170    #---------------------------------------------------------------------------
171
172    # Remap and add the isomorphs
173    for g_idx, s in enumerate( states ) :
174        
175        group_id = global_idx_to_group[ g_idx ]
176        c_state  = global_idx_to_component_state[ g_idx ]
177        
178        isomorphs = set()
179        pseudo_isomorphs = set()
180
181        for c_iso in c_state.isomorphs :
182            # Components are duplicated, which means isomorphs are too 
183            g_isos = cstate_name_to_global_idxs[ group_id ][ c_iso ]
184            for g_iso in g_isos :
185                isomorphs.add( states[ g_iso ].name )
186
187        for c_iso in c_state.pseudo_isomorphs :
188            g_isos = cstate_name_to_global_idxs[ group_id ][ c_iso ]
189            for g_iso in g_isos :
190                pseudo_isomorphs.add( states[ g_iso ].name )
191
192        states[ g_idx ] = states[ g_idx ].modified_deep_copy(
193            isomorphs=isomorphs,
194            pseudo_isomorphs=pseudo_isomorphs
195        )
196
197    ###############################################################################################
198    # Construct a cohesive component graph skeleton before adding the actual compononent connections
199
200    G = nx.DiGraph()
201    G.add_nodes_from( [ i for i in range( total_subgraphs ) ] )
202
203    # minus self edges
204    all_edges = [
205        (u, v) 
206        for u in range(total_subgraphs) 
207        for v in range(total_subgraphs) 
208        if u != v
209    ]
210
211    np_rng.shuffle( all_edges )
212
213    n_edges_from = defaultdict(int)
214    n_edges_to = defaultdict(int)
215    edges_added = 0
216
217    cycle_edges = []
218    extra_edges = []
219
220    max_edges = ( 1.0 + composition_connectivity_factor ) * total_subgraphs
221
222    # start with hamiltonian cycle
223    for u, v in all_edges:
224
225        if n_edges_from[u] >= 1 :
226            continue
227
228        if n_edges_to[v] >= 1:
229            continue
230
231        if nx.has_path( G, v, u ) and edges_added < total_subgraphs - 1: 
232            continue
233
234        G.add_edge(u, v)
235        edges_added += 1
236        
237        n_edges_from[ u ] += 1
238        n_edges_to[v] += 1
239
240        cycle_edges.append( ( u,v ) )
241
242        if edges_added >= total_subgraphs :
243            if nx.is_strongly_connected(G):
244                break
245
246    # Get some candidate extra edges
247    for u, v in all_edges :
248
249        if G.has_edge( u , v ) :
250            continue
251
252        if end_of_component_symbols is not None and n_edges_from[u] >= len( end_of_component_symbols ) :
253            continue
254
255        edges_added += 1
256
257        n_edges_from[ u ] += 1
258        n_edges_to[ v ] += 1
259
260        extra_edges.append( (u,v) )
261
262        if edges_added >= max_edges :
263            break
264
265    ##############################################################
266    # calculate best entry and exit nodes for each subgraph
267
268    # lower is better
269    def score_as_from( m, s_idx ) :
270        n_outgoing = 0
271        has_self_loop = False
272        for tr in m.transitions :
273            if tr.origin_state_idx == s_idx :
274                n_outgoing += 1
275            if tr.origin_state_idx == s_idx and tr.target_state_idx == s_idx :
276                has_self_loop = True
277        return ( n_outgoing, not has_self_loop, s_idx )
278
279    # higher is better
280    def score_as_to( m, s_idx, path_length ) :
281        n_incoming = 0
282        has_self_loop = False
283        for tr in m.transitions :
284            if tr.target_state_idx == s_idx :
285                n_incoming += 1
286            if tr.origin_state_idx == s_idx and tr.target_state_idx == s_idx :
287                has_self_loop = True
288
289        return ( path_length, 1.0/n_incoming, has_self_loop, s_idx )
290
291    best_entry_and_exits = []
292
293    if connector_state_pools is None :
294        
295        for cm, _, _ in component_pool:
296
297            n_states = len(cm.states)
298            cm_dg = cm.as_digraph()
299
300            best_from = min(range(n_states), key=lambda s: score_as_from(cm, s))
301
302            lengths_to_from = nx.single_source_shortest_path_length(
303                cm_dg.reverse(), best_from)
304
305            best_to = max(
306                range(n_states),
307                key=lambda s: score_as_to(cm, s, lengths_to_from.get(s, 0)))
308
309            best_entry_and_exits.append( ( best_to, best_from ) )
310    else :
311
312        for cm, g_id, cmi in component_pool:
313          
314            if g_id not in connector_state_pools :
315                raise ValueError( "connector_states missing group id" )
316          
317            if cmi < 0 or cmi >= len( connector_state_pools[ g_id ] ) :
318                raise ValueError( f"Invalid connector state list for group {g_id}." )
319
320            n_states = len(cm.states)
321            pool = connector_state_pools[ g_id ][ cmi ]
322
323            if len( pool ) == 0:
324                raise ValueError( f"Empty connector pool for group {g_id}, component {cmi}." )
325
326            invalid  = [ 
327                ( entry_state, exit_state ) for entry_state, exit_state in pool 
328                if ( entry_state >= n_states 
329                  or exit_state  >= n_states 
330                  or entry_state < 0 
331                  or exit_state  < 0 )
332            ]
333
334            if invalid:
335                raise ValueError( f"connector_state_pools['{g_id}'][{cmi}] has out-of-range indices" )
336
337            entry_options = [ p[ 0 ] for p in pool ]
338            exit_options  = [ p[ 1 ] for p in pool ]
339
340            cm_dg = cm.as_digraph()
341
342            best_from = min( exit_options, key=lambda s: score_as_from(cm, s) )
343
344            lengths_to_from = nx.single_source_shortest_path_length(
345                cm_dg.reverse(), best_from)
346
347            best_to = max(
348                entry_options,
349                key=lambda s: score_as_to(cm, s, lengths_to_from.get(s, 0)))
350
351            best_entry_and_exits.append( ( best_to, best_from ) )
352
353    # The states which are entry or exit states between components are no longer true 
354    # ismorphs of the states they used to be isomorphic to
355    
356    instance_idx = 0
357    for m_idx, (cm, g_idx, g_offset ) in enumerate(component_pool):
358        
359        best_to_local, best_from_local = best_entry_and_exits[m_idx]
360        n_instances = instances_per_component[ g_idx ][ g_offset ]
361        
362        for _ in range(n_instances):
363
364            g_best_to   = cstate_instance_to_global_idx[instance_idx][best_to_local]
365            g_best_from = cstate_instance_to_global_idx[instance_idx][best_from_local]
366            
367            states[g_best_to ] = states[g_best_to].modified_deep_copy(
368                isomorphs = set(),
369                pseudo_isomorphs = states[g_best_to].pseudo_isomorphs | states[g_best_to].isomorphs
370            )
371
372            states[g_best_from ] = states[g_best_from].modified_deep_copy(
373                isomorphs = set(),
374                pseudo_isomorphs = states[g_best_from].pseudo_isomorphs | states[g_best_from].isomorphs
375            )
376
377            instance_idx += 1
378
379    ##############################################################
380
381    # Expand component to instances and apply noise
382
383    nodes = list(G.nodes())
384    N = len(nodes)
385    node_to_idx = {node: i for i, node in enumerate(nodes)}
386
387    assert sum(component_instance_counts) == N
388
389    pos_to_comp = [c for c, n in enumerate(component_instance_counts) for _ in range(n)]
390    instance_cohesion = component_cohesion[np.ix_(pos_to_comp, pos_to_comp)]
391    
392    lo, hi = component_cohesion.min(), component_cohesion.max()
393    noise = np_rng.uniform( lo, hi, ( N, N ) )
394    instance_cohesion = (1 - instance_cohesion_randomness) * instance_cohesion + instance_cohesion_randomness * noise
395
396    ##########################################################################
397
398    # Note: CP-SAT is about on par with efficiently implemented brute force for this problem
399    # and ILS tends to find optimal soltion anyways when the search space is this small
400    # probably either always use ILS, or consider implementing pure brute force for small cases
401    # would be better, however there is a chance that a better cp-sat formulation chould change the 
402    # calculus. 
403
404    # Search space is N!
405    if N < 12 : 
406        node_permutation = optimize_node_permutation_cpsat( 
407            G=G, 
408            cohesion_matrix=instance_cohesion,
409            max_search_time=60,
410            np_rng=np_rng
411        )
412    else :
413        node_permutation = optimize_node_permutation_ils( 
414            G=G, 
415            cohesion_matrix=instance_cohesion,
416            np_rng=np_rng
417        )
418
419    nodes = list( G.nodes() )
420    mapping = { i : node_permutation[i] for i in range(total_subgraphs) }
421
422    cycle_edges  = [ ( mapping[u], mapping[v] ) for u, v in cycle_edges ]
423    extra_edges  = [ ( mapping[u], mapping[v] ) for u, v in extra_edges ]
424
425    # connectors is a list over the component instances of the global enter / exit state of the instances  
426    connectors : list[tuple[int,int]] = []
427
428    for i in range( total_subgraphs ) :
429        m_idx = instance_idx_to_component_idx[i]
430        best_to_local, best_from_local = best_entry_and_exits[ m_idx ]
431        connectors.append((
432            cstate_instance_to_global_idx[i][ best_to_local   ], # entry
433            cstate_instance_to_global_idx[i][ best_from_local ]  # exit
434        ) )
435
436    ##############################################################
437    # Connect the component instances
438
439    # In theory could score to and from nodes also based on what kind of cohesion scores are possible
440    # Given what free symbols they would have, but this is a bit complicated, considering adding it later.
441
442    extra_edges_added = 0
443    for edge_idx, ( u, v ) in enumerate( cycle_edges + extra_edges ) :
444
445        is_cycle_edge = edge_idx < len( cycle_edges )
446
447        m_u_idx = instance_idx_to_component_idx[ u ]
448        m_v_idx = instance_idx_to_component_idx[ v ]
449
450        best_from = best_entry_and_exits[ m_u_idx ][ 1 ]
451        best_to   = best_entry_and_exits[ m_v_idx ][ 0 ]
452
453        origin = cstate_instance_to_global_idx[ u ][ best_from ]
454        target = cstate_instance_to_global_idx[ v ][ best_to   ]
455
456        # if end_of_component_symbols is None we use the whole alphabet, otherwise eoc only
457        free_eoc_symbols    = end_of_component_symbols.copy() if end_of_component_symbols is not None else set()
458        free_normal_symbols = set( alphabet ) if end_of_component_symbols is None else set()
459
460        # reweight existing transitions and figure out which eoc symbols are available
461        for tr_idx, tr in enumerate( transitions ) :
462            if tr.origin_state_idx == origin :
463
464                transitions[ tr_idx  ] = transitions[ tr_idx  ].modified_deep_copy(
465                    prob = transitions[ tr_idx  ].prob * component_residency_factor
466                )
467
468                symbol = alphabet[ tr.symbol_idx ]
469                
470                if symbol in free_eoc_symbols :
471                    free_eoc_symbols.remove( symbol )
472                
473                if symbol in free_normal_symbols :
474                    free_normal_symbols.remove( symbol )
475
476        # Shouldn't happen on a cycle edge
477        if len( free_eoc_symbols ) == 0 and end_of_component_symbols is not None :
478            if is_cycle_edge :
479                raise ValueError( "Ran out of eoc symbols attempting to connect components." )
480            else :
481                continue
482
483        # This could happen on a cycle edge
484        if len( free_normal_symbols ) == 0 and end_of_component_symbols is None :
485            if is_cycle_edge :
486                raise ValueError( "Ran out of symbols attempting to connect components." )
487            else :
488                continue
489
490        if not is_cycle_edge :
491            extra_edges_added += 1
492
493        if len( free_eoc_symbols ) > 0 :
494            symbol = np_rng.choice( sorted( free_eoc_symbols ) )
495
496        elif len( free_normal_symbols ) > 0 :
497            avail = list( free_normal_symbols )
498            cohesion_scores = [ 0 for _ in avail ]
499            for i, s in enumerate( avail ) : 
500                for tr in transitions :
501                    if tr.target_state_idx == origin :
502                        cohesion_scores[ i ] += symbol_cohesion[ tr.symbol_idx, symbol_idx_map[ s ] ]
503                    elif tr.origin_state_idx == target :
504                        cohesion_scores[ i ] += symbol_cohesion[ symbol_idx_map[ s ], tr.symbol_idx ]
505            best_i = np.argmax( np.array( cohesion_scores ) )
506            symbol = avail[ best_i ]
507        else :
508            continue
509
510        transitions.append(
511            Transition(
512                origin_state_idx=origin,
513                target_state_idx=target,
514                prob=( 1.0 - component_residency_factor ),
515                symbol_idx=symbol_idx_map[ symbol ],
516                pq=None,
517                composition_depth=composition_depth
518        ) )
519
520    print( f"Was able to add {extra_edges_added}/{len(extra_edges)} extra cross-component transitions beyond a Hamiltonian cycle" )
521
522    res = HMM( 
523        states=states,
524        transitions=transitions,
525        alphabet=alphabet,
526        start_state=0,
527        composition_depth=composition_depth
528    )
529
530    if not res.is_row_stochastic() :
531        raise Exception( "Structured composition is not row stochastic" )
532
533    if not  res.is_unifilar() :
534        raise Exception( "Structured composition is not unifilar" )
535
536    return res, connectors