knm.py 12.6 KB
 Daniel Brown committed Jul 22, 2014 1 2 3 4 from itertools import combinations_with_replacement as combinations from pykat.utilities.optics.gaussian_beams import beam_param, HG_beam from pykat.exceptions import BasePyKatException  Daniel Brown committed Aug 12, 2014 5 import time  Daniel Brown committed Jul 22, 2014 6 7 8 9 10 11 import maps import os.path import numpy as np import pykat import collections import math  Daniel Brown committed Aug 12, 2014 12 import cmath  Daniel Brown committed Aug 01, 2014 13   Daniel Brown committed Jul 25, 2014 14 from romhom import u_star_u  Daniel Brown committed Aug 01, 2014 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 from progressbar import ProgressBar, ETA, Percentage, Bar def makeCouplingMatrix(max_order): 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) row.append(e) M.append(row)  Daniel Brown committed Aug 12, 2014 37  return np.array(M)  Daniel Brown committed Jul 22, 2014 38   Daniel Brown committed Aug 12, 2014 39 def riemann_HG_knm(x, y, mode_in, mode_out, q1, q2, q1y=None, q2y=None, Axy=None, cache=None, delta=(0,0)):  Daniel Brown committed Jul 22, 2014 40 41 42 43 44 45 46 47 48 49 50 51  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]")  Daniel Brown committed Aug 01, 2014 52   Daniel Brown committed Jul 22, 2014 53  dx = abs(x[1] - x[0])  Daniel Brown committed Aug 01, 2014 54 55 56 57 58  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])  Daniel Brown committed Jul 22, 2014 59   Daniel Brown committed Aug 12, 2014 60  U1 = Hg_in.Unm(x+delta[0], y+delta[1])  Daniel Brown committed Aug 01, 2014 61 62 63 64 65 66 67 68 69 70  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]))  Daniel Brown committed Jul 22, 2014 71   Daniel Brown committed Aug 01, 2014 72   Daniel Brown committed Aug 12, 2014 73 74  def __gen_riemann_knm_cache(x, y, couplings, q1, q2, q1y=None, q2y=None, delta=(0,0)):  Daniel Brown committed Aug 01, 2014 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92  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])  Daniel Brown committed Aug 12, 2014 93 94  #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])  Daniel Brown committed Aug 01, 2014 95 96  if strx not in cache:  Daniel Brown committed Aug 12, 2014 97 98  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()  Daniel Brown committed Aug 01, 2014 99 100  if stry not in cache:  Daniel Brown committed Aug 12, 2014 101 102  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()  Daniel Brown committed Aug 01, 2014 103 104 105 106 107 108  except StopIteration: break return cache  Daniel Brown committed Aug 12, 2014 109 110   Daniel Brown committed Jul 25, 2014 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 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  Daniel Brown committed Jul 22, 2014 130   Daniel Brown committed Jul 25, 2014 131 132 133 134 135 136 137 138 139  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:  Daniel Brown committed Aug 01, 2014 140  cache[strx] = u_star_u(q1.z, q2.z, q1.w0, q2.w0, mode_in[0], mode_out[0], weights.EI["xm"].nodes)  Daniel Brown committed Jul 25, 2014 141   Daniel Brown committed Aug 01, 2014 142 143  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)  Daniel Brown committed Jul 25, 2014 144 145 146 147 148  except StopIteration: break return cache  Daniel Brown committed Aug 12, 2014 149 150 151   Daniel Brown committed Jul 25, 2014 152 153 154 def ROM_HG_knm(weights, mode_in, mode_out, q1, q2, q1y=None, q2y=None, cache=None): if q1y == None: q1y = q1  Daniel Brown committed Jul 22, 2014 155   Daniel Brown committed Jul 25, 2014 156 157 158 159 160 161 162 163 164 165 166  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]  Daniel Brown committed Aug 12, 2014 167 168 169  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)  Daniel Brown committed Aug 01, 2014 170   Daniel Brown committed Aug 12, 2014 171 172 173 174 175 176 177 178 179 180 181  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 committed Jul 31, 2014 182   Daniel Brown committed Aug 12, 2014 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209  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: 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 committed Jul 31, 2014 210   Daniel Brown committed Aug 12, 2014 211 212 213 214 215  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)  Daniel Brown committed Jul 30, 2014 216   Daniel Brown committed Aug 12, 2014 217 218  return k_ROQ  Daniel Brown committed Jul 31, 2014 219   Daniel Brown committed Aug 12, 2014 220 __fac_cache = []  Daniel Brown committed Jul 31, 2014 221   Daniel Brown committed Aug 12, 2014 222 223 224 225 def fac(n): global __fac_cache return __fac_cache[n] #return math.factorial(int(n))  Daniel Brown committed Jul 31, 2014 226   Daniel Brown committed Aug 12, 2014 227 228 229 def m_1_pow(n): if n % 2 == 0: return 1  Daniel Brown committed Jul 31, 2014 230  else:  Daniel Brown committed Aug 12, 2014 231 232 233 234 235 236 237 238  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 committed Jul 31, 2014 239   Daniel Brown committed Aug 12, 2014 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280  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  Daniel Brown committed Jul 25, 2014 281   Daniel Brown committed Aug 12, 2014 282 283  _F = K / (2.0 * (1.0+K0)) F = K.conjugate() / 2.0  Daniel Brown committed Jul 25, 2014 284   Daniel Brown committed Aug 12, 2014 285  Sg = __S(n, _n, X, _X, F, _F)  Daniel Brown committed Jul 25, 2014 286   Daniel Brown committed Aug 12, 2014 287 288 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 314  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]) def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riemann", verbose=False, profile=False, gamma=(0,0), delta=(0,0)):  Daniel Brown committed Jul 25, 2014 315 316 317 318 319 320 321  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  Daniel Brown committed Jul 22, 2014 322 323 324 325 326 327 328 329  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']")  Daniel Brown committed Aug 12, 2014 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345  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: Axy = surface_map.z_xy(q1.wavelength) x = surface_map.x y = surface_map.y  Daniel Brown committed Jul 22, 2014 346 347 348 349 350 351  K = np.zeros((couplings.size/4,), dtype=np.complex128) it = np.nditer(couplings, flags=['refs_ok','f_index']) i = 0  Daniel Brown committed Aug 12, 2014 352 353 354  if profile: t0 = time.time()  Daniel Brown committed Aug 01, 2014 355   Daniel Brown committed Jul 25, 2014 356 357 358 359 360 361 362 363  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.")  Daniel Brown committed Aug 01, 2014 364 365  cache = __gen_ROM_HG_knm_cache(weights, couplings, q1=q1, q2=q2, q1y=q1y, q2y=q2y)  Daniel Brown committed Jul 25, 2014 366   Daniel Brown committed Aug 01, 2014 367  elif method == "riemann":  Daniel Brown committed Aug 12, 2014 368 369 370 371  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)  Daniel Brown committed Jul 25, 2014 372  else:  Daniel Brown committed Aug 01, 2014 373  cache = None  Daniel Brown committed Jul 25, 2014 374  weights = None  Daniel Brown committed Aug 01, 2014 375   Daniel Brown committed Aug 12, 2014 376 377 378 379  if profile: cache_time = time.time() - t0 Ktime = np.zeros((couplings.size/4,), dtype=np.float64)  Daniel Brown committed Aug 01, 2014 380 381 382  if verbose: p = ProgressBar(maxval=couplings.size, widgets=["Knm (%s): " % method, Percentage(), Bar(), ETA()])  Daniel Brown committed Jul 22, 2014 383 384  while not it.finished: try:  Daniel Brown committed Aug 12, 2014 385 386 387  if profile: t0 = time.time()  Daniel Brown committed Jul 22, 2014 388 389  mode_in = [int(it.next()), int(it.next())] mode_out = [int(it.next()), int(it.next())]  Daniel Brown committed Jul 25, 2014 390   Daniel Brown committed Jul 22, 2014 391  if method == "riemann":  Daniel Brown committed Aug 12, 2014 392  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)  Daniel Brown committed Jul 25, 2014 393  elif method == "romhom":  Daniel Brown committed Aug 01, 2014 394  K[i] = ROM_HG_knm(weights, mode_in, mode_out, q1=q1, q2=q2, q1y=q1y, q2y=q2y, cache=cache)  Daniel Brown committed Aug 12, 2014 395 396  elif method == "bayerhelms": K[i] = bayerhelms_HG_knm(mode_in, mode_out, q1=q1, q2=q2, q1y=q1y, q2y=q2y, gamma=gamma)  Daniel Brown committed Jul 22, 2014 397  else:  Daniel Brown committed Jul 25, 2014 398  raise BasePyKatException("method value '%s' not accepted" % method)  Daniel Brown committed Aug 12, 2014 399 400 401 402  if profile: Ktime[i] = time.time() - t0  Daniel Brown committed Jul 22, 2014 403  i +=1  Daniel Brown committed Aug 01, 2014 404 405 406 407  if verbose: p.update(i*4)  Daniel Brown committed Jul 22, 2014 408 409 410 411  except StopIteration: break  Daniel Brown committed Aug 12, 2014 412 413 414 415  if profile: return K.reshape(couplings.shape[:-1]), Ktime.reshape(couplings.shape[:-1]), cache_time else: return K.reshape(couplings.shape[:-1])  Daniel Brown committed Aug 01, 2014 416 417   Daniel Brown committed Aug 12, 2014 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 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=0 and row