# Importing libraries:
from pylab import arange, zeros, plot, savefig, clf, grid, legend, array, xscale, yscale, title, xlabel, ylabel, axis, fabs
from pylab import pi, exp, sin, cos, loglog
# Explicit Euler (RK1) method
def RK1(F, y0, t_array, n, dt):
y_array = zeros(n+1)
y_array[0] = y0
for i in range(n):
y_array[i+1] = y_array[i] + dt*F(y_array[i], t_array[i])
return y_array
# RK2 method
def RK2(F, y0, t_array, n, dt):
b1 = 1./4
b2 = 3./4
c2 = 2./3
a21 = 2./3
y_array = zeros(n+1)
y_array[0] = y0
for i in range(n):
k1 = F(y_array[i], t_array[i])
k2 = F(y_array[i] + a21*dt*k1, t_array[i] + c2*dt)
y_array[i+1] = y_array[i] + dt*(b1*k1 + b2*k2)
return y_array
# RK4 method
def RK4(F, y0, t_array, n, dt):
b1 = 1./6
b2 = 1./3
b3 = 1./3
b4 = 1./6
c2 = 1./2
c3 = 1./2
c4 = 1.
a21 = 1./2
a31 = 0.
a32 = 1./2
a41 = 0.
a42 = 0.
a43 = 1.
y_array = zeros(n+1)
y_array[0] = y0
for i in range(n):
k1 = F(y_array[i], t_array[i])
k2 = F(y_array[i] + dt*(a21*k1), t_array[i] + c2*dt)
k3 = F(y_array[i] + dt*(a31*k1 + a32*k2), t_array[i] + c3*dt)
k4 = F(y_array[i] + dt*(a41*k1 + a42*k2 + a43*k3), t_array[i] + c4*dt)
y_array[i+1] = y_array[i] + dt*(b1*k1 + b2*k2 + b3*k3 + b4*k4)
return y_array
# Plotting exact solution
def plot_exact(E, t_min, t_max, n_exact):
dt_exact = (t_max - t_min)/n_exact
t_array_exact = arange(t_min, t_max + dt_exact/2., dt_exact)
E_array = array([E(t) for t in t_array_exact])
label = "Exact solution E(t)"
plot(t_array_exact, E_array, "g-", label=label)
# Error calculation: Finds maximum of |E(t) - y_approx(t)|
# over all discrete time steps
def calc_error(E, t_array, y_array, n):
max = 0.
for i in range(n):
if fabs(E(t_array[i]) - y_array[i]) > max:
max = fabs(E(t_array[i]) - y_array[i])
return max
# Go through div_list, always calculate and plot the three
# solutions, and store the maximum errors for all three
# methods
def browse_list(t_min, t_max, E, F, y0, n_array):
RK1_error = []
RK2_error = []
RK4_error = []
for n in n_array:
print "Subdivision = ", n
# Time step
dt = (t_max - t_min)/n
# Array of discrete time points
t_array = arange(t_min, t_max + dt/2., dt)
clf()
name = "Convergence for division n = " + str(n)
title(name)
# Explicit Euler (RK1) method
Y_RK1_array = RK1(F, y0, t_array, n, dt)
RK1_error.append(calc_error(E, t_array, Y_RK1_array, n))
plot(t_array, Y_RK1_array, "b-", label = "RK1")
# Second-order Runge-Kutta (RK2) method
Y_RK2_array = RK2(F, y0, t_array, n, dt)
RK2_error.append(calc_error(E, t_array, Y_RK2_array, n))
plot(t_array, Y_RK2_array, "y-", label = "RK2")
# Fourth-order Runge-Kutta (RK4) method
Y_RK4_array = RK4(F, y0, t_array, n, dt)
RK4_error.append(calc_error(E, t_array, Y_RK4_array, n))
plot(t_array, Y_RK4_array, "r-", label = "RK4")
# Plot exact solution if known
if exact_known == True: plot_exact(E, t_min, t_max, 500)
legend()
grid(True)
name = "a" + str(1000000+n) + ".png"
savefig(name)
# Print errors
print "RK1_error:"
print RK1_error
print "RK2_error:"
print RK2_error
print "RK4_error:"
print RK4_error
# Plot convergence
clf()
#axis('equal')
title("Convergence")
xlabel("Number of time steps (log scale)")
ylabel("Error (log scale)")
loglog(n_array, RK1_error, "g-o", label="RK1")
loglog(n_array, RK2_error, "b-o", label="RK2")
loglog(n_array, RK4_error, "y-o", label="RK4")
legend()
grid(True)
savefig("graph.png")
