gaussian_beams.py 13.2 KB
Newer Older
1
import pykat.exceptions as pkex
2
import numpy as np
3
import math
4
import copy
5
6
7
import warnings
import cmath
from scipy.special import hermite
8
from pykat.SIfloat import SIfloat
9

Daniel Brown's avatar
Daniel Brown committed
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def fit_beam_size(z_data, w_data, w0guess, z0guess, lam=1064e-9, show_plot=True,
                  plotpts=100, title='Beam scan fit', w2zero=True, filename=None):
    """
    Function to fit a beam size (w0 and z) to a 2D intensity array. Plotting can
    be switched on or off.

    Returns best fit of (w0, z)
    
    Contributed by Paul Fulda on 15/01/16
    """
    
    import scipy.optimize
    
    def zw0z02w(z, w0, z0, lam):
        z_R = pl.pi*w0**2/lam
        w = w0*pl.sqrt(1+((z-z0)/z_R)**2)
        return w
    
    popt,pcov = scipy.optimize.curve_fit(lambda z_data, w0, z0: zw0z02w(z_data, w0, z0, lam), z_data, w_data, p0=[w0guess,z0guess])
    w0out=popt[0]
    z0out=popt[1]
    
    z_fit = pl.linspace(min(z_data),max(z_data),plotpts)
    w_fit = zw0z02w(z_fit, w0out, z0out, lam)


    if show_plot or filename is not None:
        import pylab as pl
        
        um=1e6
        
        pl.figure()
        pl.plot(z_data,w_data*um,'bo', label = 'Data')
        pl.plot(z_fit,w_fit*um,'b', label = 'Fit')
        pl.tight_layout
        pl.grid()
        pl.xlabel('Position [m]')
        pl.ylabel('Beam size [$\mu$m]')
        pl.xlim((min(z_data),max(z_data)))
        
        if w2zero:
            pl.ylim((0,max(w_data)*um))
            
        pl.legend(loc=0)
        pl.title(title)
        
        if filename is not None:
            pl.savefig(filename)
            
    return w0out, z0out      
60
    
Daniel Brown's avatar
Daniel Brown committed
61
62
    

63
64
class gauss_param(object):
    """
65
    Use beam_param instead, will be the future name of this object.
66
    
67
68
    Gaussian beam complex parameter
    
Daniel Brown's avatar
Daniel Brown committed
69
    beam_param is effectively a complex number with extra
70
71
    functionality to determine beam parameters.
    
Daniel Brown's avatar
Daniel Brown committed
72
    Defaults to 1064e-9m for wavelength and refractive index 1
73
74
75
    usage:
        q = gauss_param(w0=w0, z=z)
        q = gauss_param(z=z, zr=zr)
76
        q = gauss_param(w=w, rc=rc)
77
        q = gauss_param(q=a) # where a is a complex number
78
79
80
81
82
83
84
        
        or change default wavelength and refractive index with:
        
        q = gauss_param(wavelength, nr, w0=w0, zr=zr)
    """
    
    def __init__(self, wavelength=1064e-9, nr=1, *args, **kwargs):
85
86
87
        if self.__class__ != beam_param:
            warnings.warn("Name changed. Use beam_param instead of gauss_param.")
            
88
        self.__q = None
89
90
        self.__lambda = SIfloat(wavelength)
        self.__nr = SIfloat(nr)
91
92
        
        if len(args) == 1:
Daniel Brown's avatar
Daniel Brown committed
93
            self.__q = complex(args[0])
94
95
96
97
98
        
        elif len(kwargs) == 1:
            if "q" in kwargs:
                self.__q = complex(kwargs["q"])        
            else:
99
                raise pkex.BasePyKatException("Must specify: z and w0 or z and zr or rc and w or q, to define the beam parameter")
100
                
101
102
103
        elif len(kwargs) == 2:        
            
            if "w0" in kwargs and "z" in kwargs:
104
                q = SIfloat(kwargs["z"]) + 1j * math.pi*SIfloat(kwargs["w0"])**2/(self.__lambda/self.__nr)
105
            elif "z" in kwargs and "zr" in kwargs:
106
                q = SIfloat(kwargs["z"]) + 1j * SIfloat(kwargs["zr"]) 
107
            elif "rc" in kwargs and "w" in kwargs:
108
                one_q = 1 / SIfloat(kwargs["rc"]) - 1j * SIfloat(wavelength) / (math.pi * SIfloat(nr) * SIfloat(kwargs["w"])**2)
109
110
                q = 1/one_q
            else:
111
                raise pkex.BasePyKatException("Must specify: z and w0 or z and zr or rc and w or q, to define the beam parameter")
112
113
114
115
116
117
118
                
            self.__q = q
        else:
            raise pkex.BasePyKatException("Incorrect usage for gauss_param constructor")
    
    @property
    def wavelength(self): return self.__lambda
119
120
    @wavelength.setter
    def wavelength(self,value): self.__lambda = SIfloat(value)
121
122
123
124
125
126
127
128
129
130
131
132
133
134
    
    @property
    def nr(self): return self.__nr
    
    @property
    def q(self): return self.__q
    
    @property
    def z(self): return self.__q.real
    
    @property
    def zr(self): return self.__q.imag
    
    @property
135
    def w(self):
136
        return np.abs(self.__q)* np.sqrt(self.__lambda / (self.__nr * math.pi * self.__q.imag))
137
    
138
    def beamsize(self, z=None, wavelength=None, nr=None, w0=None):
139

140
        if z is None:
141
142
143
144
            z = self.z
        else:
            z = np.array(z)
                
145
        if wavelength is None:
146
147
148
149
            wavelength = self.wavelength
        else:
            wavelength = np.array(wavelength)
            
150
        if nr is None:
151
152
153
154
            nr = self.nr
        else:
            nr = np.array(nr)
            
155
        if w0 is None:
156
157
158
159
160
161
162
163
164
            w0 = self.w0
        else:
            w0 = np.array(w0)
        
        q = z + 1j * math.pi * w0 **2 / wavelength
        
        return np.abs(q)*np.sqrt(wavelength / (nr * math.pi * q.imag))
    
    def gouy(self, z=None, wavelength=None, nr=None, w0=None):
165
        if z is None:
166
167
168
169
            z = self.z
        else:
            z = np.array(z)
                
170
        if wavelength is None:
171
172
173
174
            wavelength = self.wavelength
        else:
            wavelength = np.array(wavelength)
            
175
        if nr is None:
176
177
178
179
            nr = self.nr
        else:
            nr = np.array(nr)
            
180
        if w0 is None:
181
182
183
184
185
186
187
188
            w0 = self.w0
        else:
            w0 = np.array(w0)
        
        q = z + 1j * math.pi * w0 **2 / wavelength
        
        return np.arctan2(q.real, q.imag)
        
189
190
    @property
    def w0(self):
191
        return np.sqrt(self.__q.imag * self.__lambda / (self.__nr * math.pi))    
192
193
194

    @property
    def Rc(self):
195
196
197
198
199
200
201
202
203
        def __rc(z, zr):
            if z != 0:
                return z * (1 + (zr/z)**2)
            else:
                return float("inf")
                
        v = np.vectorize(__rc)
        
        return v(self.z, self.zr)
204
    
205
    def curvature(self, z=None, wavelength=None, nr=None, w0=None):
206
        if z is None:
207
208
209
210
            z = self.z
        else:
            z = np.array(z)
                
211
        if wavelength is None:
212
213
214
215
            wavelength = self.wavelength
        else:
            wavelength = np.array(wavelength)
            
216
        if nr is None:
217
218
219
220
            nr = self.nr
        else:
            nr = np.array(nr)
            
221
        if w0 is None:
222
223
224
225
226
227
228
229
            w0 = self.w0
        else:
            w0 = np.array(w0)
        
        q = z + 1j * math.pi * w0 **2 / wavelength
        
        return q.real * (1+ (q.imag/q.real)**2)
        
230
231
232
233
234
235
236
237
238
239
240
241
242
    @staticmethod
    def overlap(q1, q2):
        """
        Computes the projection from one beam parameter to another to give a measure of the
        overlap between the two beam parameters.
        
        This function was provided by Paul Fulda and Antonio Perreca, which came originally
        from Chris Mueller.
        
        Added on 20/4/2015
        """
        return abs(4*q1.imag * q2.imag)/abs(q1.conjugate()-q2)**2
        
243
    def conjugate(self):
Daniel Brown's avatar
Daniel Brown committed
244
        return beam_param(self.__lambda, self.__nr, self.__q.conjugate())
245
    
246
247
248
    def __abs__(self):
        return abs(complex(self.__q))
        
249
250
251
252
253
254
255
    def __complex__(self):
        return self.__q
    
    def __str__(self):
        return str(self.__q)
    
    def __mul__(self, a):
Daniel Brown's avatar
Daniel Brown committed
256
        return beam_param(self.__lambda, self.__nr, self.__q * complex(a))
257
258
    
    def __imul__(self, a):
259
        self.__q *= complex(a)
260
261
262
263
264
        return self
        
    __rmul__ = __mul__
    
    def __add__(self, a):
Daniel Brown's avatar
Daniel Brown committed
265
        return beam_param(self.__lambda, self.__nr, self.__q + complex(a))
266
267
268
269
270
271
272
273
    
    def __iadd__(self, a):
        self.__q += complex(a)
        return self
        
    __radd__ = __add__
    
    def __sub__(self, a):
Daniel Brown's avatar
Daniel Brown committed
274
        return beam_param(self.__lambda, self.__nr, self.__q - complex(a))
275
276
277
278
279
    
    def __isub__(self, a):
        self.__q -= complex(a)
        return self
        
280
    def __rsub__(self, a):
Daniel Brown's avatar
Daniel Brown committed
281
        return beam_param(self.__lambda, self.__nr, complex(a) - self.__q)
282
283
    
    def __div__(self, a):
Daniel Brown's avatar
Daniel Brown committed
284
        return beam_param(self.__lambda, self.__nr, self.__q / complex(a))
285
286
287
288
289
290
    
    def __idiv__(self, a):
        self.__q /= complex(a)
        return self
    
    def __pow__(self, q):
Daniel Brown's avatar
Daniel Brown committed
291
        return beam_param(self.__lambda, self.__nr, self.__q**q)
292
293

    def __neg__(self, q):
Daniel Brown's avatar
Daniel Brown committed
294
        return beam_param(self.__lambda, self.__nr, -self.__q)
295
296
        
    def __eq__(self, q):
297
        if q is None:
298
299
            return False
            
300
301
302
303
304
305
306
307
308
309
        return complex(q) == self.__q
        
    @property
    def real(self): return self.__q.real
    @real.setter
    def real(self, value): self.__q.real = SIfloat(value)
    
    @property
    def imag(self): return self.__q.imag
    @imag.setter
310
    def imag(self, value): self.__q.imag = SIfloat(value)
311
312
313
314

    # reverse beam direction 
    def reverse(self):
        self.__q = -1.0 * self.__q.real + 1j * self.__q.imag
315

316

317
318
class beam_param(gauss_param):
    pass
319
    
320
class HG_beam(object):
321
322
323
324
325
    
    def __init__(self, qx, qy=None, n=0, m=0):
        self._qx = copy.deepcopy(qx)
        self._2pi_qrt = math.pow(2.0/math.pi, 0.25)
        
326
        if qy is None:
327
            self._qy = copy.deepcopy(qx)
328
        else:
329
            self._qy = copy.deepcopy(qy)
330
    
331
332
333
334
        self._n = int(n)
        self._m = int(m)
        self._hn = hermite(self._n)
        self._hm = hermite(self._m)
335
336
337
338
339
340
        self._calc_constants()
        
    @property
    def n(self): return self._n
    @n.setter
    def n(self,value): 
341
        self._n = int(value)
342
        self._calc_constants()
343
        self._hn = hermite(self._n)
344
345
346
347
348

    @property
    def m(self): return self._m
    @m.setter
    def m(self,value): 
349
        self._m = int(value)
350
        self._calc_constants()
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
        self._hm = hermite(self._m)
            
    @property
    def q(self):
        if self._qx.q == self._qy.q:
            return self._qx.q
        else:
            return (self._qx.q, self._qy.q)
    @q.setter
    def q(self, value):
        if value.__class__ == beam_param:
            self._qx = copy.deepcopy(value)
            self._qy = copy.deepcopy(value)
        else:
            self._qx = beam_param(q=complex(value))
            self._qy = beam_param(q=complex(value))
    
    @property
    def qx(self):
        return self._qx.q
        
    @qx.setter
    def qx(self, value):
        if value.__class__ == beam_param:
            self._qx = copy.deepcopy(value)
        else:
            self._qx = beam_param(q=complex(value))
    
    @property
    def qy(self):
        return self._qy.q
382
        
383
384
385
386
387
388
    @qy.setter
    def qy(self, value):
        if value.__class__ == beam_param:
            self._qy = copy.deepcopy(value)
        else:
            self._qy = beam_param(q=complex(value))
389
390
391
392
393
394
395
396
397
    
    @property
    def constant_x(self):
        return self.__xpre_const
        
    @property
    def constant_y(self):
        return self.__ypre_const
        
398
399
    def _calc_constants(self):
        self.__xpre_const = math.pow(2.0/math.pi, 0.25)
400
        self.__xpre_const *= np.sqrt(1.0/(self._qx.w0 * 2**(self._n) * np.math.factorial(self._n)))
401
        self.__xpre_const *= np.sqrt(1j*self._qx.imag / self._qx.q)
402
        self.__xpre_const *= ((1j*self._qx.imag * self._qx.q.conjugate())/(-1j*self._qx.imag * self._qx.q)) ** ( self._n/2.0)
403
404
        
        self.__ypre_const = math.pow(2.0/math.pi, 0.25)
405
        self.__ypre_const *= np.sqrt(1.0/(self._qy.w0 * 2**(self._m) * np.math.factorial(self._m)))
406
        self.__ypre_const *= np.sqrt(1j*self._qy.imag / self._qy.q)
Daniel Brown's avatar
Daniel Brown committed
407
        self.__ypre_const *= ((1j*self._qy.imag * self._qy.q.conjugate())/(-1j*self._qy.imag * self._qy.q)) **(self._m/2.0)
408
409
410
411
412
413
414
415
416
    
        self.__sqrt2_wxz = math.sqrt(2) / self._qx.w
        self.__sqrt2_wyz = math.sqrt(2) / self._qy.w
        
        self.__kx =  2*math.pi / self._qx.wavelength
        self.__ky =  2*math.pi / self._qy.wavelength
        
        self.__invqx = 1/ self._qx.q
        self.__invqy = 1/ self._qy.q
417
    
Daniel Brown's avatar
Daniel Brown committed
418
    def Un(self, x):
419
420
        return self.__xpre_const * self._hn(self.__sqrt2_wxz * x) * np.exp(-0.5j * self.__kx * x*x * self.__invqx)
    
Daniel Brown's avatar
Daniel Brown committed
421
422
    def Um(self, y):
        return self.__ypre_const * self._hm(self.__sqrt2_wyz * y) * np.exp(-0.5j * self.__ky * y*y * self.__invqy)
423
        
424
425
426
427
    def Unm(self, x, y):
        _un = self.Un(x)  
        _um = self.Um(y)
        return np.outer(_un, _um)
428
429
430
431
432
433
434
435
436
437
        
    def plot(self, ndx=100, ndy=100, xscale=4, yscale=4):
        import pylab
        
        xrange = xscale * np.linspace(-self._qx.w, self._qx.w, ndx)
        yrange = yscale * np.linspace(-self._qy.w, self._qy.w, ndy)

        dx = xrange[1]-xrange[0]
        dy = yrange[1]-yrange[0]

438
        data = self.Unm(xrange,yrange)
439
440
441
442
443
444
445

        fig = pylab.figure()
        axes = pylab.imshow(np.abs(data), aspect=dx/dy, extent=[min(xrange),max(xrange),min(yrange),max(yrange)])
        pylab.xlabel('x [m]')
        pylab.ylabel('y [m]')
        cbar = fig.colorbar(axes)
        pylab.show()
446
        
447