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

updating plotting command, adding example notebook

parent 1da582c5
import pykat import pykat
pykat.init_pykat_plotting(mode="display", dpi=100) pykat.init_pykat_plotting()
from pykat import finesse from pykat import finesse
from pykat.detectors import * from pykat.detectors import *
...@@ -14,7 +14,7 @@ l l1 1 0 0 n1 ...@@ -14,7 +14,7 @@ l l1 1 0 0 n1
s s1 10 1 n1 n2 s s1 10 1 n1 n2
m m1 0.9 0.1 0 n2 n3 m m1 0.9 0.1 0 n2 n3
s s2 10 1 n3 n4 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 s s3 10 1 n5 n6
yaxis abs:deg yaxis abs:deg
...@@ -24,8 +24,6 @@ ad circ 0 0 0 n4 ...@@ -24,8 +24,6 @@ ad circ 0 0 0 n4
ad tran 0 0 0 n5 ad tran 0 0 0 n5
pd pd_cav n3 pd pd_cav n3
cav c1 m1 n3 m2 n4 cav c1 m1 n3 m2 n4
attr m1 Rc 1 attr m1 Rc 1
...@@ -44,4 +42,6 @@ kat.m2.Rcy = 1000.0 ...@@ -44,4 +42,6 @@ kat.m2.Rcy = 1000.0
kat.maxtem = 0 kat.maxtem = 0
out = kat.run() 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): ...@@ -240,7 +240,23 @@ class katRun(object):
self.katVersion = None self.katVersion = None
self.yaxis = 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 matplotlib.pyplot as pyplot
import pykat.plotting as plt import pykat.plotting as plt
...@@ -248,7 +264,8 @@ class katRun(object): ...@@ -248,7 +264,8 @@ class katRun(object):
kat.verbose = False kat.verbose = False
kat.parseCommands(self.katScript) kat.parseCommands(self.katScript)
plot_cmd = None if yaxis is not None:
kat.yaxis = yaxis
if "log" in kat.yaxis: if "log" in kat.yaxis:
if kat.xaxis.scale == "log": if kat.xaxis.scale == "log":
...@@ -265,46 +282,69 @@ class katRun(object): ...@@ -265,46 +282,69 @@ class katRun(object):
_func1 = np.abs _func1 = np.abs
_func2 = None _func2 = None
plot_cmd1 = None
plot_cmd2 = None
if "re:im" in kat.yaxis: if "re:im" in kat.yaxis:
_func1 = np.real _func1 = np.real
_func2 = np.imag _func2 = np.imag
plot_cmd1 = plot_cmd2 = plot_cmd
dual_plot = True dual_plot = True
elif "abs:deg" in kat.yaxis: elif "abs:deg" in kat.yaxis:
_func1 = np.abs _func1 = np.abs
_func2 = lambda x: np.rad2deg(np.angle(x)) _func2 = lambda x: np.rad2deg(np.angle(x))
plot_cmd1 = plot_cmd
plot_cmd2 = pyplot.plot
dual_plot = True dual_plot = True
elif "db:deg" in kat.yaxis: elif "db:deg" in kat.yaxis:
_func1 = lambda x: 10*np.log10(x) _func1 = lambda x: 10*np.log10(x)
_func2 = lambda x: np.rad2deg(np.angle(x)) _func2 = lambda x: np.rad2deg(np.angle(x))
plot_cmd1 = plot_cmd
plot_cmd2 = pyplot.plot
dual_plot = True dual_plot = True
elif "abs" in kat.yaxis: elif "abs" in kat.yaxis:
_func1 = np.abs _func1 = np.abs
plot_cmd1 = plot_cmd
elif "db" in kat.yaxis: elif "db" in kat.yaxis:
_func1 = lambda x: 10*np.log10(x) _func1 = lambda x: 10*np.log10(x)
plot_cmd1 = plot_cmd
elif "deg" in kat.yaxis: elif "deg" in kat.yaxis:
_func1 = lambda x: np.rad2deg(np.angle(x)) _func1 = lambda x: np.rad2deg(np.angle(x))
plot_cmd1 = plot_cmd
if dual_plot: if dual_plot:
fig = plt.figure(width="full", height=1) fig = plt.figure(width="full", height=1)
else: else:
fig = plt.figure(width="full") fig = plt.figure(width="full")
for lbl in self.ylabels: if detectors is None:
lbl = lbl.split()[0] detectors = [lbl.split()[0] for lbl in self.ylabels]
if not dual_plot:
plot_cmd(self.x, _func1(self[lbl])) detectors = list(set(detectors))
else: detectors.sort()
pyplot.subplot(2,1,1)
l, = plot_cmd(self.x, _func1(self[lbl])) 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.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: if dual_plot:
pyplot.subplot(2,1,1) pyplot.subplot(2,1,1)
pyplot.xlabel(self.xlabel, fontsize=pyplot.rcParams["font.size"]) pyplot.xlabel(self.xlabel, fontsize=pyplot.rcParams["font.size"])
pyplot.xlim(self.x.min(), self.x.max()) 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 "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"]) if "re" in kat.yaxis: pyplot.ylabel("Real part [au]", fontsize=pyplot.rcParams["font.size"])
...@@ -318,8 +358,12 @@ class katRun(object): ...@@ -318,8 +358,12 @@ class katRun(object):
pyplot.xlabel(self.xlabel, fontsize=pyplot.rcParams["font.size"]) pyplot.xlabel(self.xlabel, fontsize=pyplot.rcParams["font.size"])
pyplot.ylabel(" [au]") pyplot.ylabel(" [au]")
pyplot.xlim(self.x.min(), self.x.max()) 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: if filename is not None:
fig.savefig(filename) fig.savefig(filename)
......
...@@ -27,10 +27,13 @@ def init_pykat_plotting(mode="display", dpi=100): ...@@ -27,10 +27,13 @@ def init_pykat_plotting(mode="display", dpi=100):
__DPI__ = int(dpi) __DPI__ = int(dpi)
if in_ipython(): if in_ipython():
from IPython.display import set_matplotlib_formats try:
set_matplotlib_formats('pdf', 'svg') from IPython.display import set_matplotlib_formats
ipy = get_ipython() set_matplotlib_formats('pdf', 'svg')
ipy.magic("matplotlib inline") ipy = get_ipython()
ipy.magic("matplotlib inline")
except:
pass
if mode == "display": if mode == "display":
__mode__ = mode __mode__ = mode
...@@ -92,4 +95,5 @@ def figure(width="full", height=0.618, textwidth=6, **kwargs): ...@@ -92,4 +95,5 @@ def figure(width="full", height=0.618, textwidth=6, **kwargs):
return fig return fig
\ No newline at end of file
Supports Markdown
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