GitLab Repo

amachine.am_vis 

  1from pathlib import Path
  2import graphviz
  3from matplotlib import colormaps
  4from matplotlib.colors import to_hex
  5import matplotlib.pyplot as plt
  6from IPython.display import Image, display
  7import matplotlib.image as mpimg
  8from io import BytesIO
  9import imageio
 10from pathlib import Path
 11import numpy as np
 12from scipy.spatial.distance import pdist, squareform
 13
 14from matplotlib.colors import Normalize
 15from matplotlib import colormaps
 16from matplotlib.colors import Colormap
 17
 18import ternary
 19
 20def create_digraph(engine="dot"):
 21
 22    engine_configs = {
 23        "dot": {
 24            'rankdir': 'LR',
 25            # 'rankdir': 'TB',
 26            'ranksep': '0.85',
 27            'nodesep': '1.0',
 28            'splines': 'spline',
 29            'constraint' : 'true',
 30            'concentrate': 'false',
 31            'ratio': 'auto',
 32            # 'nslimit' : '0',
 33            # 'nslimit1' : '2'
 34        },
 35        "neato": {
 36            'overlap': 'scale',
 37            'overlap_scaling': '-4',  
 38            'esep': '+2.5',           
 39            'sep': '+1.75',           
 40            'model': 'shortpath',
 41            'damping': '0.85',
 42            'epsilon': '0.00001',
 43            'maxiter': '10000',
 44            'start': '5',
 45        },
 46        "fdp": {
 47            'overlap': 'prism',
 48            'sep': '+1.5',
 49            'K': '1.0',
 50            'splines': 'true',
 51            'len' : '3.0',
 52            'maxiter': '5000'
 53        }
 54    }
 55
 56    graph_attr = engine_configs.get(engine, {})
 57    graph_attr['outputorder'] = 'edgesfirst'
 58
 59    node_attr = {
 60        'shape': 'box',
 61        'style': 'rounded, filled',
 62        'fillcolor': 'lightblue',
 63        'fontname': 'Helvetica'
 64    }
 65    
 66    edge_attr = {
 67        'penwidth': '12.0',
 68        'color': 'gray40',
 69        'arrowsize' : '1.75'
 70    }
 71
 72    return graphviz.Digraph(
 73        engine=engine,
 74        graph_attr=graph_attr,
 75        node_attr=node_attr,
 76        edge_attr=edge_attr
 77    )
 78
 79
 80def draw_graph( 
 81    aM, 
 82    output_dir : Path | str | None = None,
 83    title : str ="am_graph", 
 84    view : bool = True, 
 85    format : str = "png",
 86    strongly_connected_components : list[list[int]] | None = None,
 87    engine : str ="dot",
 88    symbols_only: bool = False,
 89    no_labels : bool = False,
 90    color_nodes_by : str | None = "strongly_connected_components",
 91    color_borders_by : str | None = None,
 92    color_edges_by : str | None = None,
 93    symbol_categories : str | None = None ):
 94
 95    if color_edges_by == "category" and symbol_categories is None :
 96        raise ValueError( "Coloring edges by category requires passing symbol categories" )
 97
 98    #########################################################################################
 99
100    GV = create_digraph(engine=engine)
101
102    #########################################################################################
103
104    node_cmap = colormaps['tab20_r']
105    n_node_colors = 8
106    node_colors = None
107    node_classes = None
108
109    if strongly_connected_components and color_nodes_by=="strongly_connected_components" :
110        node_colors = [to_hex(node_cmap(i % n_node_colors)) for i in range(len(strongly_connected_components))]
111
112    elif color_nodes_by  :
113        class_code = color_nodes_by.split( "_" )[ 0 ]
114        node_classes = set()
115        for state in aM.states :
116            for cl in state.classes :
117                code = cl.split( "_" )[ 0 ]
118                if code == class_code :
119                    node_classes.add( cl )
120
121        node_classes = list( node_classes )
122        node_colors = [to_hex(node_cmap(i % n_node_colors)) for i in range(len(node_classes))]
123
124    #########################################################################################
125
126    edge_colors = None
127
128    if color_edges_by == "category" :
129        edge_cmap = colormaps['Dark2']
130        n_edge_colors = 8
131        categories  = list( set( symbol_categories.values() ) )
132        edge_colors = [to_hex(edge_cmap(i % n_edge_colors)) for i in range(len(categories))]
133
134    elif color_edges_by == "symbol" :
135        edge_cmap = colormaps['tab20']
136        n_edge_colors = 20
137        edge_colors = [to_hex(edge_cmap(i % n_edge_colors)) for i in range(len(aM.alphabet))]
138
139    #########################################################################################
140
141    border_cmap = colormaps['Set1']
142    n_border_colors = 8
143    border_colors = None
144    border_classes = None
145
146    if strongly_connected_components and color_borders_by=="strongly_connected_components" :
147        border_colors = [to_hex(border_cmap(i % n_border_colors)) for i in range(len(strongly_connected_components))]
148
149    elif color_borders_by  :
150        class_code = color_borders_by.split( "_" )[ 0 ]
151        border_classes = set()
152        for state in aM.states :
153            for cl in state.classes :
154                code = cl.split( "_" )[ 0 ]
155                if code == class_code :
156                    border_classes.add( cl )
157
158        border_classes = list( border_classes )
159        border_colors = [to_hex(border_cmap(i % n_border_colors)) for i in range(len(border_classes))]
160
161    #########################################################################################
162
163    for node_idx in range( len(aM.states) ) :
164
165        #-------------------------------------------------------------------------------------
166
167        nb_color   = [ 'gray', 'black' ]
168
169        nb_colors  = (    node_colors,    border_colors )
170        nb_cby     = ( color_nodes_by, color_borders_by )
171        nb_classes = (   node_classes,   border_classes )
172
173        for nb_i, c_by in enumerate( nb_cby ) :
174
175            if nb_classes[ nb_i ] is None :
176                continue
177
178            if strongly_connected_components and c_by=="strongly_connected_components" :
179                
180                node_component = -1
181                for i, sg in enumerate( strongly_connected_components ) :
182                    if node_idx in sg :
183                        node_component = i
184
185                if node_component >= 0 :
186                    nb_color[ nb_i ] = nb_colors[ nb_i ][ node_component % len(nb_colors[nb_i]) ]
187
188            elif c_by :
189                state_classes = aM.states[ node_idx ].classes
190                for cl in state_classes : 
191                    try : 
192                        color_index = nb_classes[ nb_i ].index( cl )
193                        nb_color[ nb_i ] = nb_colors[ nb_i ][ color_index % len(nb_colors[nb_i]) ]
194                    except :
195                        continue
196
197        #-------------------------------------------------------------------------------------
198
199        node_tex = aM.states[ node_idx ].name
200        if not node_tex :
201            node_tex = str(node_idx)
202
203        GV.node(
204            str(node_idx), 
205            label=node_tex, 
206            shape='circle',
207            style='bold,filled',
208            fillcolor=nb_color[ 0 ],
209            color=nb_color[ 1 ],
210            width='2.0',
211            penwidth='26.0' )
212
213    for tr in aM.transitions :
214
215        u = tr.origin_state_idx
216        v = tr.target_state_idx
217
218        pr_str = str( tr.pq ) if aM.is_q_weighted else str( round( tr.prob, 4 ) )
219        
220        fontsize= '80' if symbols_only else '12'
221
222        if symbols_only :
223            label_text = f"{aM.alphabet[tr.symbol_idx]}"
224        else :
225            label_text = f"{aM.alphabet[tr.symbol_idx]}({pr_str})"
226
227        html_label = (
228            f'<<TABLE BORDER="0" CELLBORDER="0" CELLSPACING="0" CELLPADDING="2">'
229            f'<TR><TD>{label_text}</TD></TR>'
230            f'</TABLE>>'
231        )
232
233        edge_color = "black"
234
235        if color_edges_by == "category" :
236            symbol = aM.alphabet[ tr.symbol_idx ]
237            if symbol in symbol_categories :
238                cat = symbol_categories[ symbol ]
239                try :
240                    cl_idx = categories.index( cat )
241                    edge_color = edge_colors[ cl_idx % n_edge_colors ]
242                except :
243                    edge_color = "black"
244
245        elif color_edges_by == "symbol" :
246            edge_color = edge_colors[ tr.symbol_idx % n_edge_colors ]
247
248        if no_labels :
249            html_label = None
250
251        GV.edge(
252            str(u), str(v), 
253            label=html_label,
254            headlabel=" ",
255            fontsize=fontsize,
256            fontname='Times-Italic',
257            #labelfloat="true",
258            fontcolor="#1f1f1f",
259            color=edge_color,
260            arrowhead='nonenonenormal'
261        )
262
263    if view == True :
264
265        GV.graph_attr.update(dpi='150')
266        png_bytes = GV.pipe(format="png")
267        buf = BytesIO( png_bytes )
268        img = mpimg.imread( buf )
269
270        plt.figure(figsize=(8, 6))
271        plt.imshow(img)
272        plt.axis('off')
273        plt.tight_layout()
274        plt.show()
275
276    if output_dir is not None :		
277        GV.attr(bgcolor='transparent')
278        GV.render( title, directory=output_dir, view=False, format=format, cleanup=True)
279
280def animated_distance_plot( vectors, max_res=1440, stride=7, apply_threshold=False, threshold=0.01 ) :
281
282    use_cuda = False # cpx is not None
283
284    n_vectors = len( vectors )
285    max_res = min( n_vectors, max_res )
286
287    if use_cuda :
288        state_vectors = cp.asarray( vectors[ 0:max_states, : ] )
289    else :
290        state_vectors = vectors
291
292    imgs = []
293    for start in range( 0, n_vectors, stride ) :
294        
295        if start % (stride*10) == 0 :
296            print( f"offset: {start}" )
297
298        end = min( start+max_res, n_vectors )
299        
300        if use_cuda :
301            distance_plot = squareform( 
302                cpx.scipy.spatial.distance.pdist( state_vectors[ start:end, : ], metric='jensenshannon' ).get()
303            )
304        else :
305            distance_plot = squareform( 
306                pdist( state_vectors[ start:end, : ], metric='jensenshannon' )
307            )
308
309        if len( distance_plot ) < max_res :
310            distance_plot = np.pad( 
311                distance_plot, 
312                ( ( 0, max_res-len( distance_plot ) ), (0, max_res-len( distance_plot )) ), 
313                mode='constant' )
314
315        if apply_threshold : 
316            distance_plot = np.where( distance_plot > threshold, 1, 0 ).astype( np.uint8 )
317            imgs.append( distance_plot )
318        else :
319            imgs.append( distance_plot.astype( np.float16 ) )
320
321    print( f"Rendering {len(imgs)} frames." )
322    to_video( imgs, fps=60 )
323
324def to_video(
325    images      : list[np.ndarray],
326    output_path : str = "output.mp4",
327    fps         : int = 3,
328    cmap        : str | Colormap = "plasma_r",
329) -> None:
330    if not images:
331        raise ValueError("images list is empty")
332
333    colormap = colormaps[cmap] if isinstance(cmap, str) else cmap
334
335    # Global min/max so colormapping is consistent across frames
336    global_min = min(f.min() for f in images)
337    global_max = max(f.max() for f in images)
338
339    if global_min == global_max:
340        raise ValueError("All frames are constant, cannot normalize")
341
342    norm = Normalize( vmin=global_min, vmax=global_max )
343
344    with imageio.get_writer(output_path, fps=fps) as writer:
345        for frame in images:
346            rgba   = colormap(norm(frame))
347            rgb    = (rgba[..., :3] * 255).astype(np.uint8)
348            writer.append_data(rgb)
349
350def triangular_barycentric_plot(
351    coords : np.ndarray,  
352    coord_classes : np.ndarray,
353    axis_labels : tuple[str, str, str],
354    title : str="" ) :
355
356    figure, tax = ternary.figure(scale=40)
357
358    tax.boundary(linewidth=2.0)
359
360    fontsize = 20
361    tax.set_title("Simplex Boundary and Gridlines", fontsize=fontsize)
362
363    tax.left_axis_label( axis_labels[0], fontsize=fontsize)
364    tax.right_axis_label( axis_labels[1], fontsize=fontsize)
365    tax.bottom_axis_label( axis_labels[2], fontsize=fontsize)
366
367    # Remove default Matplotlib Axes
368    tax.clear_matplotlib_ticks()
369
370    tax.scatter( coords, marker='*' )
371
372    ternary.plt.show()
def create_digraph(engine='dot'): 
22def create_digraph(engine="dot"):
23
24    engine_configs = {
25        "dot": {
26            'rankdir': 'LR',
27            # 'rankdir': 'TB',
28            'ranksep': '0.85',
29            'nodesep': '1.0',
30            'splines': 'spline',
31            'constraint' : 'true',
32            'concentrate': 'false',
33            'ratio': 'auto',
34            # 'nslimit' : '0',
35            # 'nslimit1' : '2'
36        },
37        "neato": {
38            'overlap': 'scale',
39            'overlap_scaling': '-4',  
40            'esep': '+2.5',           
41            'sep': '+1.75',           
42            'model': 'shortpath',
43            'damping': '0.85',
44            'epsilon': '0.00001',
45            'maxiter': '10000',
46            'start': '5',
47        },
48        "fdp": {
49            'overlap': 'prism',
50            'sep': '+1.5',
51            'K': '1.0',
52            'splines': 'true',
53            'len' : '3.0',
54            'maxiter': '5000'
55        }
56    }
57
58    graph_attr = engine_configs.get(engine, {})
59    graph_attr['outputorder'] = 'edgesfirst'
60
61    node_attr = {
62        'shape': 'box',
63        'style': 'rounded, filled',
64        'fillcolor': 'lightblue',
65        'fontname': 'Helvetica'
66    }
67    
68    edge_attr = {
69        'penwidth': '12.0',
70        'color': 'gray40',
71        'arrowsize' : '1.75'
72    }
73
74    return graphviz.Digraph(
75        engine=engine,
76        graph_attr=graph_attr,
77        node_attr=node_attr,
78        edge_attr=edge_attr
79    )
def draw_graph( aM, output_dir: pathlib.Path | str | None = None, title: str = 'am_graph', view: bool = True, format: str = 'png', strongly_connected_components: list[list[int]] | None = None, engine: str = 'dot', symbols_only: bool = False, no_labels: bool = False, color_nodes_by: str | None = 'strongly_connected_components', color_borders_by: str | None = None, color_edges_by: str | None = None, symbol_categories: str | None = None): 
 82def draw_graph( 
 83    aM, 
 84    output_dir : Path | str | None = None,
 85    title : str ="am_graph", 
 86    view : bool = True, 
 87    format : str = "png",
 88    strongly_connected_components : list[list[int]] | None = None,
 89    engine : str ="dot",
 90    symbols_only: bool = False,
 91    no_labels : bool = False,
 92    color_nodes_by : str | None = "strongly_connected_components",
 93    color_borders_by : str | None = None,
 94    color_edges_by : str | None = None,
 95    symbol_categories : str | None = None ):
 96
 97    if color_edges_by == "category" and symbol_categories is None :
 98        raise ValueError( "Coloring edges by category requires passing symbol categories" )
 99
100    #########################################################################################
101
102    GV = create_digraph(engine=engine)
103
104    #########################################################################################
105
106    node_cmap = colormaps['tab20_r']
107    n_node_colors = 8
108    node_colors = None
109    node_classes = None
110
111    if strongly_connected_components and color_nodes_by=="strongly_connected_components" :
112        node_colors = [to_hex(node_cmap(i % n_node_colors)) for i in range(len(strongly_connected_components))]
113
114    elif color_nodes_by  :
115        class_code = color_nodes_by.split( "_" )[ 0 ]
116        node_classes = set()
117        for state in aM.states :
118            for cl in state.classes :
119                code = cl.split( "_" )[ 0 ]
120                if code == class_code :
121                    node_classes.add( cl )
122
123        node_classes = list( node_classes )
124        node_colors = [to_hex(node_cmap(i % n_node_colors)) for i in range(len(node_classes))]
125
126    #########################################################################################
127
128    edge_colors = None
129
130    if color_edges_by == "category" :
131        edge_cmap = colormaps['Dark2']
132        n_edge_colors = 8
133        categories  = list( set( symbol_categories.values() ) )
134        edge_colors = [to_hex(edge_cmap(i % n_edge_colors)) for i in range(len(categories))]
135
136    elif color_edges_by == "symbol" :
137        edge_cmap = colormaps['tab20']
138        n_edge_colors = 20
139        edge_colors = [to_hex(edge_cmap(i % n_edge_colors)) for i in range(len(aM.alphabet))]
140
141    #########################################################################################
142
143    border_cmap = colormaps['Set1']
144    n_border_colors = 8
145    border_colors = None
146    border_classes = None
147
148    if strongly_connected_components and color_borders_by=="strongly_connected_components" :
149        border_colors = [to_hex(border_cmap(i % n_border_colors)) for i in range(len(strongly_connected_components))]
150
151    elif color_borders_by  :
152        class_code = color_borders_by.split( "_" )[ 0 ]
153        border_classes = set()
154        for state in aM.states :
155            for cl in state.classes :
156                code = cl.split( "_" )[ 0 ]
157                if code == class_code :
158                    border_classes.add( cl )
159
160        border_classes = list( border_classes )
161        border_colors = [to_hex(border_cmap(i % n_border_colors)) for i in range(len(border_classes))]
162
163    #########################################################################################
164
165    for node_idx in range( len(aM.states) ) :
166
167        #-------------------------------------------------------------------------------------
168
169        nb_color   = [ 'gray', 'black' ]
170
171        nb_colors  = (    node_colors,    border_colors )
172        nb_cby     = ( color_nodes_by, color_borders_by )
173        nb_classes = (   node_classes,   border_classes )
174
175        for nb_i, c_by in enumerate( nb_cby ) :
176
177            if nb_classes[ nb_i ] is None :
178                continue
179
180            if strongly_connected_components and c_by=="strongly_connected_components" :
181                
182                node_component = -1
183                for i, sg in enumerate( strongly_connected_components ) :
184                    if node_idx in sg :
185                        node_component = i
186
187                if node_component >= 0 :
188                    nb_color[ nb_i ] = nb_colors[ nb_i ][ node_component % len(nb_colors[nb_i]) ]
189
190            elif c_by :
191                state_classes = aM.states[ node_idx ].classes
192                for cl in state_classes : 
193                    try : 
194                        color_index = nb_classes[ nb_i ].index( cl )
195                        nb_color[ nb_i ] = nb_colors[ nb_i ][ color_index % len(nb_colors[nb_i]) ]
196                    except :
197                        continue
198
199        #-------------------------------------------------------------------------------------
200
201        node_tex = aM.states[ node_idx ].name
202        if not node_tex :
203            node_tex = str(node_idx)
204
205        GV.node(
206            str(node_idx), 
207            label=node_tex, 
208            shape='circle',
209            style='bold,filled',
210            fillcolor=nb_color[ 0 ],
211            color=nb_color[ 1 ],
212            width='2.0',
213            penwidth='26.0' )
214
215    for tr in aM.transitions :
216
217        u = tr.origin_state_idx
218        v = tr.target_state_idx
219
220        pr_str = str( tr.pq ) if aM.is_q_weighted else str( round( tr.prob, 4 ) )
221        
222        fontsize= '80' if symbols_only else '12'
223
224        if symbols_only :
225            label_text = f"{aM.alphabet[tr.symbol_idx]}"
226        else :
227            label_text = f"{aM.alphabet[tr.symbol_idx]}({pr_str})"
228
229        html_label = (
230            f'<<TABLE BORDER="0" CELLBORDER="0" CELLSPACING="0" CELLPADDING="2">'
231            f'<TR><TD>{label_text}</TD></TR>'
232            f'</TABLE>>'
233        )
234
235        edge_color = "black"
236
237        if color_edges_by == "category" :
238            symbol = aM.alphabet[ tr.symbol_idx ]
239            if symbol in symbol_categories :
240                cat = symbol_categories[ symbol ]
241                try :
242                    cl_idx = categories.index( cat )
243                    edge_color = edge_colors[ cl_idx % n_edge_colors ]
244                except :
245                    edge_color = "black"
246
247        elif color_edges_by == "symbol" :
248            edge_color = edge_colors[ tr.symbol_idx % n_edge_colors ]
249
250        if no_labels :
251            html_label = None
252
253        GV.edge(
254            str(u), str(v), 
255            label=html_label,
256            headlabel=" ",
257            fontsize=fontsize,
258            fontname='Times-Italic',
259            #labelfloat="true",
260            fontcolor="#1f1f1f",
261            color=edge_color,
262            arrowhead='nonenonenormal'
263        )
264
265    if view == True :
266
267        GV.graph_attr.update(dpi='150')
268        png_bytes = GV.pipe(format="png")
269        buf = BytesIO( png_bytes )
270        img = mpimg.imread( buf )
271
272        plt.figure(figsize=(8, 6))
273        plt.imshow(img)
274        plt.axis('off')
275        plt.tight_layout()
276        plt.show()
277
278    if output_dir is not None :		
279        GV.attr(bgcolor='transparent')
280        GV.render( title, directory=output_dir, view=False, format=format, cleanup=True)
def animated_distance_plot( vectors, max_res=1440, stride=7, apply_threshold=False, threshold=0.01): 
282def animated_distance_plot( vectors, max_res=1440, stride=7, apply_threshold=False, threshold=0.01 ) :
283
284    use_cuda = False # cpx is not None
285
286    n_vectors = len( vectors )
287    max_res = min( n_vectors, max_res )
288
289    if use_cuda :
290        state_vectors = cp.asarray( vectors[ 0:max_states, : ] )
291    else :
292        state_vectors = vectors
293
294    imgs = []
295    for start in range( 0, n_vectors, stride ) :
296        
297        if start % (stride*10) == 0 :
298            print( f"offset: {start}" )
299
300        end = min( start+max_res, n_vectors )
301        
302        if use_cuda :
303            distance_plot = squareform( 
304                cpx.scipy.spatial.distance.pdist( state_vectors[ start:end, : ], metric='jensenshannon' ).get()
305            )
306        else :
307            distance_plot = squareform( 
308                pdist( state_vectors[ start:end, : ], metric='jensenshannon' )
309            )
310
311        if len( distance_plot ) < max_res :
312            distance_plot = np.pad( 
313                distance_plot, 
314                ( ( 0, max_res-len( distance_plot ) ), (0, max_res-len( distance_plot )) ), 
315                mode='constant' )
316
317        if apply_threshold : 
318            distance_plot = np.where( distance_plot > threshold, 1, 0 ).astype( np.uint8 )
319            imgs.append( distance_plot )
320        else :
321            imgs.append( distance_plot.astype( np.float16 ) )
322
323    print( f"Rendering {len(imgs)} frames." )
324    to_video( imgs, fps=60 )
def to_video( images: list[numpy.ndarray], output_path: str = 'output.mp4', fps: int = 3, cmap: str | matplotlib.colors.Colormap = 'plasma_r') -> None: 
326def to_video(
327    images      : list[np.ndarray],
328    output_path : str = "output.mp4",
329    fps         : int = 3,
330    cmap        : str | Colormap = "plasma_r",
331) -> None:
332    if not images:
333        raise ValueError("images list is empty")
334
335    colormap = colormaps[cmap] if isinstance(cmap, str) else cmap
336
337    # Global min/max so colormapping is consistent across frames
338    global_min = min(f.min() for f in images)
339    global_max = max(f.max() for f in images)
340
341    if global_min == global_max:
342        raise ValueError("All frames are constant, cannot normalize")
343
344    norm = Normalize( vmin=global_min, vmax=global_max )
345
346    with imageio.get_writer(output_path, fps=fps) as writer:
347        for frame in images:
348            rgba   = colormap(norm(frame))
349            rgb    = (rgba[..., :3] * 255).astype(np.uint8)
350            writer.append_data(rgb)
def triangular_barycentric_plot( coords: numpy.ndarray, coord_classes: numpy.ndarray, axis_labels: tuple[str, str, str], title: str = ''): 
352def triangular_barycentric_plot(
353    coords : np.ndarray,  
354    coord_classes : np.ndarray,
355    axis_labels : tuple[str, str, str],
356    title : str="" ) :
357
358    figure, tax = ternary.figure(scale=40)
359
360    tax.boundary(linewidth=2.0)
361
362    fontsize = 20
363    tax.set_title("Simplex Boundary and Gridlines", fontsize=fontsize)
364
365    tax.left_axis_label( axis_labels[0], fontsize=fontsize)
366    tax.right_axis_label( axis_labels[1], fontsize=fontsize)
367    tax.bottom_axis_label( axis_labels[2], fontsize=fontsize)
368
369    # Remove default Matplotlib Axes
370    tax.clear_matplotlib_ticks()
371
372    tax.scatter( coords, marker='*' )
373
374    ternary.plt.show()