updating plotting command, adding example notebook

bin/test_plot.py

 import pykat
-pykat.init_pykat_plotting(mode="display", dpi=100)
+pykat.init_pykat_plotting()
 
 from pykat import finesse
 from pykat.detectors import *
@@ -14,7 +14,7 @@ l l1 1 0 0 n1
 s s1 10 1 n1 n2
 m m1 0.9 0.1 0 n2 n3
 s s2 10 1 n3 n4
-m m2 0.5 0.5 0 n4 n5
+m m2 0.91 0.09 0 n4 n5
 s s3 10 1 n5 n6
 
 yaxis abs:deg
@@ -24,8 +24,6 @@ ad circ 0 0 0 n4
 ad tran 0 0 0 n5
 pd pd_cav n3
 
-
-
 cav c1 m1 n3 m2 n4
 
 attr m1 Rc 1
@@ -44,4 +42,6 @@ kat.m2.Rcy =  1000.0
 kat.maxtem = 0
 
 out = kat.run()
-fig = out.plot("test_plot.pdf")
+fig = out.plot(yaxis="log abs:deg")
+
+#fig.savefig("test_plot.pdf")

pykat/finesse.py

@@ -240,7 +240,23 @@ class katRun(object):
         self.katVersion = None
         self.yaxis = None
         
-    def plot(self, filename=None, show=True):
+    def plot(self, detectors=None, filename=None, show=True, yaxis=None, legend=True, loc=0, title=None):
+        """
+        This will generate a plot for the output data of this particular pykat run.
+        It will attempt to generate a plot that shows all the various traces and plots
+        by default for quick viewing of the data. Similar to that which would be
+        generated by running the Finesse file from the command line.
+        
+        There are some additional keyword options to customise the plot output slightly:
+        
+            detectors: a list of detectors that you want to plot
+            filename: providing a filename here will save the plot to a file. The format is given by the extension provided.
+            show: True | False - whether to display the plot or not
+            yaxis: Set the Finesse yaxis command to base the plot on. By default the original one will be used.
+            legend: True | False - whether to include a legend
+            loc: Location value for the legend, the usual matplotlib one.
+            title: Provide a title for the plot if required.
+        """
         import matplotlib.pyplot as pyplot
         import pykat.plotting as plt
 
@@ -248,7 +264,8 @@ class katRun(object):
         kat.verbose = False
         kat.parseCommands(self.katScript)
 
-        plot_cmd = None
+        if yaxis is not None:
+            kat.yaxis = yaxis
 
         if "log" in kat.yaxis:
             if kat.xaxis.scale == "log":
@@ -265,45 +282,68 @@ class katRun(object):
         _func1 = np.abs
         _func2 = None
 
+        plot_cmd1 = None
+        plot_cmd2 = None
+        
         if "re:im" in kat.yaxis:
             _func1 = np.real
             _func2 = np.imag
+            plot_cmd1 = plot_cmd2 = plot_cmd
+            
             dual_plot = True
         elif "abs:deg" in kat.yaxis:
             _func1 = np.abs
             _func2 = lambda x: np.rad2deg(np.angle(x))
+
+            plot_cmd1 = plot_cmd
+            plot_cmd2 = pyplot.plot
+            
             dual_plot = True
         elif "db:deg" in kat.yaxis:
             _func1 = lambda x: 10*np.log10(x)
             _func2 = lambda x: np.rad2deg(np.angle(x))
+
+            plot_cmd1 = plot_cmd
+            plot_cmd2 = pyplot.plot
+            
             dual_plot = True
         elif "abs" in kat.yaxis:
             _func1 = np.abs
+            plot_cmd1 = plot_cmd
         elif "db" in kat.yaxis:
             _func1 = lambda x: 10*np.log10(x)
+            plot_cmd1 = plot_cmd
         elif "deg" in kat.yaxis:
             _func1 = lambda x: np.rad2deg(np.angle(x))
+            plot_cmd1 = plot_cmd
             
         if dual_plot:
             fig = plt.figure(width="full", height=1)
         else:
             fig = plt.figure(width="full")   
 
-        for lbl in self.ylabels:
-            lbl = lbl.split()[0]
+        if detectors is None:
+            detectors = [lbl.split()[0] for lbl in self.ylabels]
+            
+        detectors = list(set(detectors))
+        detectors.sort()
+        
+        for det in detectors:
+            if not getattr(kat, det).noplot:
                 if not dual_plot:
-                plot_cmd(self.x, _func1(self[lbl]))
+                    plot_cmd1(self.x, _func1(self[det]), label=det)
                 else:
                     pyplot.subplot(2,1,1)
-                l, = plot_cmd(self.x, _func1(self[lbl]))
+                    l, = plot_cmd1(self.x, _func1(self[det]), label=det)
         
                     pyplot.subplot(2,1,2)
-                pyplot.plot(self.x, _func2(self[lbl]), color=l.get_color(), ls=l.get_linestyle())
+                    plot_cmd2(self.x, _func2(self[det]), color=l.get_color(), ls=l.get_linestyle(), label=det)
 
         if dual_plot:
             pyplot.subplot(2,1,1)
             pyplot.xlabel(self.xlabel, fontsize=pyplot.rcParams["font.size"])
             pyplot.xlim(self.x.min(), self.x.max())
+            if title is not None: pyplot.title(title, fontsize=pyplot.rcParams["font.size"]-1)
             
             if "abs" in kat.yaxis: pyplot.ylabel("Absolute [au]", fontsize=pyplot.rcParams["font.size"])
             if "re" in kat.yaxis: pyplot.ylabel("Real part [au]", fontsize=pyplot.rcParams["font.size"])
@@ -318,8 +358,12 @@ class katRun(object):
             pyplot.xlabel(self.xlabel, fontsize=pyplot.rcParams["font.size"])
             pyplot.ylabel(" [au]")
             pyplot.xlim(self.x.min(), self.x.max())
+            if title is not None: pyplot.title(title, fontsize=pyplot.rcParams["font.size"]-1)
+    
+        pyplot.tight_layout()
     
-        fig.tight_layout()
+        if legend:
+            fig.axes[0].legend(loc=loc, fontsize=pyplot.rcParams["font.size"]-1)
         
         if filename is not None:
             fig.savefig(filename)

pykat/plotting.py

@@ -27,10 +27,13 @@ def init_pykat_plotting(mode="display", dpi=100):
     __DPI__ = int(dpi)
     
     if in_ipython():
+        try:
             from IPython.display import set_matplotlib_formats
             set_matplotlib_formats('pdf', 'svg')
             ipy = get_ipython()
             ipy.magic("matplotlib inline")
+        except:
+            pass
     
     if mode == "display":
         __mode__ = mode
@@ -93,3 +96,4 @@ def figure(width="full", height=0.618, textwidth=6, **kwargs):
     
     
     
+    
\ No newline at end of file