from itertools import combinations_with_replacement as combinations from pykat.optics.gaussian_beams import beam_param, HG_beam from pykat.exceptions import BasePyKatException from pykat.optics.romhom import u_star_u from pykat.external.progressbar import ProgressBar, ETA, Percentage, Bar from scipy.interpolate import interp2d from scipy.integrate import dblquad from pykat.optics.romhom import ROMWeights import time import pykat.optics.maps import os.path import numpy as np import pykat import collections import math import cmath def makeCouplingMatrix(max_order, Neven=True, Nodd=True, Meven=True, Modd=True): 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) 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 row.append(e) M.append(np.array(row).squeeze()) return np.array(M) def adaptive_knm(mode_in, mode_out, q1, q2, q1y=None, q2y=None, smap=None, delta=(0,0), params={}): if q1y == None: q1y = q1 if q2y == None: 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]") 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]) Nfuncs = [] Nfuncs.append(0) if smap != None: 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 def riemann_HG_knm(x, y, mode_in, mode_out, q1, q2, q1y=None, q2y=None, Axy=None, cache=None, delta=(0,0), params={}): 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]") dx = abs(x[1] - x[0]) 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]) U1 = Hg_in.Unm(x+delta[0], y+delta[1]) 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])) def __gen_riemann_knm_cache(x, y, couplings, q1, q2, q1y=None, q2y=None, delta=(0,0), params={}): 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]) #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() except StopIteration: break return cache 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 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 return cache def ROM_HG_knm(weights, mode_in, mode_out, q1, q2, q1y=None, q2y=None, cache=None): if q1y == None: q1y = q1 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] if isinstance(weights, ROMWeights): if cache == None: u_x_nodes = u_star_u(q1.z, q2.z, q1.w0, q2.w0, n, m, weights.EI["x"].nodes) u_y_nodes = u_star_u(q1y.z, q2y.z, q1y.w0, q2y.w0, npr, mpr, weights.EI["y"].nodes) 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]) 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) 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: if cache == None: u_x_nodes = u_star_u(q1.z, q2.z, q1.w0, q2.w0, n, m, weights.EI["x"].nodes) u_y_nodes = u_star_u(q1y.z, q2y.z, q1y.w0, q2y.w0, npr, mpr, weights.EI["y"].nodes) w_ij = weights.w_ij else: 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] u_xy_nodes = np.outer(u_x_nodes, u_y_nodes) k_ROQ = np.einsum('ij,ij', u_xy_nodes, w_ij) return k_ROQ __fac_cache = [] def fac(n): global __fac_cache if len(__fac_cache) == 0: return math.factorial(int(n)) else: return __fac_cache[n] def m_1_pow(n): if n % 2 == 0: return 1 else: 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)) 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 _F = K / (2.0 * (1.0+K0)) F = K.conjugate() / 2.0 Sg = __S(n, _n, X, _X, F, _F) 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), params={}): 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 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']") 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(wavelength=q1.wavelength) x = surface_map.x y = surface_map.y K = np.zeros((couplings.size/4,), dtype=np.complex128) it = np.nditer(couplings, flags=['refs_ok','f_index']) i = 0 if profile: t0 = time.time() 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.") cache = __gen_ROM_HG_knm_cache(weights, couplings, q1=q1, q2=q2, q1y=q1y, q2y=q2y) elif method == "riemann": 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) else: cache = None weights = None if profile: cache_time = time.time() - t0 Ktime = np.zeros((couplings.size/4,), dtype=np.float64) if verbose: p = ProgressBar(maxval=couplings.size, widgets=["Knm (%s): " % method, Percentage(), Bar(), ETA()]) while not it.finished: try: if profile: t0 = time.time() mode_in = [int(it.next()), int(it.next())] mode_out = [int(it.next()), int(it.next())] 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) except StopIteration: break if profile: return K.reshape(couplings.shape[:-1]), Ktime.reshape(couplings.shape[:-1]), cache_time else: return K.reshape(couplings.shape[:-1]) 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