Commit 36f3b074 authored by Daniel Brown's avatar Daniel Brown

updating plotting command, adding example notebook

parent 1da582c5
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")
This source diff could not be displayed because it is too large. You can view the blob instead.
This source diff could not be displayed because it is too large. You can view the blob instead.
......@@ -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,46 +282,69 @@ 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 not dual_plot:
plot_cmd(self.x, _func1(self[lbl]))
else:
pyplot.subplot(2,1,1)
l, = plot_cmd(self.x, _func1(self[lbl]))
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_cmd1(self.x, _func1(self[det]), label=det)
else:
pyplot.subplot(2,1,1)
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())
pyplot.subplot(2,1,2)
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)
......
......@@ -27,10 +27,13 @@ def init_pykat_plotting(mode="display", dpi=100):
__DPI__ = int(dpi)
if in_ipython():
from IPython.display import set_matplotlib_formats
set_matplotlib_formats('pdf', 'svg')
ipy = get_ipython()
ipy.magic("matplotlib inline")
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
......@@ -92,4 +95,5 @@ def figure(width="full", height=0.618, textwidth=6, **kwargs):
return fig
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment