finesse.py 133 KB
Newer Older
Daniel Brown's avatar
Daniel Brown committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
# -*- coding: utf-8 -*-
"""
Created on Sun Jan 27 09:56:53 2013

PyKat - Python interface and wrapper for FINESSE
Copyright (C) 2013 Daniel David Brown

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

Contact at ddb@star.sr.bham.ac.uk

@author: Daniel Brown
"""
26 27 28 29 30
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

31
import traceback
Daniel Brown's avatar
Daniel Brown committed
32
import warnings
33
import codecs
34
import uuid
Daniel Brown's avatar
Daniel Brown committed
35 36 37 38 39 40
import sys
import os
import subprocess
import tempfile
import numpy as np
import datetime
41
import time
Daniel Brown's avatar
Daniel Brown committed
42
import pickle
43 44
import base64
import zlib
Daniel Brown's avatar
Daniel Brown committed
45 46
import pykat
import warnings
47
import re
48
import math
Daniel Brown's avatar
Daniel Brown committed
49
import itertools
Daniel Brown's avatar
Daniel Brown committed
50
import ctypes
Daniel Brown's avatar
Daniel Brown committed
51
import ctypes.util
Daniel Brown's avatar
Daniel Brown committed
52
import collections
53
import re
54
import copy
55

56 57
from subprocess import Popen, PIPE

58 59 60 61 62 63
try:
    # Python 2
    from itertools import izip_longest
except ImportError:
    # Python 3
    from itertools import zip_longest as izip_longest
Daniel Brown's avatar
Daniel Brown committed
64 65 66 67 68 69 70


try:
    # Add exception in Python 2
    FileNotFoundError
except NameError:
    FileNotFoundError = IOError
71

72 73

from math import erfc, pi
74
from collections import namedtuple, OrderedDict
75

Daniel Brown's avatar
Daniel Brown committed
76
from pykat.node_network import NodeNetwork
77
from pykat.detectors import BaseDetector as Detector
Daniel Brown's avatar
Daniel Brown committed
78
from pykat.components import Component
79
from pykat.commands import Command, xaxis, Constant
80
from pykat.SIfloat import *
81
from pykat.param import Param, AttrParam
82
from pykat.external import progressbar
83
from pykat.freeze import canFreeze
84

85
import pykat.external.six as six
Daniel Brown's avatar
Daniel Brown committed
86 87
import pykat.exceptions as pkex

88
from pykat import USE_GUI, HAS_OPTIVIS, NoGUIException, isContainer
89 90 91 92 93

if HAS_OPTIVIS:
    from optivis.bench.labels import Label as optivis_label
    from optivis.geometry import Coordinates as optivis_coord
    import PyQt4
94 95 96 97 98

if USE_GUI:
    from pykat.gui.gui import pyKatGUI
    from PyQt4.QtCore import QCoreApplication
    from PyQt4.QtGui import QApplication
Daniel Brown's avatar
Daniel Brown committed
99

Daniel Brown's avatar
Daniel Brown committed
100 101
from multiprocessing import Process, Manager

102
PYKAT_DATA = "#PYKAT_DATA="
Daniel Brown's avatar
Daniel Brown committed
103
NO_BLOCK = "NO_BLOCK"
104 105
pykat_web = "www.gwoptics.org/pykat"

Daniel Brown's avatar
Daniel Brown committed
106 107 108 109
# containers used in the trace routine
space_trace = namedtuple("space_trace", ['gouyx','gouyy'])
node_trace = namedtuple("node_trace", ['qx','qy'])
cav_trace = namedtuple("cav_trace", ['isStable','gx','gy','qx','qy','finesse','loss','length','FSR','FWHM','pole'])
110

Daniel Brown's avatar
Daniel Brown committed
111
lkat_location = ctypes.util.find_library("kat")
112 113


Daniel Brown's avatar
Daniel Brown committed
114 115 116
def f__lkat_process(callback, cmd, kwargs):
    """
    """
117

118
    if lkat_location is None:
Daniel Brown's avatar
Daniel Brown committed
119
        raise RuntimeError("Could not find shared library 'libkat', please install to a system location or copy to the same directory as this script")
120

Daniel Brown's avatar
Daniel Brown committed
121
    lkat = ctypes.PyDLL(lkat_location)
Daniel Brown's avatar
Daniel Brown committed
122 123 124 125 126 127 128 129 130

    try:
        lkat._pykat_preInit() # must always be called, sets up
                        # exception handling and such no simulation
                        # specifc code here

        # reads in the kat.ini and setups up other parts
        lkat._pykat_init()
        lkat._pykat_setup(cmd)
131

Daniel Brown's avatar
Daniel Brown committed
132
        callback(lkat, **kwargs)
133 134

    except Exception as ex:
135
        print ("Exception caught in python: ", ex.message)
Daniel Brown's avatar
Daniel Brown committed
136 137 138
    finally:
        # This should always be called no matter what
        lkat._pykat_finish(0)
Daniel Brown's avatar
Daniel Brown committed
139 140


141
def f__lkat_trace_callback(lkat, trace_info, getCavities, getNodes, getSpaces):
Daniel Brown's avatar
Daniel Brown committed
142 143 144 145 146 147 148 149 150 151 152 153
    """
    lkat callback for computing the beam traces through a setup.
    Returns a dictionary of nodes, spaces and cavities and the
    various outputs of the tracing algorithm.
    """
    import pylibkat

    # first we need to get a handle on the internals of Finesse
    inter = pylibkat.interferometer.in_dll(lkat, "inter")

    lkat._pykat_step()

154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180
    if getNodes:
        for n in range(0, inter.num_nodes):
            node = inter.node_list[n]

            node_info = node_trace(
                                    qx = complex(node.qx.re, node.qx.im),
                                    qy = complex(node.qy.re, node.qy.im)
                                    )

            trace_info[node.name] = node_info

    if getCavities:
        for c in range(0, inter.num_cavities):
            cav = inter.cavity_list[c]

            cav_info = cav_trace(
                                isStable = (cav.stable == 1),
                                gx = cav.stability_x,
                                gy = cav.stability_y,
                                qx = complex(cav.qx.re, cav.qx.im),
                                qy = complex(cav.qy.re, cav.qy.im),
                                finesse = cav.finesse,
                                FSR = cav.FSR,
                                FWHM = cav.FWHM,
                                loss = cav.loss,
                                length = cav.length,
                                pole = cav.pole
Daniel Brown's avatar
Daniel Brown committed
181 182
                                )

183
            trace_info[cav.name] = cav_info
Daniel Brown's avatar
Daniel Brown committed
184

185 186 187
    if getSpaces:
        for s in range(0, inter.num_spaces):
            space = inter.space_list[s]
Daniel Brown's avatar
Daniel Brown committed
188

189 190
            trace_info[space.name] = space_trace(gouyx = space.gouy_x,
                                                 gouyy = space.gouy_y)
191

Daniel Brown's avatar
Daniel Brown committed
192

193 194 195 196 197 198 199 200 201 202 203
@canFreeze
class GaussCommands(object):
    """
    Container class to store Gauss commands in
    """
    def __init__(self):
        self._freeze()

    def _add(self, cmd):
        if hasattr(self, cmd.name):
            raise pkex.BasePyKatException("Gauss command called %s already added" % cmd.name)
204

205 206 207
        self._unfreeze()
        setattr(self, cmd.name, self)
        self._freeze()
208

209 210 211
    def _remove(self, cmd):
        if hasattr(self, cmd.name):
            raise pkex.BasePyKatException("Gauss command called %s not found" % cmd.name)
212

213 214 215
        self._unfreeze()
        delattr(self, cmd.name)
        self._freeze()
Daniel Brown's avatar
Daniel Brown committed
216 217

class BlockedKatFile(object):
218
    """
Daniel Brown's avatar
Daniel Brown committed
219
    Allows manipulation of blocked kat file.
220

Daniel Brown's avatar
Daniel Brown committed
221 222
    Example:
        bkf = BlockedKatFile()
223

Daniel Brown's avatar
Daniel Brown committed
224 225 226 227
        bkf.read(katfile)
        bkf.add('tester', "blah\nblah", addAfter="Tunings")
        bkf.remove("Laser")
        bkf.write("mytest.kat")
228
    """
229

230 231
    def __str__(self):
         rtn = ""
232

233 234 235 236
         for block in self.ordering:
             rtn += "\n%%% FTblock " + block + "\n"
             rtn += self.blocks[block]
             rtn += "%%% FTend " + block + "\n"
237

238
         return rtn
239

Daniel Brown's avatar
Daniel Brown committed
240 241 242 243 244
    def __init__(self, NO_BLOCK="NO_BLOCK"):
        self.__NO_BLOCK = NO_BLOCK
        self.ordering = [self.__NO_BLOCK]
        self.blocks = {self.__NO_BLOCK:""}
        self.__currentBlock = self.__NO_BLOCK
245

246 247 248 249 250
    def remove(self, *blocks):
        if len(blocks[0]) > 1 and not isinstance(blocks[0], six.string_types):
            # if we've got an iterable thing that isn't a string, eg list or tuple
            # just use that
            blocks = blocks[0]
251

252 253 254 255 256 257
        for block in blocks:
            if block not in self.ordering or block not in self.blocks:
               raise Exception("%s block not found")

            self.ordering.remove(block)
            self.blocks.pop(block)
258

Daniel Brown's avatar
Daniel Brown committed
259
    def add(self, block, contents, addAfter=None):
Daniel Brown's avatar
Daniel Brown committed
260

Daniel Brown's avatar
Daniel Brown committed
261 262
        if block in self.ordering or block in self.blocks:
            raise Exception("%s block already present")
263

Daniel Brown's avatar
Daniel Brown committed
264 265 266 267
        if addAfter is not None:
            self.ordering.insert(self.ordering.index(addAfter)+1, block)
        else:
            self.ordering.append(block)
268

Daniel Brown's avatar
Daniel Brown committed
269
        self.blocks[block] = contents + "\n"
270

Daniel Brown's avatar
Daniel Brown committed
271 272 273 274 275 276
    def write(self, katfile):
        with open(katfile, "w") as f:
            for block in self.ordering:
                f.write("\n%%% FTblock " + block + "\n")
                f.write(self.blocks[block])
                f.write("%%% FTend " + block + "\n")
277

Daniel Brown's avatar
Daniel Brown committed
278 279 280 281
    def read(self, katfile):
        """
        For a given kat file, the blocks are parsed into dictionary as raw strings.
        """
282

Daniel Brown's avatar
Daniel Brown committed
283 284
        with open(katfile, "r") as f:
            commands = f.readlines()
285

Daniel Brown's avatar
Daniel Brown committed
286 287
        for line in commands:
            line = line.strip()
Daniel Brown's avatar
Daniel Brown committed
288

Daniel Brown's avatar
Daniel Brown committed
289 290
            # Looking for block start or end
            values = line.split()
291

Daniel Brown's avatar
Daniel Brown committed
292 293 294
            if len(values) >= 3 and values[0] == "%%%":
                if values[1] == "FTblock":
                    newTag = values[2]
295 296 297 298

                    if self.__currentBlock != None and self.__currentBlock != self.__NO_BLOCK:
                        warnings.warn("found block {0} before block {1} ended".format(newTag, self.__currentBlock))

Daniel Brown's avatar
Daniel Brown committed
299
                    if newTag in self.blocks:
300 301 302
                        #raise pkex.BasePyKatException("Block `{0}` has already been read".format(newTag))
                        self.__currentBlock = newTag
                        continue
303

Daniel Brown's avatar
Daniel Brown committed
304 305 306
                    self.blocks[newTag] = ""
                    self.__currentBlock = newTag
                    self.ordering.append(newTag)
307

Daniel Brown's avatar
Daniel Brown committed
308 309
                if values[1] == "FTend":
                    self.__currentBlock = self.__NO_BLOCK
310

Daniel Brown's avatar
Daniel Brown committed
311
                continue
312

Daniel Brown's avatar
Daniel Brown committed
313 314
            if(len(line) == 0 and (self.__currentBlock == self.__NO_BLOCK)):
                continue
315

Daniel Brown's avatar
Daniel Brown committed
316
            self.blocks[self.__currentBlock] += line + "\n"
317 318


Daniel Brown's avatar
Daniel Brown committed
319
class KatBatch(object):
320

321 322 323 324 325 326
    def __init__(self):
        from IPython.parallel import Client

        self._c = Client()
        self._lb = c.load_balanced_view()
        self.lb.block = False
327

328
        self._todo = []
329

330 331 332 333
    def _run(dir, commands, **kwargs):
        import pykat
        kat = pykat.finesse.kat()
        kat.verbose = False
334
        kat.parse(commands)
335

336
        kw = dict()
337

338 339
        if "cmd_args" in kwargs:
            kw["cmd_args"] = kwargs["cmd_args"]
340

341
        return kat.run(**kw)
342

343 344 345 346 347 348
    def addKat(self, kat, **kwargs):
        import os
        cdir = os.getcwd()
        script = "\n".join(kat.generateKatScript())
        self.todo.append(self.lb.apply_async(self._run, script, **kwargs))
        return self.todo[-1]
349

350 351
    def wait(self):
        return self.lb.wait(self.todo)
352

353 354 355
    def results(self):
        return self.todo

356

357 358 359 360 361
def GUILength(L):
    """
    Should scale the lengths in some way to handle km and mm for time being
    """
    return L # * ( 40 * erfc(L/400.0) + 0.01)
362 363

@canFreeze
364
class KatRun(object):
Daniel Brown's avatar
Daniel Brown committed
365
    def __init__(self):
366
        self._unfreeze()
367 368
        self.runtime = None
        self.StartDateTime = datetime.datetime.now()
Daniel Brown's avatar
Daniel Brown committed
369
        self.x = None
370 371 372
        self.stdout = None
        self.stderr = None
        self.runDateTime = None
Daniel Brown's avatar
Daniel Brown committed
373 374 375 376 377
        self.y = None
        self.xlabel = None
        self.ylabels = None
        self.katScript = None
        self.katVersion = None
378
        self.yaxis = None
379
        self._freeze()
380

381 382 383 384 385
    def __contains__(self, key):
        """
        The Finesse output file will output headers with re:im or abs:deg in the name.
        This confuses some python scripts when checking whether some output is inside
        this katrun or not.
386

387
        This tidies up the ylabels list and allows us to use:
388

389
            out = kat.run()
390

391 392
            if "output" in out:
                print("it is")
393

394 395
        """
        return key in set(_.split()[0] for _ in self.ylabels)
396

Daniel Brown's avatar
Daniel Brown committed
397
    def info(self):
398

Daniel Brown's avatar
Daniel Brown committed
399 400
        kat = pykat.finesse.kat()
        kat.verbose = False
401
        kat.parse(self.katScript)
402

Daniel Brown's avatar
Daniel Brown committed
403 404
        detectors = list(set([lbl.split()[0] for lbl in self.ylabels]))
        detectors.sort()
405

406
        print("")
Daniel Brown's avatar
Daniel Brown committed
407
        print("--- Output info ---")
408
        print("")
Daniel Brown's avatar
Daniel Brown committed
409 410
        print("Run date and time: %s" % self.StartDateTime)
        print("Detectors used: %s" % (", ".join(detectors)))
411 412
        print("")

Daniel Brown's avatar
Daniel Brown committed
413 414 415 416
        if kat.noxaxis:
            print("No xaxis used")
        else:
            print("One xaxis used: %s" % kat.xaxis.getFinesseText())
417

418 419 420 421
        import numpy as np

        maxs = np.max(self.y, 0)
        mins = np.min(self.y, 0)
422 423 424

        maxlbl = max([len(lbl) for lbl in self.ylabels])

425 426 427
        for i, lbl in enumerate(self.ylabels):
            a = "{0:" + str(maxlbl) + "} : min = {1:.15e} max = {2:.15e}"
            print(a.format(lbl, mins[i], maxs[i]))
428 429


Daniel Brown's avatar
Daniel Brown committed
430 431
    def plot(self, detectors=None, filename=None, show=True,
                   yaxis=None, legend=True, loc=0, title=None, styles=None,
432
                   ylabel=None, y2label=None, xlabel=None, x2label=None,
433 434
                   xlim=None, x2lim=None, ylim=None, y2lim=None, return_fig=False,
                   xscale=None, y1scale=None, y2scale=None):
435 436 437 438 439
        """
        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.
440

441
        There are some additional keyword options to customise the plot output slightly:
442

Daniel Brown's avatar
Daniel Brown committed
443 444 445 446 447 448 449 450 451 452 453 454
            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.
            styles:             A dictionary which keys being the detector names and the
                                value being a colour and linestyle of the sort 'k:'
            ylabel, xlabel:     Text for the first plot x and y labels
455 456 457 458 459 460
            y2label, x2label:   Text for the second plot x and y labels

            xlim, ylim:         Limits of x- and y-axes of the first plot. List or tuple
                                of length 2.
            x2lim, y2lim:       Limits of x- and y-axes of the second plot. List or tuple
                                of length 2.
461 462
            xscale, y1scale     SI scale either axes. See pykat.SI dictionary for possible
            y2scale:            values
463
        """
464 465
        import matplotlib.pyplot as pyplot
        import pykat.plotting as plt
466

Daniel Brown's avatar
Daniel Brown committed
467 468 469
        if not show:
            pyplot.ioff()

470 471
        kat = pykat.finesse.kat()
        kat.verbose = False
472
        kat.parse(self.katScript)
473

474 475 476
        if kat.noxaxis == True:
            raise  pkex.BasePyKatException("This kat object has noxaxis=True, so there is nothing to plot.")

477 478
        original_yaxis = kat.yaxis

479 480
        if yaxis is not None:
            kat.yaxis = yaxis
481 482

        dual_plot = False
Daniel Brown's avatar
Daniel Brown committed
483 484
        _func1 = np.abs
        _func2 = None
485

486 487
        plot_cmd1 = None
        plot_cmd2 = None
488

489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515
        if "re:im" in kat.yaxis or "abs:deg" in kat.yaxis or "db:deg" in kat.yaxis:
            dual_plot = True

        if dual_plot:
            fig = plt.figure(width="full", height=1)
            ax1 = fig.add_subplot(2, 1, 1)
            ax2 = fig.add_subplot(2, 1, 2)
        else:
            fig = plt.figure(width="full")
            ax1 = fig.add_subplot(1, 1, 1)
            ax2 = ax1

        if "log" in kat.yaxis:
            if kat.xaxis.scale == "log":
                plot_cmd1 = ax1.loglog
                plot_cmd2 = ax2.loglog
            else:
                plot_cmd1 = ax1.semilogy
                plot_cmd2 = ax2.semilogy
        else:
            if kat.xaxis.scale == "log":
                plot_cmd1 = ax1.semilogx
                plot_cmd2 = ax2.semilogx
            else:
                plot_cmd1 = ax1.plot
                plot_cmd2 = ax2.plot

Daniel Brown's avatar
Daniel Brown committed
516 517 518
        if "re:im" in kat.yaxis:
            _func1 = np.real
            _func2 = np.imag
519

Daniel Brown's avatar
Daniel Brown committed
520 521 522
        elif "abs:deg" in kat.yaxis:
            _func1 = np.abs
            _func2 = lambda x: np.rad2deg(np.angle(x))
523

524
            plot_cmd2 = ax2.plot if kat.xaxis.scale == "lin" else ax2.semilogx
525

Daniel Brown's avatar
Daniel Brown committed
526
        elif "db:deg" in kat.yaxis:
527 528 529 530
            if "db" not in original_yaxis:
                _func1 = lambda x: 10*np.log10(x)
            else:
                _func1 = lambda x: x
531

Daniel Brown's avatar
Daniel Brown committed
532
            _func2 = lambda x: np.rad2deg(np.angle(x))
533

534
            plot_cmd2 = ax2.plot if kat.xaxis.scale == "lin" else ax2.semilogx
535

Daniel Brown's avatar
Daniel Brown committed
536
        elif "abs" in kat.yaxis:
537 538
            # _func1 = np.abs
            _func1 = np.real
Daniel Brown's avatar
Daniel Brown committed
539
        elif "db" in kat.yaxis:
540 541 542 543
            if "db" not in original_yaxis:
                _func1 = lambda x: 10*np.log10(x)
            else:
                _func1 = lambda x: x
544

Daniel Brown's avatar
Daniel Brown committed
545 546
        elif "deg" in kat.yaxis:
            _func1 = lambda x: np.rad2deg(np.angle(x))
547

548 549
        if detectors is None:
            detectors = [lbl.split()[0] for lbl in self.ylabels]
550

551

552 553
        detectors = list(set(detectors))
        detectors.sort()
554

555 556
        _x = self.x/pykat.SI[xscale]
        
557
        for det in detectors:
Daniel Brown's avatar
Daniel Brown committed
558
            if not hasattr(kat, det) or (hasattr(kat, det) and not getattr(kat, det).noplot):
Daniel Brown's avatar
Daniel Brown committed
559
                if styles is not None and det in styles:
560
                    l, = plot_cmd1(_x, _func1(self[det])/pykat.SI[y1scale], styles[det], label=det)
Daniel Brown's avatar
Daniel Brown committed
561
                else:
562
                    l, = plot_cmd1(_x, _func1(self[det])/pykat.SI[y1scale], label=det)
563 564

                if dual_plot:
565
                    plot_cmd2(_x, _func2(self[det])/pykat.SI[y2scale], color=l.get_color(), ls=l.get_linestyle(), label=det)
566

Daniel Brown's avatar
Daniel Brown committed
567 568 569 570
        if dual_plot:
            if ylabel is None:
                if "abs" in kat.yaxis: ylabel = "Absolute [au]"
                if "re" in kat.yaxis:  ylabel = "Real part [au]"
571

Daniel Brown's avatar
Daniel Brown committed
572 573 574
            if y2label is None:
                if "deg" in kat.yaxis: y2label = "Phase [deg]"
                if "im" in kat.yaxis:  y2label = "Imaginary part [au]"
575

576
            if xlim is None:
577
                xlim = (_x.min(), _x.max())
578

579
            if x2lim is None:
580
                 x2lim = (_x.min(), _x.max())
581

Daniel Brown's avatar
Daniel Brown committed
582 583 584
        else:
            if ylabel is None:
                ylabel = "[au]"
585

586
            if xlim is None:
587
                xlim = (_x.min(), _x.max())
588 589


Daniel Brown's avatar
Daniel Brown committed
590 591
        if xlabel is None:
            xlabel = self.xlabel
592

Daniel Brown's avatar
Daniel Brown committed
593 594
        if x2label is None:
            x2label = self.xlabel
595

Daniel Brown's avatar
Daniel Brown committed
596
        font_label_size = pyplot.rcParams["font.size"]-1
597

598
        if dual_plot:
599 600 601 602
            #ax = plt.subplot(2,1,1)
            ax1.set_xlabel(xlabel, fontsize=font_label_size)
            ax1.set_ylabel(ylabel, fontsize=font_label_size)
            ax1.set_xlim(xlim[0], xlim[1])
603
            if ylim is not None:
604
                ax1.set_ylim(ylim[0],ylim[1])
605

Daniel Brown's avatar
Daniel Brown committed
606
            if title is not None:
607
                ax1.set_title(title, fontsize=font_label_size)
608

609 610 611
            #plt.subplot(2,1,2)
            ax2.set_xlabel(x2label, fontsize=font_label_size)
            ax2.set_ylabel(y2label, fontsize=font_label_size)
612

613
            ax2.set_xlim(x2lim[0], x2lim[1])
614
            if y2lim is not None:
615
                ax2.set_ylim(y2lim[0],y2lim[1])
616

617
        else:
618 619 620
            ax1.set_xlabel(xlabel, fontsize=font_label_size)
            ax1.set_ylabel(ylabel)
            ax1.set_xlim(xlim[0], xlim[1])
621

Daniel Brown's avatar
Daniel Brown committed
622
            if title is not None:
623
                ax1.set_title(title, fontsize=font_label_size)
624
            if ylim is not None:
625
                ax1.set_ylim(ylim[0],ylim[1])
626

627
        pyplot.margins(0, 0.05)
628
        pyplot.tight_layout()
629

630
        if legend:
Daniel Brown's avatar
Daniel Brown committed
631
            fig.axes[0].legend(loc=loc, fontsize=font_label_size)
632

633 634
        if filename is not None:
            fig.savefig(filename)
635

636 637
        if show:
            pyplot.show(fig)
Daniel Brown's avatar
Daniel Brown committed
638
            pyplot.ion()
639

640 641
        if return_fig:
            return fig
642

643
    def saveKatRun(self, filename):
Daniel Brown's avatar
Daniel Brown committed
644 645
        with open(filename,'w') as outfile:
            pickle.dump(self, outfile)
646

Daniel Brown's avatar
Daniel Brown committed
647 648 649 650
    @staticmethod
    def loadKatRun(filename):
        with open(filename,'r') as infile:
            return pickle.load(infile)
651

Daniel Brown's avatar
Daniel Brown committed
652
    def get(self, value): return self[value]
653

Daniel Brown's avatar
Daniel Brown committed
654
    def __getitem__(self, value):
655 656 657 658
        if isContainer(value):
            results = []
            for _ in value:
                results.append(self[_])
659

660 661 662 663
            return np.array(results).squeeze()
        else:
            idx = [i for i in range(len(self.ylabels)) if self.ylabels[i].split()[0] == str(value)]
            out = None
664

665 666
            if len(idx) > 0:
                #out = self.y[:, idx]
667

668 669 670 671 672
                if len(idx) == 1:
                    if "abs:deg" in self.yaxis:
                        out = self.y[:, idx[0]]
                    elif "re:im" in self.yaxis:
                        out = self.y[:, idx[0]]
673
                else:
674 675 676 677 678 679 680 681 682 683 684 685
                    if "abs:deg" in self.yaxis:
                        out = self.y[:, idx[0]] * np.exp(1j*math.pi*self.y[:, idx[1]]/180.0)
                    elif "re:im" in self.yaxis :
                        out = self.y[:, idx[0]] + 1j*self.y[:, idx[1]]

                if out is None:
                    out = self.y[:, idx]

                if out.size == 1:
                    return out[0].squeeze()
                else:
                    return out.squeeze()
686
            else:
687
                raise  pkex.BasePyKatException("No output by the name '{0}' found in the output".format(str(value)))
688 689

@canFreeze
690
class KatRun2D(object):
Daniel Brown's avatar
Daniel Brown committed
691
    def __init__(self):
692
        self._unfreeze()
693
        self.runtime = None
694
        self.runDateTime = None
695
        self.startDateTime = datetime.datetime.now()
Daniel Brown's avatar
Daniel Brown committed
696 697 698
        self.x = None
        self.y = None
        self.z = None
699
        self.yaxis = None
Daniel Brown's avatar
Daniel Brown committed
700 701 702 703 704
        self.xlabel = None
        self.ylabel = None
        self.zlabels = None
        self.katScript = None
        self.katVersion = None
705 706
        self.stderr = None
        self.stdout = None
707
        self._freeze()
708

Daniel Brown's avatar
Daniel Brown committed
709 710 711
    def saveKatRun(self, filename):
        with open(filename,'w') as outfile:
            pickle.dump(self, outfile)
712

Daniel Brown's avatar
Daniel Brown committed
713 714 715 716
    @staticmethod
    def loadKatRun(filename):
        with open(filename,'r') as infile:
            return pickle.load(infile)
717

718 719 720 721 722 723 724 725 726
    # def get(self, value): return self[value].squeeze()
    #
    # def __getitem__(self, value):
    #     idx = [i for i in range(len(self.zlabels)) if self.zlabels[i].split()[0] == str(value)]
    #
    #     if len(idx) > 0:
    #         return self.z[idx].squeeze()
    #     else:
    #         raise  pkex.BasePyKatException("No output by the name {0} found".format(str(value)))
727

728
    def get(self, value): return self[value]
729

Daniel Brown's avatar
Daniel Brown committed
730
    def __getitem__(self, value):
731 732 733 734
        if isContainer(value):
            results = []
            for _ in value:
                results.append(self[_])
735

736
            return np.array(results).squeeze()
Daniel Brown's avatar
Daniel Brown committed
737
        else:
738 739
            idx = [i for i in range(len(self.zlabels)) if self.zlabels[i].split()[0] == str(value)]
            out = None
740

741 742
            if len(idx) > 0:
                #out = self.y[:, idx]
743

744 745 746 747 748
                if len(idx) == 1:
                    if "abs:deg" in self.yaxis:
                        out = self.z[idx[0], :, :]
                    elif "re:im" in self.yaxis:
                        out = self.z[idx[0], :, :]
749
                else:
750 751 752 753 754 755
                    if "abs:deg" in self.yaxis:
                        out = self.z[idx[0], :, :] * np.exp(1j*math.pi*self.z[idx[1], :, :]/180.0)
                    elif "re:im" in self.yaxis :
                        out = self.z[idx[0], :, :] + 1j*self.z[idx[1], :, :]

                if out is None:
756
                    out = self.z[idx]
757 758 759 760 761 762 763

                if out.size == 1:
                    return out[0].squeeze()
                else:
                    return out.squeeze()
            else:
                raise  pkex.BasePyKatException("No output by the name '{0}' found in the output".format(str(value)))
764 765

@canFreeze
766
class Signals(object):
767 768

    @canFreeze
769
    class _fsig(object):
Daniel Brown's avatar
Daniel Brown committed
770
        def __init__(self, param, name, amplitude, phase, signal):
771
            self._unfreeze()
772 773 774
            self._params = []
            self.__target = param
            self.__name = name
775
            self.__removed = False
Daniel Brown's avatar
Daniel Brown committed
776
            self.__signal = signal
777 778 779

            self.__amplitude = Param("amp", self, SIfloat(amplitude))
            self.__phase = Param("phase", self, SIfloat(phase))
780

781
            self._freeze()
782
            # unfortunatenly the target names for fsig are not the same as the
783
            # various parameter names of the components, e.g. mirror xbeta is x
784 785 786
            # for fsig. So we need to check here what type of component we are targetting
            # and then based on the parameter specfied get the name
            if not param.canFsig:
787
                raise  pkex.BasePyKatException("Cannot fsig parameter {1} on component {0}".format(str(param._owner().name), param.name))
788

789 790
        def __deepcopy__(self, memo):
            # Here we need to update the params with new owners
Daniel Brown's avatar
Daniel Brown committed
791
            result = self.__class__.__new__(self.__class__)
792 793
            memo[id(self)] = result
            result.__dict__ = copy.deepcopy(self.__dict__, memo)
794

795 796
            for _ in result._params:
                _._updateOwner(result)
797

798
            return result
799

800 801
        def _register_param(self, param):
            self._params.append(param)
802

803
        @property
804 805
        def removed(self):
            return self.__removed
806

807
        def remove(self):
808
            self.__signal._kat.remove(self)
809

810
        def _on_remove(self):
811 812 813
            if self.__removed:
                raise pkex.BasePyKatException("Signal {0} has already been marked as removed".format(self.name))
            else:
Daniel Brown's avatar
Daniel Brown committed
814
                self.__signal.targets.remove(self)
Daniel Brown's avatar
Daniel Brown committed
815
                self.__removed = True
816

817 818 819 820 821
        @property
        def name(self): return self.__name

        @property
        def amplitude(self): return self.__amplitude
822 823 824
        @amplitude.setter
        def amplitude(self,value): self.__amplitude.value = SIfloat(value)

825 826 827

        @property
        def phase(self): return self.__phase
828 829
        @phase.setter
        def phase(self,value): self.__phase.value = SIfloat(value)
830 831 832 833 834

        @property
        def target(self): return self.__target.fsig_name

        @property
835
        def owner(self): return self.__target._owner().name
836 837


838 839
        def getFinesseText(self):
            rtn = []
840

841 842
            for p in self._params:
                rtn.extend(p.getFinesseText())
843

844
            return rtn
845

846 847 848 849 850 851 852 853
    def __init__(self, kat):
        self._unfreeze()
        self._default_name = "fsignal"
        self.targets = []
        self._params = []
        self.__f = Param("f", self, None)
        self._kat = kat
        self._freeze()
854

855 856 857 858 859
    def __deepcopy__(self, memo):
        # Here we need to update the params with new owners
        result = self.__class__.__new__(self.__class__)
        memo[id(self)] = result
        result.__dict__ = copy.deepcopy(self.__dict__, memo)
860

861 862
        for _ in result._params:
            <