adding function for finite aperture clipping loss for LG modes

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
-         
-          
+