finesse.py 129 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 48
import re
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 71


try:
    # Add exception in Python 2
    FileNotFoundError
except NameError:
    FileNotFoundError = IOError
    
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 110 111
# 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'])
         
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):
    """
    """
Daniel Brown's avatar
Daniel Brown committed
117
    
118
    if lkat_location is None:
Daniel Brown's avatar
Daniel Brown committed
119 120 121
        raise RuntimeError("Could not find shared library 'libkat', please install to a system location or copy to the same directory as this script")
        
    lkat = ctypes.PyDLL(lkat_location)
Daniel Brown's avatar
Daniel Brown committed
122 123 124 125 126 127 128 129 130 131 132 133 134

    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)
    
        callback(lkat, **kwargs)
    
    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)
Daniel Brown's avatar
Daniel Brown committed
191 192
                                                 

Daniel Brown's avatar
Daniel Brown committed
193 194

class BlockedKatFile(object):
195
    """
Daniel Brown's avatar
Daniel Brown committed
196 197 198 199 200 201 202 203 204
    Allows manipulation of blocked kat file.
    
    Example:
        bkf = BlockedKatFile()
        
        bkf.read(katfile)
        bkf.add('tester', "blah\nblah", addAfter="Tunings")
        bkf.remove("Laser")
        bkf.write("mytest.kat")
205
    """
Daniel Brown's avatar
Daniel Brown committed
206
    
207 208 209 210 211 212 213 214 215 216
    def __str__(self):
         rtn = ""
         
         for block in self.ordering:
             rtn += "\n%%% FTblock " + block + "\n"
             rtn += self.blocks[block]
             rtn += "%%% FTend " + block + "\n"
         
         return rtn
         
Daniel Brown's avatar
Daniel Brown committed
217 218 219 220 221
    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
222
             
223 224 225 226 227
    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]
Daniel Brown's avatar
Daniel Brown committed
228
        
229 230 231 232 233 234 235
        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)
            
Daniel Brown's avatar
Daniel Brown committed
236
    def add(self, block, contents, addAfter=None):
Daniel Brown's avatar
Daniel Brown committed
237

Daniel Brown's avatar
Daniel Brown committed
238 239
        if block in self.ordering or block in self.blocks:
            raise Exception("%s block already present")
Daniel Brown's avatar
Daniel Brown committed
240
    
Daniel Brown's avatar
Daniel Brown committed
241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264
        if addAfter is not None:
            self.ordering.insert(self.ordering.index(addAfter)+1, block)
        else:
            self.ordering.append(block)
            
        self.blocks[block] = contents + "\n"
        
    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")
    
    def read(self, katfile):
        """
        For a given kat file, the blocks are parsed into dictionary as raw strings.
        """
    
        with open(katfile, "r") as f:
            commands = f.readlines()
    
        for line in commands:
            line = line.strip()
Daniel Brown's avatar
Daniel Brown committed
265

Daniel Brown's avatar
Daniel Brown committed
266 267
            # Looking for block start or end
            values = line.split()
Daniel Brown's avatar
Daniel Brown committed
268
    
Daniel Brown's avatar
Daniel Brown committed
269 270 271
            if len(values) >= 3 and values[0] == "%%%":
                if values[1] == "FTblock":
                    newTag = values[2]
272
                    
Daniel Brown's avatar
Daniel Brown committed
273 274
                    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
275
    
Daniel Brown's avatar
Daniel Brown committed
276
                    if newTag in self.blocks:
277 278 279
                        #raise pkex.BasePyKatException("Block `{0}` has already been read".format(newTag))
                        self.__currentBlock = newTag
                        continue
Daniel Brown's avatar
Daniel Brown committed
280
    
Daniel Brown's avatar
Daniel Brown committed
281 282 283 284 285 286
                    self.blocks[newTag] = ""
                    self.__currentBlock = newTag
                    self.ordering.append(newTag)
                
                if values[1] == "FTend":
                    self.__currentBlock = self.__NO_BLOCK
Daniel Brown's avatar
Daniel Brown committed
287
        
Daniel Brown's avatar
Daniel Brown committed
288 289 290 291
                continue
                
            if(len(line) == 0 and (self.__currentBlock == self.__NO_BLOCK)):
                continue
Daniel Brown's avatar
Daniel Brown committed
292
        
Daniel Brown's avatar
Daniel Brown committed
293
            self.blocks[self.__currentBlock] += line + "\n"
Daniel Brown's avatar
Daniel Brown committed
294 295 296
    
    
class KatBatch(object):
297 298 299 300 301 302 303 304 305 306 307 308 309 310
    
    def __init__(self):
        from IPython.parallel import Client

        self._c = Client()
        self._lb = c.load_balanced_view()
        self.lb.block = False
        
        self._todo = []
    
    def _run(dir, commands, **kwargs):
        import pykat
        kat = pykat.finesse.kat()
        kat.verbose = False
311
        kat.parse(commands)
312 313 314 315 316 317
        
        kw = dict()
        
        if "cmd_args" in kwargs:
            kw["cmd_args"] = kwargs["cmd_args"]
        
318
        return kat.run(**kw)
319 320 321 322 323 324 325 326 327 328 329 330 331 332 333
    
    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]
    
    def wait(self):
        return self.lb.wait(self.todo)
        
    def results(self):
        return self.todo

                                  
334 335 336 337 338
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)
339 340

@canFreeze
341
class KatRun(object):
Daniel Brown's avatar
Daniel Brown committed
342
    def __init__(self):
343
        self._unfreeze()
344 345
        self.runtime = None
        self.StartDateTime = datetime.datetime.now()
Daniel Brown's avatar
Daniel Brown committed
346
        self.x = None
347 348 349
        self.stdout = None
        self.stderr = None
        self.runDateTime = None
Daniel Brown's avatar
Daniel Brown committed
350 351 352 353 354
        self.y = None
        self.xlabel = None
        self.ylabels = None
        self.katScript = None
        self.katVersion = None
355
        self.yaxis = None
356
        self._freeze()
357
        
Daniel Brown's avatar
Daniel Brown committed
358 359 360 361
    def info(self):
        
        kat = pykat.finesse.kat()
        kat.verbose = False
362
        kat.parse(self.katScript)
Daniel Brown's avatar
Daniel Brown committed
363 364 365 366
        
        detectors = list(set([lbl.split()[0] for lbl in self.ylabels]))
        detectors.sort()
        
367
        print("")
Daniel Brown's avatar
Daniel Brown committed
368
        print("--- Output info ---")
369
        print("")
Daniel Brown's avatar
Daniel Brown committed
370 371
        print("Run date and time: %s" % self.StartDateTime)
        print("Detectors used: %s" % (", ".join(detectors)))
372 373
        print("")

Daniel Brown's avatar
Daniel Brown committed
374 375 376 377
        if kat.noxaxis:
            print("No xaxis used")
        else:
            print("One xaxis used: %s" % kat.xaxis.getFinesseText())
378 379 380 381 382 383 384 385 386 387 388
            
        import numpy as np

        maxs = np.max(self.y, 0)
        mins = np.min(self.y, 0)
        
        maxlbl = max([len(lbl) for lbl in self.ylabels])    
        
        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]))
Daniel Brown's avatar
Daniel Brown committed
389 390
        
        
Daniel Brown's avatar
Daniel Brown committed
391 392
    def plot(self, detectors=None, filename=None, show=True,
                   yaxis=None, legend=True, loc=0, title=None, styles=None,
393
                   ylabel=None, y2label=None, xlabel=None, x2label=None,
394
                   xlim=None, x2lim=None, ylim=None, y2lim=None, return_fig=False):
395 396 397 398 399 400 401 402
        """
        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:
        
Daniel Brown's avatar
Daniel Brown committed
403 404 405 406 407 408 409 410 411 412 413 414
            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
415 416 417 418 419 420
            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.
421
        """
422 423
        import matplotlib.pyplot as pyplot
        import pykat.plotting as plt
424
            
Daniel Brown's avatar
Daniel Brown committed
425 426 427
        if not show:
            pyplot.ioff()

428 429
        kat = pykat.finesse.kat()
        kat.verbose = False
430
        kat.parse(self.katScript)
431

432 433 434
        if kat.noxaxis == True:
            raise  pkex.BasePyKatException("This kat object has noxaxis=True, so there is nothing to plot.")

435 436
        original_yaxis = kat.yaxis

437 438
        if yaxis is not None:
            kat.yaxis = yaxis
439 440 441 442 443 444 445 446 447 448 449 450 451

        if "log" in kat.yaxis:
            if kat.xaxis.scale == "log":
                plot_cmd = pyplot.loglog
            else:
                plot_cmd = pyplot.semilogy
        else:
            if kat.xaxis.scale == "log":
                plot_cmd = pyplot.semilogx
            else:
                plot_cmd = pyplot.plot

        dual_plot = False
Daniel Brown's avatar
Daniel Brown committed
452 453
        _func1 = np.abs
        _func2 = None
454

455 456 457
        plot_cmd1 = None
        plot_cmd2 = None
        
Daniel Brown's avatar
Daniel Brown committed
458 459 460
        if "re:im" in kat.yaxis:
            _func1 = np.real
            _func2 = np.imag
461 462
            plot_cmd1 = plot_cmd2 = plot_cmd
            
Daniel Brown's avatar
Daniel Brown committed
463 464 465 466
            dual_plot = True
        elif "abs:deg" in kat.yaxis:
            _func1 = np.abs
            _func2 = lambda x: np.rad2deg(np.angle(x))
467 468

            plot_cmd1 = plot_cmd
Daniel Brown's avatar
Daniel Brown committed
469
            plot_cmd2 = pyplot.plot if kat.xaxis.scale == "lin" else pyplot.semilogx
470
            
Daniel Brown's avatar
Daniel Brown committed
471 472
            dual_plot = True
        elif "db:deg" in kat.yaxis:
473 474 475 476 477
            if "db" not in original_yaxis:
                _func1 = lambda x: 10*np.log10(x)
            else:
                _func1 = lambda x: x
                
Daniel Brown's avatar
Daniel Brown committed
478
            _func2 = lambda x: np.rad2deg(np.angle(x))
479 480

            plot_cmd1 = plot_cmd
Daniel Brown's avatar
Daniel Brown committed
481
            plot_cmd2 = pyplot.plot if kat.xaxis.scale == "lin" else pyplot.semilogx
482
            
483
            dual_plot = True
Daniel Brown's avatar
Daniel Brown committed
484
        elif "abs" in kat.yaxis:
485 486
            # _func1 = np.abs
            _func1 = np.real
487
            plot_cmd1 = plot_cmd
Daniel Brown's avatar
Daniel Brown committed
488
        elif "db" in kat.yaxis:
489 490 491 492 493
            if "db" not in original_yaxis:
                _func1 = lambda x: 10*np.log10(x)
            else:
                _func1 = lambda x: x
                
494
            plot_cmd1 = plot_cmd
Daniel Brown's avatar
Daniel Brown committed
495 496
        elif "deg" in kat.yaxis:
            _func1 = lambda x: np.rad2deg(np.angle(x))
497
            plot_cmd1 = plot_cmd
Daniel Brown's avatar
Daniel Brown committed
498 499 500
            
        if dual_plot:
            fig = plt.figure(width="full", height=1)
501
        else:
Daniel Brown's avatar
Daniel Brown committed
502
            fig = plt.figure(width="full")   
503

504 505 506 507 508 509 510
        if detectors is None:
            detectors = [lbl.split()[0] for lbl in self.ylabels]
            
        detectors = list(set(detectors))
        detectors.sort()
        
        for det in detectors:
Daniel Brown's avatar
Daniel Brown committed
511
            if not hasattr(kat, det) or (hasattr(kat, det) and not getattr(kat, det).noplot):
512
                
Daniel Brown's avatar
Daniel Brown committed
513
                if dual_plot:
Daniel Brown's avatar
Daniel Brown committed
514
                    ax = pyplot.subplot(2,1,1)
Daniel Brown's avatar
Daniel Brown committed
515 516 517 518
                    
                if styles is not None and det in styles:
                    l, = plot_cmd1(self.x, _func1(self[det]), styles[det], label=det)
                else:
519
                    l, = plot_cmd1(self.x, _func1(self[det]), label=det)
Daniel Brown's avatar
Daniel Brown committed
520 521
                
                if dual_plot: 
522 523
                    pyplot.subplot(2,1,2)
                    plot_cmd2(self.x, _func2(self[det]), color=l.get_color(), ls=l.get_linestyle(), label=det)
524

Daniel Brown's avatar
Daniel Brown committed
525 526 527 528
        if dual_plot:
            if ylabel is None:
                if "abs" in kat.yaxis: ylabel = "Absolute [au]"
                if "re" in kat.yaxis:  ylabel = "Real part [au]"
529
                    
Daniel Brown's avatar
Daniel Brown committed
530 531 532
            if y2label is None:
                if "deg" in kat.yaxis: y2label = "Phase [deg]"
                if "im" in kat.yaxis:  y2label = "Imaginary part [au]"
533 534 535 536 537 538 539
                    
            if xlim is None:
                xlim = (self.x.min(), self.x.max())
                
            if x2lim is None:
                 x2lim = (self.x.min(), self.x.max())
                 
Daniel Brown's avatar
Daniel Brown committed
540 541 542
        else:
            if ylabel is None:
                ylabel = "[au]"
543 544 545 546
                
            if xlim is None:
                xlim = (self.x.min(), self.x.max())
            
Daniel Brown's avatar
Daniel Brown committed
547
        
Daniel Brown's avatar
Daniel Brown committed
548 549 550 551 552 553
        if xlabel is None:
            xlabel = self.xlabel
            
        if x2label is None:
            x2label = self.xlabel
            
Daniel Brown's avatar
Daniel Brown committed
554 555
        font_label_size = pyplot.rcParams["font.size"]-1
        
556
        if dual_plot:
Daniel Brown's avatar
Daniel Brown committed
557
            ax = pyplot.subplot(2,1,1)
Daniel Brown's avatar
Daniel Brown committed
558 559
            pyplot.xlabel(xlabel, fontsize=font_label_size)
            pyplot.ylabel(ylabel, fontsize=font_label_size)
560 561 562 563
            pyplot.xlim(xlim[0], xlim[1])
            if ylim is not None:
                pyplot.ylim(ylim[0],ylim[1])

Daniel Brown's avatar
Daniel Brown committed
564 565
            if title is not None:
                pyplot.title(title, fontsize=font_label_size)
566 567
    
            pyplot.subplot(2,1,2)
Daniel Brown's avatar
Daniel Brown committed
568 569
            pyplot.xlabel(x2label, fontsize=font_label_size)
            pyplot.ylabel(y2label, fontsize=font_label_size)
570 571 572 573 574
        
            pyplot.xlim(x2lim[0], x2lim[1])
            if y2lim is not None:
                pyplot.ylim(y2lim[0],y2lim[1])
            
575
        else:
Daniel Brown's avatar
Daniel Brown committed
576 577
            pyplot.xlabel(xlabel, fontsize=font_label_size)
            pyplot.ylabel(ylabel)
578
            pyplot.xlim(self.x.min(), self.x.max())
Daniel Brown's avatar
Daniel Brown committed
579 580 581
            
            if title is not None:
                pyplot.title(title, fontsize=font_label_size)
582 583
            if ylim is not None:
                pyplot.ylim(ylim[0],ylim[1])
584
    
585
        pyplot.margins(0, 0.05)
586
        pyplot.tight_layout()
587
    
588
        if legend:
Daniel Brown's avatar
Daniel Brown committed
589
            fig.axes[0].legend(loc=loc, fontsize=font_label_size)
590 591 592 593 594 595
        
        if filename is not None:
            fig.savefig(filename)
            
        if show:
            pyplot.show(fig)
Daniel Brown's avatar
Daniel Brown committed
596
            pyplot.ion()
597
        
598 599
        if return_fig:
            return fig
Daniel Brown's avatar
Daniel Brown committed
600
        
601
    def saveKatRun(self, filename):
Daniel Brown's avatar
Daniel Brown committed
602 603 604 605 606 607 608 609 610 611 612
        with open(filename,'w') as outfile:
            pickle.dump(self, outfile)
    
    @staticmethod
    def loadKatRun(filename):
        with open(filename,'r') as infile:
            return pickle.load(infile)
    
    def get(self, value): return self[value]
    
    def __getitem__(self, value):
613 614 615 616 617 618 619 620 621
        if isContainer(value):
            results = []
            for _ in value:
                results.append(self[_])
                
            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
622
        
623 624
            if len(idx) > 0:
                #out = self.y[:, idx]
625
            
626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643
                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]]
                else: 
                    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()
644
            else:
645
                raise  pkex.BasePyKatException("No output by the name '{0}' found in the output".format(str(value)))
646 647
            
@canFreeze 
648
class KatRun2D(object):
Daniel Brown's avatar
Daniel Brown committed
649
    def __init__(self):
650
        self._unfreeze()
651
        self.runtime = None
652
        self.runDateTime = None
653
        self.startDateTime = datetime.datetime.now()
Daniel Brown's avatar
Daniel Brown committed
654 655 656
        self.x = None
        self.y = None
        self.z = None
657
        self.yaxis = None
Daniel Brown's avatar
Daniel Brown committed
658 659 660 661 662
        self.xlabel = None
        self.ylabel = None
        self.zlabels = None
        self.katScript = None
        self.katVersion = None
663 664
        self.stderr = None
        self.stdout = None
665
        self._freeze()
Daniel Brown's avatar
Daniel Brown committed
666 667 668 669 670 671 672 673 674 675
        
    def saveKatRun(self, filename):
        with open(filename,'w') as outfile:
            pickle.dump(self, outfile)
    
    @staticmethod
    def loadKatRun(filename):
        with open(filename,'r') as infile:
            return pickle.load(infile)
    
676 677 678 679 680 681 682 683 684 685 686
    # 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)))
    
    def get(self, value): return self[value]
Daniel Brown's avatar
Daniel Brown committed
687 688
    
    def __getitem__(self, value):
689 690 691 692 693 694
        if isContainer(value):
            results = []
            for _ in value:
                results.append(self[_])
                
            return np.array(results).squeeze()
Daniel Brown's avatar
Daniel Brown committed
695
        else:
696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713
            idx = [i for i in range(len(self.zlabels)) if self.zlabels[i].split()[0] == str(value)]
            out = None
            
            if len(idx) > 0:
                #out = self.y[:, idx]
            
                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], :, :]
                else: 
                    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:
714
                    out = self.z[idx]
715 716 717 718 719 720 721 722

                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)))
                
723
@canFreeze    
724
class Signals(object):
725 726
    
    @canFreeze 
727
    class fsig(object):
Daniel Brown's avatar
Daniel Brown committed
728
        def __init__(self, param, name, amplitude, phase, signal):
729
            self._unfreeze()
730 731 732
            self._params = []
            self.__target = param
            self.__name = name
733
            self.__removed = False
Daniel Brown's avatar
Daniel Brown committed
734
            self.__signal = signal
735 736 737

            self.__amplitude = Param("amp", self, SIfloat(amplitude))
            self.__phase = Param("phase", self, SIfloat(phase))
738
            
739
            self._freeze()
740
            # unfortunatenly the target names for fsig are not the same as the
Daniel Brown's avatar
Daniel Brown committed
741
            # various parameter names of the components, e.g. mirror xbeta is x 
742 743 744
            # 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:
745
                raise  pkex.BasePyKatException("Cannot fsig parameter {1} on component {0}".format(str(param._owner().name), param.name))
746 747 748
        
        def __deepcopy__(self, memo):
            # Here we need to update the params with new owners
Daniel Brown's avatar
Daniel Brown committed
749
            result = self.__class__.__new__(self.__class__)
750 751 752 753 754 755 756
            memo[id(self)] = result
            result.__dict__ = copy.deepcopy(self.__dict__, memo)
        
            for _ in result._params:
                _._updateOwner(result)
            
            return result
757 758 759
            
        def _register_param(self, param):
            self._params.append(param)
760 761
        
        @property
762 763
        def removed(self):
            return self.__removed
764
  
765
        def remove(self):
766 767 768
            self.__signal._kat.remove(self)
            
        def _on_remove(self):
769 770 771
            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
772
                self.__signal.targets.remove(self)
Daniel Brown's avatar
Daniel Brown committed
773
                self.__removed = True
774
        
775 776 777 778 779
        @property
        def name(self): return self.__name

        @property
        def amplitude(self): return self.__amplitude
780 781 782
        @amplitude.setter
        def amplitude(self,value): self.__amplitude.value = SIfloat(value)

783 784 785

        @property
        def phase(self): return self.__phase
786 787
        @phase.setter
        def phase(self,value): self.__phase.value = SIfloat(value)
788 789 790 791 792

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

        @property
793
        def owner(self): return self.__target._owner().name
794
    
795
        
796 797 798 799 800 801 802
        def getFinesseText(self):
            rtn = []
    
            for p in self._params:
                rtn.extend(p.getFinesseText())
        
            return rtn
803 804 805 806 807 808 809 810 811
            
    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()
812
    
813 814 815 816 817
    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)
818
    
819 820 821 822 823
        for _ in result._params:
            _._updateOwner(result)
        
        return result
        
824 825 826 827 828 829 830
    @property
    def name(self):
        # if we don't have any signals yet then use a dummy name
        # however we need to always tune a real fsig command
        # so need to get the name of at least one of them
        # as if you tune one you tune them all
        if len(self.targets) == 0:
831
            return self._default_name
832 833
        else:
            return self.targets[0].name
834 835
            
    @property
Daniel Brown's avatar
Daniel Brown committed
836
    def removed(self): return False # we can never remove the Signal object altogethr just the individual fsig targets
837

838 839
    def remove(self):
        for t in self.targets:
Daniel Brown's avatar
Daniel Brown committed
840
            t.remove()
841 842 843
        
        del self.targets[:]
        
Daniel Brown's avatar
Daniel Brown committed
844 845
        self.f = None
        
846 847
    @property
    def f(self): return self.__f
848
    
849
    @f.setter
850
    def f(self,value):
Daniel Brown's avatar
Daniel Brown committed
851 852 853 854
        if value is None:
            self.__f.value = None
            return
            
855 856
        v = SIfloat(value)
        
Daniel Brown's avatar
Daniel Brown committed
857
        if v <= 0:
858 859
            raise pkex.BasePyKatException("Signal frequency must be greater than 0.")
            
Daniel Brown's avatar
Daniel Brown committed
860
        self.__f.value = SIfloat(value)
Daniel Brown's avatar
Daniel Brown committed
861
        
862 863 864 865
    def _register_param(self, param):
        self._params.append(param)
        
    def apply(self, target, amplitude, phase, name=None):
866
        if target is None:
867 868
            raise  pkex.BasePyKatException("No target was specified for signal to be applied")
        
869
        if name is None:
870
            name = "sig_" + target._owner().name + "_" + target.name
871
        
Daniel Brown's avatar
Daniel Brown committed
872
        self.targets.append(Signals.fsig(target, name, amplitude, phase, self))
873 874 875 876
        
    def getFinesseText(self):
        rtn = []
        
877 878 879 880 881 882 883 884
        if self.f.value is not None and self.f is not None:
            if len(self.targets) == 0:
                rtn.append("fsig {name} {frequency}"
                                .format(name = self.name,
                                        frequency=str(self.f.value)))
            else:
                for t in self.targets:
                    rtn.extend(t.getFinesseText())
Daniel Brown's avatar
Daniel Brown committed
885
            
886 887 888 889 890 891 892
                    rtn.append("fsig {name} {comp} {target} {frequency} {phase} {amplitude}"
                                    .format(name = t.name,
                                            comp=t.owner,
                                            target=t.target,
                                            frequency=str(self.f.value),
                                            phase=str(t.phase),
                                            amplitude=str(t.amplitude if t.amplitude != None else "")))
Daniel Brown's avatar
Daniel Brown committed
893

894 895 896 897
        for p in self._params:
            rtn.extend(p.getFinesseText())
        
        return rtn
898

Daniel Brown's avatar
Daniel Brown committed
899
class Block:
900 901 902
    class Placeholder(object):
        pass
        
Daniel Brown's avatar
Daniel Brown committed
903 904 905 906 907 908 909 910
    def __init__(self, name):
        self.__name = name
        self.contents = [] # List of objects and strings of finesse code
        self.enabled = True 
        
    @property
    def name(self): return self.__name
    
911 912 913 914 915 916 917 918 919 920 921 922 923 924 925
    def __str__(self):
        objs = []
        
        for _ in self.contents:
            if isinstance(_, six.string_types):
                objs.append(_)
            else:
                __ = _.getFinesseText()

                if isinstance(__, six.string_types):
                    objs.append(__)
                else:
                    objs += __
        
        return "\n".join(objs)
926 927 928

id___ = 0

929
@canFreeze
930
class kat(object):  
Daniel Brown's avatar
Daniel Brown committed
931 932 933
    #def __del__(self):
    #    print ("\33[101m__del__() called: object %08x destroyed (%s)\33[0m" % (id(self), self.__class__))
        
934 935 936 937 938 939 940 941 942
    def __new__(cls, *args, **kwargs):
        # This may seem like an arbitrary step but here we are creating a
        # new class that is a base class of itself. This is because when
        # the kat object adds new components it also adds properties for
        # each of these. There properties are unique to each kat object,
        # but properties are part of the class definition. Thus if two
        # kat objects share the same class definition they also have the
        # same properties regardless of whether they have the actual
        # object added to it. So we create an instance specific class.
943 944 945
        global id___
        id___ += 1
        cnew = type(pykat.finesse.kat.__name__ + str("_") + str(id___), (pykat.finesse.kat,), {})
Daniel Brown's avatar
Daniel Brown committed
946 947 948 949 950 951
        
        obj = object.__new__(cnew)
        
        #print ("\33[104m__del__() called: object %08x created (%s)\33[0m" % (id(obj), obj.__class__))
        
        return obj
952
    
953
    def __init__(self, kat_file=None, kat_code=None, katdir="", katname="", tempdir=None, tempname=None):
Daniel Brown's avatar
Daniel Brown committed
954
        
955
        self._unfreeze()
Daniel Brown's avatar
Daniel Brown committed
956
        self.__looking = False
Daniel Brown's avatar
Daniel Brown committed
957 958
        self.scene = None # scene object for GUI
        self.verbose = True
959
        self.__blocks = OrderedDict() # dictionary of blocks that are used
Daniel Brown's avatar
Daniel Brown committed
960 961 962 963 964 965 966 967 968 969
        self.__components = {}  # dictionary of optical components      
        self.__detectors = {}   # dictionary of detectors
        self.__commands = {}    # dictionary of commands
        self.__gui = None
        self.nodes = NodeNetwork(self)  
        self.__katdir = katdir
        self.__katname = katname
        self.__tempdir = tempdir
        self.__tempname = tempname
        self.pykatgui = None
Daniel Brown's avatar
Daniel Brown committed
970
        self.__signals = Signals(self)
971
        self.constants = {}
972
        self.vacuum = []
973 974
        self.__prevrunfilename = None
        self.printmatrix = None
975
        self.__variables = {}
Daniel Brown's avatar
Daniel Brown committed
976
        self.IFO = None
977
        self.mf = []
978
        
979 980
        self.data = {}
        
981 982 983
        # initialise default block
        self.__currentTag= NO_BLOCK
        self.__blocks[NO_BLOCK] = Block(NO_BLOCK)
Daniel Brown's avatar
Daniel Brown committed
984 985 986
        
        # Various options for running finesse, typicaly the commands with just 1 input
        # and have no name attached to them.
987 988
        self.retrace = None
        self.deriv_h = None
989
        self.scale = None
ssl's avatar
ssl committed
990
        self.__trace = None
Daniel Brown's avatar
Daniel Brown committed
991 992
        self.__phase = None
        self.__maxtem = None
993
        self.__noxaxis = False
Daniel Brown's avatar
Daniel Brown committed
994
        self.__time_code = None
995
        self.__yaxis = "abs" # default yaxis
996
        self.__lambda0 = 1064e-9
997
        self.__finesse_dir = None
Daniel Brown's avatar
Daniel Brown committed
998 999 1000 1001 1002
        
        if kat_code != None and kat_file != None:
            raise pkex.BasePyKatException("Specify either a Kat file or some Kat code, not both.")
        
        if kat_code != None:
1003
            self.parse(kat_code)
Daniel Brown's avatar
Daniel Brown committed
1004 1005
        
        if kat_file != None:
1006
            self.load(kat_file)
1007 1008
    
        self._freeze()
1009 1010 1011 1012 1013 1014 1015 1016 1017 1018
    
    def __and__(self, other):
        """
        Quick syntax for returning a string of the block if present.
        
        Example:
            print(kat & 'locks')
        """
        return self.getBlockString(other)
        
1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039
    def _data2str(self):
        """
        From the pykat data object we serialise, compress and convert into a base64 string.
        """
        if not isinstance(self.data, dict):
            raise pkex.BasePyKatException("Data object is not a dictionary")
            
        return base64.b64encode(zlib.compress(pickle.dumps(self.data))).decode('ascii')
    
    def _str2data(self, data_str):
        """
        Takes a compressed string of _data2str and converts it back into a data object
        """
        dic = pickle.loads(zlib.decompress(base64.decodebytes(bytes(data_str, 'ascii'))))
        
        if not isinstance(dic, dict):
            raise pkex.BasePyKatException("Read data object is not a dictionary")
        
        if not isinstance(self.data, dict):
            raise pkex.BasePyKatException("Data object is not a dictionary")
             
1040
        self.data.update(dic)
1041
        
Daniel Brown's avatar
Daniel Brown committed
1042 1043 1044 1045 1046 1047 1048 1049 1050 1051
    @property
    def binaryDirectory(self):
        """
        Returns the directory of the FINESSE binary that will be called when running the simulation. 
        """
        if len(self.__katdir) == 0:
            return os.environ.get('FINESSE_DIR')
        else:
            return self.__katdir
      
Daniel Brown's avatar
updates  
Daniel Brown committed
1052 1053 1054
    def deepcopy(self):
        return copy.deepcopy(self)
    
Daniel Brown's avatar
Daniel Brown committed
1055
    def getAll(self, type, parameter=None):
1056 1057 1058 1059 1060 1061 1062
        """
        Returns a collection of all objects of the type argument that are
        part of this kat object.
        
        Example:
            # returns all cav commands that are present in this kat object
            cavs = kat.getAll(pykat.commands.cavity)