romhom.py 10.4 KB
Newer Older
1
import math
2
3
4
import os.path
import pykat
import collections
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
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
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
        self.limits = limits
        
    def writeToFile(self, filename):
        f = open(filename, 'w+')
        
        f.write(repr(self.limits) + "\n")
        
        f.write(repr(self.w_ij_Q1.shape) + "\n")
        
        for v in self.w_ij_Q1.flatten():
            f.write("%s " % repr(complex(v))[1:-1])
        
        f.write("\n")
        
        f.write(repr(self.w_ij_Q2.shape) + "\n")
        
        for v in self.w_ij_Q2.flatten():
            f.write("%s " % repr(complex(v))[1:-1])
        
        f.write("\n")
        
        f.write(repr(self.w_ij_Q3.shape) + "\n")
        
        for v in self.w_ij_Q3.flatten():
            f.write("%s " % repr(complex(v))[1:-1])
        
        f.write("\n")
        
        f.write(repr(self.w_ij_Q4.shape) + "\n")
        
        for v in self.w_ij_Q4.flatten():
            f.write("%s " % repr(complex(v))[1:-1])
        
        f.write("\n")
            
        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)

108
    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)))
109
110


111
112
113
114
115
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()
116
117

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

    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
250
    B = np.array(B_matrix(invV, e_x))
251
    
Daniel Brown's avatar
Daniel Brown committed
252
253
254
255
256
257
258
259
260
261
262
    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)
263
    
Daniel Brown's avatar
Daniel Brown committed
264
265

def makeWeights(smap, EI, verbose=True):
266
267
    
    # get full A_xy
268
    A_xy = smap.z_xy()
269
    
Daniel Brown's avatar
Daniel Brown committed
270
271
272
273
274
275
276
277
278
279
    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)
280
281
    
    # get A_xy in the four quadrants of the x-y plane
Daniel Brown's avatar
Daniel Brown committed
282
283
284
285
286
287
288
289
    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()])
    
290
291
292
293
294
295
    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
296
    if verbose:
297
        count  = 4*len(EI["xm"].B) * len(EI["ym"].B)
Daniel Brown's avatar
Daniel Brown committed
298
299
300
301
        p = ProgressBar(maxval=count, widgets=["Computing weights: ", Percentage(), Bar(), ETA()])
    
    n = 0
    
302
    # make integration weights
Daniel Brown's avatar
Daniel Brown committed
303
    Bx = EI["xm"].B
304
    By = EI["ym"].B[:,::-1]
Daniel Brown's avatar
Daniel Brown committed
305
306
307
308
309
    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])
310
            w_ij_Q1[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q1, A_xy_Q1)	
311
            
Daniel Brown's avatar
Daniel Brown committed
312
313
314
315
            if verbose:
                p.update(n)
                n+=1

316
317
    Bx = EI["xm"].B[:,::-1]
    By = EI["ym"].B[:,::-1]
Daniel Brown's avatar
Daniel Brown committed
318
319
320
321
322
    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])
323
            w_ij_Q2[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q2, A_xy_Q2)
324
            
Daniel Brown's avatar
Daniel Brown committed
325
326
327
328
            if verbose:
                p.update(n)
                n+=1

329
    Bx = EI["xm"].B[:,::-1]
Daniel Brown's avatar
Daniel Brown committed
330
331
332
333
334
335
    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])
336
337
            w_ij_Q3[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q3, A_xy_Q3)

Daniel Brown's avatar
Daniel Brown committed
338
339
340
341
342
343
344
345
346
347
348
            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])
349
            w_ij_Q4[i][j] = dx*dy*np.einsum('ij,ij', B_ij_Q4, A_xy_Q4)
Daniel Brown's avatar
Daniel Brown committed
350
351
352
353
354
355
356
357
358
359

            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)
    
    
360
    
361