added sensible size defaults to Chimera, Pegasus and Zephyr plots

boothby · boothby · commit 11a2defc64a3 · 2021-12-10T10:45:05.000-08:00
diff --git a/dwave_networkx/drawing/chimera_layout.py b/dwave_networkx/drawing/chimera_layout.py
@@ -20,14 +20,14 @@
 import networkx as nx
 from networkx import draw
 
-from dwave_networkx.drawing.qubit_layout import draw_qubit_graph, draw_embedding, draw_yield
+from dwave_networkx.drawing.qubit_layout import draw_qubit_graph, draw_embedding, draw_yield, normalize_size_and_aspect
 from dwave_networkx.generators.chimera import chimera_graph, find_chimera_indices, chimera_coordinates
 
 
 __all__ = ['chimera_layout', 'draw_chimera', 'draw_chimera_embedding', 'draw_chimera_yield']
 
 
-def chimera_layout(G, scale=1., center=None, dim=2):
+def chimera_layout(G, scale=1., center=None, dim=2, normalize_kwargs = None):
     """Positions the nodes of graph G in a Chimera cross topology.
 
     NumPy (https://scipy.org) is required for this function.
@@ -51,6 +51,11 @@ def chimera_layout(G, scale=1., center=None, dim=2):
         Number of dimensions. When dim > 2, all extra dimensions are
         set to 0.
 
+    normalize_kwargs : None or dict (default None)
+        A dict of keyword arguments to be used in a plotting function.  If not
+        None, we will populate the "ax" keyword with matplotlib axes, and the
+        "node_size" and "width" keywords with defaults if they are not set.
+
     Returns
     -------
     pos : dict
@@ -75,9 +80,9 @@ def chimera_layout(G, scale=1., center=None, dim=2):
         n = G.graph['columns']
         t = G.graph['tile']
         # get a node placement function
-        xy_coords = chimera_node_placer_2d(m, n, t, scale, center, dim)
+        xy_coords = chimera_node_placer_2d(m, n, t, scale, center, dim, normalize_kwargs = normalize_kwargs)
 
-        if G.graph.get('labels') == 'coordinate':
+        if G.graph['labels'] == 'coordinate':
             pos = {v: xy_coords(*v) for v in G.nodes()}
         elif G.graph.get('data'):
             pos = {v: xy_coords(*dat['chimera_index']) for v, dat in G.nodes(data=True)}
@@ -97,15 +102,15 @@ def chimera_layout(G, scale=1., center=None, dim=2):
         m = max(idx[0] for idx in chimera_indices.values()) + 1
         n = max(idx[1] for idx in chimera_indices.values()) + 1
         t = max(idx[3] for idx in chimera_indices.values()) + 1
-        xy_coords = chimera_node_placer_2d(m, n, t, scale, center, dim)
+        xy_coords = chimera_node_placer_2d(m, n, t, scale, center, dim, normalize_kwargs = normalize_kwargs)
 
         # compute our coordinates
         pos = {v: xy_coords(i, j, u, k) for v, (i, j, u, k) in chimera_indices.items()}
 
     return pos
 
 
-def chimera_node_placer_2d(m, n, t, scale=1., center=None, dim=2):
+def chimera_node_placer_2d(m, n, t, scale=1., center=None, dim=2, normalize_kwargs = None):
     """Generates a function that converts Chimera indices to x, y
     coordinates for a plot.
 
@@ -130,6 +135,11 @@ def chimera_node_placer_2d(m, n, t, scale=1., center=None, dim=2):
     dim : int (default 2)
         Number of dimensions. When dim > 2, all extra dimensions are
         set to 0.
+        
+    normalize_kwargs : None or dict (default None)
+        A dict of keyword arguments to be used in a plotting function.  If not
+        None, we will populate the "ax" keyword with matplotlib axes, and the
+        "node_size" and "width" keywords with defaults if they are not set.
 
     Returns
     -------
@@ -144,8 +154,13 @@ def chimera_node_placer_2d(m, n, t, scale=1., center=None, dim=2):
     center_pad = 1
     tile_center = t // 2
     tile_length = t + 2 + center_pad  # 2 for spacing between tiles
+
     # want the enter plot to fill in [0, 1] when scale=1
-    scale /= max(m, n) * tile_length - 2 - center_pad
+    fabric_scale = max(m, n) * tile_length - 2 - center_pad
+    scale /= fabric_scale
+
+    if normalize_kwargs is not None:
+        normalize_size_and_aspect(fabric_scale, 200, normalize_kwargs)
 
     grid_offsets = {}
 
@@ -228,8 +243,8 @@ def draw_chimera(G, **kwargs):
     >>> plt.show()  # doctest: +SKIP
 
     """
-    draw_qubit_graph(G, chimera_layout(G), **kwargs)
-
+    layout = chimera_layout(G, normalize_kwargs = kwargs)
+    draw_qubit_graph(G, layout, **kwargs)
 
 def draw_chimera_embedding(G, *args, **kwargs):
     """Draws an embedding onto the chimera graph G, according to layout.
@@ -282,7 +297,8 @@ def draw_chimera_embedding(G, *args, **kwargs):
        function. If `linear_biases` or `quadratic_biases` are provided,
        any provided `node_color` or `edge_color` arguments are ignored.
     """
-    draw_embedding(G, chimera_layout(G), *args, **kwargs)
+    layout = chimera_layout(G, normalize_kwargs = kwargs)
+    draw_embedding(G, layout, *args, **kwargs)
 
 
 def draw_chimera_yield(G, **kwargs):
@@ -325,5 +341,6 @@ def draw_chimera_yield(G, **kwargs):
         tile, and label attributes to be able to identify faulty qubits.")
 
     perfect_graph = chimera_graph(m,n,t, coordinates=coordinates)
+    layout = chimera_layout(perfect_graph, normalize_kwargs = kwargs)
+    draw_yield(G, layout, perfect_graph, **kwargs)
 
-    draw_yield(G, chimera_layout(perfect_graph), perfect_graph, **kwargs)
diff --git a/dwave_networkx/drawing/pegasus_layout.py b/dwave_networkx/drawing/pegasus_layout.py
@@ -19,7 +19,7 @@
 import networkx as nx
 from networkx import draw
 
-from dwave_networkx.drawing.qubit_layout import draw_qubit_graph, draw_embedding, draw_yield
+from dwave_networkx.drawing.qubit_layout import draw_qubit_graph, draw_embedding, draw_yield, normalize_size_and_aspect
 from dwave_networkx.generators.pegasus import pegasus_graph, pegasus_coordinates
 from dwave_networkx.drawing.chimera_layout import chimera_node_placer_2d
 
@@ -31,7 +31,7 @@
            ]
 
 
-def pegasus_layout(G, scale=1., center=None, dim=2, crosses=False):
+def pegasus_layout(G, scale=1., center=None, dim=2, crosses=False, normalize_kwargs=None):
     """Positions the nodes of graph G in a Pegasus topology.
 
     `NumPy <https://scipy.org>`_ is required for this function.
@@ -58,6 +58,11 @@ def pegasus_layout(G, scale=1., center=None, dim=2, crosses=False):
         rather than L configuration. Ignored if G is defined with
         ``nice_coordinates=True``.
 
+    normalize_kwargs : None or dict (default None)
+        A dict of keyword arguments to be used in a plotting function.  If not
+        None, we will populate the "ax" keyword with matplotlib axes, and the
+        "node_size" and "width" keywords with defaults if they are not set.
+
     Returns
     -------
     pos : dict
@@ -75,28 +80,27 @@ def pegasus_layout(G, scale=1., center=None, dim=2, crosses=False):
     if not isinstance(G, nx.Graph) or G.graph.get("family") != "pegasus":
         raise ValueError("G must be generated by dwave_networkx.pegasus_graph")
 
-    if G.graph.get('labels') == 'nice':
-        m = 3*(G.graph['rows']-1)
-        c_coords = chimera_node_placer_2d(m, m, 4, scale=scale, center=center, dim=dim)
-        def xy_coords(t, y, x, u, k):
-            return c_coords(3*y+2-t, 3*x+t, u, k)
+    labels = G.graph.get('labels')
+    xy_coords = pegasus_node_placer_2d(G, scale, center, dim, crosses=(crosses or labels == 'nice_coordinates'), normalize_kwargs = normalize_kwargs)
+
+    if labels == 'coordinate':
         pos = {v: xy_coords(*v) for v in G.nodes()}
+    elif G.graph['data']:
+        pos = {v: xy_coords(*dat['pegasus_index']) for v, dat in G.nodes(data=True)}
     else:
-        xy_coords = pegasus_node_placer_2d(G, scale, center, dim, crosses=crosses)
-
-        if G.graph.get('labels') == 'coordinate':
-            pos = {v: xy_coords(*v) for v in G.nodes()}
-        elif G.graph.get('data'):
-            pos = {v: xy_coords(*dat['pegasus_index']) for v, dat in G.nodes(data=True)}
-        else:
-            m = G.graph.get('rows')
-            coord = pegasus_coordinates(m)
+        m = G.graph['rows']
+        coord = pegasus_coordinates(m)
+        if labels == 'int':
             pos = {v: xy_coords(*coord.linear_to_pegasus(v)) for v in G.nodes()}
+        elif labels == 'nice':
+            pos = {v: xy_coords(*coord.nice_to_pegasus(v)) for v in G.nodes()}
+        else:
+            raise ValueError(f"Pegasus labeling {labels} not recognized")
 
     return pos
 
 
-def pegasus_node_placer_2d(G, scale=1., center=None, dim=2, crosses=False):
+def pegasus_node_placer_2d(G, scale=1., center=None, dim=2, crosses=False, normalize_kwargs=None):
     """Generates a function to convert Pegasus indices to plottable coordinates.
 
     Parameters
@@ -138,7 +142,11 @@ def pegasus_node_placer_2d(G, scale=1., center=None, dim=2, crosses=False):
     cross_shift = 2 if crosses else 0
 
     # want the enter plot to fill in [0, 1] when scale=1
-    scale /= m*tile_width - 2*odd_k_wobble - 1
+    fabric_scale = m*tile_width - 2*odd_k_wobble - 1
+    scale /= fabric_scale
+
+    if normalize_kwargs is not None:
+        normalize_size_and_aspect(fabric_scale, 150, normalize_kwargs)
 
     if center is None:
         center = np.zeros(dim)
@@ -219,11 +227,10 @@ def draw_pegasus(G, crosses=False, **kwargs):
     >>> plt.show()    # doctest: +SKIP
 
     """
+    layout = pegasus_layout(G, normalize_kwargs = kwargs)
+    draw_qubit_graph(G, layout, **kwargs)
 
-    draw_qubit_graph(G, pegasus_layout(G, crosses=crosses), **kwargs)
-
-
-def draw_pegasus_embedding(G, *args, **kwargs):
+def draw_pegasus_embedding(G, *args, crosses=False, **kwargs):
     """Draws an embedding onto Pegasus graph G.
 
     Parameters
@@ -276,10 +283,10 @@ def draw_pegasus_embedding(G, *args, **kwargs):
        function. If ``linear_biases`` or ``quadratic_biases`` are provided,
        any provided ``node_color`` or ``edge_color`` arguments are ignored.
     """
-    crosses = kwargs.pop("crosses", False)
-    draw_embedding(G, pegasus_layout(G, crosses=crosses), *args, **kwargs)
+    layout = pegasus_layout(G, crosses=crosses, normalize_kwargs=kwargs)
+    draw_embedding(G, layout, *args, **kwargs)
 
-def draw_pegasus_yield(G, **kwargs):
+def draw_pegasus_yield(G, crosses=False, **kwargs):
     """Draws the given graph G with highlighted faults, according to layout.
 
     Parameters
@@ -302,6 +309,11 @@ def draw_pegasus_yield(G, **kwargs):
     fault_style : string, optional (default='dashed')
         Edge fault line style (solid|dashed|dotted|dashdot)
 
+    crosses: boolean (optional, default False)
+        If True, :math:`K_{4,4}` subgraphs are shown in a cross
+        rather than L configuration. Ignored if G is defined with
+        ``nice_coordinates=True``.
+
     kwargs : optional keywords
        See networkx.draw_networkx() for a description of optional keywords,
        with the exception of the `pos` parameter which is not used by this
@@ -321,5 +333,5 @@ def draw_pegasus_yield(G, **kwargs):
 
 
     perfect_graph = pegasus_graph(m, offset_lists=offset_lists, coordinates=coordinates, nice_coordinates=nice)
-
-    draw_yield(G, pegasus_layout(perfect_graph), perfect_graph, **kwargs)
+    layout = pegasus_layout(perfect_graph, crosses=crosses, normalize_kwargs=kwargs)
+    draw_yield(G, layout, perfect_graph, **kwargs)
diff --git a/dwave_networkx/drawing/qubit_layout.py b/dwave_networkx/drawing/qubit_layout.py
@@ -499,3 +499,37 @@ def draw_yield(G, layout, perfect_graph, unused_color=(0.9,0.9,0.9,1.0),
         draw(perfect_graph, layout, nodelist=nodelist, edgelist=edgelist,
             node_color=unused_node_color, edge_color=unused_edge_color,
             **kwargs)
+
+def normalize_size_and_aspect(scale, node_scale, kwargs):
+    ax = kwargs.get('ax')
+    if ax is None:
+        try:
+            import matplotlib.pyplot as plt
+        except ImportError:
+            raise ImportError("Matplotlib required for graph drawing")
+        cf = plt.gcf()
+    else:
+        cf = ax.get_figure()
+    cf.set_facecolor("w")
+    if ax is None:
+        if cf._axstack() is None:
+            ax = cf.add_axes((0, 0, 1, 1))
+        else:
+            ax = cf.gca()
+    kwargs['ax'] = ax
+    ax.set_aspect(1)
+    fig_scale = min(cf.get_figheight(), cf.get_figwidth())
+    ax.set_axis_off()
+
+    if kwargs.get('node_size') is None:
+        if kwargs.get("line_plot"):
+            kwargs['node_size'] = 50*(fig_scale/scale)**2
+        else:
+            kwargs['node_size'] = node_scale*(fig_scale/scale)**2
+
+    if kwargs.get('width') is None:
+        if kwargs.get("line_plot"):
+            kwargs['width'] = 5*(fig_scale / scale)
+        else:
+            kwargs['width'] = 2*(fig_scale / scale)
+
diff --git a/dwave_networkx/drawing/zephyr_layout.py b/dwave_networkx/drawing/zephyr_layout.py
@@ -19,7 +19,7 @@
 import networkx as nx
 from networkx import draw
 
-from dwave_networkx.drawing.qubit_layout import draw_qubit_graph, draw_embedding, draw_yield
+from dwave_networkx.drawing.qubit_layout import draw_qubit_graph, draw_embedding, draw_yield, normalize_size_and_aspect
 from dwave_networkx.generators.zephyr import zephyr_graph, zephyr_coordinates
 
 
@@ -30,7 +30,7 @@
            ]
 
 
-def zephyr_layout(G, scale=1., center=None, dim=2):
+def zephyr_layout(G, scale=1., center=None, dim=2, normalize_kwargs=None):
     """Positions the nodes of graph G in a Zephyr topology.
 
     `NumPy <https://scipy.org>`_ is required for this function.
@@ -52,6 +52,11 @@ def zephyr_layout(G, scale=1., center=None, dim=2):
         Number of dimensions. When dim > 2, all extra dimensions are
         set to 0.
 
+    normalize_kwargs : None or dict (default None)
+        A dict of keyword arguments to be used in a plotting function.  If not
+        None, we will populate the "ax" keyword with matplotlib axes, and the
+        "node_size" and "width" keywords with defaults if they are not set.
+
     Returns
     -------
     pos : dict
@@ -69,22 +74,22 @@ def zephyr_layout(G, scale=1., center=None, dim=2):
     if not isinstance(G, nx.Graph) or G.graph.get("family") != "zephyr":
         raise ValueError("G must be generated by dwave_networkx.zephyr_graph")
 
-    xy_coords = zephyr_node_placer_2d(G, scale, center, dim)
+    xy_coords = zephyr_node_placer_2d(G, scale, center, dim, normalize_kwargs)
 
-    if G.graph.get('labels') == 'coordinate':
+    if G.graph['labels'] == 'coordinate':
         pos = {v: xy_coords(*v) for v in G.nodes()}
-    elif G.graph.get('data'):
+    elif G.graph['data']:
         pos = {v: xy_coords(*dat['zephyr_index']) for v, dat in G.nodes(data=True)}
     else:
-        m = G.graph.get('rows')
-        t = G.graph.get('tile')
+        m = G.graph['rows']
+        t = G.graph['tile']
         coord = zephyr_coordinates(m, t)
         pos = {v: xy_coords(*coord.linear_to_zephyr(v)) for v in G.nodes()}
 
     return pos
 
 
-def zephyr_node_placer_2d(G, scale=1., center=None, dim=2):
+def zephyr_node_placer_2d(G, scale=1., center=None, dim=2, normalize_kwargs = None):
     """Generates a function to convert Zephyr indices to plottable coordinates.
 
     Parameters
@@ -104,6 +109,11 @@ def zephyr_node_placer_2d(G, scale=1., center=None, dim=2):
         Number of dimensions. When dim > 2, all extra dimensions are
         set to 0.
 
+    normalize_kwargs : None or dict (default None)
+        A dict of keyword arguments to be used in a plotting function.  If not
+        None, we will populate the "ax" keyword with matplotlib axes, and the
+        "node_size" and "width" keywords with defaults if they are not set.
+
     Returns
     -------
     xy_coords : function
@@ -141,7 +151,12 @@ def _xy_coords(u, w, k, j, z):
         return np.hstack((xy * scale, paddims)) + center
 
     # want the enter plot to fill in [0, 1] when scale=1
-    scale /= max(map(abs, _xy_coords(0, 2*m, G.graph["tile"]-1, 1, m-1)))
+    fabric_scale = max(map(abs, _xy_coords(0, 2*m, G.graph["tile"]-1, 1, m-1)))
+    scale /= fabric_scale
+
+    if normalize_kwargs is not None:
+        normalize_size_and_aspect(fabric_scale, 200, normalize_kwargs)
+
     return _xy_coords
 
 def draw_zephyr(G, **kwargs):
