maps.py 84.3 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
Andreas Freise's avatar
Andreas Freise committed
19
from pykat.math.zernike import *        
20
from pykat.exceptions import BasePyKatException
21
from copy import deepcopy
22

Daniel Brown's avatar
Daniel Brown committed
23
import numpy as np
24
import math
25
import pickle
26

Daniel Toyra's avatar
Daniel Toyra committed
27

28
29
30
31
32
33
34
35
36
37
38
39
40
41
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)
42
            
Daniel Brown's avatar
Daniel Brown committed
43
class surfacemap(object):
44
    
45
    def __init__(self, name, maptype, size=None, center=None, step_size=1.0, scaling=1.0e-9, data=None,
46
                 notNan=None, zOffset=None, xyOffset=(.0,.0)):
47
48
49
50
        '''
        size, center, step_size, xyOffset are all tuples of the form (x, y),
        i.e., (col, row).
        '''
51
52
53
        
        self.name = name
        self.type = maptype
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
        if data is None:
            if size is None:
                raise BasePyKatException("One of the parameters data or size needs to be specified to create map")
            elif isinstance(size, tuple) or isinstance(size, list):
                if len(size) == 2:
                    self.data = np.zeros(size[::-1])
                else:
                    raise BasePyKatException("Parameter size must have length 2")
            elif isinstance(size, int):
                self.data = np.zeros((size, size))
        else:
            self.data = data

        self.notNan = notNan

69
        # Currently "beam center", i.e., mirror_center + xyOffset. 
70
71
72
        self.center = center
        self.step_size = step_size
        self.scaling = scaling
73
        self._RcRemoved = None
74
        # Offset of fitted sphere. Proably unnecessary to have here.
75
        self.zOffset = zOffset
76
        self.__interp = None
Daniel Toyra's avatar
Daniel Toyra committed
77
        self._zernikeRemoved = {}
Daniel Toyra's avatar
Daniel Toyra committed
78
        self._betaRemoved = None
Daniel Toyra's avatar
Daniel Toyra committed
79
        self._xyOffset = xyOffset
Daniel Brown's avatar
Daniel Brown committed
80

81
        self._rom_weights = None
82
83
84
85
86
87
88
        
    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))
89
            mapfile.write("% Size: {0} {1}\n".format(self.data.shape[0], self.data.shape[1]))
90
91
92
93
94
95
96
97
98
            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")
99

100
101
102
103
104
105
106
107
108
109
110
111
112
    @property
    def xyOffset(self):
        '''
        The offset from the mirror center measured in meters.
        '''
        return self._xyOffset

    @xyOffset.setter
    def xyOffset(self,offset):
        '''
        Give 'offset' is in meters.
        '''
        
113
114
115
116
117
118
119
        # These two lines aren't compatible with recenter(). Changed.
        # x0new = self.center[0] + (-self.xyOffset[0] + offset[0])/self.step_size[0]
        # y0new = self.center[1] + (-self.xyOffset[1] + offset[1])/self.step_size[1]
        
        x0new = self.center[0] + offset[0]/self.step_size[0]
        y0new = self.center[1] + offset[1]/self.step_size[1]
        
120
121
        # Checking so new suggested "center" is within the the mirror surface
        if (x0new>0 and x0new<self.size[0]-1 and y0new>0 and y0new<self.size[1]-1 and
122
            self.notNan[round(y0new), round(x0new)]):
123
124
125
126
127
            
            self.center = (x0new, y0new)
            self._xyOffset = (offset[0], offset[1])
        
        else:
128
            raise BasePyKatException( ('Error in xyOffset: ({:.2e}, {:.2e}) m --> pos = ({:.2f}, {:.2f}) '+
129
                  'is outside mirror surface.').format(offset[0],offset[1],x0new,y0new))
130
                    
Daniel Toyra's avatar
Daniel Toyra committed
131
132
133
134
135
136
137
138
    @property
    def betaRemoved(self):
        return self._betaRemoved
    
    @betaRemoved.setter
    def betaRemoved(self, v):
        self._betaRemoved = v
        
139
    @property
Daniel Toyra's avatar
Daniel Toyra committed
140
141
    def zernikeRemoved(self):
        return self._zernikeRemoved
142
    
Daniel Toyra's avatar
Daniel Toyra committed
143
144
    @zernikeRemoved.setter
    def zernikeRemoved(self, v):
145
146
147
        '''
        v = tuple(m,n,amplitude)
        '''
Daniel Toyra's avatar
Daniel Toyra committed
148
        self._zernikeRemoved["%i%i" % (v[0],v[1])] = v
149

150
151
152
153
154
155
156
157
    @property
    def data(self):
        return self.__data
    
    @data.setter
    def data(self, value):
        self.__data = value
        self.__interp = None
158
159
160
161
162
163
164
165
166
167
168
169
170

    @property
    def notNan(self):
        return self.__notNan

    @notNan.setter
    def notNan(self,value):
        if value is None:
            if not hasattr(self,'notNan'):
                self.__notNan = np.ones(self.size[::-1], dtype=bool)
        else:
            self.__notNan = value
            
171
172
173
174
175
176
    @property
    def center(self):
        return self.__center
    
    @center.setter
    def center(self, value):
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
        if value is None:
            if not hasattr(self,'center'):
                self.__center = self.recenter()
        elif isinstance(value, tuple):
            if len(value) == 2:
                self.__center = value
        elif isinstance(value, list):
            if len(value) == 2:
                self.__center = (value[0],value[1])
        elif isinstance(value, np.ndarray):
            if len(value) == 2:
                self.__center = (value[0],value[1])
        elif isinstance(value, float) or isinstance(value, int):
            self.__center = (value, value)
        else:
            raise BasePyKatException("Invalid format of center, array of length 2 wanted.")
            
194
195
196
197
198
199
200
201
        self.__interp = None
    
    @property
    def step_size(self):
        return self.__step_size
    
    @step_size.setter
    def step_size(self, value):
202
203
204
205
206
207
        if isinstance(value, tuple):
            self.__step_size = value
        elif isinstance(value, list):
            self.__step_size = (value[0],value[1])
        elif isinstance(value, float) or isinstance(value, int):
            self.__step_size = (value,value)
208
209
210
211
212
213
214
215
216
217
218
        self.__interp = None

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

219
220
221
222
    # 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
223
224
    @property
    def x(self):
225
226
        return self.step_size[0] * (np.array(range(0, self.data.shape[1])) - self.center[0])
    
Daniel Brown's avatar
Daniel Brown committed
227
228
    @property
    def y(self):
229
230
231
232
        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
233
234
    @property
    def size(self):
235
236
        return self.data.shape[::-1]
        
237
238
    @property
    def offset(self):
239
        return np.array(self.step_size)*(np.array(self.center) - (np.array(self.size)-1)/2.0)
240
241
242
243
    
    @property
    def ROMWeights(self):
        return self._rom_weights
244
    
245
    def z_xy(self, x=None, y=None, wavelength=1064e-9, direction="reflection_front", nr1=1.0, nr2=1.0):
246
247
248
249
250
251
252
        """
        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.
253
            
254
            wavelength: Wavelength of light in vacuum [m]
255
256
257
258
259
260
261
262
263
            
            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"
                                
264
            nr1       : refractive index on front side
265
            
266
            nr2       : refractive index on back side
267
            
268
269
270
271
        """
        
        assert(nr1 >= 1)
        assert(nr2 >= 1)
272
        
273
        if x is None and y is None:
274
275
            data = self.scaling * self.data
        else:
276
            if self.__interp is None:
277
278
279
                self.__interp = interp2d(self.x, self.y, self.data * self.scaling)
                
            data = self.__interp(x, y)
280
281
        
        if direction == "reflection_front" or direction == "reflection_back":
282
            if "phase" in self.type:
283
                k = math.pi * 2 / wavelength
284
                
285
286
                if direction == "reflection_front":
                    return np.exp(-2j * nr1 * k * data)
287
                else:
288
                    return np.exp(2j * nr2 * k * data[:,::-1])
289
                
290
            elif "absorption" in self.type:
291
                if direction == "reflection_front":
292
293
294
                    return np.sqrt(1.0 - data)
                else:
                    return np.sqrt(1.0 - data[:, ::-1])
295
296
297
298
299
            elif "surface_motion" in self.type:
                if direction == "reflection_front":
                    return data
                else:
                    return data[:, ::-1]
300
301
            else:
                raise BasePyKatException("Map type needs handling")
302
                
303
        elif direction == "transmission_front" or direction == "transmission_back":
304
305
            if "phase" in self.type:
                k = math.pi * 2 / wavelength
306
                
307
308
                if direction == "transmission_front":
                    return np.exp((nr1-nr2) * k * data)
309
                else:
310
                    return np.exp((nr2-nr1) * k * data[:, ::-1])
311
                
312
            elif "absorption" in self.type:
313
                if direction == "transmission_front":
314
315
316
                    return np.sqrt(1.0 - data)
                else:
                    return np.sqrt(1.0 - data[:, ::-1])
317
318
319
                    
            elif "surface_motion" in self.type:
                return np.ones(data.shape)
320
321
            else:
                raise BasePyKatException("Map type needs handling")
322
                
323
        else:
324
            raise ValueError("Direction not valid")
325
        
Daniel Brown's avatar
Daniel Brown committed
326

327
    
328
329
    def generateROMWeights(self, EIxFilename, EIyFilename=None, nr1=1.0, nr2=1.0, verbose=False, interpolate=False, newtonCotesOrder=8):
        
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
        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)
        
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
        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
367

368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
        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
            
387
388
389
    def interpolate(self, nx, ny, **kwargs):
        """
        Interpolates the map for some new x and y values.
390
        
391
392
393
394
395
396
397
398
399
400
        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)
        
401
        data = D(nx-self.offset[0], ny-self.offset[1])
402
        
403
404
        Dx = interp1d(nx, np.arange(0,len(nx)))
        Dy = interp1d(ny, np.arange(0,len(ny)))
405
406
407
408
409
        
        self.center = (Dx(0), Dy(0))
        self.step_size = (nx[1]-nx[0], ny[1]-ny[0])
        self.data = data

410
    # xlim and ylim given in centimeters
411
    def plot(self, show=True, clabel=None, xlim=None, ylim=None, isBlock=False):
412
413
414

        import matplotlib
        import matplotlib.pyplot as plt
415
        
416
        if xlim is not None:
417
            # Sorts out the x-values within xlim
418
419
420
421
            _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:
422
            # Uses the whole available x-range
423
424
425
426
            xmin = 0
            xmax = len(self.x)-1
            xlim = [self.x.min()*100, self.x.max()*100]
    
427
        if ylim is not None:
428
            # Sorts out the y-values within ylim
429
430
431
432
            _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:
433
            # Uses the whole available y-range
434
435
436
            ymin = 0
            ymax = len(self.y)-1
            ylim = [self.y.min()*100, self.y.max()*100]
437
            
438
439
        # CHANGED! y corresponds to rows and x to columns, so this should be
        # correct now. Switch the commented/uncommented things to revert. /DT
440
        # ------------------------------------------------------
441
442
443
444
445
        # 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()
446
        # ------------------------------------------------------
447
        
448
        # 100 factor for scaling to cm
449
450
        xRange = 100*self.x
        yRange = 100*self.y
451
        # xRange, yRange = np.meshgrid(xRange,yRange)
452
        
453
        fig = plt.figure()
454
455
456

        # Important to remember here is that xRange corresponds to column and
        # yRange to row indices of the matrix self.data.
457
        pcm = plt.pcolormesh(xRange, yRange, self.data)
458
        pcm.set_rasterized(True)
459
        
460
461
        plt.xlabel('x [cm]')
        plt.ylabel('y [cm]')
462

463
464
        if xlim is not None: plt.xlim(xlim)
        if ylim is not None: plt.ylim(ylim)
465

466
        plt.title('Surface map {0}, type {1}'.format(self.name, self.type))
467
        
468
        cbar = plt.colorbar()
469
470
        cbar.set_clim(zmin, zmax)
        # cbar.set_clim(-1.86, 1.04)
471
        
472
        if clabel is not None:
473
            cbar.set_label(clabel)
474
    
475
        if show:
476
            plt.show(block=isBlock)
477
        
478
        return fig
479

480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
    def rescaleData(self,scaling):
        '''
        Rescaling data. Assumes that the current map is scaled correctly, i.e.,
        surfacemap.scaling*surfacemap.data would give the data in meters.

        Input: scaling
        scaling  - float number that representing the new scaling. E.g.,
                   scaling = 1.0e-9 means that surfacemap.data is converted to
                   being given in nano meters.
        '''
        if isinstance(scaling, float):
            self.data[self.notNan] = self.data[self.notNan]*self.scaling/scaling
            self.scaling = scaling
        else:
            print('Error: Scaling factor must be of type float, {:s} found'.
                  format(type(scaling)))
            

Daniel Toyra's avatar
Daniel Toyra committed
498
    def rms(self, w=None):
499
500
501
502
503
504
        '''
        Returns the root mean square value of the mirror map, expressed in the unit
        given by self.scaling. If w [m] is specified, only the area inside the radius
        is considered. The area is r<w is centered around the self.center, i.e., the
        beam center (not mirror center).
        '''
Daniel Toyra's avatar
Daniel Toyra committed
505
506
507
508
509
510
511
512
513
514
515
516
517
518
        if w is None:
            return math.sqrt((self.data[self.notNan]**2).sum())/self.notNan.sum()
        else:
            R = self.find_radius(unit='meters')
            if w>=R:
                return math.sqrt((self.data[self.notNan]**2).sum())/self.notNan.sum()
            else:
                rho = self.createPolarGrid()[0]
                inside = np.zeros(self.data.shape,dtype=bool)
                tmp = rho<w
                inside[tmp] = self.notNan[tmp]
                return math.sqrt((self.data[inside]**2).sum())/inside.sum()
            
    def avg(self, w=None):
519
520
521
522
523
524
        '''
        Returns the average height of the mirror map, expressed in the unit
        given by self.scaling. If w [m] is specified, only the area inside the radius
        is considered. The area is r<w is centered around the self.center, i.e., the
        beam center (not mirror center).
        '''
Daniel Toyra's avatar
Daniel Toyra committed
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
        if w is None:
            tot = self.data[self.notNan].sum()
            sgn = np.sign(tot)
            return sgn*math.sqrt(sgn*tot)/self.notNan.sum()
        else:
            R = self.find_radius(unit='meters')
            if w>=R:
                tot = self.data[self.notNan].sum()
                sgn = np.sign(tot)
                return sgn*math.sqrt(sgn*tot)/self.notNan.sum()
            else:
                rho = self.createPolarGrid()[0]
                inside = np.zeros(self.data.shape,dtype=bool)
                tmp = rho<w
                inside[tmp] = self.notNan[tmp]
                tot = self.data[inside].sum()
                sgn = np.sign(tot)
                return sgn*math.sqrt(sgn*tot)/inside.sum()
            
            
545
    def find_radius(self, method = 'max', unit='points'):
546
        '''
547
548
549
550
        Estimates the radius of the mirror in the xy-plane. If xyOffset is different from
        (0,0), then the radius will be the minimum distance from the beam center to the
        nearest mirror edge.
        
551
552
553
554
555
        method   - 'max' gives maximal distance from beam centre to the edge.
                 - 'xy' returns the distance from centre to edge along x- and y-directions.
                   Doesn't make sense if the beam center is offset from the mirror center.
                 - 'area' calculates the area and uses this to estimate the mean radius.
                 - 'min' gives minimal distance from beam center to the edge
556
557
        unit     - 'points' gives radius in data points.
                   'meters' gives radius in meters.
Daniel Toyra's avatar
Daniel Toyra committed
558
559

        Based on Simtools function 'FT_find_map_radius.m' by Charlotte Bond.
560
        '''
561

562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
        # 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':
585
                r = self.step_size[0]*math.sqrt(len(cIndx)/math.pi)
586
587
            else:
                r = math.sqrt(len(cIndx)/math.pi)
Daniel Toyra's avatar
Daniel Toyra committed
588
589
590
591
592
593
594
595
596
597
598
        elif method=='min':
            # Arrays of row and column indices where data is NaN.
            r,c = np.where(self.notNan == False)
            # Sets the map center to index [0,0]
            x = c-self.center[0]
            y = r-self.center[1]
            if unit == 'meters' or unit == 'metres' or unit=='Meters' or unit=='Metres':
                x = x*self.step_size[0]
                y = y*self.step_size[1]
            # Minimum radius squared
            r = math.sqrt((x**2+y**2).min())
599
        return r
600

601

Daniel Toyra's avatar
Daniel Toyra committed
602
    def zernikeConvol(self,n,m=None):
603
        '''
Daniel Toyra's avatar
Daniel Toyra committed
604
605
606
607
608
609
610
611
612
613
614
615
616
        Performs a convolution with the map surface and Zernike polynomials.

        Input: n, m
        n  -  If m (see below) is set to None, n is the maximum order of Zernike
              polynomials that will be convolved with the map. If m is an integer
              or a list, only the exact order n will be convolved.
        m  -  = None. All Zernike polynomials up to and equal to order n will
              be convolved.
              = value [int]. The Zernike polynomial characterised by (n,m) will
              be convolved witht the map.
              = list of integers. The Zernike polynomials of order n with m-values
              in m will be convolved with the map.
             
617
        Returns: A
Daniel Toyra's avatar
Daniel Toyra committed
618
619
620
621
622
623
624
625
        A  -  If m is None:
                List with lists of amplitude coefficients. E.g. A[n] is a list of
                amplitudes for the n:th order with m ranging from -n to n, i.e.,
                A[n][0] is the amplitude of the Zernike polynomial Z(m,n) = Z(-n,n).
              If m is list or int:
                Returns list with with Zernike amplitudes ordered in the same way as
                in the input list m. If m == integer, there is only one value in the
                list A.
626
627
628
629
        
        Based on the simtool function 'FT_zernike_map_convolution.m' by Charlotte
        Bond.
        '''
Daniel Toyra's avatar
Daniel Toyra committed
630
631
632
633
634
635
        
        mVals = []
        if m is None:
            nVals = range(0,n+1)
            for k in nVals:
                mVals.append(range(-k,k+1,2))
636
637
638
639
        elif isinstance(m, list):
            nVals = range(n,n+1)
            mVals.append(m)
        elif isinstance(m, range):
Daniel Toyra's avatar
Daniel Toyra committed
640
641
642
643
644
645
            nVals = range(n,n+1)
            mVals.append(m)
        elif isinstance(m,int):
            nVals = range(n,n+1)
            mVals.append(range(m,m+1))
            
646
647
648
649
650
651
        # Amplitudes
        A = []
        # Radius
        R = self.find_radius(unit='meters')
        
        # Grid for for Zernike-polynomials (same size as map).
Daniel Toyra's avatar
Daniel Toyra committed
652
653
        rho, phi = self.createPolarGrid()
        rho = rho/R
654
        # Loops through all n-values up to n_max
Daniel Toyra's avatar
Daniel Toyra committed
655
656
        for k in range(len(nVals)):
            n = nVals[k]
657
658
            # Loops through all possible m-values for each n.
            tmp = []
Daniel Toyra's avatar
Daniel Toyra committed
659
            for m in mVals[k]:
660
661
662
663
664
665
666
667
668
                # (n,m)-Zernike polynomial for this grid.
                Z = zernike(m, n, rho, phi)
                # Z = znm(n, m, rho, phi)
                # Convolution (unnecessary conjugate? map.data is real numbers.)
                c = (Z[self.notNan]*np.conjugate(self.data[self.notNan])).sum()
                cNorm = (Z[self.notNan]*np.conjugate(Z[self.notNan])).sum()
                c = c/cNorm
                tmp.append(c)
            A.append(tmp)
Daniel Toyra's avatar
Daniel Toyra committed
669
670
671
672
673
        if len(A) == 1:
            A = A[0]
            if len(A)==1:
                A = A[0]
                
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
        return A

    def rmZernikeCurvs(self, zModes='all',interpol=False):
        '''
        Removes curvature from mirror map by fitting Zernike polynomials to the 
        mirror surface. Also stores the estimated radius of curvature
        and the Zernike polynomials and amplitudes.
        
        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).
                   'all' uses both defocus and the astigmatic modes.
        interpol - Boolean where 'true' means that the surface map is interpolated to get
                   more data points.

        Returns [Rc,A]
        Rc       - Estimated radius of curvature removed [m].
        A        - Amplitudes of Zernike polynomials removed [surfacemap.scaling].

        Based on the Simtools functions 'FT_remove_zernike_curvatures_from_map.m',
        'FT_zernike_map_convolution.m', etc. by Charlotte Bond.
        '''
        
        tmp = deepcopy(self)
        ny,nx = self.data.shape
699
700
        # Interpolating for more precise convolution, in case size of map is small. Not sure
        # this is necessary. It isn't used in the latest Simtools tutorials.
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
        if interpol and (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+1),tmp.step_size[1]/N)

            '''
            # --------------------------------------------------------------
            # Interpolating data
            tmp.interpolate(x,y)
            tmp.plot()
            # 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
            # --------------------------------------------------------------
            '''
            # Later when compatible with interpolation function, switch code
            # below for code above (but need to change 
            # --------------------------------------------------------------
            # 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))
            # --------------------------------------------------------------

        # Convolution between the surface map and Zernike polynomials
        # to get amplitudes of the polynomials. Only interested in n=2.
        A = tmp.zernikeConvol(2)[2]
        
        # Radius of mirror.
        R = self.find_radius(unit='meters')
Daniel Toyra's avatar
Daniel Toyra committed
747
748
749
750
751
752
        # Grid in polar coordinates
        rho,phi = self.createPolarGrid()
        rho = rho/R
        
        # Creates the choosen Zernike polynomials and removes them from the
        # mirror map.
753
754
755
756
757
        if zModes=='all' or zModes=='All':
            for k in range(0,3):
                m = (k-1)*2
                Z = A[k]*zernike(m, 2, rho, phi)
                self.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
Daniel Toyra's avatar
Daniel Toyra committed
758
                self.zernikeRemoved = (m, 2, A[k])
759
760
761
762
763
764
            # Estimating radius of curvature
            Rc = znm2Rc([a*self.scaling for a in A], R)
        elif zModes=='astigmatism' or zModes=='Astigmatism':
            for k in range(0,3,2):
                m = (k-1)*2
                Z = A[k]*zernike(m, 2, rho, phi)
765
                self.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
Daniel Toyra's avatar
Daniel Toyra committed
766
                self.zernikeRemoved = (m, 2, A[k])
767
768
769
770
            Rc = znm2Rc([a*self.scaling for a in A[::2]], R)
        elif zModes=='defocus' or zModes=='Defocus':
            Z = A[1]*zernike(0, 2, rho, phi)
            self.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
Daniel Toyra's avatar
Daniel Toyra committed
771
            self.zernikeRemoved = (0, 2, A[1])
772
773
            Rc = znm2Rc(A[1]*self.scaling, R)
        
774
        self.RcRemoved = Rc
775
        
776
        return self.RcRemoved, self.zernikeRemoved
777
778


Daniel Toyra's avatar
Daniel Toyra committed
779
    def rmTilt(self, method='fitSurf', w=None, xbeta=None, ybeta=None, zOff=None):
Daniel Toyra's avatar
Daniel Toyra committed
780
        
Daniel Toyra's avatar
Daniel Toyra committed
781
782
783
784
785
786
787
788
789
        R = self.find_radius(unit='meters')
        
        if method=='zernike' or method=='Zernike':
            A1 = self.rmZernike(1, m = 'all')
            # Calculating equivalent tilts in radians
            xbeta = np.arctan(A1[1]*self.scaling/R)
            ybeta = np.arctan(A1[0]*self.scaling/R)
            self.betaRemoved = (xbeta,ybeta)
            return A1, xbeta, ybeta
Daniel Toyra's avatar
Daniel Toyra committed
790

Daniel Toyra's avatar
Daniel Toyra committed
791
792
793
794
795
796
797
798
        else:
            X,Y = np.meshgrid(self.x,self.y)
            r2 = X**2 + Y**2

            if w is not None:
                weight = 2/(math.pi*w**2) * np.exp(-2*r2[self.notNan]/w**2)

            def f(p):
Daniel Toyra's avatar
Daniel Toyra committed
799
800
801
802
803
804
805
                # This is used in simtools, but gives less good results from what I have seen, also in simtools.
                # I would guess the problem is that the initial guess is 0, and the minimising algorithm would
                # need to go quite far away from zero to find a minumum even slightly away way from zero.
                # I think the idea is to make the algorithm to search for small angles. Of course, if used, this
                # must be converted back after minimisation.
                # p[0] = p[0]*1.0e-9
                # p[1] = p[1]*1.0e-9
Daniel Toyra's avatar
Daniel Toyra committed
806
807
808
809
810
811
                Z = self.createSurface(0,X,Y,p[2],0,0,p[0],p[1])
                if w is None:
                    res = math.sqrt(((self.data[self.notNan]-Z[self.notNan])**2).sum())/self.notNan.sum()
                else:
                    # weight = 2/(math.pi*w**2) * np.exp(-2*r2[self.notNan]/w**2)
                    res = math.sqrt((weight*(self.data[self.notNan]-Z[self.notNan])**2).sum())/weight.sum()
Daniel Toyra's avatar
Daniel Toyra committed
812

Daniel Toyra's avatar
Daniel Toyra committed
813
                return res
Daniel Toyra's avatar
Daniel Toyra committed
814

Daniel Toyra's avatar
Daniel Toyra committed
815
816
817
818
819
820
821
822
            if xbeta is None:
                xbeta = 0
            if ybeta is None:
                ybeta = 0
            if zOff is None:
                zOff = 0

            params = [xbeta,ybeta,zOff]
Daniel Toyra's avatar
Daniel Toyra committed
823

Daniel Toyra's avatar
Daniel Toyra committed
824
            opts = {'xtol': 1.0e-9, 'ftol': 1.0e-9, 'maxiter': 1000, 'maxfev': 1000, 'disp': False}
Daniel Toyra's avatar
Daniel Toyra committed
825
            out = minimize(f, params, method='Nelder-Mead', options=opts)
Daniel Toyra's avatar
Daniel Toyra committed
826
827
828
            if not out['success']:
                msg = '  Warning: ' + out['message'].split('.')[0] + ' (nfev={:d}).'.format(out['nfev'])
                print(msg)
Daniel Toyra's avatar
Daniel Toyra committed
829

Daniel Toyra's avatar
Daniel Toyra committed
830
831
832
            xbeta = out['x'][0]
            ybeta = out['x'][1]
            zOff = out['x'][2]
Daniel Toyra's avatar
Daniel Toyra committed
833

Daniel Toyra's avatar
Daniel Toyra committed
834
835
836
            Z = self.createSurface(0,X,Y,zOff,0,0,xbeta,ybeta)
            self.data[self.notNan] = self.data[self.notNan] - Z[self.notNan]
            self.betaRemoved = (xbeta, ybeta)
Daniel Toyra's avatar
Daniel Toyra committed
837

Daniel Toyra's avatar
Daniel Toyra committed
838
839
840
841
842
843
            # Equivalent Zernike amplitude
            A1 = R*np.tan(np.array([ybeta,xbeta]))/self.scaling
            self.zernikeRemoved = (-1,1,A1[0])
            self.zernikeRemoved = (1,1,A1[1])
            
            return A1,xbeta,ybeta,zOff
Daniel Toyra's avatar
Daniel Toyra committed
844
845
        

846
    def rmSphericalSurf(self, Rc0, w=None, zOff=None, isCenter=[False,False], maxfev=2000):
Daniel Toyra's avatar
Daniel Toyra committed
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
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
        '''
        Fits spherical surface to the mirror map and removes it.

        Inputs: Rc0, w, zOff, isCenter
        Rc       - Initial guess of the Radius of curvature. [m]
        w        - Gaussian weighting parameter. The distance from the center where the
                   weigths have decreased by a factor of exp(-2) compared to the center.
                   Should preferrably be equal to the beam radius at the mirror. [m]
        zOff     - Initial guess of the z-offset. Generally unnecessary. [surfacemap.scaling]
        isCenter - 2D-list with booleans. isCenter[0] Determines if the center of the
                   sphere is to be fitted (True) or not (False, recommended). If the center is
                   fitted, then isCenter[1] determines if the weights (in case w!=None) are
                   centered around the fitted center (True) or centered around the center of
                   the data-grid (False, highly recommended).
                   
        if isCenter[0] == False
        Returns: Rc, zOff
        Rc       - Radius of curvature of the sphere removed from the mirror map. [m]
        zOff     - z-offset of the sphere removed. [surfacemap.scaling]
        isCenter[0] == True
        Returns Rc, zOff, center
        center   - [x0, y0] giving the coordinates of the center of the fitted sphere.
        
        Based on the file 'FT_remove_curvature_from_mirror_map.m' by Charlotte Bond.
        '''
        
        # 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])]

        # If fitting center of the sphere, four variables are fitted. Initial guess
        # of deviation from notNan-data-grid-center: (x0,y0) = (0,0).
        if isCenter[0]:
            p = [Rc0, zOff, 0.0, 0.0]
        # Otherwise two.
        else:
            p = [Rc0, zOff]

        # Grids with X,Y and r2 values. X and Y crosses zero in the center
        # of the xy-plane.
        X,Y = np.meshgrid(self.x, self.y)
        r2 = X**2 + Y**2
Daniel Toyra's avatar
Daniel Toyra committed
890
        
Daniel Toyra's avatar
Daniel Toyra committed
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
        # 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.
            else:
                Rc = p[0]
                zOff = p[1]
                x0 = 0
                y0 = 0

Daniel Toyra's avatar
Daniel Toyra committed
906
            Z = self.createSurface(Rc,X,Y,zOff,x0,y0)
Daniel Toyra's avatar
Daniel Toyra committed
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929

            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.
930
        opts = {'xtol': 1.0e-5, 'ftol': 1.0e-9, 'maxiter': 10000, 'maxfev': maxfev, 'disp': False}
Daniel Toyra's avatar
Daniel Toyra committed
931
932
933
934
935
        out = minimize(costFunc, p, method='Nelder-Mead', options=opts)
        if not out['success']:
            msg = '  Warning: ' + out['message'].split('.')[0] + ' (nfev={:d}).'.format(out['nfev'])
            print(msg)
            
Daniel Toyra's avatar
Daniel Toyra committed
936
        # Assigning values to the instance variables
937
        self.RcRemoved = out['x'][0]
Daniel Toyra's avatar
Daniel Toyra committed
938
939
940
941
942
943
        if self.zOffset is None:
            self.zOffset = 0
        self.zOffset = self.zOffset + out['x'][1]

        # Equivalent Zernike (n=2,m=0) amplitude.
        R = self.find_radius(unit='meters')
944
        A20 = Rc2znm(self.RcRemoved,R)/self.scaling
Daniel Toyra's avatar
Daniel Toyra committed
945
        self.zernikeRemoved = (0,2,A20)
Daniel Toyra's avatar
Daniel Toyra committed
946
947
948
949
950
951
952
953
954
955

        # 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
956
            Z = self.createSurface(self.RcRemoved, X, Y, self.zOffset, x0, y0)
Daniel Toyra's avatar
Daniel Toyra committed
957
958
            # Subtracting sphere from map
            self.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
959
            return self.RcRemoved, self.zOffset, self.center, A20
Daniel Toyra's avatar
Daniel Toyra committed
960
961
962
        # Subtracting fitted sphere from mirror map.
        else:
            # Creating fitted sphere
963
            Z = self.createSurface(self.RcRemoved,X,Y,self.zOffset)
Daniel Toyra's avatar
Daniel Toyra committed
964
965
            # Subtracting sphere from map
            self.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
966
            return self.RcRemoved, self.zOffset, A20
967
968

    def remove_curvature(self, method='zernike', w=None, zOff=None,
969
                         isCenter=[False,False], zModes = 'all'):
970
        '''
Daniel Toyra's avatar
Daniel Toyra committed
971
972
973
974
975
976
977
978
979
980
981
982
983
984
        Removes curvature from mirror map by fitting a sphere to the mirror map, or
        by convolving with second order Zernike polynomials and then extracting
        these polynomials with the obtained amplitudes. 

        Inputs: method, w, zOff, isCenter, zModes
        method   - 'zernike' (default) 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 second order modes that are fitted is
                   determined by zMods (see below).
                 - 'sphere' fits a sphere to the mirror map. The fitting is Gaussian
                   weighted if w (see below) is specifed. 
        w        - Gaussian weighting parameter. The distance from the center where the
                   weigths have decreased by a factor of exp(-2) compared to the center.
                   Should preferrably be equal to the beam radius at the mirror. [m]
985
        zOff     - Initial guess of the z-offset. Only used together with method='sphere'.
Daniel Toyra's avatar
Daniel Toyra committed
986
                   Generally unnecessary. [surfacemap.scaling]
987
        isCenter - 2D-list with booleans. isCenter[0] Determines if the center of the
988
                   sphere is to be fitted (True) or not (False, recommended). If the center is
Daniel Toyra's avatar
Daniel Toyra committed
989
990
991
992
993
994
995
996
                   fitted, then isCenter[1] determines if the weights (in case w!=None) are
                   centered around the fitted center (True) or centered around the center of
                   the data-grid (False, highly recommended).
        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).
                   'all' uses both the defocus and the astigmatic modes.

997
        if method == 'zernike':
Daniel Toyra's avatar
Daniel Toyra committed
998
999
1000
        Returns: Rc, znm
        Rc       - Approximate spherical radius of curvature removed.
        znm      - Dictionary specifying which Zernike polynomials that have been removed
For faster browsing, not all history is shown. View entire blame