33"""
44
55import random
6+ import math
67import networkx as nx
78
89
@@ -58,6 +59,9 @@ def format_output(obj, expr, format=None):
5859 return boxes
5960
6061
62+ cached_pair = None
63+
64+
6165def hierarchy_pos (
6266 G , root = None , width = 1.0 , vert_gap = 0.2 , vert_loc = 0 , leaf_vs_root_factor = 0.5
6367):
@@ -117,6 +121,10 @@ def hierarchy_pos(
117121 if not nx .is_tree (G ):
118122 raise TypeError ("cannot use hierarchy_pos on a graph that is not a tree" )
119123
124+ global cached_pair
125+ if cached_pair is not None :
126+ return cached_pair
127+
120128 # These get swapped if tree edge directions point to the root.
121129 decendants = nx .descendants
122130 out_degree = G .out_degree if hasattr (G , "out_degree" ) else G .degree
@@ -247,7 +255,8 @@ def _hierarchy_pos(
247255 if n <= 0 :
248256 continue
249257 min_sep = min ([x_list [i + 1 ] - x_list [i ] for i in range (n )] + [min_sep ])
250- return pos , min_sep
258+ cached_pair = pos , min_sep
259+ return cached_pair
251260
252261
253262node_size = 300 # this is networkx's default size
@@ -273,6 +282,34 @@ def tree_layout(G):
273282}
274283
275284
285+ def clamp (value , min = - math .inf , max = math .inf ):
286+ if value <= min :
287+ return min
288+ if value >= max :
289+ return max
290+ return value
291+
292+
293+ DEFAULT_NODE_SIZE = 300.0
294+ DEFAULT_POINT_SIZE = 12
295+
296+
297+ def harmonize_parameters (draw_options : dict ):
298+ if draw_options .get ("with_labels" , False ):
299+ draw_options ["edgecolors" ] = draw_options .get ("edgecolors" , "black" )
300+ draw_options ["node_color" ] = draw_options .get ("node_color" , "white" )
301+
302+ if "width" not in draw_options :
303+ width = clamp (node_size / DEFAULT_NODE_SIZE , min = 0.15 )
304+ draw_options ["width" ] = width
305+
306+ if "font_size" not in draw_options :
307+ font_size = clamp (
308+ int ((node_size * DEFAULT_POINT_SIZE ) / DEFAULT_NODE_SIZE ), min = 1 , max = 18
309+ )
310+ draw_options ["font_size" ] = font_size
311+
312+
276313def format_graph (G ):
277314 """
278315 Format a Graph
@@ -281,12 +318,25 @@ def format_graph(G):
281318 import matplotlib .pyplot as plt
282319
283320 global node_size
284- node_size = 300 # This is networkx's default
321+ global cached_pair
322+
323+ cached_pair = None
285324
286325 graph_layout = G .graph_layout if hasattr (G , "graph_layout" ) else None
326+
327+ node_size = DEFAULT_NODE_SIZE
328+ draw_options = {
329+ "node_size" : node_size ,
330+ # "with_labels": vertex_labels # Set below
331+ # "font_size": 12, # Harmonized
332+ # "node_color": "white", # Set below
333+ # "edgecolors": "black", # Set below
334+ # "width": 5, # Marmonized
335+ }
336+
287337 vertex_labels = G .vertex_labels if hasattr (G , "vertex_labels" ) else False
288338 if vertex_labels :
289- vertex_labels = vertex_labels .to_python () or False
339+ draw_options [ "with_labels" ] = vertex_labels .to_python () or False
290340
291341 if hasattr (G , "title" ) and G .title .get_string_value ():
292342 fig , ax = plt .subplots () # Create a figure and an axes
@@ -296,26 +346,19 @@ def format_graph(G):
296346 if not isinstance (graph_layout , str ):
297347 graph_layout = graph_layout .get_string_value ()
298348 layout_fn = NETWORKX_LAYOUTS .get (graph_layout , None )
349+ if layout_fn == tree_layout :
350+ # Call this to compute node_size. Cache the
351+ # results
352+ tree_layout (G )
353+ draw_options ["node_size" ] = node_size
299354 else :
300355 layout_fn = None
301356
302- options = {
303- # "font_size": 36,
304- "node_size" : node_size ,
305- # "node_color": "white", # Set below
306- # "edgecolors": "black", # Set below
307- # "linewidths": 5,
308- # "width": 5,
309- "with_labels" : vertex_labels ,
310- }
311-
312- if vertex_labels :
313- options ["node_color" ] = "white"
314- options ["edgecolors" ] = "black"
357+ harmonize_parameters (draw_options )
315358
316359 if layout_fn :
317- nx .draw (G , pos = layout_fn (G ), ** options )
360+ nx .draw (G , pos = layout_fn (G ), ** draw_options )
318361 else :
319- nx .draw_shell (G , ** options )
362+ nx .draw_shell (G , ** draw_options )
320363 plt .show ()
321364 return None
0 commit comments