Removed scitools imports

Author: EdCaunt

src/vib/vib.py

@@ -1,8 +1,8 @@
 import numpy as np
 import sympy as sp
 from devito import Dimension, Constant, TimeFunction, Eq, solve, Operator
-#import matplotlib.pyplot as plt
-import scitools.std as plt
+import matplotlib.pyplot as plt
+# import scitools.std as plt
 
 def solver(I, V, m, b, s, F, dt, T, damping='linear'):
     """
@@ -13,33 +13,27 @@ def solver(I, V, m, b, s, F, dt, T, damping='linear'):
     'quadratic', f(u')=b*u'*abs(u').
     F(t) and s(u) are Python functions.
     """
-    dt = float(dt)
-    b = float(b)
-    m = float(m)
+    dt = float(dt); b = float(b); m = float(m) # avoid integer div.
     Nt = int(round(T/dt))
-    t = Dimension('t', spacing=Constant('h_t'))
-
-    u = TimeFunction(name='u', dimensions=(t,),
-                     shape=(Nt+1,), space_order=2)
-
-    u.data[0] = I
+    u = np.zeros(Nt+1)
+    t = np.linspace(0, Nt*dt, Nt+1)
 
+    u[0] = I
     if damping == 'linear':
-        # dtc for central difference (default for time is forward, 1st order)
-        eqn = m*u.dt2 + b*u.dtc + s(u) - F(u)
-        stencil = Eq(u.forward, solve(eqn, u.forward))
+        u[1] = u[0] + dt*V + dt**2/(2*m)*(-b*V - s(u[0]) + F(t[0]))
     elif damping == 'quadratic':
-        # fd_order set as backward derivative used is 1st order
-        eqn = m*u.dt2 + b*u.dt*sp.Abs(u.dtl(fd_order=1)) + s(u) - F(u)
-        stencil = Eq(u.forward, solve(eqn, u.forward))
-    # First timestep needs to have the backward timestep substituted
-    stencil_init = stencil.subs(u.backward, u.forward-2*t.spacing*V)
-    op_init = Operator(stencil_init, name='first_timestep')
-    op = Operator(stencil, name='main_loop')
-    op_init.apply(h_t=dt, t_M=1)
-    op.apply(h_t=dt, t_m=1, t_M=Nt-1)
+        u[1] = u[0] + dt*V + \
+               dt**2/(2*m)*(-b*V*abs(V) - s(u[0]) + F(t[0]))
 
-    return u.data, np.linspace(0, Nt*dt, Nt+1)
+    for n in range(1, Nt):
+        if damping == 'linear':
+            u[n+1] = (2*m*u[n] + (b*dt/2 - m)*u[n-1] +
+                      dt**2*(F(t[n]) - s(u[n])))/(m + b*dt/2)
+        elif damping == 'quadratic':
+            u[n+1] = (2*m*u[n] - m*u[n-1] + b*u[n]*abs(u[n] - u[n-1])
+                      + dt**2*(F(t[n]) - s(u[n])))/\
+                      (m + b*abs(u[n] - u[n-1]))
+    return u, t
 
 def visualize(u, t, title='', filename='tmp'):
     plt.plot(t, u, 'b-')
@@ -212,9 +206,9 @@ def plot_empirical_freq_and_amplitude(u, t):
     plt.figure()
     from math import pi
     w = 2*pi/p    
-    plt.plot(range(len(p)), w, 'r-')
+    plt.plot(list(range(len(p))), w, 'r-')
     plt.hold('on')
-    plt.plot(range(len(a)), a, 'b-')
+    plt.plot(list(range(len(a))), a, 'b-')
     ymax = 1.1*max(w.max(), a.max())
     ymin = 0.9*min(w.min(), a.min())
     plt.axis([0, max(len(p), len(a)), ymin, ymax])
@@ -247,7 +241,7 @@ def visualize_front(u, t, window_width, savefig=False):
                     axis=plot_manager.axis(),
                     show=not savefig) # drop window if savefig
             if savefig:
-                print 't=%g' % t[n]
+                print('t=%g' % t[n])
                 st.savefig('tmp_vib%04d.png' % n)
         plot_manager.update(n)
 
@@ -263,7 +257,7 @@ def visualize_front_ascii(u, t, fps=10):
 
     p = Plotter(ymin=umin, ymax=umax, width=60, symbols='+o')
     for n in range(len(u)):
-        print p.plot(t[n], u[n]), '%.2f' % (t[n])
+        print(p.plot(t[n], u[n]), '%.2f' % (t[n]))
         time.sleep(1/float(fps))
 
 def minmax(t, u):