diff --git a/pykat/optics/knm.py b/pykat/optics/knm.py index 0ffd761eb9e881a9e2842cb64ad4a57c43023103..8c919880163c2707596455b40f96f1088cb9045b 100644 --- a/pykat/optics/knm.py +++ b/pykat/optics/knm.py @@ -10,6 +10,7 @@ from math import factorial from pykat.math.hermite import hermite from scipy.misc import comb from scipy.integrate import newton_cotes +from scipy.special import gammainc from pykat.math import newton_weights import time @@ -33,237 +34,237 @@ def makeCouplingMatrix(max_order, Neven=True, Nodd=True, Meven=True, Modd=True): 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 is None: q1y = q1 - + if q2y is 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_mode(qx=q1, qy=q1y, n=mode_in[0], m=mode_in[1]) Hg_out = HG_mode(qx=q2, qy=q2y, n=mode_out[0], m=mode_out[1]) - + Nfuncs = [] Nfuncs.append(0) - + if smap is not 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(phi) 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={}, newtonCotesOrder=0): if Axy is None: Axy == np.ones((len(x), len(y))) - + if q1y is None: q1y = q1 - + if q2y is 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]") + 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]) - + dy = abs(y[1] - y[0]) + if cache is None: Hg_in = HG_mode(qx=q1, qy=q1y, n=mode_in[0], m=mode_in[1]) Hg_out = HG_mode(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() if newtonCotesOrder > 0: - + W = newton_cotes(newtonCotesOrder, 1)[0] - + if newtonCotesOrder > 1: if (len(x) - len(W)) % newtonCotesOrder != 0: raise ValueError("To use Newton-Cotes order {0} the number of data points in x must ensure: (N_x - ({0}+1)) mod {0} == 0".format(newtonCotesOrder) ) if (len(y) - len(W)) % newtonCotesOrder != 0: raise ValueError("To use Newton-Cotes order {0} the number of data points in y must ensure: (N_y - ({0}+1)) mod {0} == 0".format(newtonCotesOrder) ) - - wx = np.zeros(x.shape, dtype=np.float64) + + wx = np.zeros(x.shape, dtype=np.float64) wy = np.zeros(y.shape, dtype=np.float64) - + N = len(W) for i in range(0, (len(wx)-1)/newtonCotesOrder): wx[(i*(N-1)):(i*(N-1)+N)] += W for i in range(0, (len(wy)-1)/newtonCotesOrder): wy[(i*(N-1)):(i*(N-1)+N)] += W - + Wxy = np.outer(wx, wy) - + if newtonCotesOrder == 0: return dx * dy * np.einsum('ij,ij', Axy, U1*U2) else: return dx * dy * np.einsum('ij,ij', Axy, U1*U2*Wxy) 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 is None: q1y = q1 - + if q2y is 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())] - + couplings = couplings.copy() couplings.resize(int(couplings.size/4), 4) - + for _ in couplings: mode_in = _[:2] - mode_out = _[2:] + mode_out = _[2:] 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() - + 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]) + 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 is None: q1y = q1 - + if q2y is None: q2y = q2 - - - + + + 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 @@ -274,7 +275,7 @@ def __gen_ROM_HG_knm_cache(weights, couplings, q1, q2, q1y=None, q2y=None): couplings = couplings.copy() couplings.resize(int(couplings.size/4), 4) - + for _ in couplings: mode_in = _[:2] mode_out = _[2:] @@ -287,7 +288,7 @@ def __gen_ROM_HG_knm_cache(weights, couplings, q1, q2, q1y=None, q2y=None): 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) - + # it = np.nditer(couplings, flags=['refs_ok','f_index']) # while not it.finished: # try: @@ -305,7 +306,7 @@ def __gen_ROM_HG_knm_cache(weights, couplings, q1, q2, q1y=None, q2y=None): # # except StopIteration: # break - + return cache @@ -316,7 +317,7 @@ def ROM_HG_knm(weights, mode_in, mode_out, q1, q2, q1y=None, q2y=None, cache=Non if q2y is None: q2y = q2 - + # x modes n = mode_in[0] m = mode_out[0] @@ -324,12 +325,12 @@ def ROM_HG_knm(weights, mode_in, mode_out, q1, q2, q1y=None, q2y=None, cache=Non # y modes npr = mode_in[1] mpr = mode_out[1] - + if isinstance(weights, ROMWeights): if cache is None: u_x_nodes = u_star_u(q1.z, q2.z, q1.w0, q2.w0, n, m, weights.EIx.nodes) u_y_nodes = u_star_u(q1y.z, q2y.z, q1y.w0, q2y.w0, npr, mpr, weights.EIy.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 @@ -337,14 +338,14 @@ def ROM_HG_knm(weights, mode_in, mode_out, q1, q2, q1y=None, q2y=None, cache=Non 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"] @@ -352,7 +353,7 @@ def ROM_HG_knm(weights, mode_in, mode_out, q1, q2, q1y=None, q2y=None, cache=Non 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 @@ -374,12 +375,12 @@ def ROM_HG_knm(weights, mode_in, mode_out, q1, q2, q1y=None, q2y=None, cache=Non 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 is None: u_x_nodes = u_star_u(q1.z, q2.z, q1.w0, q2.w0, n, m, weights.EIx.nodes) u_y_nodes = u_star_u(q1y.z, q2y.z, q1y.w0, q2y.w0, npr, mpr, weights.EIy.nodes) - + w_ij = weights.w_ij else: strx = "x[%i,%i]" % (mode_in[0], mode_out[0]) @@ -387,11 +388,11 @@ def ROM_HG_knm(weights, mode_in, mode_out, q1, q2, q1y=None, q2y=None, cache=Non 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 = [] @@ -418,56 +419,56 @@ def __sum_lim(m): def __Sg(m, _m, X, _X, F, _F): lim1 = __sum_lim(m) - lim2 = __sum_lim(_m) + lim2 = __sum_lim(_m) r = 0 - + def __Ss(u, _u, F, _F): 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)*fac(u-s)*fac(_u-s)) - + return r - - + + 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) - + return r - + def __Su(m, _m, X, _X, F, _F): lim1 = __sum_lim(m-1) lim2 = __sum_lim(_m-1) - + r = 0 - + def __Ss(u, _u, F, _F): 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+1)*fac(u-s)*fac(_u-s)) - + return r - - + + for u in range(0, lim1+1): for _u in range(0, lim2+1): r += m_1_pow(u) * _X**(m-2*u-1) * X**(_m-2*_u-1) / ( fac(m-2*u-1)*fac(_m-2*_u-1) ) * __Ss(u, _u, F, _F) - + 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: @@ -480,9 +481,9 @@ def __bayerhelms_kn(n, _n, q1, q2, gamma=0.0): _X = 0.0 X = 0.0 Ex = 1.0 - + _F = K / (2.0 * (1.0+K0)) - F = K.conjugate() / 2.0 + F = K.conjugate() / 2.0 Sg = __Sg(n, _n, X, _X, F, _F) @@ -490,20 +491,20 @@ def __bayerhelms_kn(n, _n, q1, q2, gamma=0.0): Su = 0 else: Su = __Su(n, _n, X, _X, F, _F) - - b = ((-1)**(_n) * + + b = ((-1)**(_n) * np.sqrt(fac(n) * fac(_n)) * (1.0 + K0)**(n/2.0 + 1/4.0) * (1.0 + K.conjugate()) ** (-(n+_n+1)/2.0)) - + #print(K, F, _F, b, Ex, Sg, Su) - + return b * Ex * (Sg - Su) def bayerhelms_HG_knm(mode_in, mode_out, q1, q2, q1y=None, q2y=None, gamma=(0,0)): """ - + """ if q1y is None: q1y = q1 @@ -527,23 +528,23 @@ def __sqaure_knm_int(n, _n, R): # thus making it easier to solve expR = math.exp(-(R**2)) S = 0 - + for j in range(0, min(n, _n)+1): _j1 = _n + n - 2*j - 1 - + if _j1+1 == 0: # for the zeroth order we just have the gaussian integral to solve - L = math.sqrt(math.pi) * math.erf(R) + L = math.sqrt(math.pi) * math.erf(R) elif (_j1+1) % 2 == 1: # if the Hermite is odd then the integral is always 0 as its anti-symmetric L = 0 else: L = 2 * hermite(_j1, 0) - expR * (hermite(_j1, R) - hermite(_j1, -R)) - + I = 2**j * factorial(j) * comb(n, j) * comb(_n, j) - + S += I * L - + return S @@ -558,22 +559,22 @@ def square_aperture_HG_knm(mode_in, mode_out, q, R): # y modes m = mode_in[1] _m = mode_out[1] - + hg1 = HG_mode(q, n=n, m=m) hg2 = HG_mode(q, n=_n, m=_m) - + kx = hg1.constant_x * hg2.constant_x.conjugate() ky = hg1.constant_y * hg2.constant_y.conjugate() - + f = q.w / math.sqrt(2) R = R / (q.w / math.sqrt(2)) - + kx *= f kx *= __sqaure_knm_int(n, _n, R) - + ky *= f ky *= __sqaure_knm_int(m, _m, R) - + return kx * ky @@ -586,15 +587,15 @@ def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riema - Surface maps (b) - Tilts (c) - Displacements (d) - + Mathematically the 2D overlap integral between an incoming and outgoing beam is solved in the Hermite-Gaussian mode basis. - + Some of these effects aren't supported in all solvers. The markers a,b,c, and d are listed below next to each method. - + Parameters: - couplings - Matrix of coupling + couplings - Matrix of coupling q1 - Incoming mode parameter (x and y direction if q1y not specified) q2 - Outgoing mode parameter (x and y direction if q2y not specified) q1y - Incoming mode parameter (y direction) @@ -605,11 +606,11 @@ def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riema adaptive - Numerical adaptive quadrature solver (a,b,c,d) romhom - Reduced order modelling solver (a,b,c) verbose - If true outputs more debug and timing information - gamma - Tuple of misalignment tilts of incoming and outgoing mode, e.g. (x_tilt, y_tilt) - delta - Tuple of displacements between incoming and outgoing mode, e.g. (dx, dy) + gamma - Tuple of misalignment tilts of incoming and outgoing mode, e.g. (x_tilt, y_tilt) + delta - Tuple of displacements between incoming and outgoing mode, e.g. (dx, dy) cache - If True caching of certain results are used to speed up calculations - - + + Example using Maps: import pykat from pykat.optics.knm import knmHG, makeCouplingMatrix, plot_knm_matrix @@ -633,7 +634,7 @@ def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riema plot_knm_matrix(C, abs(Kmap)) - + Example using Bayer-Helms: import pykat from pykat.optics.knm import knmHG, makeCouplingMatrix, plot_knm_matrix @@ -650,82 +651,82 @@ def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riema """ if q1y is None: q1y = q1 - + if q2y is 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, int(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 is not None: + + if surface_map is not None: Axy = surface_map.z_xy(wavelength=q1.wavelength) - + x = surface_map.x y = surface_map.y - + K = np.zeros((int(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 is None: raise BasePyKatException("Using 'romhom' method requires a surface map to be specified") - + weights = surface_map.ROMWeights - + if weights is 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" and cache: if surface_map is 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()]) - + _couplings = couplings.copy() _couplings.resize(int(_couplings.size/4), 4) - + for _ in _couplings: mode_in = _[:2] mode_out = _[2:] if profile: t0 = time.time() - + 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=None, delta=delta) elif method == "romhom": @@ -736,12 +737,12 @@ def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riema K[i] = adaptive_knm(mode_in, mode_out, q1=q1, q2=q2, q1y=q1y, q2y=q2y, smap=surface_map, delta=delta, params=kwargs) else: raise BasePyKatException("method value '%s' not accepted" % method) - + if profile: Ktime[i] = time.time() - t0 - + i +=1 - + if verbose: p.update(i*4) @@ -753,43 +754,57 @@ def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riema def plot_knm_matrix(couplings, knm, cmap=None, show=True, fig=None): import matplotlib.pyplot as plt - + if fig is None: fig = plt.figure() - + ax = fig.add_subplot(111) cax = ax.pcolormesh(abs(knm), cmap=cmap) fig.colorbar(cax) - + numrows, numcols = knm.shape - + c = couplings[:, 0, :2] c_ = [] for d in c: c_.append("[%i,%i]"%(d[0], d[1])) - + A = np.arange(1, len(c)+1)-0.5 ax.set_xticks(A) ax.set_yticks(A) ax.set_xticklabels(c_) ax.set_yticklabels(c_) - + ax.set_xlim(None, max(A)+0.5) ax.set_ylim(None, max(A)+0.5) - + def format_coord(x, y): col = int(np.floor(x)) row = int(np.floor(y)) - + if col>=0 and col=0 and row