init_mathjax(); Pig in the Python: 2012

Friday, September 28, 2012

Using the IPython Notebook and _repr_svg_ with networkx

If you use the excellent networkx package from working with graphs, here's a cool way to use the _repr_svg_ in the IPython Notebook or qtconsole.

In [2]:
import networkx as nx
import matplotlib.pyplot as plt
import StringIO
from matplotlib.figure import Figure

class MyGraph(nx.Graph):
  def _repr_svg_(self):
     plt.ioff() # turn off interactive mode
     fig=plt.figure(figsize=(2,2))
     ax = fig.add_subplot(111)
     nx.draw_shell(self, ax=ax)
     output = StringIO.StringIO()
     fig.savefig(output,format='svg')
     plt.ion() # turn on interactive mode
     return output.getvalue()

Basically, I'm creating a class to use as a harness for networkx graphs so that I don't have to keep plotting them to see what changes have been made to the graph as a result of various operations on the graph (e.g. removing nodes, adding edges). For example,here's a binomial graph:

In [3]:
g=nx.random_graphs.binomial_graph(10,.5)
f=MyGraph(g)
f
Out [3]:

Now, if I remove some edges as shown:

In [4]:
f.remove_nodes_from([0,3,7,5] )
f # this triggers the _repr_svg function in the notebook
Out [4]:

Pretty cool, huh? Now, you can interact with networkx graphs and get quick SVG's the represent those changes. Because I've defined a subclass, then I can do everything I normally do with network graphs.

I'm sure there's a neater way to do this and I'd love to hear it! You can download this notebook on my github site.

Addendum by Min RK in the comments

Min has come up with a better way to do this without subclassing. I'm including his work below to keep everything in one place.

Let's define a function to render a graph to SVG:

In [5]:
import matplotlib.pyplot as plt
from io import BytesIO

def graph_to_svg(g):
    """return the SVG of a matplotlib figure generated from a graph"""
    fig=plt.figure(figsize=(2,2))
    ax = fig.add_subplot(111)
    nx.draw_shell(g, ax=ax)
    output = BytesIO()
    fig.savefig(output, format='svg')
    plt.close(fig)
    return output.getvalue()

And call it:

In [6]:
from IPython.display import display, SVG
display(SVG(graph_to_svg(g)))

Now we can register graph_to_svg as the formatter to be called any time an nx.Graph object is seen:

In [7]:
formatter = get_ipython().display_formatter.formatters['image/svg+xml']
formatter.for_type(nx.Graph, graph_to_svg)

From now on, any time we would return a graph, we will see it as an SVG:

In [8]:
g
Out [8]:

Now, if I remove some edges as shown:

In [9]:
g.remove_nodes_from([0,3,7,5])
g # this triggers the graph_to_svg function in the notebook
Out [9]:

Pretty cool, huh? Now, you can interact with networkx graphs and get quick SVG's the represent those changes, without any subclasses.

Thanks for the input, Min! This is a much cleaner way to accomplish this rendering.