romhom.py 11.1 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
        
Daniel Brown's avatar
Daniel Brown committed
44
45
46
47
48
49
50
51
52
53
        f.write("xnodes=%i\n" % len(self.EI["xm"].nodes))
        
        for v in self.EI["xm"].nodes.flatten():
            f.write("%s\n" % repr(float(v)))
        
        f.write("ynodes=%i\n" % len(self.EI["ym"].nodes))
        
        for v in self.EI["ym"].nodes.flatten():
            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
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
            
        f.close()

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)

119
    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)))
120
121


122
123
124
125
126
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()
127
128

    
129
def makeReducedBasis(x, isModeMatched=True, tolerance = 1e-12, sigma = 1, greedyfile=None):
130
    
131
132
    if greedyfile != None:
        greedypts = str(greedyfile)
133
    else:
134
135
136
137
        if isModeMatched:
            greedypts = 'matched20.txt'
        else:
            greedypts = 'mismatched20.txt'
138
    
139
    greedyfile = os.path.join(pykat.__path__[0],'optics','greedypoints',greedypts)
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
    
    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
161
    
162
163
164
165
166
167
168
169
    #### 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:
170
            q1 = beam_param(w0=float(params[i][0]), z=float(params[i][1]))
171
172
173
174
            q2 = q1
            n = int(params[i][2])
            m = int(params[i][3])
        else:
175
            q1 = beam_param(w0=float(params[i][0]), z=float(params[i][2]))
176
            q2 = beam_param(w0=float(params[i][1]), z=float(params[i][3]))            
177
178
179
            n = int(params[i][4])
            m = int(params[i][5])
            
180
181
182
183
184
        w0_1 = q1.w0
        w0_2 = q2.w0
        re_q1 = q1.z
        re_q2 = q2.z
            
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
        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
204
    
205
206
207
208
209
    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
210
        
211
212
213
    	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) 
214
        
215
216
217
218
219
    	#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)		
220
221
222
223
        
    	iter += 1
        
        
224
225
    return ReducedBasis(RB=np.array(RB_matrix), limits=romlimits, x=x)
    
Daniel Brown's avatar
Daniel Brown committed
226
def makeEmpiricalInterpolant(RB, sort=False):
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260

    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
261
    B = np.array(B_matrix(invV, e_x))
262
    
Daniel Brown's avatar
Daniel Brown committed
263
264
265
266
267
268
269
270
271
272
273
    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)
274
    
Daniel Brown's avatar
Daniel Brown committed
275
276

def makeWeights(smap, EI, verbose=True):
277
278
    
    # get full A_xy
279
    A_xy = smap.z_xy()
280
    
Daniel Brown's avatar
Daniel Brown committed
281
282
283
284
285
286
287
288
289
290
    xm = smap.x[smap.x < 0]
    xp = smap.x[smap.x > 0]
    
    ym = smap.y[smap.y < 0]
    yp = smap.y[smap.y > 0]
    
    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)
291
292
    
    # get A_xy in the four quadrants of the x-y plane
Daniel Brown's avatar
Daniel Brown committed
293
294
295
296
297
298
299
300
    A_xy_Q1 = A_xy[Q1xy]
    A_xy_Q2 = A_xy[Q2xy]
    A_xy_Q3 = A_xy[Q3xy]
    A_xy_Q4 = A_xy[Q4xy]
    
    # 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()])
    
301
302
303
304
305
306
    full_x = smap.x
    full_y = smap.y

    dx = full_x[1] - full_x[0]
    dy = full_y[1] - full_y[0]

Daniel Brown's avatar
Daniel Brown committed
307
    if verbose:
308
        count  = 4*len(EI["xm"].B) * len(EI["ym"].B)
Daniel Brown's avatar
Daniel Brown committed
309
310
311
312
        p = ProgressBar(maxval=count, widgets=["Computing weights: ", Percentage(), Bar(), ETA()])
    
    n = 0
    
313
    # make integration weights
Daniel Brown's avatar
Daniel Brown committed
314
    Bx = EI["xm"].B
315
    By = EI["ym"].B[:,::-1]
Daniel Brown's avatar
Daniel Brown committed
316
317
318
319
320
    w_ij_Q1 = np.zeros((len(Bx),len(By)), dtype = complex)
    
    for i in range(len(Bx)):
        for j in range(len(By)):
            B_ij_Q1 = np.outer(Bx[i], By[j])
321
            w_ij_Q1[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q1, A_xy_Q1)	
322
            
Daniel Brown's avatar
Daniel Brown committed
323
324
325
326
            if verbose:
                p.update(n)
                n+=1

327
328
    Bx = EI["xm"].B[:,::-1]
    By = EI["ym"].B[:,::-1]
Daniel Brown's avatar
Daniel Brown committed
329
330
331
332
333
    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])
334
            w_ij_Q2[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q2, A_xy_Q2)
335
            
Daniel Brown's avatar
Daniel Brown committed
336
337
338
339
            if verbose:
                p.update(n)
                n+=1

340
    Bx = EI["xm"].B[:,::-1]
Daniel Brown's avatar
Daniel Brown committed
341
342
343
344
345
346
    By = EI["ym"].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])
347
348
            w_ij_Q3[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q3, A_xy_Q3)

Daniel Brown's avatar
Daniel Brown committed
349
350
351
352
353
354
355
356
357
358
359
            if verbose:
                p.update(n)
                n+=1

    Bx = EI["xm"].B
    By = EI["ym"].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])
360
            w_ij_Q4[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q4, A_xy_Q4)
Daniel Brown's avatar
Daniel Brown committed
361
362
363
364
365
366
367
368
369
370

            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["xm"].limits)
    
    
371
    
372