finesse.py 103 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

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

53
54
from subprocess import Popen, PIPE

55
56
57
58
59
60
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
61
62
63
64
65
66
67
68


try:
    # Add exception in Python 2
    FileNotFoundError
except NameError:
    FileNotFoundError = IOError
    
69
70
71
72
73
74
75
76
77
78
"""
try:
    from future_builtins import zip_longest
except ImportError: # not 2.6+ or is 3.x
    try:
        from itertools import izip_longest as zip_longest # < 2.5 or 3.x
    except ImportError:
        print("boom")
        pass
"""
79
80

from math import erfc, pi
81
from collections import namedtuple, OrderedDict
82

Daniel Brown's avatar
Daniel Brown committed
83
from pykat.node_network import NodeNetwork
84
from pykat.detectors import BaseDetector as Detector
Daniel Brown's avatar
Daniel Brown committed
85
86
from pykat.components import Component
from pykat.commands import Command, xaxis
87
from pykat.SIfloat import *
88
from pykat.param import Param, AttrParam
89
from pykat.external import progressbar
90
from pykat.freeze import canFreeze
91
92
import pykat.external.six as six

Daniel Brown's avatar
Daniel Brown committed
93
94
import pykat.exceptions as pkex

95
96
97
98
99
100
from pykat import USE_GUI, HAS_OPTIVIS, NoGUIException

if HAS_OPTIVIS:
    from optivis.bench.labels import Label as optivis_label
    from optivis.geometry import Coordinates as optivis_coord
    import PyQt4
101
102
103
104
105

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
106

Daniel Brown's avatar
Daniel Brown committed
107
108
from multiprocessing import Process, Manager

Daniel Brown's avatar
Daniel Brown committed
109
NO_BLOCK = "NO_BLOCK"
110
111
pykat_web = "www.gwoptics.org/pykat"

Daniel Brown's avatar
Daniel Brown committed
112
113
114
115
116
117
118
# 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")
                                     
Daniel Brown's avatar
Daniel Brown committed
119
120
121
def f__lkat_process(callback, cmd, kwargs):
    """
    """
Daniel Brown's avatar
Daniel Brown committed
122
    
123
    if lkat_location is None:
Daniel Brown's avatar
Daniel Brown committed
124
125
126
        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
127
128
129
130
131
132
133
134
135
136
137
138
139

    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: 
140
        print ("Exception caught in python: ", ex.message)
Daniel Brown's avatar
Daniel Brown committed
141
142
143
    finally:
        # This should always be called no matter what
        lkat._pykat_finish(0)
Daniel Brown's avatar
Daniel Brown committed
144
145


146
def f__lkat_trace_callback(lkat, trace_info, getCavities, getNodes, getSpaces):
Daniel Brown's avatar
Daniel Brown committed
147
148
149
150
151
152
153
154
155
156
157
158
    """
    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()

159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
    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
186
187
                                )

188
            trace_info[cav.name] = cav_info
Daniel Brown's avatar
Daniel Brown committed
189

190
191
192
    if getSpaces:
        for s in range(0, inter.num_spaces):
            space = inter.space_list[s]
Daniel Brown's avatar
Daniel Brown committed
193

194
195
            trace_info[space.name] = space_trace(gouyx = space.gouy_x,
                                                 gouyy = space.gouy_y)
Daniel Brown's avatar
Daniel Brown committed
196
197
                                                 

Daniel Brown's avatar
Daniel Brown committed
198
199

class BlockedKatFile(object):
200
    """
Daniel Brown's avatar
Daniel Brown committed
201
202
203
204
205
206
207
208
209
    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")
210
    """
Daniel Brown's avatar
Daniel Brown committed
211
    
212
213
214
215
216
217
218
219
220
221
    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
222
223
224
225
226
227
    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
        
228
229
230
231
232
    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
233
        
234
235
236
237
238
239
240
        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
241
    def add(self, block, contents, addAfter=None):
Daniel Brown's avatar
Daniel Brown committed
242

Daniel Brown's avatar
Daniel Brown committed
243
244
        if block in self.ordering or block in self.blocks:
            raise Exception("%s block already present")
Daniel Brown's avatar
Daniel Brown committed
245
    
Daniel Brown's avatar
Daniel Brown committed
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
        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
270

Daniel Brown's avatar
Daniel Brown committed
271
272
            # Looking for block start or end
            values = line.split()
Daniel Brown's avatar
Daniel Brown committed
273
    
Daniel Brown's avatar
Daniel Brown committed
274
275
276
            if len(values) >= 3 and values[0] == "%%%":
                if values[1] == "FTblock":
                    newTag = values[2]
277
                    
Daniel Brown's avatar
Daniel Brown committed
278
279
                    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
280
    
Daniel Brown's avatar
Daniel Brown committed
281
                    if newTag in self.blocks:
282
283
284
                        #raise pkex.BasePyKatException("Block `{0}` has already been read".format(newTag))
                        self.__currentBlock = newTag
                        continue
Daniel Brown's avatar
Daniel Brown committed
285
    
Daniel Brown's avatar
Daniel Brown committed
286
287
288
289
290
291
                    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
292
        
Daniel Brown's avatar
Daniel Brown committed
293
294
295
296
                continue
                
            if(len(line) == 0 and (self.__currentBlock == self.__NO_BLOCK)):
                continue
Daniel Brown's avatar
Daniel Brown committed
297
        
Daniel Brown's avatar
Daniel Brown committed
298
            self.blocks[self.__currentBlock] += line + "\n"
Daniel Brown's avatar
Daniel Brown committed
299
300
301
    
    
class KatBatch(object):
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
    
    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
        kat.parseCommands(commands)
        
        kw = dict()
        
        if "cmd_args" in kwargs:
            kw["cmd_args"] = kwargs["cmd_args"]
        
323
        return kat.run(**kw)
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
    
    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

                                  
339
340
341
342
343
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)
344
345

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

Daniel Brown's avatar
Daniel Brown committed
379
380
381
382
        if kat.noxaxis:
            print("No xaxis used")
        else:
            print("One xaxis used: %s" % kat.xaxis.getFinesseText())
383
384
385
386
387
388
389
390
391
392
393
            
        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
394
395
        
        
Daniel Brown's avatar
Daniel Brown committed
396
397
    def plot(self, detectors=None, filename=None, show=True,
                   yaxis=None, legend=True, loc=0, title=None, styles=None,
398
399
                   ylabel=None, y2label=None, xlabel=None, x2label=None,
                   xlim=None, x2lim=None, ylim=None, y2lim=None):
400
401
402
403
404
405
406
407
        """
        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
408
409
410
411
412
413
414
415
416
417
418
419
            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
420
421
422
423
424
425
            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.
426
        """
427
428
        import matplotlib.pyplot as pyplot
        import pykat.plotting as plt
429
            
Daniel Brown's avatar
Daniel Brown committed
430
431
432
        if not show:
            pyplot.ioff()

433
434
435
436
        kat = pykat.finesse.kat()
        kat.verbose = False
        kat.parseCommands(self.katScript)

437
438
439
        if kat.noxaxis == True:
            raise  pkex.BasePyKatException("This kat object has noxaxis=True, so there is nothing to plot.")

440
441
        original_yaxis = kat.yaxis

442
443
        if yaxis is not None:
            kat.yaxis = yaxis
444
445
446
447
448
449
450
451
452
453
454
455
456

        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
457
458
        _func1 = np.abs
        _func2 = None
459

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

            plot_cmd1 = plot_cmd
Daniel Brown's avatar
Daniel Brown committed
474
            plot_cmd2 = pyplot.plot if kat.xaxis.scale == "lin" else pyplot.semilogx
475
            
Daniel Brown's avatar
Daniel Brown committed
476
477
            dual_plot = True
        elif "db:deg" in kat.yaxis:
478
479
480
481
482
            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
483
            _func2 = lambda x: np.rad2deg(np.angle(x))
484
485

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

509
510
511
512
513
514
515
        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
516
            if not hasattr(kat, det) or (hasattr(kat, det) and not getattr(kat, det).noplot):
517
                
Daniel Brown's avatar
Daniel Brown committed
518
                if dual_plot:
Daniel Brown's avatar
Daniel Brown committed
519
                    ax = pyplot.subplot(2,1,1)
Daniel Brown's avatar
Daniel Brown committed
520
521
522
523
                    
                if styles is not None and det in styles:
                    l, = plot_cmd1(self.x, _func1(self[det]), styles[det], label=det)
                else:
524
                    l, = plot_cmd1(self.x, _func1(self[det]), label=det)
Daniel Brown's avatar
Daniel Brown committed
525
526
                
                if dual_plot: 
527
528
                    pyplot.subplot(2,1,2)
                    plot_cmd2(self.x, _func2(self[det]), color=l.get_color(), ls=l.get_linestyle(), label=det)
529

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

Daniel Brown's avatar
Daniel Brown committed
569
570
            if title is not None:
                pyplot.title(title, fontsize=font_label_size)
571
572
    
            pyplot.subplot(2,1,2)
Daniel Brown's avatar
Daniel Brown committed
573
574
            pyplot.xlabel(x2label, fontsize=font_label_size)
            pyplot.ylabel(y2label, fontsize=font_label_size)
575
576
577
578
579
        
            pyplot.xlim(x2lim[0], x2lim[1])
            if y2lim is not None:
                pyplot.ylim(y2lim[0],y2lim[1])
            
580
        else:
Daniel Brown's avatar
Daniel Brown committed
581
582
            pyplot.xlabel(xlabel, fontsize=font_label_size)
            pyplot.ylabel(ylabel)
583
            pyplot.xlim(self.x.min(), self.x.max())
Daniel Brown's avatar
Daniel Brown committed
584
585
586
            
            if title is not None:
                pyplot.title(title, fontsize=font_label_size)
587
588
            if ylim is not None:
                pyplot.ylim(ylim[0],ylim[1])
589
    
590
        pyplot.margins(0, 0.05)
591
        pyplot.tight_layout()
592
    
593
        if legend:
Daniel Brown's avatar
Daniel Brown committed
594
            fig.axes[0].legend(loc=loc, fontsize=font_label_size)
595
596
597
598
599
600
        
        if filename is not None:
            fig.savefig(filename)
            
        if show:
            pyplot.show(fig)
Daniel Brown's avatar
Daniel Brown committed
601
            pyplot.ion()
602
603
        
        return fig
Daniel Brown's avatar
Daniel Brown committed
604
        
605
    def saveKatRun(self, filename):
Daniel Brown's avatar
Daniel Brown committed
606
607
608
609
610
611
612
613
614
615
616
        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):
617
        idx = [i for i in range(len(self.ylabels)) if self.ylabels[i].split()[0] == str(value)]
618
619
        out = None
        
620
        if len(idx) > 0:
621
            #out = self.y[:, idx]
622
            
623
            if len(idx) == 1:
Daniel Brown's avatar
Daniel Brown committed
624
                if "abs:deg" in self.yaxis:
625
                    out = self.y[:, idx[0]]
Daniel Brown's avatar
Daniel Brown committed
626
                elif "re:im" in self.yaxis:
627
628
                    out = self.y[:, idx[0]]
            else: 
Daniel Brown's avatar
Daniel Brown committed
629
                if "abs:deg" in self.yaxis:
630
                    out = self.y[:, idx[0]] * np.exp(1j*math.pi*self.y[:, idx[1]]/180.0)
Daniel Brown's avatar
Daniel Brown committed
631
                elif "re:im" in self.yaxis :
632
                    out = self.y[:, idx[0]] + 1j*self.y[:, idx[1]]
633

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

637
            if out.size == 1:
638
                return out[0].squeeze()
639
            else:
640
                return out.squeeze()
Daniel Brown's avatar
Daniel Brown committed
641
        else:
642
            raise  pkex.BasePyKatException("No output by the name '{0}' found in the output".format(str(value)))
643
644
            
@canFreeze 
645
class KatRun2D(object):
Daniel Brown's avatar
Daniel Brown committed
646
    def __init__(self):
647
        self._unfreeze()
648
        self.runtime = None
649
        self.startDateTime = datetime.datetime.now()
Daniel Brown's avatar
Daniel Brown committed
650
651
652
653
654
655
656
657
        self.x = None
        self.y = None
        self.z = None
        self.xlabel = None
        self.ylabel = None
        self.zlabels = None
        self.katScript = None
        self.katVersion = None
658
659
        self.stderr = None
        self.stdout = None
660
        self._freeze()
Daniel Brown's avatar
Daniel Brown committed
661
662
663
664
665
666
667
668
669
670
        
    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)
    
671
    def get(self, value): return self[value].squeeze()
Daniel Brown's avatar
Daniel Brown committed
672
673
    
    def __getitem__(self, value):
674
        idx = [i for i in range(len(self.zlabels)) if self.zlabels[i].split()[0] == str(value)]
Daniel Brown's avatar
Daniel Brown committed
675
676
677
678
679
        
        if len(idx) > 0:
            return self.z[idx].squeeze()
        else:
            raise  pkex.BasePyKatException("No output by the name {0} found".format(str(value)))
680
    
681
@canFreeze    
682
class Signals(object):
683
684
    
    @canFreeze 
685
    class fsig(object):
Daniel Brown's avatar
Daniel Brown committed
686
        def __init__(self, param, name, amplitude, phase, signal):
687
            self._unfreeze()
688
689
690
691
692
            self._params = []
            self.__target = param
            self.__name = name
            self.__amplitude = Param("amp", self, SIfloat(amplitude))
            self.__phase = Param("phase", self, SIfloat(phase))
693
            self.__removed = False
Daniel Brown's avatar
Daniel Brown committed
694
            self.__signal = signal
695
            self._freeze()
696
            
697
            # unfortunatenly the target names for fsig are not the same as the
Daniel Brown's avatar
Daniel Brown committed
698
            # various parameter names of the components, e.g. mirror xbeta is x 
699
700
701
            # 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:
702
                raise  pkex.BasePyKatException("Cannot fsig parameter {1} on component {0}".format(str(param._owner().name), param.name))
703
704
705
            
        def _register_param(self, param):
            self._params.append(param)
706
707
708
        
        @property
        def removed(self): return self.__removed
709
  
710
        def remove(self):
711
712
713
            self.__signal._kat.remove(self)
            
        def _on_remove(self):
714
715
716
            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
717
718
                self.__signal.targets.remove(self)
                self.__remove = True
719
        
720
721
722
723
724
        @property
        def name(self): return self.__name

        @property
        def amplitude(self): return self.__amplitude
725
726
727
        @amplitude.setter
        def amplitude(self,value): self.__amplitude.value = SIfloat(value)

728
729
730

        @property
        def phase(self): return self.__phase
731
732
        @phase.setter
        def phase(self,value): self.__phase.value = SIfloat(value)
733
734
735
736
737

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

        @property
738
        def owner(self): return self.__target._owner().name
739
740
741
742
743
744
745
746
747
    
        def getFinesseText(self):
            rtn = []
    
            for p in self._params:
                rtn.extend(p.getFinesseText())
        
            return rtn
    
748
    
749
750
751
752
753
754
755
    @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:
756
            return self._default_name
757
758
        else:
            return self.targets[0].name
759
760
            
    @property
Daniel Brown's avatar
Daniel Brown committed
761
    def removed(self): return False # we can never remove the Signal object altogethr just the individual fsig targets
762

763
764
    def remove(self):
        for t in self.targets:
Daniel Brown's avatar
Daniel Brown committed
765
            t.remove()
766
767
768
        
        del self.targets[:]
        
769
770
    @property
    def f(self): return self.__f
771
    @f.setter
772
773
774
775
776
777
778
    def f(self,value):
        v = SIfloat(value)
        
        if v <= 0:
            raise pkex.BasePyKatException("Signal frequency must be greater than 0.")
            
        self.__f.value = SIfloat(value)
779
    
Daniel Brown's avatar
Daniel Brown committed
780
    def __init__(self, kat):
781
        self._unfreeze()
782
        self._default_name = "fsignal"
783
784
        self.targets = []
        self._params = []
785
        self.__f = Param("f", self, None)
786
        self._kat = kat
787
        self._freeze()
788
        
789
790
791
792
793
    def _register_param(self, param):
        self._params.append(param)
        
    def apply(self, target, amplitude, phase, name=None):
        
794
        if target is None:
795
796
            raise  pkex.BasePyKatException("No target was specified for signal to be applied")
        
797
        if name is None:
798
            name = "sig_" + target._owner().name + "_" + target.name
799
        
Daniel Brown's avatar
Daniel Brown committed
800
        self.targets.append(Signals.fsig(target, name, amplitude, phase, self))
801
802
803
804
        
    def getFinesseText(self):
        rtn = []
        
805
806
807
808
809
810
811
812
        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
813
            
814
815
816
817
818
819
820
                    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
821

822
823
824
825
826
        for p in self._params:
            rtn.extend(p.getFinesseText())
        
        return rtn
        
Daniel Brown's avatar
Daniel Brown committed
827
828
829
830
831
832
833
834
835
class Block:
    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
    
836
Constant = namedtuple('Constant', 'name, value, usedBy')
837
838
839

id___ = 0

840
@canFreeze
841
class kat(object):  
842

843
844
845
846
847
848
849
850
851
    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.
852
853
854
        global id___
        id___ += 1
        cnew = type(pykat.finesse.kat.__name__ + str("_") + str(id___), (pykat.finesse.kat,), {})
855
856
        return object.__new__(cnew)
    
857
    def __init__(self, kat_file=None, kat_code=None, katdir="", katname="", tempdir=None, tempname=None):
858
        self._unfreeze()
Daniel Brown's avatar
Daniel Brown committed
859
        self.__looking = False
Daniel Brown's avatar
Daniel Brown committed
860
861
        self.scene = None # scene object for GUI
        self.verbose = True
862
        self.__blocks = OrderedDict() # dictionary of blocks that are used
Daniel Brown's avatar
Daniel Brown committed
863
864
865
866
867
868
869
870
871
872
        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
873
        self.__signals = Signals(self)
874
        self.constants = {}
875
        self.vacuum = []
876
877
        self.__prevrunfilename = None
        self.printmatrix = None
878
        self.__variables = {}
879
880
881
882
        
        # initialise default block
        self.__currentTag= NO_BLOCK
        self.__blocks[NO_BLOCK] = Block(NO_BLOCK)
Daniel Brown's avatar
Daniel Brown committed
883
884
885
        
        # Various options for running finesse, typicaly the commands with just 1 input
        # and have no name attached to them.
886
887
        self.retrace = None
        self.deriv_h = None
888
        self.scale = None
Sean Leavey's avatar
Sean Leavey committed
889
        self.__trace = None
Daniel Brown's avatar
Daniel Brown committed
890
891
        self.__phase = None
        self.__maxtem = None
892
        self.__noxaxis = False
Daniel Brown's avatar
Daniel Brown committed
893
        self.__time_code = None
894
        self.__yaxis = "abs" # default yaxis
895
        self.__lambda0 = 1064e-9
896
        self.__finesse_dir = None
Daniel Brown's avatar
Daniel Brown committed
897
898
899
900
901
902
903
904
905
        
        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:
            self.parseCommands(kat_code)
        
        if kat_file != None:
            self.loadKatFile(kat_file)
906
907
908
    
        self._freeze()
        
Daniel Brown's avatar
updates    
Daniel Brown committed
909
910
911
    def deepcopy(self):
        return copy.deepcopy(self)
    
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
    def getAll(self, type):
        """
        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)
        """
        items = []
        
        for a in (item for item in self.__class__.__dict__):
            b = getattr(self, a)
            
            if isinstance(b, type):
                items.append(b)

        return tuple(items)
        
931
    def __deepcopy__(self, memo):
932
933
934
935
936
937
938
939
940
941
        """
        When deep copying a kat object we need to take into account
        the instance specific properties. 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.
        """
Daniel Brown's avatar
updates    
Daniel Brown committed
942
        result = self.__class__.__new__(self.__class__.__base__)
943
        memo[id(self)] = result
944
945
946
947
948
949
950
951
952
953
        result.__dict__ = copy.deepcopy(self.__dict__, memo)

        # Find all properties in class we are copying
        # and deep copy these to the new class instance
        for x in self.__class__.__dict__.items():
            if isinstance(x[1], property):
                setattr(result.__class__, x[0], x[1])
    
        result.nodes._NodeNetwork__update_nodes_properties()
                
954
955
956
        # Update any weakrefs
        for c in result.components:
            result.components[c]._Component__update_node_setters()
957
        
958
959
        return result
    
960
961
962
    @property
    def signals(self): return self.__signals

963
    yaxis_options = ["abs:deg","db:deg","re:im","abs","db","deg"]
Daniel Brown's avatar
Daniel Brown committed
964
    
965
966
967
968
    @property
    def yaxis(self): return self.__yaxis
    @yaxis.setter
    def yaxis(self, value):
Daniel Brown's avatar
Daniel Brown committed
969
970
971
972
973
974
975
976
        values = value.split()
        
        if len(values) == 2:
            scale = values[0]
            mode = values[1]
        else:
            scale = "lin"
            mode = value
977
        
Daniel Brown's avatar
Daniel Brown committed
978
979
980
981
982
        if not str(scale) in ["lin", "log"]:
            raise pkex.BasePyKatException("yaxis value '{0}' is not a valid option. Valid options are: lin or log".format(str(mode)))
                
        if not str(mode) in self.yaxis_options:
            raise pkex.BasePyKatException("yaxis value '{0}' is not a valid option. Valid options are: {1}".format(str(mode), ",".join(self.yaxis_options) ))
983
            
Daniel Brown's avatar
Daniel Brown committed
984
        self.__yaxis = str(scale + " " + mode).strip()
985

Sean Leavey's avatar
Sean Leavey committed
986
987
988
989
990
991
992
993
994
995
996
    @property
    def trace(self): return self.__trace
    @trace.setter
    def trace(self, value):
        value = int(value)

        if value < 0 or value > 255:
            raise pkex.BasePyKatException('trace command only accepts values in the range 0-255.')
        else:
            self.__trace = value

997
998
999
1000
1001
1002
1003
1004
1005
1006
    @property
    def lambda0(self): return self.__lambda0
    @lambda0.setter
    def lambda0(self, value):
        self.__lambda0 = SIfloat(value)
        
        for node in self.nodes.getNodes():
            if self.nodes[node].q != None:
                self.nodes[node].q.wavelength = self.__lambda0

Daniel Brown's avatar
Daniel Brown committed
1007
1008
1009
    @property
    def maxtem(self): return self.__maxtem
    @maxtem.setter
1010
1011
1012
1013
1014
    def maxtem(self,value):
        if value == "off":
            self.__maxtem = -1
        else:
            self.__maxtem = int(value)
Daniel Brown's avatar
Daniel Brown committed
1015
1016
1017
1018
1019
    
    @property
    def phase(self): return self.__phase
    @phase.setter
    def phase(self,value): self.__phase = int(value)
1020
        
Daniel Brown's avatar
Daniel Brown committed
1021
    @property
1022
1023
1024
    def timeCode(self): return self.__time_code
    @timeCode.setter
    def timeCode(self,value): self.__time_code = bool(value)
Daniel Brown's avatar
Daniel Brown committed
1025
    
Daniel Brown's avatar
Daniel Brown committed
1026
1027
1028
1029
1030
1031
1032
1033
    @property
    def components(self):
        return self.__components.copy()
    
    @property
    def detectors(self):
        return self.__detectors.copy()
        
1034
1035
1036
1037
    @property
    def commands(self):
        return self.__commands.copy()
        
Daniel Brown's avatar
Daniel Brown committed
1038
1039
1040
    @property
    def noxaxis(self): return self.__noxaxis
    @noxaxis.setter
1041
    def noxaxis(self,value): self.__noxaxis = bool(value) 
Daniel Brown's avatar
Daniel Brown committed
1042

1043
1044
    @staticmethod
    def logo():
1045
        print ("""                                              ..-
1046
    PyKat {0:7}         _                  '(
Daniel Brown's avatar
Daniel Brown committed
1047
1048
1049
1050
1051
1052
                          \\`.|\\.__...-\"\"""-_." )
       ..+-----.._        /  ' `            .-'
   . '            `:      7/* _/._\\    \\   (
  (        '::;;+;;:      `-"' =" /,`"" `) /
  L.        \\`:::a:f            c_/     n_'
  ..`--...___`.  .    ,  
1053
   `^-....____:   +.      {1}\n""".format(pykat.__version__, pykat_web))
Daniel Brown's avatar
Daniel Brown committed
1054
    
1055
    def loadKatFile(self, katfile, blocks=None):
1056
1057
1058
        with open(katfile) as f:
            commands= f.read()
            
1059
        self.parseCommands(commands, blocks=blocks)
Daniel Brown's avatar
Daniel Brown committed
1060
    
1061
    def parseKatCode(self, code, blocks=None):
Daniel Brown's avatar
Daniel Brown committed
1062
        warnings.warn('parseKatCode depreciated, use parseCommands.', stacklevel=2)
Daniel Brown's avatar
typo    
Daniel Brown committed
1063
        self.parseCommands(code, blocks=blocks)
1064
1065
1066
1067
1068
1069
1070

    def processConstants(self, commands):
        """
        Before fully parsing a bunch of commands firstly any constants or variables
        to be recorded and replaced.
        """
        
Daniel Brown's avatar
Daniel Brown committed
1071
1072
1073
1074
1075
        try:
            constants = self.constants
        
            for line in commands:
                values = line.split()
1076
            
Daniel Brown's avatar
Daniel Brown committed
1077
                if len(values)>0 and values[0] == 'const':
1078
                
Daniel Brown's avatar
Daniel Brown committed
1079
1080
1081
1082
1083
                    if len(values) >= 3:
                        if values[1] in constants:
                            raise pkex.BasePyKatException('const command with the name "{0}" already used'.format(values[1]))
                        else:
                            constants[str(values[1])] = Constant(values[1], values[2], [])
1084
                    else:
Daniel Brown's avatar
Daniel Brown committed
1085
                        raise pkex.BasePyKatException('const command "{0}" was not the correct format'.format(line))
1086
        
Daniel Brown's avatar
Daniel Brown committed
1087
            commands_new = []
1088
        
Daniel Brown's avatar
Daniel Brown committed
1089
1090
            for line in commands:
                values = line.split()
1091
            
Daniel Brown's avatar
Daniel Brown committed
1092
1093
1094
1095
1096
                if len(values) > 0 and values[0] != 'const':
                    # check if we have a var/constant in this line
                    if line.find('$') >= 0:
                        for key in constants.keys():
                            # TODO: need to fix this for checking mulitple instances of const in a single line
1097
                        
Daniel Brown's avatar
Daniel Brown committed
1098
                            chars = [' ', '+', '-', '*', '/', ')']
1099

Daniel Brown's avatar
Daniel Brown committed
1100
1101
                            for c in chars:
                                none_found = False
1102
                            
Daniel Brown's avatar
Daniel Brown committed
1103
1104
1105
1106
1107
1108
                                while not none_found:
                                    if line.find('$'+key+c) > -1:
                                        constants[key].usedBy.append(line)
                                        line = line.replace('$'+key+c, str(constants[key].value)+ c)
                                    else:
                                        none_found = True
1109
                        
Daniel Brown's avatar
Daniel Brown committed
1110
1111
1112
                            if line.endswith('$'+key):
                                constants[key].usedBy.append(line)
                                line = line.replace('$'+key, str(constants[key].value))
1113
                        
Daniel Brown's avatar
Daniel Brown committed
1114
                    commands_new.append(line)
1115
    
Daniel Brown's avatar
Daniel Brown committed
1116
            self.constants = constants
1117
        
Daniel Brown's avatar
Daniel Brown committed
1118
1119
1120
1121
1122
            return commands_new
            
        except pkex.BasePyKatException as ex:
            pkex.PrintError("Error processing constants:", ex)
            sys.exit(1)
1123
    
Daniel Brown's avatar
updates    
Daniel Brown committed
1124
1125
1126
1127
    def getBlocks(self):
        return self.__blocks.keys()
    
    def removeBlock(self, name, failOnBlockNotFound=True):
1128
1129
        
        if name not in self.__blocks:
Daniel Brown's avatar
updates    
Daniel Brown committed
1130
1131
1132
1133
1134
            if failOnBlockNotFound:
                pkex.PrintError("Error removing block:", pkex.BasePyKatException('Block "{0}" was not found'.format(name)))
                sys.exit(1)
            else:
                return
1135
        
1136
        for o in list(self.__blocks[name].contents):
1137
1138
1139
            self.remove(o)
        
        del self.__blocks[name]
1140
1141
1142
1143
    
    def __str__(self):
         return "".join(self.generateKatScript())
         
1144
1145
1146
1147
1148