diff --git a/blflow.py b/blflow.py
old mode 100755
new mode 100644
index 491f56f..5cdb65f
--- a/blflow.py
+++ b/blflow.py
@@ -1,72 +1,83 @@
 #!/usr/bin/python
-
 # Boundary layer flow
-from numpy import *
-import time 
-import matplotlib
-matplotlib.use('TkAgg')
-# from matplotlib.pylab import *
-import pylab as p
-
-# import matplotlib.animation as animation
-def K(t):                                 #Forcing function
-    return (1-exp(-0.1*t))*cos(t)
+import numpy as np
+import matplotlib.pyplot as plt
 
 
-s=10
-#Define domain
-n=50                                      #Number of gridpoints
-y=linspace(0,1,n)
+def K(t):
+    """
+    Forcing function as a function of time.
+    """
+    return (1-np.exp(-0.1*t))*np.cos(t)
 
-dy=y[1]-y[0]
-dt=0.0005
 
-l=(dt/(s**2*dy**2))
+s = 10  # Shear wave number
+n = 50   # Number of gridpoints in vertical direction
 
-hnu=exp(-sqrt(1j)*s*y)
+# Grid
+y = np.linspace(0, 1, n)
+
+dy = y[1]-y[0]
+dt = 0.0005
+
+l = (dt/(s**2*dy**2))
+
+hnu = np.exp(-np.sqrt(1j)*s*y)
 # fnu=(1-1j)/s
-fnu=0
+fnu = 0
+
 
 def u_ex(tn):
-    return (((1-hnu)/(1-fnu))*exp(1j*(tn))/1j).real
-    
+    """
+    Exact harmonic solution of the velocity field
+    """
+    return (((1-hnu)/(1-fnu))*np.exp(1j*(tn))/1j).real
 
-def u_np1(un,tn,dt):
-    Kn=K(tn)
-    unp1=un
-    unp1[0]=0                             #Velocity zero ver here
-    unp1[1:-1]=un[1:-1]+dt*Kn+l*(un[0:-2]-2*un[1:-1]+un[2:])
-    unp1[-1]=unp1[-2]                     #Approximate 'infinity' bc
+
+def u_np1(un, tn, dt):
+    Kn = K(tn)
+    unp1 = un
+    unp1[0] = 0  # Velocity zero ver here
+    unp1[1:-1] = un[1:-1]+dt*Kn+l*(un[0:-2]-2*un[1:-1]+un[2:])
+    unp1[-1] = unp1[-2]  # Approximate 'infinity' bc
     return unp1
 
-un0=zeros(n,float)
-t=0
-un=un0
+
+un0 = np.zeros(n, float)
+t = 0
+un = un0
 # un.append(un0)
 
 # Make the plot
-p.ion()
-linefd, = p.plot(un0,y)
-linee, = p.plot(un0,y)
-p.legend(('Finite difference','Periodic exact'))
-p.ylim(0,1)
-p.xlim(-1.5,1.5)
-p.ylabel('y')
-p.xlabel('u')
-p.grid('on')
-i=0
-uold=un
-while(True):
-    t+=dt
-    uold=un
-    un=u_np1(uold,t,dt)
-    if(i%100==0):
-        linefd.set_xdata(un)
-        linee.set_xdata(u_ex(t))
-        p.draw()
-        # print("Time:",t)
-    i+=1
-
+# plt.ioff()
+f = plt.figure(1, figsize=(12, 6))
 
+linefd, = plt.plot(un0, y)
+linee, = plt.plot(un0, y)
+plt.legend(('Finite difference', 'Periodic exact'))
+plt.ylim(0, 1)
+plt.xlim(-1.5, 1.5)
+plt.ylabel('Distance from wall [m]')
+plt.xlabel('Velocity [m/s]')
+plt.grid(True)
+plt.show(block=False)
+i = 0
+uold = un
 
+try:
+    while(True):
+        t += dt
+        uold = un
+        un = u_np1(uold, t, dt)
+        if(i % 100 == 0):
+            linefd.set_xdata(un)
+            linee.set_xdata(u_ex(t))
+            f.canvas.draw_idle()
+            plt.pause(.00001)
+            if not plt.fignum_exists(1):
+                break
+        i += 1
+except Exception:
+    pass
 
+exit(0)
diff --git a/parplateflow.py b/parplateflow.py
old mode 100755
new mode 100644
index 5897ae5..d2c1b94
--- a/parplateflow.py
+++ b/parplateflow.py
@@ -1,74 +1,66 @@
 #!/usr/bin/python
 
 # Boundary layer flow
-from numpy import *
-import time 
-import matplotlib
-matplotlib.use('TkAgg')
-# from matplotlib.pylab import *
-import pylab as p
-
+from numpy import zeros, linspace, exp, cos, cosh, sqrt, tanh
+import matplotlib.pyplot as plt
 
 
 # import matplotlib.animation as animation
-def K(t):                                 #Forcing function
+def K(t):  # Forcing function
     return (1-exp(-0.1*t))*cos(t)
 
 
-s=10
-#Define domain
-n=50                                      #Number of gridpoints
-y=linspace(-1,1,n)
+s = 10
 
-dy=y[1]-y[0]
-dt=0.0005
+# Define domain
+n = 50  # Number of gridpoints
+y = linspace(-1, 1, n)
 
-l=(dt/(s**2*dy**2))
+dy = y[1]-y[0]
+dt = 0.0005
 
-hnu=cosh(sqrt(1j)*s*y)/cosh(sqrt(1j)*s)
-fnu=tanh(sqrt(1j)*s)/(sqrt(1j)*s)
+l = (dt/(s**2*dy**2))
+
+hnu = cosh(sqrt(1j)*s*y)/cosh(sqrt(1j)*s)
+fnu = tanh(sqrt(1j)*s)/(sqrt(1j)*s)
 
 
 def u_ex(tn):
     return (((1-hnu)/(1-fnu))*exp(1j*(tn))/1j).real
-    
 
-def u_np1(un,tn,dt):
-    Kn=K(tn)
-    unp1=un
-    unp1[0]=0                                  #Velocity zero ver here
-    unp1[1:-1]=un[1:-1]+dt*Kn+l*(un[0:-2]-2*un[1:-1]+un[2:])
-    unp1[-1]=0                     #Boundary other side
+
+def u_np1(un, tn, dt):
+    Kn = K(tn)
+    unp1 = un
+    unp1[0] = 0  # Velocity zero ver here
+    unp1[1:-1] = un[1:-1]+dt*Kn+l*(un[0:-2]-2*un[1:-1]+un[2:])
+    unp1[-1] = 0  # Boundary other side
     return unp1
 
-un0=zeros(n,float)
-t=0
-un=un0
+
+un0 = zeros(n, float)
+t = 0
+un = un0
 # un.append(un0)
 
 # Make the plot
-p.ion()
-linefd, = p.plot(un0,y)
-linee, = p.plot(un0,y)
-p.legend(('Finite difference','Periodic exact'))
-p.ylim(-1,1)
-p.xlim(-1.5,1.5)
-p.ylabel('y')
-p.xlabel('u')
-p.grid('on')
-i=0
-uold=un
+linefd, = plt.plot(un0, y)
+linee, = plt.plot(un0, y)
+plt.legend(('Finite difference', 'Periodic exact'))
+plt.ylim(-1, 1)
+plt.xlim(-1.5, 1.5)
+plt.ylabel('y')
+plt.xlabel('u')
+plt.grid('on')
+i = 0
+uold = un
 while(True):
-    t+=dt
-    uold=un
-    un=u_np1(uold,t,dt)
-    if(i%100==0):
+    t += dt
+    uold = un
+    un = u_np1(uold, t, dt)
+    if(i % 100 == 0):
         linefd.set_xdata(un)
         linee.set_xdata(u_ex(t))
-        p.draw()
+        plt.draw()
         # print("Time:",t)
-    i+=1
-
-
-
-
+    i += 1