Merge pull request #597 from viktorashi/fix/save_diag_to_file

added savefig option to .draw()

tellurium/tests/test_tellurium.py

@@ -13,6 +13,7 @@
 import numpy as np
 import matplotlib
 import antimony
+from matplotlib.pyplot import savefig
 
 CELLML_SUPPORT = hasattr(antimony, "loadCellMLString")
 
@@ -356,6 +357,20 @@ def test_draw(self):
         except ImportError:
             pass
 
+    def test_draw2file(self):
+        r = te.loada("""
+            S1 -> S2; k1*S1;
+            k1 = 0.1; S1 = 40; S2 = 0.0;
+        """)
+        try:
+            import pygraphviz
+            f, temp_file_path = tempfile.mkstemp(suffix='.png')
+            r.draw(savefig=temp_file_path)
+            os.close(f)
+            os.remove(temp_file_path)
+        except ImportError:
+            pass
+
     def test_plot(self):
         """ Regression tests for plotting.
         The following calls should work. """

tellurium/visualization/sbmldiagram.py

@@ -134,19 +134,25 @@ def _createGraph(model, species={}, reactions={}, reactants={}, products={}, mod
                 g.add_edge(s.getSpecies(), r.getId(), **modifiers)
         return g
 
-    def draw(self, layout='neato', **kwargs):
+    def draw(self, layout='neato', savefig=None, **kwargs):
         """ Draw the graph.
         Optional layout=['neato'|'dot'|'twopi'|'circo'|'fdp'|'nop']
         will use specified graphviz layout method.
 
         :param layout: pygraphviz layout algorithm (default: 'neato')
+        :param savefig: .png file in which to save the figure (displays to the front-end without saving by default)
         :type layout: str
         """
-        f, filePath = tempfile.mkstemp(suffix='.png')
-        self.g.layout(prog=layout)
-        self.g.draw(filePath)
-
-        i = Image(filename=filePath)
-        display(i)
-        os.close(f)
-        os.remove(filePath)
+        if not savefig:
+            f, temp_file_path = tempfile.mkstemp(suffix='.png')
+            self.g.layout(prog=layout)
+            self.g.draw(temp_file_path)
+
+            i = Image(filename=temp_file_path)
+            display(i)
+            os.close(f)
+            os.remove(temp_file_path)
+        else:
+            self.g.layout(prog=layout)
+            self.g.draw(savefig)
+