knm.py 21.2 KB
Newer Older
1
from itertools import combinations_with_replacement as combinations
2
from pykat.optics.gaussian_beams import BeamParam, HG_mode
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
from pykat.optics.romhom import ROMWeights
9
from math import factorial
10
from pykat.math.hermite import hermite
11
12
from scipy.misc import comb
from scipy.integrate import newton_cotes
13
from pykat.math import newton_weights
14

15
import time
16
import pykat.optics.maps
17
18
19
20
21
import os.path
import numpy as np
import pykat
import collections
import math
22
import cmath
23

24
def makeCouplingMatrix(max_order, Neven=True, Nodd=True, Meven=True, Modd=True):
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
    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)
40
41
42
43
44
45
            
            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
            
46
47
            row.append(e)
        
48
49
        
        M.append(np.array(row).squeeze())
50
    
51
    return np.array(M)
52

53
54
def adaptive_knm(mode_in, mode_out, q1, q2, q1y=None, q2y=None, smap=None, delta=(0,0), params={}):
    
55
    if q1y is None:
56
57
        q1y = q1
        
58
    if q2y is None:
59
60
61
62
63
64
65
66
67
        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]")
    
68
69
    Hg_in  = HG_mode(qx=q1, qy=q1y, n=mode_in[0], m=mode_in[1])
    Hg_out = HG_mode(qx=q2, qy=q2y, n=mode_out[0], m=mode_out[1])
70
71
72
73
    
    Nfuncs = []
    Nfuncs.append(0)
    
74
    if smap is not None:
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
147
148
149
150
151
        
        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
    
152
153
def riemann_HG_knm(x, y, mode_in, mode_out, q1, q2, q1y=None, q2y=None,
                     Axy=None, cache=None, delta=(0,0), params={}, newtonCotesOrder=0):
154

155
    if Axy is None:
156
157
        Axy == np.ones((len(x), len(y)))
    
158
    if q1y is None:
159
160
        q1y = q1
        
161
    if q2y is None:
162
163
164
165
        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]")        
166

167
    dx = abs(x[1] - x[0])
168
169
    dy = abs(y[1] - y[0])    
        
170
    if cache is None:
171
172
        Hg_in  = HG_mode(qx=q1, qy=q1y, n=mode_in[0], m=mode_in[1])
        Hg_out = HG_mode(qx=q2, qy=q2y, n=mode_out[0], m=mode_out[1])
173
        
174
        U1 = Hg_in.Unm(x+delta[0], y+delta[1])
175
        U2 = Hg_out.Unm(x,y).conjugate()
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201

        if newtonCotesOrder > 0:
                
            W = newton_cotes(newtonCotesOrder, 1)[0]
            
            if newtonCotesOrder > 1:
                if (len(x) - len(W)) % newtonCotesOrder != 0:
                    raise ValueError("To use Newton-Cotes order {0} the number of data points in x must ensure: (N_x - ({0}+1)) mod {0} == 0".format(newtonCotesOrder) )

                if (len(y) - len(W)) % newtonCotesOrder != 0:
                    raise ValueError("To use Newton-Cotes order {0} the number of data points in y must ensure: (N_y - ({0}+1)) mod {0} == 0".format(newtonCotesOrder) )
                
            wx = np.zeros(x.shape, dtype=np.float64)    
            wy = np.zeros(y.shape, dtype=np.float64)
    
            N = len(W)

            for i in range(0, (len(wx)-1)/newtonCotesOrder): wx[(i*(N-1)):(i*(N-1)+N)] += W
            for i in range(0, (len(wy)-1)/newtonCotesOrder): wy[(i*(N-1)):(i*(N-1)+N)] += W
            
            Wxy = np.outer(wx, wy)
            
        if newtonCotesOrder == 0:
            return dx * dy * np.einsum('ij,ij', Axy, U1*U2)
        else:
            return dx * dy * np.einsum('ij,ij', Axy, U1*U2*Wxy)
202
203
204
205
206
207
208
    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]))

209

210

211
    
212
def __gen_riemann_knm_cache(x, y, couplings, q1, q2, q1y=None, q2y=None, delta=(0,0), params={}):
213
    if q1y is None:
214
215
        q1y = q1
        
216
    if q2y is None:
217
218
        q2y = q2
        
Daniel Brown's avatar
Daniel Brown committed
219
    #it = np.nditer(couplings, flags=['refs_ok','f_index'])
220
221
222
    
    cache = {}
    
Daniel Brown's avatar
Daniel Brown committed
223
224
225
226
    #while not it.finished:
    #    try:
    #        mode_in = [int(it.next()), int(it.next())]
    #        mode_out = [int(it.next()), int(it.next())]
227
    
Daniel Brown's avatar
Daniel Brown committed
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
    couplings = couplings.copy()
    couplings.resize(int(couplings.size/4), 4)
    
    for _ in couplings:
        mode_in = _[:2]
        mode_out = _[2:]        
        strx = "u1[%i,%i]" % (mode_in[0], mode_out[0])
        stry = "u2[%i,%i]" % (mode_in[1], mode_out[1])
        
        #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])

        if strx not in cache:
            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()   
        
        if stry not in cache:
            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()
247
            
Daniel Brown's avatar
Daniel Brown committed
248
249
    #    except StopIteration:
    #        break
250
251
252
    
    return cache
    
253
254
    
    
255
256
def __gen_ROM_HG_knm_cache(weights, couplings, q1, q2, q1y=None, q2y=None):

257
    if q1y is None:
258
259
        q1y = q1
        
260
    if q2y is None:
261
262
        q2y = q2
        
Daniel Brown's avatar
Daniel Brown committed
263
    
264
265
266
267
268
269
270
271
272
273
    
    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
Daniel Brown's avatar
Daniel Brown committed
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307

    couplings = couplings.copy()
    couplings.resize(int(couplings.size/4), 4)
    
    for _ in couplings:
        mode_in = _[:2]
        mode_out = _[2:]

        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:
            cache[strx] = u_star_u(q1.z,   q2.z,  q1.w0,  q2.w0, mode_in[0], mode_out[0], weights.EI["xm"].nodes)

        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)
    
    # it = np.nditer(couplings, flags=['refs_ok','f_index'])
    # 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:
    #             cache[strx] = u_star_u(q1.z,   q2.z,  q1.w0,  q2.w0, mode_in[0], mode_out[0], weights.EI["xm"].nodes)
    #
    #         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)
    #
    #     except StopIteration:
    #         break
308
309
    
    return cache
310
311
312



313
def ROM_HG_knm(weights, mode_in, mode_out, q1, q2, q1y=None, q2y=None, cache=None):
314
    if q1y is None:
315
        q1y = q1
316

317
    if q2y is None:
318
319
320
321
322
323
324
325
326
327
        q2y = q2
    
    # x modes
    n = mode_in[0]
    m = mode_out[0]

    # y modes
    npr = mode_in[1]
    mpr = mode_out[1]
    
328
    if isinstance(weights, ROMWeights):
329
        if cache is None:
330
331
            u_x_nodes = u_star_u(q1.z,   q2.z,  q1.w0,  q2.w0, n,     m,   weights.EIx.nodes)
            u_y_nodes = u_star_u(q1y.z,   q2y.z,  q1y.w0,  q2y.w0, npr, mpr,   weights.EIy.nodes)
332
        
333
334
335
336
337
338
339
            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
340
        
341
342
343
        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
344

345
346
            u_x_nodes = cache[strx]
            u_y_nodes = cache[stry]
347
        
348
349
350
351
352
353
354
            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"]
355
        
356
        u_xy_nodes = np.outer(u_x_nodes, u_y_nodes)
357

358
359
360
361
        n_mod_2 = n % 2
        m_mod_2 = m % 2
        npr_mod_2 = npr % 2
        mpr_mod_2 = mpr % 2
362

363
364
        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)
365

366
367
368
369
370
        elif n_mod_2 != m_mod_2:
            if npr_mod_2 == mpr_mod_2:
                k_ROQ = np.einsum('ij,ij', u_xy_nodes, w_ij_Q1Q4) - np.einsum('ij,ij', u_xy_nodes, w_ij_Q2Q3)
            else:
                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
371

372
373
374
375
376
377
378
        elif npr_mod_2 != mpr_mod_2:
            if n_mod_2 == m_mod_2:
                k_ROQ = np.einsum('ij,ij', u_xy_nodes, w_ij_Q3Q4) - np.einsum('ij,ij', u_xy_nodes,  w_ij_Q1Q2)
            else:
                k_ROQ = np.einsum('ij,ij', u_xy_nodes, w_ij_Q2Q4) - np.einsum('ij,ij', u_xy_nodes, w_ij_Q1Q3)
    
    else:
379
        if cache is None:
380
381
            u_x_nodes = u_star_u(q1.z,   q2.z,  q1.w0,  q2.w0, n,     m,   weights.EIx.nodes)
            u_y_nodes = u_star_u(q1y.z,   q2y.z,  q1y.w0,  q2y.w0, npr, mpr,   weights.EIy.nodes)
382
383
    
            w_ij = weights.w_ij
384
        else:
385
386
387
388
389
            strx = "x[%i,%i]" % (mode_in[0], mode_out[0])
            stry = "y[%i,%i]" % (mode_in[1], mode_out[1])

            u_x_nodes = cache[strx]
            u_y_nodes = cache[stry]
Daniel Brown's avatar
Daniel Brown committed
390
    
391
392
393
394
        u_xy_nodes = np.outer(u_x_nodes, u_y_nodes)

        k_ROQ = np.einsum('ij,ij', u_xy_nodes, w_ij)
         
395
    return k_ROQ
Daniel Brown's avatar
Daniel Brown committed
396

397
__fac_cache = []
Daniel Brown's avatar
Daniel Brown committed
398

399
400
def fac(n):
    global __fac_cache
401
402
403
404
    if len(__fac_cache) == 0:
        return math.factorial(int(n))
    else:
        return __fac_cache[n]
Daniel Brown's avatar
Daniel Brown committed
405

406
407
408
def m_1_pow(n):
    if n % 2 == 0:
        return 1
Daniel Brown's avatar
Daniel Brown committed
409
    else:
410
411
412
413
414
415
416
417
        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
418
        
419
420
421
422
423
    return r


def __S(m, _m, X, _X, F, _F, d=0):
    if m % 2 == 1:
Daniel Brown's avatar
Daniel Brown committed
424
        lim1 = int((m-1)/2)
425
    else:
Daniel Brown's avatar
Daniel Brown committed
426
        lim1 = int(m/2 )
427
428

    if _m % 2 == 1:
Daniel Brown's avatar
Daniel Brown committed
429
        lim2 = int((_m-1)/2)
430
    else:
Daniel Brown's avatar
Daniel Brown committed
431
        lim2 = int(_m/2)
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
    
    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
460
    
461
462
    _F  = K / (2.0 * (1.0+K0))
    F = K.conjugate() / 2.0 
463

464
    Sg = __S(n, _n, X, _X, F, _F)
465

466
467
468
469
470
471
472
473
474
475
476
    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)):
477
    if q1y is None:
478
479
        q1y = q1

480
    if q2y is None:
481
482
483
484
485
486
487
488
489
490
491
492
        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])

493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
def __sqaure_knm_int(n, _n, R):
    # This uses the H_n(x) * H_m(x) product identity to reduce the overlap into
    # a sum of factorial and an integral of a single Hermite with a gaussian function
    # thus making it easier to solve
    expR = math.exp(-(R**2))
    S = 0
    
    for j in range(0, min(n, _n)+1):
        _j1 = _n + n - 2*j - 1
        
        if _j1+1 == 0:
            # for the zeroth order we just have the gaussian integral to solve
            L = math.sqrt(math.pi) * math.erf(R)    
        elif (_j1+1) % 2 == 1:
            # if the Hermite is odd then the integral is always 0 as its anti-symmetric
            L = 0
        else:
            L = 2 * hermite(_j1, 0) - expR * (hermite(_j1, R) - hermite(_j1, -R))
        
        I = 2**j * factorial(j) * comb(n, j) * comb(_n, j)
        
        S += I * L
                
    return S


def square_aperture_HG_knm(mode_in, mode_out, q, R):
    """
    Computes the coupling coefficients for a square aperture.
    """
    # x modes
    n = mode_in[0]
    _n = mode_out[0]

    # y modes
    m = mode_in[1]
    _m = mode_out[1]
    
531
532
    hg1 = HG_mode(q, n=n, m=m)
    hg2 = HG_mode(q, n=_n, m=_m)
533
534
535
        
    kx = hg1.constant_x * hg2.constant_x.conjugate()
    ky = hg1.constant_y * hg2.constant_y.conjugate()
536
     
537
538
539
540
541
542
543
544
545
546
547
    f = q.w / math.sqrt(2)
    R = R / (q.w / math.sqrt(2))
    
    kx *= f
    kx *= __sqaure_knm_int(n, _n, R)
    
    ky *= f
    ky *= __sqaure_knm_int(m, _m, R)
    
    return kx * ky

548
549
550


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={}):
551
    if q1y is None:
552
553
        q1y = q1
        
554
    if q2y is None:
555
556
557
        q2y = q2
        
    assert q1.wavelength == q2.wavelength and q1y.wavelength == q2y.wavelength and q1y.wavelength == q1.wavelength
558
559
560
561
562
563
564
565
    
    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']")
    
566
567
568
    maxtem = 0
    c = couplings.flatten()
    
Daniel Brown's avatar
Daniel Brown committed
569
    for i in range(0, int(c.size/2)):
570
571
572
573
574
575
576
        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))
    
577
    if surface_map is not None:  
578
        Axy = surface_map.z_xy(wavelength=q1.wavelength)
579
580
581
    
        x = surface_map.x
        y = surface_map.y
582
583
584
    
    K = np.zeros((couplings.size/4,), dtype=np.complex128)
    
Daniel Brown's avatar
Daniel Brown committed
585
    #it = np.nditer(couplings, flags=['refs_ok','f_index'])
586
587
    
    i = 0
588
589
590
    
    if profile:
        t0 = time.time()
591
        
592
    if method == "romhom":
593
        if surface_map is None:
594
595
596
597
            raise BasePyKatException("Using 'romhom' method requires a surface map to be specified")
            
        weights = surface_map.ROMWeights
        
598
        if weights is None:
599
            raise BasePyKatException("The ROM weights need to be generated for this map before use.")
600
601

        cache = __gen_ROM_HG_knm_cache(weights, couplings, q1=q1, q2=q2, q1y=q1y, q2y=q2y)
602
        
603
    elif method == "riemann":
604
        if surface_map is None:
605
606
607
            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)
608
    else:
609
        cache = None
610
        weights = None
611
    
612
613
614
615
    if profile:
        cache_time = time.time() - t0
        Ktime = np.zeros((couplings.size/4,), dtype=np.float64)
    
616
617
618
    if verbose:
        p = ProgressBar(maxval=couplings.size, widgets=["Knm (%s): " % method, Percentage(), Bar(), ETA()])
    
Daniel Brown's avatar
Daniel Brown committed
619
620
621
622
623
624
625
626
627
    _couplings = couplings.copy()
    _couplings.resize(int(_couplings.size/4), 4)
    
    for _ in _couplings:
        mode_in = _[:2]
        mode_out = _[2:]

        if profile:
            t0 = time.time()
628
                
Daniel Brown's avatar
Daniel Brown committed
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
        
        
        if method == "riemann":
            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)
        elif method == "romhom":
            K[i] = ROM_HG_knm(weights, mode_in, mode_out, q1=q1, q2=q2, q1y=q1y, q2y=q2y, cache=cache)
        elif method == "bayerhelms":
            K[i] = bayerhelms_HG_knm(mode_in, mode_out, q1=q1, q2=q2, q1y=q1y, q2y=q2y, gamma=gamma)
        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)
        else:
            raise BasePyKatException("method value '%s' not accepted" % method)
        
        if profile:
            Ktime[i] = time.time() - t0
        
        i +=1
        
        if verbose:
            p.update(i*4)
649

650
651
652
653
    if profile:
        return K.reshape(couplings.shape[:-1]), Ktime.reshape(couplings.shape[:-1]), cache_time
    else:
        return K.reshape(couplings.shape[:-1])
654
655


Daniel Brown's avatar
Daniel Brown committed
656
def plot_knm_matrix(couplings, knm, cmap=None, show=True):
657
658
659
660
    import pylab as plt
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
Daniel Brown's avatar
Daniel Brown committed
661
    cax = ax.pcolormesh(abs(knm), cmap=cmap)
662
663
664
665
666
667
    fig.colorbar(cax)
    
    numrows, numcols = knm.shape
    
    c = couplings[:, 0, :2]
    c_ = []
Daniel Brown's avatar
Daniel Brown committed
668

669
    for d in c:
Daniel Brown's avatar
Daniel Brown committed
670
        c_.append("[%i,%i]"%(d[0], d[1]))
671
    
Daniel Brown's avatar
Daniel Brown committed
672
673
674
    A = np.arange(1, len(c)+1)-0.5
    ax.set_xticks(A)
    ax.set_yticks(A)
675
676
677
    ax.set_xticklabels(c_)
    ax.set_yticklabels(c_)
    
Daniel Brown's avatar
Daniel Brown committed
678
679
680
    ax.set_xlim(None, max(A)+0.5)
    ax.set_ylim(None, max(A)+0.5)
    
681
    def format_coord(x, y):
Daniel Brown's avatar
Daniel Brown committed
682
683
        col = int(np.floor(x))
        row = int(np.floor(y))
684
685
686
687
        
        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)
Daniel Brown's avatar
Daniel Brown committed
688
689
        
        return None
690

691
    ax.format_coord = format_coord
692

Daniel Brown's avatar
Daniel Brown committed
693
694
695
696
697
    fig.tight_layout()
    
    if show: plt.show()
    
    return fig