knm.py 16.9 KB
Newer Older
1
from itertools import combinations_with_replacement as combinations
2
from pykat.optics.gaussian_beams import beam_param, HG_beam
3
from pykat.exceptions import BasePyKatException
4
from pykat.optics.romhom import u_star_u
5
from pykat.external.progressbar import ProgressBar, ETA, Percentage, Bar
6
7
from scipy.interpolate import interp2d
from scipy.integrate import dblquad
8

9
import time
10
import pykat.optics.maps
11
12
13
14
15
import os.path
import numpy as np
import pykat
import collections
import math
16
import cmath
17

18
def makeCouplingMatrix(max_order, Neven=True, Nodd=True, Meven=True, Modd=True):
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
    max_order = int(max_order)
    c = []
    for n in range(0, max_order+1):
        for m in range(0, max_order+1):
            if n+m <= max_order:
                c.append([n,m])

    M = []

    for i in c:
        row = []
    
        for j in c:
            e = list(i)
            e.extend(j)
34
35
36
37
38
39
            
            if not Neven and (e[0]-e[2]) % 2 == 0: continue
            if not Nodd and (e[0]-e[2]) % 2 == 1: continue
            if not Meven and (e[1]-e[3]) % 2 == 0: continue
            if not Modd and (e[1]-e[3]) % 2 == 1: continue
            
40
41
            row.append(e)
        
42
43
        
        M.append(np.array(row).squeeze())
44
    
45
    return np.array(M)
46

47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def adaptive_knm(mode_in, mode_out, q1, q2, q1y=None, q2y=None, smap=None, delta=(0,0), params={}):
    
    if q1y == None:
        q1y = q1
        
    if q2y == None:
        q2y = q2
    
    if "epsabs" not in params: params["epsabs"] = 1e-6
    if "epsrel" not in params: params["epsrel"] = 1e-6
    if "usepolar" not in params: params["usepolar"] = False
        
    if len(mode_in) != 2 or len(mode_out) != 2:
        raise BasePyKatException("Both mode in and out should be a container with modes [n m]")
    
    Hg_in  = HG_beam(qx=q1, qy=q1y, n=mode_in[0], m=mode_in[1])
    Hg_out = HG_beam(qx=q2, qy=q2y, n=mode_out[0], m=mode_out[1])
    
    Nfuncs = []
    Nfuncs.append(0)
    
    if smap != None:
        
        if not params["usepolar"]:
            xlims = (min(smap.x), max(smap.x))
            ylims = (min(smap.y), max(smap.y))
            
            def Rfunc(y,x):
                Nfuncs[-1] += len(x)
                return (Hg_in.Unm(x+delta[0], y+delta[1]) * smap.z_xy(x=x,y=y) * Hg_out.Unm(x, y).conjugate()).real
                
            def Ifunc(y,x):
                Nfuncs[-1] += len(x)
                return (Hg_in.Unm(x+delta[0], y+delta[1]) * smap.z_xy(x=x,y=y) * Hg_out.Unm(x, y).conjugate()).imag
            
        else:
            xlims = (0, 2*math.pi)
            ylims = (0, params["aperture"])
            
            def Rfunc(r, phi):
                Nfuncs[-1] += len(x)
                x = r*np.cos(phi)
                y = r*np.sin(phi)
                return (r * Hg_in.Unm(x, y) * smap.z_xy(x=x,y=y) * Hg_out.Unm(x, y).conjugate()).real
                
            def Ifunc(r, phi):
                Nfuncs[-1] += len(x)
                x = r*np.cos(phi)
                y = r*np.sin(phi)
                return (r * Hg_in.Unm(x, y) * smap.z_xy(x=x,y=y) * Hg_out.Unm(x, y).conjugate()).imag
            
    else:
        if not params["usepolar"]:
            _x = 4 * math.sqrt(1+max(mode_in[0],mode_in[1])) * q1.w
            _y = 4 * math.sqrt(1+max(mode_in[0],mode_in[1])) * q1y.w
        
            xlims = (-_x, _x)
            ylims = (-_y, _y)
        
            def Rfunc(y, x):
                Nfuncs[-1] += len(r)
                return (Hg_in.Unm(x+delta[0], y+delta[1]) * Hg_out.Unm(x, y).conjugate()).real
                
            def Ifunc(y,x):
                Nfuncs[-1] += len(r)
                return (Hg_in.Unm(x+delta[0], y+delta[1]) * Hg_out.Unm(x, y).conjugate()).imag
        else:
            xlims = (0, 2*math.pi)
            ylims = (0, params["aperture"])
            
            def Rfunc(r, phi):
                
                if hasattr(r, "__len__"):
                    Nfuncs[-1] += len(r)
                else:
                    Nfuncs[-1] += 1
                    
                x = r*np.cos(phi)
                y = r*np.sin(phi)
                return (r * Hg_in.Unm(x, y) * Hg_out.Unm(x, y).conjugate()).real
                
            def Ifunc(r, phi):
                if hasattr(r, "__len__"):
                    Nfuncs[-1] += len(r)
                else:
                    Nfuncs[-1] += 1
                    
                x = r*np.cos(phi)
                y = r*np.sin(phi)
                return (r * Hg_in.Unm(x, y) * Hg_out.Unm(x, y).conjugate()).imag
    
    R, errR = dblquad(Rfunc, xlims[0], xlims[1], lambda y: ylims[0], lambda y: ylims[1], epsabs=params["epsabs"], epsrel=params["epsrel"])
    I, errI = dblquad(Ifunc, xlims[0], xlims[1], lambda y: ylims[0], lambda y: ylims[1], epsabs=params["epsabs"], epsrel=params["epsrel"])
    
    params["Nfuncs"] = Nfuncs[0]
    params["errors"] = (errR, errI)
    
    return R + 1j * I
    
def riemann_HG_knm(x, y, mode_in, mode_out, q1, q2, q1y=None, q2y=None, Axy=None, cache=None, delta=(0,0), params={}):
147
148
149
150
151
152
153
154
155
156
157
158

    if Axy == None:
        Axy == np.ones((len(x), len(y)))
    
    if q1y == None:
        q1y = q1
        
    if q2y == None:
        q2y = q2
        
    if len(mode_in) != 2 or len(mode_out) != 2:
        raise BasePyKatException("Both mode in and out should be a container with modes [n m]")        
159

160
    dx = abs(x[1] - x[0])
161
162
163
164
165
    dy = abs(y[1] - y[0])    
        
    if cache == None:
        Hg_in  = HG_beam(qx=q1, qy=q1y, n=mode_in[0], m=mode_in[1])
        Hg_out = HG_beam(qx=q2, qy=q2y, n=mode_out[0], m=mode_out[1])
166
    
167
        U1 = Hg_in.Unm(x+delta[0], y+delta[1])
168
169
170
171
172
173
174
175
176
177
        U2 = Hg_out.Unm(x,y).conjugate()
        
        return dx * dy * np.einsum('ij,ij', Axy, U1*U2)
    else:
        
        strx = "u1[%i,%i]" % (mode_in[0], mode_out[0])
        stry = "u2[%i,%i]" % (mode_in[1], mode_out[1])
        
        return dx * dy * np.einsum('ij,ij', Axy, np.outer(cache[strx], cache[stry]))

178

179

180
    
181
def __gen_riemann_knm_cache(x, y, couplings, q1, q2, q1y=None, q2y=None, delta=(0,0), params={}):
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
    if q1y == None:
        q1y = q1
        
    if q2y == None:
        q2y = q2
        
    it = np.nditer(couplings, flags=['refs_ok','f_index'])
    
    cache = {}
    
    while not it.finished:
        try:
            mode_in = [int(it.next()), int(it.next())]
            mode_out = [int(it.next()), int(it.next())]
            
            strx = "u1[%i,%i]" % (mode_in[0], mode_out[0])
            stry = "u2[%i,%i]" % (mode_in[1], mode_out[1])
            
200
201
            #Hg_in  = HG_beam(qx=q1, qy=q1y, n=mode_in[0], m=mode_in[1])
            #Hg_out = HG_beam(qx=q2, qy=q2y, n=mode_out[0], m=mode_out[1])
202
203
    
            if strx not in cache:
204
205
                cache[strx] = u_star_u(q1.z,   q2.z,  q1.w0,  q2.w0, mode_in[0], mode_out[0], x, x+delta[0])    
                #Hg_in.Un(x) * Hg_out.Un(x).conjugate()   
206
207
            
            if stry not in cache:
208
209
                cache[stry] = u_star_u(q1y.z,   q2y.z,  q1y.w0,  q2y.w0, mode_in[1], mode_out[1], y, y+delta[1])    
                #Hg_in.Um(y) * Hg_out.Um(y).conjugate()
210
211
212
213
214
215
            
        except StopIteration:
            break
    
    return cache
    
216
217
    
    
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
def __gen_ROM_HG_knm_cache(weights, couplings, q1, q2, q1y=None, q2y=None):

    if q1y == None:
        q1y = q1
        
    if q2y == None:
        q2y = q2
        
    it = np.nditer(couplings, flags=['refs_ok','f_index'])
    
    cache = {}
    
    cache["w_ij_Q1Q3"] = weights.w_ij_Q1 + weights.w_ij_Q3
    cache["w_ij_Q2Q4"] = weights.w_ij_Q2 + weights.w_ij_Q4
    cache["w_ij_Q1Q2"] = weights.w_ij_Q1 + weights.w_ij_Q2
    cache["w_ij_Q1Q4"] = weights.w_ij_Q1 + weights.w_ij_Q4
    cache["w_ij_Q2Q3"] = weights.w_ij_Q2 + weights.w_ij_Q3
    cache["w_ij_Q3Q4"] = weights.w_ij_Q3 + weights.w_ij_Q4
    cache["w_ij_Q1Q2Q3Q4"] = weights.w_ij_Q1 + weights.w_ij_Q3 + weights.w_ij_Q2 + weights.w_ij_Q4
237
    
238
239
240
241
242
243
244
245
246
    while not it.finished:
        try:
            mode_in = [int(it.next()), int(it.next())]
            mode_out = [int(it.next()), int(it.next())]
            
            strx = "x[%i,%i]" % (mode_in[0], mode_out[0])
            stry = "y[%i,%i]" % (mode_in[1], mode_out[1])
            
            if strx not in cache:
247
                cache[strx] = u_star_u(q1.z,   q2.z,  q1.w0,  q2.w0, mode_in[0], mode_out[0], weights.EI["xm"].nodes)   
248
            
249
250
            if stry not in cache:
                cache[stry] = u_star_u(q1y.z, q2y.z, q1y.w0, q2y.w0, mode_in[1], mode_out[1], weights.EI["ym"].nodes)
251
252
253
254
255
            
        except StopIteration:
            break
    
    return cache
256
257
258



259
260
261
def ROM_HG_knm(weights, mode_in, mode_out, q1, q2, q1y=None, q2y=None, cache=None):
    if q1y == None:
        q1y = q1
262

263
264
265
266
267
268
269
270
271
272
273
    if q2y == None:
        q2y = q2
    
    # x modes
    n = mode_in[0]
    m = mode_out[0]

    # y modes
    npr = mode_in[1]
    mpr = mode_out[1]
    
274
275
276
    if cache == None:
        u_x_nodes = u_star_u(q1.z,   q2.z,  q1.w0,  q2.w0, n,     m,   weights.EI["xm"].nodes)
        u_y_nodes = u_star_u(q1y.z,   q2y.z,  q1y.w0,  q2y.w0, npr, mpr,   weights.EI["ym"].nodes)
277
        
278
279
280
281
282
283
284
285
286
287
288
        w_ij_Q1Q3 = weights.w_ij_Q1 + weights.w_ij_Q3
        w_ij_Q2Q4 = weights.w_ij_Q2 + weights.w_ij_Q4
        w_ij_Q1Q2 = weights.w_ij_Q1 + weights.w_ij_Q2
        w_ij_Q1Q4 = weights.w_ij_Q1 + weights.w_ij_Q4
        w_ij_Q2Q3 = weights.w_ij_Q2 + weights.w_ij_Q3
        w_ij_Q3Q4 = weights.w_ij_Q3 + weights.w_ij_Q4
        w_ij_Q1Q2Q3Q4 = weights.w_ij_Q1 + weights.w_ij_Q2 + weights.w_ij_Q3 + weights.w_ij_Q4
        
    else:
        strx = "x[%i,%i]" % (mode_in[0], mode_out[0])
        stry = "y[%i,%i]" % (mode_in[1], mode_out[1])
Daniel Brown's avatar
Daniel Brown committed
289

290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
        u_x_nodes = cache[strx]
        u_y_nodes = cache[stry]
        
        w_ij_Q1Q3 = cache["w_ij_Q1Q3"]
        w_ij_Q2Q4 = cache["w_ij_Q2Q4"]
        w_ij_Q1Q2 = cache["w_ij_Q1Q2"]
        w_ij_Q1Q4 = cache["w_ij_Q1Q4"]
        w_ij_Q2Q3 = cache["w_ij_Q2Q3"]
        w_ij_Q3Q4 = cache["w_ij_Q3Q4"]
        w_ij_Q1Q2Q3Q4 = cache["w_ij_Q1Q2Q3Q4"]
        

    u_xy_nodes = np.outer(u_x_nodes, u_y_nodes)

    n_mod_2 = n % 2
    m_mod_2 = m % 2
    npr_mod_2 = npr % 2
    mpr_mod_2 = mpr % 2

    if n_mod_2 == m_mod_2 and npr_mod_2 == mpr_mod_2:
        k_ROQ = np.einsum('ij,ij', u_xy_nodes, w_ij_Q1Q2Q3Q4)

    elif n_mod_2 != m_mod_2:
        if npr_mod_2 == mpr_mod_2:
Daniel Brown's avatar
Daniel Brown committed
314
            k_ROQ = np.einsum('ij,ij', u_xy_nodes, w_ij_Q1Q4) - np.einsum('ij,ij', u_xy_nodes, w_ij_Q2Q3)
315
        else:
Daniel Brown's avatar
Daniel Brown committed
316
            k_ROQ = np.einsum('ij,ij', u_xy_nodes, w_ij_Q2Q4) - np.einsum('ij,ij', u_xy_nodes, w_ij_Q1Q3)
Daniel Brown's avatar
Daniel Brown committed
317

318
319
    elif npr_mod_2 != mpr_mod_2:
        if n_mod_2 == m_mod_2:
Daniel Brown's avatar
Daniel Brown committed
320
            k_ROQ = np.einsum('ij,ij', u_xy_nodes, w_ij_Q3Q4) - np.einsum('ij,ij', u_xy_nodes,  w_ij_Q1Q2)
321
        else:
Daniel Brown's avatar
Daniel Brown committed
322
            k_ROQ = np.einsum('ij,ij', u_xy_nodes, w_ij_Q2Q4) - np.einsum('ij,ij', u_xy_nodes, w_ij_Q1Q3)
Daniel Brown's avatar
Daniel Brown committed
323
    
324
325
     
    return k_ROQ
Daniel Brown's avatar
Daniel Brown committed
326

327
__fac_cache = []
Daniel Brown's avatar
Daniel Brown committed
328

329
330
331
332
def fac(n):
    global __fac_cache
    return __fac_cache[n]
    #return math.factorial(int(n))
Daniel Brown's avatar
Daniel Brown committed
333

334
335
336
def m_1_pow(n):
    if n % 2 == 0:
        return 1
Daniel Brown's avatar
Daniel Brown committed
337
    else:
338
339
340
341
342
343
344
345
        return -1


def __Ss(u, _u, F, _F, d=0):
    r = 0
    
    for s in range(0, min(u,_u)+1):
        r += m_1_pow(s) * _F ** (u-s) * _F ** (_u-s) / (fac(2*s+d)*fac(u-s)*fac(_u-s))
Daniel Brown's avatar
Daniel Brown committed
346
        
347
348
349
350
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
382
383
384
385
386
387
    return r


def __S(m, _m, X, _X, F, _F, d=0):
    if m % 2 == 1:
        lim1 = (m-1)/2
    else:
        lim1 = m/2 

    if _m % 2 == 1:
        lim2 = (_m-1)/2
    else:
        lim2 = _m/2
    
    r = 0
    
    for u in range(0, lim1+1):
        for _u in range(0, lim2+1):
            r += m_1_pow(u) * _X**(m-2*u) * X**(_m-2*_u) / ( fac(m-2*u)*fac(_m-2*_u) )   * __Ss(u, _u, F, _F, d=d)
    
    return r
           

def __bayerhelms_kn(n, _n, q1, q2, gamma=0.0):
    
    K0 = (q1.zr - q2.zr)/q2.zr
    K2 = (q1.z - q2.z)/q2.zr
    K = (K0 + 1j*K2)/2.0
    
    Ktilde = abs(K / (1+K))

    if gamma != 0:
        a  = q2.zr * math.sin(gamma) / (cmath.sqrt(1+K.conjugate()) * q2.w0)

        _X = - a * (q2.z/q2.zr - 1j)
        X  = - a * (q2.z/q2.zr + 1j*(1+2*K.conjugate()))
        Ex = cmath.exp(-_X*X / 2.0)
    else:
        _X = 0.0
        X  = 0.0
        Ex = 1.0
388
    
389
390
    _F  = K / (2.0 * (1.0+K0))
    F = K.conjugate() / 2.0 
391

392
    Sg = __S(n, _n, X, _X, F, _F)
393

394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
    if n > 0 and _n > 0:
        Su = __S(n-1, _n-1, X, _X, F, _F, 1)
    else:
        Su = 0
    
    b = m_1_pow(_n) * cmath.sqrt(fac(n) * fac(_n) * (1.0 + K.real)**(n+0.5) * (1.0 + K.conjugate()) ** (-(n+_n+1)))
    
    return b * Ex * (Sg - Su)


def bayerhelms_HG_knm(mode_in, mode_out, q1, q2, q1y=None, q2y=None, gamma=(0,0)):
    if q1y == None:
        q1y = q1

    if q2y == None:
        q2y = q2

    # x modes
    n = mode_in[0]
    _n = mode_out[0]

    # y modes
    m = mode_in[1]
    _m = mode_out[1]

    return __bayerhelms_kn(n,_n, q1, q2, 2*gamma[0]) * __bayerhelms_kn(m, _m, q1y, q2y, 2*gamma[1])

421
422
423


def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riemann", verbose=False, profile=False, gamma=(0,0), delta=(0,0), params={}):
424
425
426
427
428
429
430
    if q1y == None:
        q1y = q1
        
    if q2y == None:
        q2y = q2
        
    assert q1.wavelength == q2.wavelength and q1y.wavelength == q2y.wavelength and q1y.wavelength == q1.wavelength
431
432
433
434
435
436
437
438
    
    couplings = np.array(couplings)
    
    a = couplings.size / 4.0
    
    if int(a) - a != 0:
        raise BasePyKatException("Iterator should be product of 4, each element of coupling array should be [n,m,n',m']")
    
439
440
441
442
443
444
445
446
447
448
449
450
    maxtem = 0
    c = couplings.flatten()
    
    for i in range(0, c.size/2):
        maxtem = max(sum(c[i*2:(i*2+2)]), maxtem)
    
    global __fac_cache
    
    for n in range(0, maxtem+1):
        __fac_cache.append(math.factorial(n))
    
    if surface_map != None:  
451
        Axy = surface_map.z_xy(wavelength=q1.wavelength)
452
453
454
    
        x = surface_map.x
        y = surface_map.y
455
456
457
458
459
460
    
    K = np.zeros((couplings.size/4,), dtype=np.complex128)
    
    it = np.nditer(couplings, flags=['refs_ok','f_index'])
    
    i = 0
461
462
463
    
    if profile:
        t0 = time.time()
464
        
465
466
467
468
469
470
471
472
    if method == "romhom":
        if surface_map == None:
            raise BasePyKatException("Using 'romhom' method requires a surface map to be specified")
            
        weights = surface_map.ROMWeights
        
        if weights == None:
            raise BasePyKatException("The ROM weights need to be generated for this map before use.")
473
474

        cache = __gen_ROM_HG_knm_cache(weights, couplings, q1=q1, q2=q2, q1y=q1y, q2y=q2y)
475
        
476
    elif method == "riemann":
477
478
479
480
        if surface_map == None:
            raise BasePyKatException("Using 'riemann' method requires a surface map to be specified")
            
        cache = __gen_riemann_knm_cache(x, y, couplings, q1, q2, q1y=None, q2y=None, delta=delta)
481
    else:
482
        cache = None
483
        weights = None
484
    
485
486
487
488
    if profile:
        cache_time = time.time() - t0
        Ktime = np.zeros((couplings.size/4,), dtype=np.float64)
    
489
490
491
    if verbose:
        p = ProgressBar(maxval=couplings.size, widgets=["Knm (%s): " % method, Percentage(), Bar(), ETA()])
    
492
493
    while not it.finished:
        try:
494
495
496
            if profile:
                t0 = time.time()
                
497
498
            mode_in = [int(it.next()), int(it.next())]
            mode_out = [int(it.next()), int(it.next())]
499
            
500
            
501
            if method == "riemann":
502
                K[i] = riemann_HG_knm(x, y, mode_in, mode_out, q1=q1, q2=q2, q1y=q1y, q2y=q2y, Axy=Axy, cache=cache, delta=delta)
503
            elif method == "romhom":
504
                K[i] = ROM_HG_knm(weights, mode_in, mode_out, q1=q1, q2=q2, q1y=q1y, q2y=q2y, cache=cache)
505
506
            elif method == "bayerhelms":
                K[i] = bayerhelms_HG_knm(mode_in, mode_out, q1=q1, q2=q2, q1y=q1y, q2y=q2y, gamma=gamma)
507
508
            elif method == "adaptive":
                K[i] = adaptive_knm(mode_in, mode_out, q1=q1, q2=q2, q1y=q1y, q2y=q2y, smap=surface_map, delta=delta, params=params)
509
            else:
510
                raise BasePyKatException("method value '%s' not accepted" % method)
511
512
513
514
            
            if profile:
                Ktime[i] = time.time() - t0
            
515
            i +=1
516
517
518
519
            
            if verbose:
                p.update(i*4)
                
520
521
522
523
                 
        except StopIteration:
            break

524
525
526
527
    if profile:
        return K.reshape(couplings.shape[:-1]), Ktime.reshape(couplings.shape[:-1]), cache_time
    else:
        return K.reshape(couplings.shape[:-1])
528
529


530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
def plot_knm_matrix(couplings, knm):
    import pylab as plt
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    cax = ax.imshow(knm, interpolation='nearest')
    fig.colorbar(cax)
    
    numrows, numcols = knm.shape
    
    c = couplings[:, 0, :2]
    c_ = []
    
    for d in c:
        c_.append("%i,%i"%(d[0], d[1]))
    
    ax.set_xticklabels(c_)
    ax.set_yticklabels(c_)
    
    def format_coord(x, y):
        col = int(x+0.5)
        row = int(y+0.5)
        
        if col>=0 and col<numcols and row>=0 and row<numrows:
            z = knm[row,col]
            return 'x=%s, y=%s, z=%1.4f' % (c_[col], c_[row], z)
        else:
            return 'x=%1.4f, y=%1.4f'%(x, y)
558

559
    ax.format_coord = format_coord
560

561
    plt.show()