updating to work with py3

pykat/optics/knm.py

@@ -216,15 +216,21 @@ def __gen_riemann_knm_cache(x, y, couplings, q1, q2, q1y=None, q2y=None, delta=(
     if q2y is None:
         q2y = q2
         
-    it = np.nditer(couplings, flags=['refs_ok','f_index'])
+    #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())]
+    #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:]        
         strx = "u1[%i,%i]" % (mode_in[0], mode_out[0])
         stry = "u2[%i,%i]" % (mode_in[1], mode_out[1])
         
@@ -239,8 +245,8 @@ def __gen_riemann_knm_cache(x, y, couplings, q1, q2, q1y=None, q2y=None, delta=(
             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
+    #    except StopIteration:
+    #        break
     
     return cache
     
@@ -254,7 +260,7 @@ def __gen_ROM_HG_knm_cache(weights, couplings, q1, q2, q1y=None, q2y=None):
     if q2y is None:
         q2y = q2
         
-    it = np.nditer(couplings, flags=['refs_ok','f_index'])
+    
     
     cache = {}
     
@@ -266,10 +272,12 @@ def __gen_ROM_HG_knm_cache(weights, couplings, q1, q2, q1y=None, q2y=None):
     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())]
+    couplings = couplings.copy()
+    couplings.resize(int(couplings.size/4), 4)
+    
+    for _ in couplings:
+        mode_in = _[:2]
+        mode_out = _[2:]
 
         strx = "x[%i,%i]" % (mode_in[0], mode_out[0])
         stry = "y[%i,%i]" % (mode_in[1], mode_out[1])
@@ -280,8 +288,23 @@ 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)
     
-        except StopIteration:
-            break
+    # it = np.nditer(couplings, flags=['refs_ok','f_index'])
+    # 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
 
@@ -543,7 +566,7 @@ def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riema
     maxtem = 0
     c = couplings.flatten()
     
-    for i in range(0, c.size/2):
+    for i in range(0, int(c.size/2)):
         maxtem = max(sum(c[i*2:(i*2+2)]), maxtem)
     
     global __fac_cache
@@ -559,7 +582,7 @@ def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riema
     
     K = np.zeros((couplings.size/4,), dtype=np.complex128)
     
-    it = np.nditer(couplings, flags=['refs_ok','f_index'])
+    #it = np.nditer(couplings, flags=['refs_ok','f_index'])
     
     i = 0
     
@@ -593,13 +616,16 @@ def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riema
     if verbose:
         p = ProgressBar(maxval=couplings.size, widgets=["Knm (%s): " % method, Percentage(), Bar(), ETA()])
     
-    while not it.finished:
-        try:
+    _couplings = couplings.copy()
+    _couplings.resize(int(_couplings.size/4), 4)
+    
+    for _ in _couplings:
+        mode_in = _[:2]
+        mode_out = _[2:]
+
         if profile:
             t0 = time.time()
                 
-            mode_in = [int(it.next()), int(it.next())]
-            mode_out = [int(it.next()), int(it.next())]
         
         
         if method == "riemann":
@@ -621,25 +647,18 @@ def knmHG(couplings, q1, q2, surface_map=None, q1y=None, q2y=None, method="riema
         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):
+def plot_knm_matrix(couplings, knm, cmap=None, show=True):
     import pylab as plt
     
     fig = plt.figure()
     ax = fig.add_subplot(111)
-    cax = ax.imshow(knm, interpolation='nearest')
+    cax = ax.pcolormesh(abs(knm), cmap=cmap)
     fig.colorbar(cax)
     
     numrows, numcols = knm.shape
@@ -648,21 +667,31 @@ def plot_knm_matrix(couplings, knm):
     c_ = []
 
     for d in c:
-        c_.append("%i,%i"%(d[0], d[1]))
+        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(x+0.5)
-        row = int(y+0.5)
+        col = int(np.floor(x))
+        row = int(np.floor(y))
         
         if col>=0 and col<numcols and row>=0 and row<numrows:
             z = knm[row,col]
             return 'x=%s, y=%s, z=%1.4f' % (c_[col], c_[row], z)
-        else:
-            return 'x=%1.4f, y=%1.4f'%(x, y)
+        
+        return None
 
     ax.format_coord = format_coord
 
-    plt.show()
+    fig.tight_layout()
+    
+    if show: plt.show()
+    
+    return fig