diff --git a/src/vib/vib.py b/src/vib/vib.py index 5b386e90..bcf1fbb5 100644 --- a/src/vib/vib.py +++ b/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):