romhom.py 13.5 KB
Newer Older
1
import math
2
3
4
import os.path
import pykat
import collections
Daniel Brown's avatar
Daniel Brown committed
5

6
from pykat.external.progressbar import ProgressBar, ETA, Percentage, Bar
7
from itertools import combinations_with_replacement as combinations
8
from pykat.optics.gaussian_beams import beam_param, HG_beam
9
10
from scipy.linalg import inv
from math import factorial
11
from pykat.math.hermite import *
12
13

import numpy as np
14
15
16
17
18
19
20

EmpiricalInterpolant = collections.namedtuple('EmpiricalInterpolant', 'B nodes node_indices limits x')
ReducedBasis = collections.namedtuple('ReducedBasis', 'RB limits x')
ROMLimits = collections.namedtuple('ROMLimits', 'zmin zmax w0min w0max max_order')

class ROMWeights:
    
Daniel Brown's avatar
Daniel Brown committed
21
    def __init__(self, w_ij_Q1, w_ij_Q2, w_ij_Q3, w_ij_Q4, EI, limits):
22
23
24
25
        self.w_ij_Q1 = w_ij_Q1
        self.w_ij_Q2 = w_ij_Q2
        self.w_ij_Q3 = w_ij_Q3
        self.w_ij_Q4 = w_ij_Q4
Daniel Brown's avatar
Daniel Brown committed
26
27
        
        self.EI = EI
28
29
30
        self.limits = limits
        
    def writeToFile(self, filename):
Daniel Brown's avatar
Daniel Brown committed
31
32
33
34
35
36
        """
        Writes this map's ROM weights to a file
        that can be used with Finesse. The filename
        is appended with '.rom' internally.
        """
        f = open(filename + ".rom", 'w+')
37
        
Daniel Brown's avatar
Daniel Brown committed
38
39
40
41
42
        f.write("zmin=%16.16e\n" % self.limits.zmin)
        f.write("zmax=%16.16e\n" % self.limits.zmax)
        f.write("w0min=%16.16e\n" % self.limits.w0min)
        f.write("w0max=%16.16e\n" % self.limits.w0max)
        f.write("maxorder=%i\n" % self.limits.max_order)
43
        
44
        f.write("xnodes=%i\n" % len(self.EI["x"].nodes))
Daniel Brown's avatar
Daniel Brown committed
45
        
46
        for v in self.EI["x"].nodes.flatten():
Daniel Brown's avatar
Daniel Brown committed
47
48
            f.write("%s\n" % repr(float(v)))
        
49
        f.write("ynodes=%i\n" % len(self.EI["y"].nodes))
Daniel Brown's avatar
Daniel Brown committed
50
        
51
        for v in self.EI["y"].nodes.flatten():
Daniel Brown's avatar
Daniel Brown committed
52
53
            f.write("%s\n" % repr(float(v)))
            
54
55
56
        f.write(repr(self.w_ij_Q1.shape) + "\n")
        
        for v in self.w_ij_Q1.flatten():
Daniel Brown's avatar
Daniel Brown committed
57
            f.write("%s\n" % repr(complex(v))[1:-1])
58
59
60
61
        
        f.write(repr(self.w_ij_Q2.shape) + "\n")
        
        for v in self.w_ij_Q2.flatten():
Daniel Brown's avatar
Daniel Brown committed
62
            f.write("%s\n" % repr(complex(v))[1:-1])
63
64
65
66
        
        f.write(repr(self.w_ij_Q3.shape) + "\n")
        
        for v in self.w_ij_Q3.flatten():
Daniel Brown's avatar
Daniel Brown committed
67
            f.write("%s\n" % repr(complex(v))[1:-1])
68
69
70
71
        
        f.write(repr(self.w_ij_Q4.shape) + "\n")
        
        for v in self.w_ij_Q4.flatten():
Daniel Brown's avatar
Daniel Brown committed
72
            f.write("%s\n" % repr(complex(v))[1:-1])
73
74
75
            
        f.close()

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
class ROMWeightsFull:
    
    def __init__(self, w_ij, EI, limits):
        self.w_ij = w_ij
        
        self.EI = EI
        self.limits = limits
        
    def writeToFile(self, filename):
        """
        Writes this map's ROM weights to a file
        that can be used with Finesse. The filename
        is appended with '.rom' internally.
        """
        f = open(filename + ".rom", 'w+')
        
        f.write("zmin=%16.16e\n" % self.limits.zmin)
        f.write("zmax=%16.16e\n" % self.limits.zmax)
        f.write("w0min=%16.16e\n" % self.limits.w0min)
        f.write("w0max=%16.16e\n" % self.limits.w0max)
        f.write("maxorder=%i\n" % self.limits.max_order)
        
        f.write("xnodes=%i\n" % len(self.EI["x"].nodes))
        
        for v in self.EI["x"].nodes.flatten():
            f.write("%s\n" % repr(float(v)))
        
        f.write("ynodes=%i\n" % len(self.EI["y"].nodes))
        
        for v in self.EI["y"].nodes.flatten():
            f.write("%s\n" % repr(float(v)))
            
        f.write(repr(self.w_ij.shape) + "\n")
        
        for v in self.w_ij.flatten():
            f.write("%s\n" % repr(complex(v))[1:-1])
        
        f.close()
        
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
152
153
154
155
156
157
def combs(a, r):
    """
    Return successive r-length combinations of elements in the array a.
    Should produce the same output as array(list(combinations(a, r))), but 
    faster.
    """
    a = np.asarray(a)
    dt = np.dtype([('', a.dtype)]*r)
    b = np.fromiter(combinations(a, r), dt)
    return b.view(a.dtype).reshape(-1, r)

def project_onto_basis(integration_weights, e, h, projections, proj_coefficients, idx):

    for j in range(len(h)):
        proj_coefficients[idx][j] = np.vdot(integration_weights* e[idx], h[j])
        projections[j] += proj_coefficients[idx][j]*e[idx]

    return projections
    
def B_matrix(invV, e):
    return np.inner(invV.T, e[0:(invV.shape[0])].T)

def emp_interp(B_matrix, func, indices):
    # B : RB matrix
    assert B_matrix.shape[0] == len(indices)
    interpolant = np.inner(func[indices].T, B_matrix.T)

    return interpolant 
   
    
def w(w0, im_q, re_q):
    return w0 * np.sqrt( 1 + (re_q / im_q)**2. )


def u(re_q1, w0_1, n1, x):
    
    im_q1 = np.pi*w0_1**2 / 1064e-9
    q_z1 = re_q1 + 1j*im_q1

    A_n1 = (2./np.pi)**(1./4.) * (1./((2.**n1)*factorial(n1)*w0_1))**(1./2.) * (im_q1 / q_z1)**(1./2.) * ( im_q1*q_z1.conjugate() / (-im_q1*q_z1)  )**(n1/2.) 

    wz1 = w(w0_1, im_q1, re_q1)

158
    return A_n1 * hermite(n1, np.sqrt(2.)*x / wz1) * np.exp(np.array(-1j*(2*math.pi/(1064e-9))* x**2 /(2.*q_z1)))
159
160


161
162
163
164
165
def u_star_u(re_q1, re_q2, w0_1, w0_2, n1, n2, x, x2=None):
    if x2 == None:
        x2 = x
        
    return u(re_q1, w0_1, n1, x) * u(re_q2, w0_2, n2, x2).conjugate()
166
167

    
168
def makeReducedBasis(x, isModeMatched=True, tolerance = 1e-12, sigma = 1, greedyfile=None):
169
    
170
171
    if greedyfile != None:
        greedypts = str(greedyfile)
172
    else:
173
174
175
176
        if isModeMatched:
            greedypts = 'matched20.txt'
        else:
            greedypts = 'mismatched20.txt'
177
    
178
    greedyfile = os.path.join(pykat.__path__[0],'optics','greedypoints',greedypts)
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
    
    limits = np.loadtxt(greedyfile, usecols=(1,))[:5]
    
    romlimits = ROMLimits(w0min=limits[0], w0max=limits[1], zmin=limits[2], zmax=limits[3], max_order=limits[4])
    
    w_min = limits[0]
    w_max = limits[1]

    re_q_min = limits[2]
    re_q_max = limits[3]
    
    max_order = int(limits[4])
    
    indices = range(0, max_order+1)
    nm_pairs = combs(indices, 2)
    
    dx = x[1] - x[0]
    
    params = np.loadtxt(greedyfile, skiprows=5)

    TS_size = len(params) # training space of TS_size number of waveforms
200
    
201
202
203
204
205
206
207
208
    #### allocate memory for training space ####

    TS = np.zeros(TS_size*len(x), dtype = complex).reshape(TS_size, len(x)) # store training space in TS_size X len(x) array

    IDx = 0 #TS index 

    for i in range(len(params)):
        if isModeMatched:
209
            q1 = beam_param(w0=float(params[i][0]), z=float(params[i][1]))
210
211
212
213
            q2 = q1
            n = int(params[i][2])
            m = int(params[i][3])
        else:
214
            q1 = beam_param(w0=float(params[i][0]), z=float(params[i][2]))
215
            q2 = beam_param(w0=float(params[i][1]), z=float(params[i][3]))            
216
217
218
            n = int(params[i][4])
            m = int(params[i][5])
            
219
220
221
222
223
        w0_1 = q1.w0
        w0_2 = q2.w0
        re_q1 = q1.z
        re_q2 = q2.z
            
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
        TS[IDx] = u_star_u(re_q1, re_q2, w0_1, w0_2, n, m, x)
        
        # normalize
        norm = np.sqrt(abs(np.vdot(dx * TS[IDx], TS[IDx])))
    
        if norm != 0:
            TS[IDx] /= norm 
            IDx += 1
        else:
            np.delete(TS, IDx)	

    #### Begin greedy: see Field et al. arXiv:1308.3565v2 #### 

    RB_matrix = [TS[0]] # seed greedy algorithm (arbitrary)

    proj_coefficients = np.zeros(TS_size*TS_size, dtype = complex).reshape(TS_size, TS_size)
    projections = np.zeros(TS_size*len(x), dtype = complex).reshape(TS_size, len(x))

    iter = 0
243
    
244
245
246
247
248
    while sigma > tolerance:
    #for k in range(TS_size-1):
    	# go through training set and find worst projection error
    	projections = project_onto_basis(dx, RB_matrix, TS, projections, proj_coefficients, iter)
    	residual = TS - projections
249
        
250
251
252
    	projection_errors = [np.vdot(dx* residual[i], residual[i]) for i in range(len(residual))]
    	sigma = abs(max(projection_errors))
    	index = np.argmax(projection_errors) 
253
        
254
255
256
257
258
    	#Gram-Schmidt to get the next basis and normalize
    	next_basis = TS[index] - projections[index]
    	next_basis /= np.sqrt(abs(np.vdot(dx* next_basis, next_basis)))

    	RB_matrix.append(next_basis)		
259
260
261
262
        
    	iter += 1
        
        
263
264
    return ReducedBasis(RB=np.array(RB_matrix), limits=romlimits, x=x)
    
Daniel Brown's avatar
Daniel Brown committed
265
def makeEmpiricalInterpolant(RB, sort=False):
266
267
268
269
270
271
272
273
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

    e_x = RB.RB
    x = RB.x
    
    node_indices = []
    x_nodes = []
    V = np.zeros((len(e_x), len(e_x)), dtype = complex)
    
    # seed EIM algorithm
    node_indices.append( int(np.argmax( np.abs(e_x[0]) ))) 
    x_nodes.append(x[node_indices]) 
    
    for i in range(1, len(e_x)): 
        #build empirical interpolant for e_iter
        for j in range(len(node_indices)):  
            for k in range(len(node_indices)):  
                V[k][j] = e_x[j][node_indices[k]]  
            
        invV = inv(V[0:len(node_indices), 0:len(node_indices)])  
        B = B_matrix(invV, e_x) 
        
        interpolant = emp_interp(B, e_x[i], node_indices) 
        res = interpolant - e_x[i] 

        index = int(np.argmax(np.abs(res))) 
        node_indices.append(index) 
        x_nodes.append( x[index] )
    
    # make B matrix with all the indices
    for j in range(len(node_indices)):
        for k in range(len(node_indices)):
            V[k][j] = e_x[j][node_indices[k]]

    invV = inv(V[0:len(node_indices), 0:len(node_indices)])
Daniel Brown's avatar
Daniel Brown committed
300
    B = np.array(B_matrix(invV, e_x))
301
    
Daniel Brown's avatar
Daniel Brown committed
302
303
304
305
306
307
308
309
310
311
312
    node_indices = np.array(node_indices, dtype=np.int32)
    nodes = np.array(x_nodes, dtype=np.float64)
    
    if sort:
        print sort
        ix = np.argsort(nodes)
        nodes = nodes[ix]
        node_indices = node_indices[ix]
        B = B[ix, :]
    
    return EmpiricalInterpolant(B=B, nodes=nodes, node_indices=node_indices, limits=RB.limits, x=RB.x)
313
    
Daniel Brown's avatar
Daniel Brown committed
314

315
def makeWeights(smap, EI, verbose=True, useSymmetry=True):
316
    
317
318
319
    if useSymmetry:
        # get full A_xy
        A_xy = smap.z_xy().transpose()
320
    
321
322
        xm = smap.x[smap.x < 0]
        xp = smap.x[smap.x > 0]
Daniel Brown's avatar
Daniel Brown committed
323
    
324
325
        ym = smap.y[smap.y < 0]
        yp = smap.y[smap.y > 0]
Daniel Brown's avatar
Daniel Brown committed
326
    
327
328
329
330
        Q1xy = np.ix_(smap.x < 0, smap.y > 0)
        Q2xy = np.ix_(smap.x > 0, smap.y > 0)
        Q3xy = np.ix_(smap.x > 0, smap.y < 0)
        Q4xy = np.ix_(smap.x < 0, smap.y < 0)
331
    
332
333
334
335
336
        # get A_xy in the four quadrants of the x-y plane
        A_xy_Q1 = A_xy[Q1xy]
        A_xy_Q2 = A_xy[Q2xy]
        A_xy_Q3 = A_xy[Q3xy]
        A_xy_Q4 = A_xy[Q4xy]
Daniel Brown's avatar
Daniel Brown committed
337
    
338
339
        # Not sure if there is a neater way to select the x=0,y=0 cross elements
        A_xy_0  = np.hstack([A_xy[:,smap.x==0].flatten(), A_xy[smap.y==0,:].flatten()])
Daniel Brown's avatar
Daniel Brown committed
340
    
341
342
        full_x = smap.x
        full_y = smap.y
343

344
345
        dx = full_x[1] - full_x[0]
        dy = full_y[1] - full_y[0]
346

347
348
349
        if verbose:
            count  = 4*len(EI["x"].B) * len(EI["y"].B)
            p = ProgressBar(maxval=count, widgets=["Computing weights: ", Percentage(), Bar(), ETA()])
Daniel Brown's avatar
Daniel Brown committed
350
    
351
        n = 0
Daniel Brown's avatar
Daniel Brown committed
352
    
353
354
355
356
        # make integration weights
        Bx = EI["x"].B
        By = EI["y"].B[:,::-1]
        w_ij_Q1 = np.zeros((len(Bx),len(By)), dtype = complex)
Daniel Brown's avatar
Daniel Brown committed
357
    
358
359
360
361
        for i in range(len(Bx)):
            for j in range(len(By)):
                B_ij_Q1 = np.outer(Bx[i], By[j])
                w_ij_Q1[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q1, A_xy_Q1)	
362
            
363
364
365
366
367
368
369
370
371
372
373
374
                if verbose:
                    p.update(n)
                    n+=1

        Bx = EI["x"].B[:,::-1]
        By = EI["y"].B[:,::-1]
        w_ij_Q2 = np.zeros((len(Bx),len(By)), dtype = complex)
    
        for i in range(len(Bx)):
            for j in range(len(By)):
                B_ij_Q2 = np.outer(Bx[i], By[j])
                w_ij_Q2[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q2, A_xy_Q2)
375
            
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
                if verbose:
                    p.update(n)
                    n+=1

        Bx = EI["x"].B[:,::-1]
        By = EI["y"].B
        w_ij_Q3 = np.zeros((len(Bx),len(By)), dtype = complex)
    
        for i in range(len(Bx)):
            for j in range(len(By)):
                B_ij_Q3 = np.outer(Bx[i], By[j])
                w_ij_Q3[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q3, A_xy_Q3)

                if verbose:
                    p.update(n)
                    n+=1

        Bx = EI["x"].B
        By = EI["y"].B
        w_ij_Q4 = np.zeros((len(Bx),len(By)), dtype = complex)
    
        for i in range(len(Bx)):
            for j in range(len(By)):
                B_ij_Q4 = np.outer(Bx[i], By[j])
                w_ij_Q4[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q4, A_xy_Q4)

                if verbose:
                    p.update(n)
                    n+=1
        if verbose:
            p.finish()
        
        return ROMWeights(w_ij_Q1=w_ij_Q1, w_ij_Q2=w_ij_Q2, w_ij_Q3=w_ij_Q3, w_ij_Q4=w_ij_Q4, EI=EI, limits=EI["x"].limits)
        
    else:
        # get full A_xy
        A_xy = smap.z_xy()
Daniel Brown's avatar
Daniel Brown committed
413

414
415
        full_x = smap.x
        full_y = smap.y
416

417
418
        dx = full_x[1] - full_x[0]
        dy = full_y[1] - full_y[0]
Daniel Brown's avatar
Daniel Brown committed
419

420
421
422
        if verbose:
            count  = len(EI["x"].B) * len(EI["y"].B)
            p = ProgressBar(maxval=count, widgets=["Computing weights: ", Percentage(), Bar(), ETA()])
Daniel Brown's avatar
Daniel Brown committed
423

424
425
426
427
428
        n = 0

        # make integration weights
        Bx = EI["x"].B
        By = EI["y"].B
Daniel Brown's avatar
Daniel Brown committed
429
        
430
431
432
433
434
435
        w_ij = np.zeros((len(Bx), len(By)), dtype = complex)

        for i in range(len(Bx)):
            for j in range(len(By)):
                B_ij = np.outer(Bx[i], By[j])
                w_ij[i][j] = dx*dy*np.einsum('ij,ij', B_ij, A_xy)	
Daniel Brown's avatar
Daniel Brown committed
436
    
437
438
439
440
441
442
443
444
445
                if verbose:
                    p.update(n)
                    n+=1
                    
        if verbose:
            p.finish()

        return ROMWeightsFull(w_ij=w_ij, EI=EI, limits=EI["x"].limits)

446
    
447