maps.py 49.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
"""
------------------------------------------------------
Utility functions for handling mirror surface
maps. Some functions based on earlier version
in Matlab (http://www.gwoptics.org/simtools/)
Work in progress, currently these functions are
untested!

http://www.gwoptics.org/pykat/
------------------------------------------------------
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

16
from pykat.optics.romhom import makeWeightsNew
17
from scipy.interpolate import interp2d, interp1d
18
from scipy.optimize import minimize
19
from pykat.maths.zernike import *        
20
from pykat.exceptions import BasePyKatException
21

Daniel Brown's avatar
Daniel Brown committed
22
import numpy as np
23
import math
24
import pickle
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

class MirrorROQWeights:
    
    def __init__(self, rFront, rBack, tFront, tBack):
        self.rFront = rFront
        self.rBack = rBack
        self.tFront = tFront
        self.tBack = tBack
    
    def writeToFile(self, romfilename):
        with open(romfilename + ".rom", "w+") as f:
            if self.rFront is not None: self.rFront.writeToFile(f=f)
            if self.rBack  is not None: self.rBack.writeToFile(f=f)
            if self.tFront is not None: self.tFront.writeToFile(f=f)
            if self.tBack  is not None: self.tBack.writeToFile(f=f)
                    
Daniel Brown's avatar
Daniel Brown committed
41
class surfacemap(object):
42
43
    def __init__(self, name, maptype, size, center, step_size, scaling, data=None,
                 notNan=None, Rc=None, zOffset=None):
44
45
46
47
48
49
        
        self.name = name
        self.type = maptype
        self.center = center
        self.step_size = step_size
        self.scaling = scaling
50
        self.notNan = notNan
51
52
        self.Rc = Rc
        self.zOffset = zOffset
53
54
        self.__interp = None
        
55
        if data is None:
Daniel Brown's avatar
Daniel Brown committed
56
57
58
            self.data = np.zeros(size)
        else:
            self.data = data
59
60
61
62
63
            
        if notNan is None:
            self.notNan = np.ones(size)
        else:
            self.notNan = notNan
Daniel Brown's avatar
Daniel Brown committed
64

65
        self._rom_weights = None
66
67
68
69
70
71
72
        
    def write_map(self, filename):
        with open(filename,'w') as mapfile:
            
            mapfile.write("% Surface map\n")
            mapfile.write("% Name: {0}\n".format(self.name))
            mapfile.write("% Type: {0}\n".format(self.type))
73
            mapfile.write("% Size: {0} {1}\n".format(self.data.shape[0], self.data.shape[1]))
74
75
76
77
78
79
80
81
82
83
            mapfile.write("% Optical center (x,y): {0} {1}\n".format(self.center[0], self.center[1]))
            mapfile.write("% Step size (x,y): {0} {1}\n".format(self.step_size[0], self.step_size[1]))
            mapfile.write("% Scaling: {0}\n".format(float(self.scaling)))
            mapfile.write("\n\n")
            
            for i in range(0, self.data.shape[0]):
                for j in range(0, self.data.shape[1]):
                    mapfile.write("%.15g " % self.data[i,j])
                mapfile.write("\n")
    
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
    @property
    def data(self):
        return self.__data
    
    @data.setter
    def data(self, value):
        self.__data = value
        self.__interp = None
    
    @property
    def center(self):
        return self.__center
    
    @center.setter
    def center(self, value):
        self.__center = value
        self.__interp = None
    
    @property
    def step_size(self):
        return self.__step_size
    
    @step_size.setter
    def step_size(self, value):
        self.__step_size = value
        self.__interp = None

    @property
    def scaling(self):
        return self.__scaling
    
    @scaling.setter
    def scaling(self, value):
        self.__scaling = value
        self.__interp = None

120
121
122
123
    # CHANGED! I swapped: self.data.shape[0]  <----> self.data.shape[1]. Because
    # the rows/columns in the data matrix are supposed to be y/x-values(?).
    # This may mess things up in other methods though... I fixed the plot-method.
    # /DT
Daniel Brown's avatar
Daniel Brown committed
124
125
    @property
    def x(self):
126
127
        return self.step_size[0] * (np.array(range(0, self.data.shape[1])) - self.center[0])
    
Daniel Brown's avatar
Daniel Brown committed
128
129
    @property
    def y(self):
130
131
132
133
        return self.step_size[1] * (np.array(range(0, self.data.shape[0]))- self.center[1])
        
    # CHANGED! Since everything else (step_size, center) are given as (x,y), and not
    # as (row, column), I changed this to the same format. /DT
134
135
    @property
    def size(self):
136
137
        return self.data.shape[::-1]
        
138
139
    @property
    def offset(self):
140
        return np.array(self.step_size)*(np.array(self.center) - (np.array(self.size)-1)/2.0)
141
142
143
144
    
    @property
    def ROMWeights(self):
        return self._rom_weights
145
    
146
    def z_xy(self, x=None, y=None, wavelength=1064e-9, direction="reflection_front", nr1=1.0, nr2=1.0):
147
148
149
150
151
152
153
        """
        For this given map the field perturbation is computed. This data
        is used in computing the coupling coefficient. It returns a grid
        of complex values representing the change in amplitude or phase
        of the field.
        
            x, y      : Points to interpolate at, 'None' for no interpolation.
154
            
155
            wavelength: Wavelength of light in vacuum [m]
156
157
158
159
160
161
162
163
164
            
            direction : Sets which distortion to return, as beams travelling
                        in different directions will see different distortions.
                        Options are:
                                "reflection_front"
                                "transmission_front" (front to back)
                                "transmission_back" (back to front)
                                "reflection_back"
                                
165
            nr1       : refractive index on front side
166
            
167
            nr2       : refractive index on back side
168
            
169
170
171
172
        """
        
        assert(nr1 >= 1)
        assert(nr2 >= 1)
173
        
174
        if x is None and y is None:
175
176
            data = self.scaling * self.data
        else:
177
            if self.__interp is None:
178
179
180
                self.__interp = interp2d(self.x, self.y, self.data * self.scaling)
                
            data = self.__interp(x, y)
181
182
        
        if direction == "reflection_front" or direction == "reflection_back":
183
            if "phase" in self.type:
184
                k = math.pi * 2 / wavelength
185
                
186
187
                if direction == "reflection_front":
                    return np.exp(-2j * nr1 * k * data)
188
                else:
189
                    return np.exp(2j * nr2 * k * data[:,::-1])
190
                
191
            elif "absorption" in self.type:
192
                if direction == "reflection_front":
193
194
195
                    return np.sqrt(1.0 - data)
                else:
                    return np.sqrt(1.0 - data[:, ::-1])
196
197
198
199
200
            elif "surface_motion" in self.type:
                if direction == "reflection_front":
                    return data
                else:
                    return data[:, ::-1]
201
202
            else:
                raise BasePyKatException("Map type needs handling")
203
                
204
        elif direction == "transmission_front" or direction == "transmission_back":
205
206
            if "phase" in self.type:
                k = math.pi * 2 / wavelength
207
                
208
209
                if direction == "transmission_front":
                    return np.exp((nr1-nr2) * k * data)
210
                else:
211
                    return np.exp((nr2-nr1) * k * data[:, ::-1])
212
                
213
            elif "absorption" in self.type:
214
                if direction == "transmission_front":
215
216
217
                    return np.sqrt(1.0 - data)
                else:
                    return np.sqrt(1.0 - data[:, ::-1])
218
219
220
                    
            elif "surface_motion" in self.type:
                return np.ones(data.shape)
221
222
            else:
                raise BasePyKatException("Map type needs handling")
223
                
224
        else:
225
            raise ValueError("Direction not valid")
226
        
Daniel Brown's avatar
Daniel Brown committed
227

228
    
229
230
    def generateROMWeights(self, EIxFilename, EIyFilename=None, nr1=1.0, nr2=1.0, verbose=False, interpolate=False, newtonCotesOrder=8):
        
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
        if interpolate == True:
            # Use EI nodes to interpolate if we
            with open(EIxFilename, 'rb') as f:
                EIx = pickle.load(f)

            if EIyFilename is None:
                EIy = EIx
            else:
                with open(EIyFilename, 'rb') as f:
                    EIy = pickle.load(f)

            x = EIx.x
            x.sort()
            nx = np.unique(np.hstack((x, -x[::-1])))
        
            y = EIy.x
            y.sort()
            ny = np.unique(np.hstack((y, -y[::-1])))
            
            self.interpolate(nx, ny)
        
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
        w_refl_front, w_refl_back, w_tran_front, w_tran_back = (None, None, None, None)
        
        if "reflection" in self.type or "both" in self.type:
            w_refl_front = makeWeightsNew(self, EIxFilename, EIyFilename,
                                      verbose=verbose, newtonCotesOrderMapWeight=newtonCotesOrder,
                                      direction="reflection_front")
            
            w_refl_front.nr1 = nr1
            w_refl_front.nr2 = nr2
            
            w_refl_back = makeWeightsNew(self, EIxFilename, EIyFilename,
                                      verbose=verbose, newtonCotesOrderMapWeight=newtonCotesOrder,
                                      direction="reflection_back")
            
            w_refl_back.nr1 = nr1
            w_refl_back.nr2 = nr2
268

269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
        if "transmission" in self.type or "both" in self.type:                                      
            w_tran_front = makeWeightsNew(self, EIxFilename, EIyFilename,
                                      verbose=verbose, newtonCotesOrderMapWeight=newtonCotesOrder,
                                      direction="transmission_front")

            w_refl_front.nr1 = nr1
            w_refl_front.nr2 = nr2
                                            
            w_tran_back  = makeWeightsNew(self, EIxFilename, EIyFilename,
                                      verbose=verbose, newtonCotesOrderMapWeight=newtonCotesOrder,
                                      direction="transmission_back")
            
            w_refl_back.nr1 = nr1
            w_refl_back.nr2 = nr2
            
        self._rom_weights = MirrorROQWeights(w_refl_front, w_refl_back, w_tran_front, w_tran_back)
        
        return self._rom_weights
            
288
289
290
    def interpolate(self, nx, ny, **kwargs):
        """
        Interpolates the map for some new x and y values.
291
        
292
293
294
295
296
297
298
299
300
301
        Uses scipy.interpolate.interp2d and any keywords arguments are
        passed on to it, thus settings like interpolation type and
        fill values can be set.
        
        The range of nx and ny must contain the value zero so that the
        center point of the map can be set.
        """

        D = interp2d(self.x, self.y, self.data, **kwargs)
        
302
        data = D(nx-self.offset[0], ny-self.offset[1])
303
        
304
305
        Dx = interp1d(nx, np.arange(0,len(nx)))
        Dy = interp1d(ny, np.arange(0,len(ny)))
306
307
308
309
310
        
        self.center = (Dx(0), Dy(0))
        self.step_size = (nx[1]-nx[0], ny[1]-ny[0])
        self.data = data

311
    # xlim and ylim given in centimeters
312
    def plot(self, show=True, clabel=None, xlim=None, ylim=None):
313
314
        import pylab
        
315
        if xlim is not None:
316
            # Sorts out the x-values within xlim
317
318
319
320
            _x = np.logical_and(self.x<=max(xlim)/100.0, self.x>=min(xlim)/100.0)
            xmin = np.min(np.where(_x == True))
            xmax = np.max(np.where(_x == True))
        else:
321
            # Uses the whole available x-range
322
323
324
325
            xmin = 0
            xmax = len(self.x)-1
            xlim = [self.x.min()*100, self.x.max()*100]
    
326
        if ylim is not None:
327
            # Sorts out the y-values within ylim
328
329
330
331
            _y = np.logical_and(self.y<=max(ylim)/100.0, self.y>=min(ylim)/100.0)
            ymin = np.min(np.where(_y == True))
            ymax = np.max(np.where(_y == True))
        else:
332
            # Uses the whole available y-range
333
334
335
            ymin = 0
            ymax = len(self.y)-1
            ylim = [self.y.min()*100, self.y.max()*100]
336
            
337
338
        # CHANGED! y corresponds to rows and x to columns, so this should be
        # correct now. Switch the commented/uncommented things to revert. /DT
339
        # ------------------------------------------------------
340
341
342
343
344
        # min and max of z-values
        zmin = self.data[ymin:ymax,xmin:xmax].min()
        zmax = self.data[ymin:ymax,xmin:xmax].max()
        # zmin = self.data[xmin:xmax,ymin:ymax].min()
        # zmax = self.data[xmin:xmax,ymin:ymax].max()
345
        # ------------------------------------------------------
346
        
347
        # 100 factor for scaling to cm
348
349
        xRange = 100*self.x
        yRange = 100*self.y
350
        # xRange, yRange = np.meshgrid(xRange,yRange)
351
        
352
        fig = pylab.figure()
353
354
355

        # Important to remember here is that xRange corresponds to column and
        # yRange to row indices of the matrix self.data.
Daniel Brown's avatar
updates    
Daniel Brown committed
356
        pcm = pylab.pcolormesh(xRange, yRange, self.data)
357
        pcm.set_rasterized(True)
358
        
359
360
        pylab.xlabel('x [cm]')
        pylab.ylabel('y [cm]')
361

362
363
        if xlim is not None: pylab.xlim(xlim)
        if ylim is not None: pylab.ylim(ylim)
364
            
365
        pylab.title('Surface map {0}, type {1}'.format(self.name, self.type))
366

Daniel Brown's avatar
updates    
Daniel Brown committed
367
368
        cbar = pylab.colorbar()
        #cbar.set_clim(zmin, zmax)
369
        
370
        if clabel is not None:
371
            cbar.set_label(clabel)
372
    
373
374
        if show:
            pylab.show()
375
        
376
        return fig
377

378
379
380
381
382
383
384
385
386
    def find_radius(self,method = 'max', unit='points'):
        '''
        Estimates the radius of the mirror in the xy-plane. Based on FT_find_map_radius.m
        method   - 'max' gives maximal distance from centre to the edge.
                   'xy' returns the distance from centre to edge along x- and y-directions.
                   'area' calculates the area and uses this to estimate the mean radius.
        unit     - 'points' gives radius in data points.
                   'meters' gives radius in meters.
        '''
387

388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
        # Row and column indices of non-NaN
        rIndx, cIndx = self.notNan.nonzero()
        if unit == 'meters' or unit == 'metres' or unit=='Meters' or unit=='Metres':
            # Offsets so that (0,0) is in the center.
            x = self.step_size[0]*(cIndx - self.center[0])
            y = self.step_size[1]*(rIndx - self.center[1])
        else:
            # Offsets so that (0,0) is in the center.
            x = cIndx - self.center[0]
            y = rIndx - self.center[1]
            
        # Maximum distance from center to the edge of the mirror.
        if method=='max' or method=='Max' or method=='MAX':
            r = math.sqrt((x**2 + y**2).max())
        # x and y radii
        elif method=='xy' or method=='XY' or method == 'yx' or method == 'YX':
            r = []
            r.append(abs(x).max()+0.5)
            r.append(abs(y).max()+0.5)
        # Mean radius by dividing the area by pi. Should change this in case
        # x and y step sizes are different.
        elif method=='area' or method=='Area' or method=='AREA':
            if unit == 'meters' or unit == 'metres' or unit=='Meters' or unit=='Metres':
                r = step_size[0]*math.sqrt(len(cIndx)/math.pi)
            else:
                r = math.sqrt(len(cIndx)/math.pi)
                
        return r
416

417
418
419
    
    def remove_curvature(self, method='sphere', Rc0=0, w=None, zOff=None,
                         isCenter=[False,False], zModes = 'all'):
420
421
422
423
424
        '''
        Removes curvature from mirror map by fitting a sphere to
        mirror surface. Based on the file 'FT_remove_curvature_from_mirror_map.m'.
        Rc0      - Initial guess of the radius of curvature [m]
        w        - Beam radius on mirror [m], used for weighting.
425
426
427
428
429
430
431
432
433
434
435
436
        method   - 'sphere' fits a sphere to the mirror map.
                   'zernike' convolves second order Zernike polynomials with the map,
                   and then removes the polynomial with the obtained amplitude from the
                   surface map. Which of the three modes that are fitted is determined by
                   zMods (see below).
        zModes   - 'defocus' uses only Zernike polynomial (n,m) = (2,0), which has a
                   parabolic shape.
                   'astigmatism' uses Zernike polynomials (n,m) = (2,-2) and (n,m) = (2,2).
                   There are astigmatic.
                   'all' uses both defocus and the astigmatic modes.
        zOff     - Initial guess of the z-offset. Only used together with method='sphere'.
                   Generally unnecessary [surfacemap.scaling].
437
        isCenter - 2D-list with booleans. isCenter[0] Determines if the center of the
438
439
440
441
                   sphere is to be fitted (True) or not (False, recommended). If the center is
                   fitted, then isCenter[1] determines if the weights (w!=None) are centered
                   around the fitted center (True) or centered around the center of the
                   data-grid (False, highly recommended).
442
        '''
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562

        if method == 'zernike' or method == 'Zernike':
            import copy
            import pylab
            tmp = copy.deepcopy(self)
            tmp2 = copy.deepcopy(self)
            R = self.find_radius(unit='meters')
            ny,nx = self.data.shape
            # Interpolating for more precise convolution, in case size of map is small. 
            if nx<1500 or ny<1500:
                # Number of extra steps inserted between two adjacent data points.
                N = math.ceil(1500.0/min(nx,ny))
                # New arrays of x-values and y-values
                x = np.arange(tmp.x[0],tmp.x[-1]+tmp.step_size[0]/(N+1),tmp.step_size[0]/N)
                y = np.arange(tmp.y[0],tmp.y[-1]+tmp.step_size[1]/(N+2),tmp.step_size[1]/N)

                
                # Interpolating data
                tmp.interpolate(x,y)
                # Interpolating for notNan.
                g = interp2d(self.x, self.y, self.notNan, kind='linear')
                tmp.notNan = np.round(g(x,y)) # round() or floor() here?! Can't decide...
                # Converting binary to boolean
                tmp.notNan = tmp.notNan==1    
                tmp.plot()

                # Switching code below for code above.
                '''
                # Interpolation functions, for the data (f) and for notNan (g).
                f = interp2d(tmp.x,tmp.y,tmp.data,kind='linear')
                g = interp2d(tmp.x, tmp.y, tmp.notNan, kind='linear')
                # Interpolating
                tmp.data = f(x,y)
                tmp.notNan = np.round(g(x,y)) # round() or floor() here?! Can't decide...
                # Converting binary to boolean
                tmp.notNan = tmp.notNan==1    
                # Setting new step size
                tmp.step_size = (x[1]-x[0],y[1]-y[0])
                # Setting new center
                fx = interp1d(x, np.arange(0,len(x)))
                fy = interp1d(y, np.arange(0,len(y)))
                tmp.center = (fx(0.0), fy(0.0))
                '''
                
                # Radius of new temporary map
                Rnew = tmp.find_radius(unit='meters')
                # Order of Zernike polynomials
                n = 2
                # m values.
                if zModes=='all' or zModes=='All' or zModes=='ALL':
                    mc=[-2, 0, 2] 
                elif zModes=='astigmatism' or zModes=='Astigmatism':
                    mc = [-2, 2]
                elif zModes=='defocus' or zModes=='Defocus':
                    mc = [0]
                    
                # Array where amplitudes will be stored
                A = np.array([])
                # Perform convolution and remove polynomial from map for each chosen
                # zernikeMode
                for m in mc:
                    # Creating Zernike polynomial to convolve with the map
                    # ----------------------------------------------------
                    X,Y,r2 = tmp.createGrid()
                    phi_z = np.arctan2(Y,X)
                    rho_z = np.sqrt(r2)/Rnew
                    Z = znm(n,m,rho_z,phi_z)
                    # ----------------------------------------------------
                    
                    # Convolution (gives amplitude)
                    # ----------------------------------
                    c = (Z[tmp.notNan]*tmp.data[tmp.notNan]).sum()
                    cNorm = (Z[tmp.notNan]*np.conjugate(Z[tmp.notNan])).sum()
                    c = c/cNorm
                    # ----------------------------------
                    # Storing amplitude
                    A = np.append(A,c)
                    # Creating Zernike polynomial for the surface map by
                    # using the obtained amplitudes.
                    # -----------------------------------
                    X,Y,r2 = self.createGrid()
                    phi_m = np.arctan2(Y,X)
                    rho_m = np.sqrt(r2)/R
                    Z = c*znm(n,m,rho_m,phi_m)
                    # -----------------------------------
                    # Removing polynomial from map.
                    self.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
                    
                # Scaling amplitudes  
                A=A*self.scaling
                self.zernike_amp = A
                # Estimating radius of curvature
                if len(mc) !=2:
                    if len(mc)==1:
                        A_rc=A[0]
                    else:
                        A_rc = A[1]
                    self.Rc = (4.0*A_rc**2+R**2)/(4*A_rc)
                else:
                    self.Rc = 0
                
                
                # -------------------------------------------------
                '''
                ss = (tmp.step_size[0]/(N+1), tmp.step_size[1]/(N+1))
                print(ss)
                ss = (x[1]-x[0],y[1]-y[0])
                print(ss)
                map_out.interpolate(x,y)
                #tmp_map.center = (tmp_map.size[0]/2,tmp_map.size[1]/2)
                map_out.plot()
                
                pylab.figure()
                pylab.plot(tmp_map.x)
                pylab.figure()
                pylab.plot(tmp_map.y)
                pylab.show()
                '''
            return tmp

563
        else:
564
565
566
567
        # z-value at centre of the data-grid. Serves as initial guess for deviation
            # from z=0, if no other first guess is given.
            if zOff is None:
                zOff = self.data[round(self.center[1]), round(self.center[0])]
568

569
570
            # If fitting center of the sphere, four variables are fitted. Initial guess
            # of deviation from notNan-data-grid-center: (x0,y0) = (0,0).
571
            if isCenter[0]:
572
573
                p = [Rc0, zOff, 0.0, 0.0]
            # Otherwise two.
574
            else:
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
                p = [Rc0, zOff]

            # Grid with X,Y and r2 = (X^2+Y^2) values. X and Y crosses zero in the center
            # of the xy-plane.
            X,Y,r2 = self.createGrid()

            # Cost-function to minimize.
            def costFunc(p):
                # If the center of the mirror is fitted, four variables are fitted.
                if isCenter[0]:
                    Rc = p[0]
                    zOff = p[1]
                    x0 = p[2]
                    y0 = p[3]
                # Otherwise 2 variables.
590
                else:
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
                    Rc = p[0]
                    zOff = p[1]
                    x0 = 0
                    y0 = 0

                Z = self.createSphere(Rc,X,Y,zOff,x0,y0)

                if w is None:
                    # Mean squared difference between map and the created sphere.
                    res = math.sqrt( ((self.data[self.notNan] -
                                       Z[self.notNan])**2).sum() )/self.notNan.sum()
                else:
                    if isCenter[0] and isCenter[1]:
                        # Weights centered around fitting spehere center. May give weird
                        # results if the mirror deviates much from a sphere.
                        weight = (2/(math.pi*w**2))*np.exp(-2*( (X-x0)**2 + (Y-y0)**2 )/w**2)
                    else:
                        # Weights centered around the center of the mirror xy-plane.
                        weight = (2/(math.pi*w**2))*np.exp(-2*r2/w**2)
                    # Weighted mean squared difference between map and the created sphere.
                    res = math.sqrt( (weight[self.notNan]*((self.data[self.notNan] -
                                       Z[self.notNan])**2)).sum() )/weight[self.notNan].sum()
                return res

            # Using the simplex Nelder-Mead method. This is the same or very
            # similar to the method used in 'FT_remove_curvature_from_mirror_map.m',
            # but there are probably better methods to use.
            out = minimize(costFunc, p, method='Nelder-Mead',tol=1.0e-6)

            # Assigning values to the instance variables
            self.Rc = out['x'][0]
            self.zOff = out['x'][1]

            # If center was fitted, assign new values to instance variable center, and
            # subtract the fitted sphere from the mirror map.
            if isCenter[0]:
                x0 = out['x'][2]
                y0 = out['x'][3]
                # Converts the deviation into a new absolut center in data points.
                self.center = (self.center[0] + x0/self.step_size[0],
                               self.center[1] + y0/self.step_size[1])
                # Creating fitted sphere
                Z = self.createSphere(self.Rc, X, Y, self.zOff, x0, y0)
                # Subtracting sphere from map
                self.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
                return self.Rc, self.zOff, x0, y0
            # Subtracting fitted sphere from mirror map.
            else:
                # Creating fitted sphere
                Z = self.createSphere(self.Rc,X,Y,self.zOff)
                # Subtracting sphere from map
                self.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
                return self.Rc, self.zOff
644
645

    def createSphere(self,Rc,X,Y,zOffset=0,x0=0,y0=0,xTilt=0,yTilt=0,isPlot=False):
646
647
648
649
650
651
652
653
        '''
        Creating spherical surface. Based on 'FT_create_sphere_for_map.m'
        Rc      - Radius of curvature, and center of sphere on z-axis in case zOffset=0 [m].
        X, Y    - Matrices of x,y-values generated by 'X,Y = numpy.meshgrid(xVec,yVec)'.
                  Preferrably created by method 'surfacemap.createGrid()' [m].
        zOffset - Surface center offset from 0. Positive value gives the surface center
                  positive z-coordinate. [self.scaling]
        x0,y0   - The center of the sphere in the xy-plane.
654
        x/yTilt - Tilt of x/y-axis [rad]. E.g. xTilt rotates around the y-axis.
655
656
657
        isPlot  - Plot or not [boolean]. Not recommended when used within optimization
                  algorithm.
        '''
658
659
660
661
662
663
664
665
666
667
668
669
670
        
        # Adjusting for tilts and offset
        Z = zOffset + (X*np.tan(xTilt) + Y*np.tan(yTilt))/self.scaling
        # Adjusting for spherical shape.
        Z = Z + (Rc - np.sqrt(Rc**2 - (X-x0)**2 - (Y-y0)**2))/self.scaling

        if isPlot:
            import pylab
            pylab.clf()
            fig = pylab.figure()
            axes = pylab.pcolormesh(X, Y, Z) #, vmin=zmin, vmax=zmax)
            cbar = fig.colorbar(axes)
            cbar.set_clim(Z.min(), Z.max())
671
            pylab.title('Z_min = ' + str(Z.min()) + '   Z_max = ' + str(Z.max()))
672
673
674
675
676
            pylab.show()
        
        return Z
    
    def createGrid(self):
677
678
679
680
681
682
683
684
685
686
        '''
        Creating grid for fitting spherical surface to the map. Based on
        'FT_create_grid_for_map.m' and 'FT_init_grid.m'. (x,y) = (0,0) is
        placed at the centre of the mirror surface defined by the instance
        variable surfacemap.center.
        
        Returns X,Y,r2
        , where X,Y = numpy.meshgrid(x,y), and r2 = X^2 + Y^2.
        '''
        
687
688
689
690
691
692
693
694
695
696
        yPoints, xPoints = self.data.shape
        # Difference between the data point centre, and the centre of
        # the mirror surface.
        xOffset = (((xPoints-1)/2-self.center[0]) )*self.step_size[0]
        yOffset = (((yPoints-1)/2-self.center[1]) )*self.step_size[1]

        # Creating arrays with values from 0 up to (xyPoints-1).
        x = np.array([n for n in range(xPoints)])
        y = np.array([n for n in range(yPoints)])

697
        # Physical size. Shouldn't this be (points-1)*step_size?
698
699
700
701
702
703
704
705
706
        xSize = xPoints*self.step_size[0]
        ySize = yPoints*self.step_size[1]
        # ...corrected here by subracting one step. Isn't this unnecessary
        # complicated btw?
        xAxis = x*self.step_size[0] + xOffset - (xSize-self.step_size[0])/2
        yAxis = y*self.step_size[1] + yOffset - (ySize-self.step_size[1])/2
        
        X,Y = np.meshgrid(xAxis,yAxis)
        r2 = X**2 + Y**2
707
708
        # r = np.sqrt(X**2 + Y**2)
        # phi = np.arctan2(Y,X)
709
710
711
        return X,Y,r2
        # ----------------------------------------------------------------

712
713
        

714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
class mergedmap:
    """
    A merged map combines multiple surfaces map to form one. Such a map can be used
    for computations of coupling coefficients but it cannot be written to a file to 
    be used with Finesse. For this you must output each map separately.
    
    """
    
    def __init__(self, name, size, center, step_size, scaling):
        
        self.name = name
        self.center = center
        self.step_size = step_size
        self.scaling = scaling
        self.__interp = None
        self._rom_weights = None
        self.__maps = []
731
732
        self.weighting = None
        
733
734
735
736
737
738
739
740
741
742
743
744
    def addMap(self, m):
        self.__maps.append(m)
    
    @property
    def center(self):
        return self.__center
    
    @center.setter
    def center(self, value):
        self.__center = value
        self.__interp = None
    
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
    @property
    def type(self):
        hasR = False
        hasT = False
        
        _type = ""
        
        for m in self.__maps:
            if "reflection" in m.type: hasR = True
            
            if "transmission" in m.type: hasT = True
            
            if "both" in m.type:
                hasR = True
                hasT = True
        
        if hasR and not hasT: _type += "reflection "
        elif hasR and not hasT: _type += "transmission "
        elif hasR and hasT: _type += "both "
        
        return _type
        
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
    @property
    def step_size(self):
        return self.__step_size
    
    @step_size.setter
    def step_size(self, value):
        self.__step_size = value
        self.__interp = None

    @property
    def scaling(self):
        return self.__scaling
    
    @scaling.setter
    def scaling(self, value):
        self.__scaling = value
        self.__interp = None
    
    @property
    def x(self):
787
        return self.step_size[0] * (np.array(range(0, self.size[0])) - self.center[0])
788
789
790
        
    @property
    def y(self):
791
        return self.step_size[1] * (np.array(range(0, self.size[1])) - self.center[1])
792
793
794
795
796
797
798

    @property
    def size(self):
        return self.__maps[0].data.shape
            
    @property
    def offset(self):
799
        return np.array(self.step_size)*(np.array(self.center) - (np.array(self.size)-1)/2.0)
800
801
802
803
804
    
    @property
    def ROMWeights(self):
        return self._rom_weights
    
805
    def z_xy(self, wavelength=1064e-9, direction="reflection_front", nr1=1.0, nr2=1.0):
806
807
808
809
810
811
        
        z_xy = np.ones(self.size, dtype=np.complex128)
        
        for m in self.__maps:
            z_xy *= m.z_xy(wavelength=wavelength, direction=direction, nr1=nr1, nr2=nr2)
            
812
813
814
815
        if self.weighting is None:
            return z_xy
        else:
            return z_xy * self.weighting
816
        
817
    def generateROMWeights(self, EIxFilename, EIyFilename=None, verbose=False, interpolate=False, newtonCotesOrder=8, nr1=1, nr2=1):
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
        if interpolate == True:
            # Use EI nodes to interpolate if we
            with open(EIxFilename, 'rb') as f:
                EIx = pickle.load(f)

            if EIyFilename is None:
                EIy = EIx
            else:
                with open(EIyFilename, 'rb') as f:
                    EIy = pickle.load(f)

            x = EIx.x
            x.sort()
            nx = np.unique(np.hstack((x, -x[::-1])))
        
            y = EIy.x
            y.sort()
            ny = np.unique(np.hstack((y, -y[::-1])))
            
            self.interpolate(nx, ny)
        
839
        w_refl_front, w_refl_back, w_tran_front, w_tran_back = (None, None, None, None)
840
        
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
        if "reflection" in self.type or "both" in self.type:
            w_refl_front = makeWeightsNew(self, EIxFilename, EIyFilename,
                                      verbose=verbose, newtonCotesOrderMapWeight=newtonCotesOrder,
                                      direction="reflection_front")
            
            w_refl_front.nr1 = nr1
            w_refl_front.nr2 = nr2
            
            w_refl_back = makeWeightsNew(self, EIxFilename, EIyFilename,
                                      verbose=verbose, newtonCotesOrderMapWeight=newtonCotesOrder,
                                      direction="reflection_back")
            
            w_refl_back.nr1 = nr1
            w_refl_back.nr2 = nr2

        if "transmission" in self.type or "both" in self.type:                                      
            w_tran_front = makeWeightsNew(self, EIxFilename, EIyFilename,
                                      verbose=verbose, newtonCotesOrderMapWeight=newtonCotesOrder,
                                      direction="transmission_front")

            w_refl_front.nr1 = nr1
            w_refl_front.nr2 = nr2
                                            
            w_tran_back  = makeWeightsNew(self, EIxFilename, EIyFilename,
                                      verbose=verbose, newtonCotesOrderMapWeight=newtonCotesOrder,
                                      direction="transmission_back")
            
            w_refl_back.nr1 = nr1
            w_refl_back.nr2 = nr2
            
        self._rom_weights = MirrorROQWeights(w_refl_front, w_refl_back, w_tran_front, w_tran_back)
        
        return self._rom_weights
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938

    def interpolate(self, nx, ny, **kwargs):
        """
        Interpolates all the maps that are used to fc
        
        Uses scipy.interpolate.interp2d and any keywords arguments are
        passed on to it, thus settings like interpolation type and
        fill values can be set.
        
        The range of nx and ny must contain the value zero so that the
        center point of the map can be set.
        """

        for m in self.__maps:
            m.interpolate(nx, ny)

    def plot(self, mode="absorption", show=True, clabel=None, xlim=None, ylim=None, wavelength=1064e-9):
        
        import pylab
        
        if xlim is not None:
            _x = np.logical_and(self.x<=max(xlim)/100.0, self.x>=min(xlim)/100.0)
            xmin = np.min(np.where(_x == True))
            xmax = np.max(np.where(_x == True))
        else:
            xmin = 0
            xmax = len(self.x)-1
            xlim = [self.x.min()*100, self.x.max()*100]
    
        if ylim is not None:
            _y = np.logical_and(self.y<=max(ylim)/100.0, self.y>=min(ylim)/100.0)
            ymin = np.min(np.where(_y == True))
            ymax = np.max(np.where(_y == True))
        else:
            ymin = 0
            ymax = len(self.y)-1
            ylim = [self.y.min()*100, self.y.max()*100]

        if mode == "absorption":
            # plots how much of field is absorbed
            data = 1-np.abs(self.z_xy())
        elif mode == "meter":
            # plot the phase in terms of meters of displacement
            k = 2*np.pi/wavelength
            data = np.angle(self.z_xy()) / (2*k)
            
        zmin = data[xmin:xmax,ymin:ymax].min()
        zmax = data[xmin:xmax,ymin:ymax].max()

        # 100 factor for scaling to cm
        xrange = 100*self.x
        yrange = 100*self.y

        fig = pylab.figure()
        axes = pylab.pcolormesh(xrange, yrange, data, vmin=zmin, vmax=zmax)
        pylab.xlabel('x [cm]')
        pylab.ylabel('y [cm]')

        if xlim is not None: pylab.xlim(xlim)
        if ylim is not None: pylab.ylim(ylim)

        pylab.title('Merged map {0}, mode {1}'.format(self.name, mode))

        cbar = fig.colorbar(axes)
        cbar.set_clim(zmin, zmax)
939
        
940
941
942
943
944
945
946
947
        if clabel is not None:
            cbar.set_label(clabel)
    
        if show:
            pylab.show()
        
        return fig

948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
class aperturemap(surfacemap):
    
    def __init__(self, name, size, step_size, R):
        surfacemap.__init__(self, name, "absorption both", size, (np.array(size)+1)/2.0, step_size, 1)
        self.R = R
        
    @property
    def R(self):
        return self.__R
    
    @R.setter
    def R(self, value):
        self.__R = value
    
        xx, yy = np.meshgrid(self.x, self.y)
        
        radius = np.sqrt(xx**2 + yy**2)
        
        self.data = np.zeros(self.size)
        self.data[radius > self.R] = 1.0
        
        
class curvedmap(surfacemap):
    
    def __init__(self, name, size, step_size, Rc):
        surfacemap.__init__(self, name, "phase reflection", size, (np.array(size)+1)/2.0, step_size, 1e-6)
        self.Rc = Rc
        
    @property
    def Rc(self):
        return self.__Rc
    
    @Rc.setter
    def Rc(self, value):
        self.__Rc = value
    
        xx, yy = np.meshgrid(self.x, self.y)
        
        Rsq = xx**2 + yy**2
        self.data = (self.Rc - math.copysign(1.0, self.Rc) * np.sqrt(self.Rc**2 - Rsq))/ self.scaling
Daniel Brown's avatar
Daniel Brown committed
988
989

class tiltmap(surfacemap):
990
991
992
993
994
995
996
997
998
999
1000
    """
    To create a tiltmap, plot it and write it to a file to use with Finesse:
        
        tilts = (1e-6, 1e-8) # tilt in (x, y) radians\
        dx = 1e-4
        L = 0.2
        N = L/dx
        
        tmap = tiltmap("tilt", (N, N), (dx,dx), tilts)
        tmap.plot()
        tmap.write_map("mytilt.map")