maps.py 76.1 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
from copy import deepcopy
22

Daniel Brown's avatar
Daniel Brown committed
23
import numpy as np
24
import math
25
import pickle
26
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)
                    
Daniel Brown's avatar
Daniel Brown committed
42
class surfacemap(object):
43
    def __init__(self, name, maptype, size, center, step_size, scaling, data=None,
Daniel Toyra's avatar
Daniel Toyra committed
44
                 notNan=None, Rc=None, zOffset=None, xyOffset=None):
45
46
47
48
49
50
        
        self.name = name
        self.type = maptype
        self.center = center
        self.step_size = step_size
        self.scaling = scaling
51
        self.notNan = notNan
52
53
        self.Rc = Rc
        self.zOffset = zOffset
54
        self.__interp = None
Daniel Toyra's avatar
Daniel Toyra committed
55
        self._zernikeRemoved = {}
Daniel Toyra's avatar
Daniel Toyra committed
56
        self._betaRemoved = None
Daniel Toyra's avatar
Daniel Toyra committed
57
        self._xyOffset = xyOffset
58
		
59
        if data is None:
Daniel Brown's avatar
Daniel Brown committed
60
61
62
            self.data = np.zeros(size)
        else:
            self.data = data
63
64
65
66
67
            
        if notNan is None:
            self.notNan = np.ones(size)
        else:
            self.notNan = notNan
Daniel Brown's avatar
Daniel Brown committed
68

69
        self._rom_weights = None
70
71
72
73
74
75
76
        
    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))
77
            mapfile.write("% Size: {0} {1}\n".format(self.data.shape[0], self.data.shape[1]))
78
79
80
81
82
83
84
85
86
            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")
87

Daniel Toyra's avatar
Daniel Toyra committed
88

Daniel Toyra's avatar
Daniel Toyra committed
89
90
91
92
93
94
95
96
    @property
    def betaRemoved(self):
        return self._betaRemoved
    
    @betaRemoved.setter
    def betaRemoved(self, v):
        self._betaRemoved = v
        
97
    @property
Daniel Toyra's avatar
Daniel Toyra committed
98
99
    def zernikeRemoved(self):
        return self._zernikeRemoved
100
    
Daniel Toyra's avatar
Daniel Toyra committed
101
102
    @zernikeRemoved.setter
    def zernikeRemoved(self, v):
103
104
105
        '''
        v = tuple(m,n,amplitude)
        '''
Daniel Toyra's avatar
Daniel Toyra committed
106
        self._zernikeRemoved["%i%i" % (v[0],v[1])] = v
107

108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
    @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

144
145
146
147
    # 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
148
149
    @property
    def x(self):
150
151
        return self.step_size[0] * (np.array(range(0, self.data.shape[1])) - self.center[0])
    
Daniel Brown's avatar
Daniel Brown committed
152
153
    @property
    def y(self):
154
155
156
157
        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
158
159
    @property
    def size(self):
160
161
        return self.data.shape[::-1]
        
162
163
    @property
    def offset(self):
164
        return np.array(self.step_size)*(np.array(self.center) - (np.array(self.size)-1)/2.0)
165
166
167
168
    
    @property
    def ROMWeights(self):
        return self._rom_weights
169
    
170
    def z_xy(self, x=None, y=None, wavelength=1064e-9, direction="reflection_front", nr1=1.0, nr2=1.0):
171
172
173
174
175
176
177
        """
        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.
178
            
179
            wavelength: Wavelength of light in vacuum [m]
180
181
182
183
184
185
186
187
188
            
            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"
                                
189
            nr1       : refractive index on front side
190
            
191
            nr2       : refractive index on back side
192
            
193
194
195
196
        """
        
        assert(nr1 >= 1)
        assert(nr2 >= 1)
197
        
198
        if x is None and y is None:
199
200
            data = self.scaling * self.data
        else:
201
            if self.__interp is None:
202
203
204
                self.__interp = interp2d(self.x, self.y, self.data * self.scaling)
                
            data = self.__interp(x, y)
205
206
        
        if direction == "reflection_front" or direction == "reflection_back":
207
            if "phase" in self.type:
208
                k = math.pi * 2 / wavelength
209
                
210
211
                if direction == "reflection_front":
                    return np.exp(-2j * nr1 * k * data)
212
                else:
213
                    return np.exp(2j * nr2 * k * data[:,::-1])
214
                
215
            elif "absorption" in self.type:
216
                if direction == "reflection_front":
217
218
219
                    return np.sqrt(1.0 - data)
                else:
                    return np.sqrt(1.0 - data[:, ::-1])
220
221
222
223
224
            elif "surface_motion" in self.type:
                if direction == "reflection_front":
                    return data
                else:
                    return data[:, ::-1]
225
226
            else:
                raise BasePyKatException("Map type needs handling")
227
                
228
        elif direction == "transmission_front" or direction == "transmission_back":
229
230
            if "phase" in self.type:
                k = math.pi * 2 / wavelength
231
                
232
233
                if direction == "transmission_front":
                    return np.exp((nr1-nr2) * k * data)
234
                else:
235
                    return np.exp((nr2-nr1) * k * data[:, ::-1])
236
                
237
            elif "absorption" in self.type:
238
                if direction == "transmission_front":
239
240
241
                    return np.sqrt(1.0 - data)
                else:
                    return np.sqrt(1.0 - data[:, ::-1])
242
243
244
                    
            elif "surface_motion" in self.type:
                return np.ones(data.shape)
245
246
            else:
                raise BasePyKatException("Map type needs handling")
247
                
248
        else:
249
            raise ValueError("Direction not valid")
250
        
Daniel Brown's avatar
Daniel Brown committed
251

252
    
253
254
    def generateROMWeights(self, EIxFilename, EIyFilename=None, nr1=1.0, nr2=1.0, verbose=False, interpolate=False, newtonCotesOrder=8):
        
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
        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)
        
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
        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
292

293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
        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
            
312
313
314
    def interpolate(self, nx, ny, **kwargs):
        """
        Interpolates the map for some new x and y values.
315
        
316
317
318
319
320
321
322
323
324
325
        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)
        
326
        data = D(nx-self.offset[0], ny-self.offset[1])
327
        
328
329
        Dx = interp1d(nx, np.arange(0,len(nx)))
        Dy = interp1d(ny, np.arange(0,len(ny)))
330
331
332
333
334
        
        self.center = (Dx(0), Dy(0))
        self.step_size = (nx[1]-nx[0], ny[1]-ny[0])
        self.data = data

335
    # xlim and ylim given in centimeters
336
    def plot(self, show=True, clabel=None, xlim=None, ylim=None):
337
338
        import pylab
        
339
        if xlim is not None:
340
            # Sorts out the x-values within xlim
341
342
343
344
            _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:
345
            # Uses the whole available x-range
346
347
348
349
            xmin = 0
            xmax = len(self.x)-1
            xlim = [self.x.min()*100, self.x.max()*100]
    
350
        if ylim is not None:
351
            # Sorts out the y-values within ylim
352
353
354
355
            _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:
356
            # Uses the whole available y-range
357
358
359
            ymin = 0
            ymax = len(self.y)-1
            ylim = [self.y.min()*100, self.y.max()*100]
360
            
361
362
        # CHANGED! y corresponds to rows and x to columns, so this should be
        # correct now. Switch the commented/uncommented things to revert. /DT
363
        # ------------------------------------------------------
364
365
366
367
368
        # 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()
369
        # ------------------------------------------------------
370
        
371
        # 100 factor for scaling to cm
372
373
        xRange = 100*self.x
        yRange = 100*self.y
374
        # xRange, yRange = np.meshgrid(xRange,yRange)
375
        
376
        fig = pylab.figure()
377
378
379

        # 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
380
        pcm = pylab.pcolormesh(xRange, yRange, self.data)
381
        pcm.set_rasterized(True)
382
        
383
384
        pylab.xlabel('x [cm]')
        pylab.ylabel('y [cm]')
385

386
387
        if xlim is not None: pylab.xlim(xlim)
        if ylim is not None: pylab.ylim(ylim)
388
            
389
        pylab.title('Surface map {0}, type {1}'.format(self.name, self.type))
390

Daniel Brown's avatar
updates    
Daniel Brown committed
391
392
        cbar = pylab.colorbar()
        #cbar.set_clim(zmin, zmax)
393
        
394
        if clabel is not None:
395
            cbar.set_label(clabel)
396
    
397
398
        if show:
            pylab.show()
399
        
400
        return fig
401

402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
    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
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
    def rms(self, w=None):
        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):
        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()
            
            
455
    def find_radius(self, method = 'max', unit='points'):
456
        '''
Daniel Toyra's avatar
Daniel Toyra committed
457
        Estimates the radius of the mirror in the xy-plane. 
458
459
460
461
462
        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.
Daniel Toyra's avatar
Daniel Toyra committed
463
464

        Based on Simtools function 'FT_find_map_radius.m' by Charlotte Bond.
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
        # 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)
Daniel Toyra's avatar
Daniel Toyra committed
493
494
495
496
497
498
499
500
501
502
503
        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())
504
        return r
505

506

Daniel Toyra's avatar
Daniel Toyra committed
507
    def zernikeConvol(self,n,m=None):
508
        '''
Daniel Toyra's avatar
Daniel Toyra committed
509
510
511
512
513
514
515
516
517
518
519
520
521
        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.
             
522
        Returns: A
Daniel Toyra's avatar
Daniel Toyra committed
523
524
525
526
527
528
529
530
        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.
531
532
533
534
        
        Based on the simtool function 'FT_zernike_map_convolution.m' by Charlotte
        Bond.
        '''
Daniel Toyra's avatar
Daniel Toyra committed
535
536
537
538
539
540
541
542
543
544
545
546
547
        
        mVals = []
        if m is None:
            nVals = range(0,n+1)
            for k in nVals:
                mVals.append(range(-k,k+1,2))
        elif isinstance(m,list):
            nVals = range(n,n+1)
            mVals.append(m)
        elif isinstance(m,int):
            nVals = range(n,n+1)
            mVals.append(range(m,m+1))
            
548
549
550
551
552
553
        # 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
554
555
        rho, phi = self.createPolarGrid()
        rho = rho/R
556
        # Loops through all n-values up to n_max
Daniel Toyra's avatar
Daniel Toyra committed
557
558
        for k in range(len(nVals)):
            n = nVals[k]
559
560
            # Loops through all possible m-values for each n.
            tmp = []
Daniel Toyra's avatar
Daniel Toyra committed
561
            for m in mVals[k]:
562
563
564
565
566
567
568
569
570
                # (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
571
572
573
574
575
        if len(A) == 1:
            A = A[0]
            if len(A)==1:
                A = A[0]
                
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
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
644
645
646
647
        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
        # Interpolating for more precise convolution, in case size of map is small. 
        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
648
649
650
651
652
653
        # Grid in polar coordinates
        rho,phi = self.createPolarGrid()
        rho = rho/R
        
        # Creates the choosen Zernike polynomials and removes them from the
        # mirror map.
654
655
656
657
658
        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
659
                self.zernikeRemoved = (m, 2, A[k])
660
661
662
663
664
665
666
            # 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)
                smap.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
Daniel Toyra's avatar
Daniel Toyra committed
667
                self.zernikeRemoved = (m, 2, A[k])
668
669
670
671
            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
672
            self.zernikeRemoved = (0, 2, A[1])
673
674
675
676
            Rc = znm2Rc(A[1]*self.scaling, R)
        
        self.Rc = Rc
        
Daniel Toyra's avatar
Daniel Toyra committed
677
        return self.Rc, self.zernikeRemoved
678
679


Daniel Toyra's avatar
Daniel Toyra committed
680
    def rmTilt(self, method='fitSurf', w=None, xbeta=None, ybeta=None, zOff=None):
Daniel Toyra's avatar
Daniel Toyra committed
681
        
Daniel Toyra's avatar
Daniel Toyra committed
682
683
684
685
686
687
688
689
690
        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
691

Daniel Toyra's avatar
Daniel Toyra committed
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
        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):
                # This is used in simtools, why?
                #p[0] = p[0]*1.0e-9
                #p[1] = p[1]*1.0e-9
                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
709

Daniel Toyra's avatar
Daniel Toyra committed
710
                return res
Daniel Toyra's avatar
Daniel Toyra committed
711

Daniel Toyra's avatar
Daniel Toyra committed
712
713
714
715
716
717
718
719
            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
720

Daniel Toyra's avatar
Daniel Toyra committed
721
722
            opts = {'xtol': 1.0e-8, 'ftol': 1.0e-8, 'maxiter': 2000, 'disp': False}
            out = minimize(f, params, method='Nelder-Mead', options=opts)
Daniel Toyra's avatar
Daniel Toyra committed
723

Daniel Toyra's avatar
Daniel Toyra committed
724
725
726
            xbeta = out['x'][0]
            ybeta = out['x'][1]
            zOff = out['x'][2]
Daniel Toyra's avatar
Daniel Toyra committed
727

Daniel Toyra's avatar
Daniel Toyra committed
728
729
730
            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
731

Daniel Toyra's avatar
Daniel Toyra committed
732
733
734
735
736
737
            # 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
738
739
        

Daniel Toyra's avatar
Daniel Toyra committed
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
    def rmSphericalSurf(self, Rc0, w=None, zOff=None, isCenter=[False,False]):
        '''
        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

        # 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
800
            Z = self.createSurface(Rc,X,Y,zOff,x0,y0)
Daniel Toyra's avatar
Daniel Toyra committed
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827

            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]
Daniel Toyra's avatar
Daniel Toyra committed
828
829
830
831
832
833
834
835
        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')
        A20 = Rc2znm(self.Rc,R)/self.scaling
        self.zernikeRemoved = (0,2,A20)
Daniel Toyra's avatar
Daniel Toyra committed
836
837
838
839
840
841
842
843
844
845

        # 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
Daniel Toyra's avatar
Daniel Toyra committed
846
            Z = self.createSurface(self.Rc, X, Y, self.zOffset, x0, y0)
Daniel Toyra's avatar
Daniel Toyra committed
847
848
            # Subtracting sphere from map
            self.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
Daniel Toyra's avatar
Daniel Toyra committed
849
            return self.Rc, self.zOffset, self.center, A20
Daniel Toyra's avatar
Daniel Toyra committed
850
851
852
        # Subtracting fitted sphere from mirror map.
        else:
            # Creating fitted sphere
Daniel Toyra's avatar
Daniel Toyra committed
853
            Z = self.createSurface(self.Rc,X,Y,self.zOffset)
Daniel Toyra's avatar
Daniel Toyra committed
854
855
            # Subtracting sphere from map
            self.data[self.notNan] = self.data[self.notNan]-Z[self.notNan]
Daniel Toyra's avatar
Daniel Toyra committed
856
            return self.Rc, self.zOffset, A20
857
858

    def remove_curvature(self, method='zernike', w=None, zOff=None,
859
                         isCenter=[False,False], zModes = 'all'):
860
        '''
Daniel Toyra's avatar
Daniel Toyra committed
861
862
863
864
865
866
867
868
869
870
871
872
873
874
        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]
875
        zOff     - Initial guess of the z-offset. Only used together with method='sphere'.
Daniel Toyra's avatar
Daniel Toyra committed
876
                   Generally unnecessary. [surfacemap.scaling]
877
        isCenter - 2D-list with booleans. isCenter[0] Determines if the center of the
878
                   sphere is to be fitted (True) or not (False, recommended). If the center is
Daniel Toyra's avatar
Daniel Toyra committed
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
                   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.

        if method == 'zernike'
        Returns: Rc, znm
        Rc       - Approximate spherical radius of curvature removed.
        znm      - Dictionary specifying which Zernike polynomials that have been removed
                   as well as their amplitudes.
        if method == 'sphere' and 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]
        if method == 'sphere' and isCenter[0] == True
        Returns Rc, zOff, center
        center   - [x0, y0] giving the coordinates of the center of the fitted sphere. 
899
        '''
900
901
        
        
902
        if method == 'zernike' or method == 'Zernike':
903
            Rc, znm = self.rmZernikeCurvs(zModes)
Daniel Toyra's avatar
Daniel Toyra committed
904
            return Rc, znm
905
906
        elif method == 'sphere' or method == 'Sphere':
            smap = deepcopy(self)
Daniel Toyra's avatar
Daniel Toyra committed
907
            # Using Zernike polynomial (m,n) = (0,2) to estimate radius of curvature.
908
            Rc0 = smap.rmZernikeCurvs('defocus')[0]
909
            if isCenter[0]:
Daniel Toyra's avatar
Daniel Toyra committed
910
911
                Rc, zOff, center, A20 = self.rmSphericalSurf(Rc0, w, zOff, isCenter)
                return Rc, zOff, self.center, A20
912
            else:
Daniel Toyra's avatar
Daniel Toyra committed
913
914
                Rc, zOff, A20 = self.rmSphericalSurf(Rc0, w, zOff, isCenter)
                return Rc, zOff, A20
915

Daniel Toyra's avatar
Daniel Toyra committed
916
            
917
918
    def recenter(self):
        '''
Daniel Toyra's avatar
Daniel Toyra committed
919
920
921
        Setting center of mirror to be in the middle of the notNan field. 
        
        Based on 'FT_recenter_mirror_map.m' by Charlotte Bond.
922
923
924
925
926
927
928
929
930
931
932
933
934
        '''
        # Row and column indices with non-NaN elements
        rIndx, cIndx = self.notNan.nonzero()
        # Finding centres
        x0 = float(cIndx.sum())/len(cIndx)
        y0 = float(rIndx.sum())/len(rIndx)
        self.center = tuple([x0,y0])
        return self.center
        # -------------------------------------------------
        
    def removePeriphery(self):
        '''
        Finds the NaN elements closest to the map center, and sets all
Daniel Toyra's avatar
Daniel Toyra committed
935
936
937
        elements further out to NaN. Then cropping and recentering.
        
        Based on FT_remove_elements_outside_map.m by Charlotte Bond.
938
939
940
941
942
943
944
945
946
        '''
        # 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]
        # Minimum radius squared measured in data points.
        r2 = (x**2+y**2).min()
        # Grid with the same dimensions as the map.
Daniel Toyra's avatar
Daniel Toyra committed
947
948
        X,Y = np.meshgrid(self.x, self.y)
        # Grid containing radial distances squared measured from the center.
949
950
951
        R2 = (X/self.step_size[0])**2 + (Y/self.step_size[1])**2
        # Matrix with True=='outside mirror surface'
        outs = R2>=r2
Daniel Toyra's avatar
Daniel Toyra committed
952
        # Setting notNan-matrix to false outside the mirror.
953
        self.notNan[outs] = False
Daniel Toyra's avatar
Daniel Toyra committed
954
        # Setting data matrix to 0 outside the mirror.
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
        self.data[outs] = 0.0

        # Removing unnecessary data points.
        # --------------------------------------
        # Radius inside data is kept. Don't know why R is set this way,
        # but trusting Dr. Bond.
        R = round(math.ceil(2.0*math.sqrt(r2)+6.0)/2.0)
        if 2*R<self.data.shape[0] and 2*R<self.data.shape[1]:
            x0 = round(self.center[0])
            y0 = round(self.center[1])

            self.data = self.data[y0-R:y0+R+1, x0-R:x0+R+1]
            self.notNan = self.notNan[y0-R:y0+R+1, x0-R:x0+R+1]
        
        self.recenter()
        
        
    
Daniel Toyra's avatar
Daniel Toyra committed
973
    def createSurface(self,Rc,X,Y,zOffset=0,x0=0,y0=0,xTilt=0,yTilt=0,isPlot=False):
974
        '''
Daniel Toyra's avatar
Daniel Toyra committed
975
        Creating surface.
Daniel Toyra's avatar
Daniel Toyra committed
976
977
978
979
        
        Inputs: Rc, X, Y, zOffset=0, x0=0, y0=0, xTilt=0, yTilt=0, isPlot=False
        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)'. [m]
980
981
982
        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.
983
        x/yTilt - Tilt of x/y-axis [rad]. E.g. xTilt rotates around the y-axis.
984
985
        isPlot  - Plot or not [boolean]. Not recommended when used within optimization
                  algorithm.
Daniel Toyra's avatar
Daniel Toyra committed
986
987
988
989
990

        Returns: Z
        Z       - Matrix of surface height values. [self.scaling]
        
        Based on Simtools function 'FT_create_sphere_for_map.m' by Charlotte Bond.
991
        '''
992
993
994
995
        
        # Adjusting for tilts and offset
        Z = zOffset + (X*np.tan(xTilt) + Y*np.tan(yTilt))/self.scaling
        # Adjusting for spherical shape.
Daniel Toyra's avatar
Daniel Toyra committed
996
997
998
        if Rc !=0 and Rc is not None:
            Z = Z + (Rc - np.sign(Rc)*np.sqrt(Rc**2 - (X-x0)**2 - (Y-y0)**2))/self.scaling
        
999
1000
        if isPlot:
            import pylab
For faster browsing, not all history is shown. View entire blame