def comparison_plot(f, u, Omega, filename='tmp.pdf',
plot_title='', ymin=None, ymax=None,
u_legend='approximation'):
"""Compare f(x) and u(x) for x in Omega in a plot."""
x = sm.Symbol('x')
print 'f:', f
f = sm.lambdify([x], f, modules="numpy")
u = sm.lambdify([x], u, modules="numpy")
if len(Omega) != 2:
raise ValueError('Omega=%s must be an interval (2-list)' % str(Omega))
# When doing symbolics, Omega can easily contain symbolic expressions,
# assume .evalf() will work in that case to obtain numerical
# expressions, which then must be converted to float before calling
# linspace below
if not isinstance(Omega[0], (int,float)):
Omega[0] = float(Omega[0].evalf())
if not isinstance(Omega[1], (int,float)):
Omega[1] = float(Omega[1].evalf())
resolution = 401 # no of points in plot
xcoor = linspace(Omega[0], Omega[1], resolution)
# Vectorized functions expressions does not work with
# lambdify'ed functions without the modules="numpy"
exact = f(xcoor)
approx = u(xcoor)
plot(xcoor, approx, '-')
hold('on')
plot(xcoor, exact, '-')
legend([u_legend, 'exact'])
title(plot_title)
xlabel('x')
if ymin is not None and ymax is not None:
axis([xcoor[0], xcoor[-1], ymin, ymax])
savefig(filename)
def run_solvers_and_check_amplitudes(solvers,
timesteps_per_period=20,
num_periods=1,
I=1,
w=2 * np.pi):
P = 2 * np.pi / w # duration of one period
dt = P / timesteps_per_period
Nt = num_periods * timesteps_per_period
T = Nt * dt
t_mesh = np.linspace(0, T, Nt + 1)
file_name = 'Amplitudes' # initialize filename for plot
for solver in solvers:
solver.set(f_kwargs={'w': w})
solver.set_initial_condition([0, I])
u, t = solver.solve(t_mesh)
solver_name = \
'CrankNicolson' if solver.__class__.__name__ == \
'MidpointImplicit' else solver.__class__.__name__
file_name = file_name + '_' + solver_name
minima, maxima = minmax(t, u[:, 0])
a = amplitudes(minima, maxima)
plt.plot(range(len(a)), a, '-', label=solver_name)
plt.hold('on')
plt.xlabel('Number of periods')
plt.ylabel('Amplitude (absolute value)')
plt.legend(loc='upper left')
plt.savefig(file_name + '.png')
plt.savefig(file_name + '.pdf')
plt.show()
def random_walk_2D(np, ns, plot_step):
xpositions = zeros(np)
ypositions = zeros(np)
# extent of the axis in the plot:
ymin = -4
ymax = 8 * sqrt(ns)
xmax = 3 * sqrt(ns)
xmin = -xmax
for step in range(ns):
for i in range(np):
r = random.randint(0, 10)
if 0 <= r <= 5:
ypositions[i] += 1
elif r == 6:
ypositions[i] -= 1
elif 7 <= r <= 8:
xpositions[i] += 1
elif 9 <= r <= 10:
xpositions[i] -= 1
# plot just every plot_step steps:
if (step + 1) % plot_step == 0:
plot(
xpositions,
ypositions,
"ko",
axis=[xmin, xmax, ymin, ymax],
title="%d researchers after %d months, with ownership of strategy" % (np, step + 1),
hardcopy="tmp_%03d.eps" % (step + 1),
)
return xpositions, ypositions
def test_easyviz():
from scitools.std import linspace, ndgrid, plot, contour, peaks, \
quiver, surfc, backend, get_backend
n = 21
x = linspace(-3, 3, n)
xx, yy = ndgrid(x, x, sparse=False)
# a basic plot
plot(x, x**2, 'bx:')
wait()
if backend in ['gnuplot', 'vtk', 'matlab', 'dx', 'visit', 'veusz']:
# a contour plot
contour(peaks(n), title="contour plot")
wait()
# a vector plot
uu = yy
vv = xx
quiver(xx, yy, uu, vv)
wait()
# a surface plot with contours
zz = peaks(xx, yy)
surfc(xx, yy, zz, colorbar=True)
wait()
if backend == 'grace':
g = get_backend()
g('exit')
def test_easyviz():
from scitools.std import linspace, ndgrid, plot, contour, peaks, \
quiver, surfc, backend, get_backend
n = 21
x = linspace(-3, 3, n)
xx, yy = ndgrid(x, x, sparse=False)
# a basic plot
plot(x, x**2, 'bx:')
wait()
if backend in ['gnuplot', 'vtk', 'matlab', 'dx', 'visit', 'veusz']:
# a contour plot
contour(peaks(n), title="contour plot")
wait()
# a vector plot
uu = yy
vv = xx
quiver(xx, yy, uu, vv)
wait()
# a surface plot with contours
zz = peaks(xx, yy)
surfc(xx, yy, zz, colorbar=True)
wait()
if backend == 'grace':
g = get_backend()
g('exit')
def visualize_front(u, t, I, w, savefig=False):
"""
Visualize u and the exact solution vs t, using a
moving plot window and continuous drawing of the
curves as they evolve in time.
Makes it easy to plot very long time series.
"""
import scitools.std as st
from scitools.MovingPlotWindow import MovingPlotWindow
P = 2*pi/w # one period
umin = -1.2*I; umax = -umin
plot_manager = MovingPlotWindow(
window_width=8*P,
dt=t[1]-t[0],
yaxis=[umin, umax],
mode='continuous drawing')
for n in range(1,len(u)):
if plot_manager.plot(n):
s = plot_manager.first_index_in_plot
st.plot(t[s:n+1], u[s:n+1], 'r-1',
t[s:n+1], I*cos(w*t)[s:n+1], 'b-1',
title='t=%6.3f' % t[n],
axis=plot_manager.axis(),
show=not savefig) # drop window if savefig
if savefig:
st.savefig('tmp_vib%04d.png' % n)
plot_manager.update(n)
def random_walk_2D(np, ns, plot_step):
xpositions = zeros(np)
ypositions = zeros(np)
# extent of the axis in the plot:
xymax = 3 * sqrt(ns)
xymin = -xymax
NORTH = 1
SOUTH = 2
WEST = 3
EAST = 4 # constants
for step in range(ns):
for i in range(np):
direction = random.randint(1, 4)
if direction == NORTH:
ypositions[i] += 1
elif direction == SOUTH:
ypositions[i] -= 1
elif direction == EAST:
xpositions[i] += 1
elif direction == WEST:
xpositions[i] -= 1
# plot just every plot_step steps:
if (step + 1) % plot_step == 0:
plot(xpositions, ypositions, 'ko',
axis=[xymin, xymax, xymin, xymax],
title='%d researchers after %d months, without strategy' % \
(np, step+1),
hardcopy='tmp_%03d.eps' % (step+1))
return xpositions, ypositions
def test(b):
""" Runs test on SIR model, with variying beta """
beta = b #0.0005 or 0.0001 # Infection rate
v = 0.1 # Prob. of recovery per dt
S0 = 1500 # Init. No. of suspectibles
I0 = 1 # Init. No. of infected
R0 = 0. # Init. No. of resistant
U0 = [S0, I0, R0] # Initial conditions
T = 60 # Duration, days
dt = 0.5 # Time step length in days
n = T / dt # No. of solve steps
f = RHS(v, dt, T, beta) # Get right hand side of equation
solver = ForwardEuler(f) # Select ODE solver method
solver.set_initial_condition(U0)
time_points = np.linspace(0, 60, n + 1)
u, t = solver.solve(time_points)
S = u[:, 0] # S is all data in array no 0
I = u[:, 1] # I is all data in array no 1
R = u[:, 2] # R is all data in array no 2
plot(t,
S,
t,
I,
t,
R,
xlabel='Days',
ylabel='Persons',
legend=['Suspectibles', 'Infected', 'Resistant'],
hold=('on'))
def process_u(u, x, t, n):
"""
Animate u, and store arrays in plot_arrays if
t coincides with chosen times for plotting (plot_times).
"""
global ymax
if animate:
plt.plot(x,
u,
'r-',
x,
f(x, t[n]) / delta,
'b-',
axis=[0, L, ymin, ymax],
title='t=%f' % t[n],
xlabel='x',
ylabel='u and f/%g' % delta)
if t[n] == 0:
time.sleep(1)
plot_arrays.append(x)
dt = t[1] - t[0]
tol = dt / 10.0
if abs(t[n] - plot_times[0]) < tol or \
abs(t[n] - plot_times[1]) < tol:
plot_arrays.append((u.copy(), f(x, t[n]) / delta))
if u.max() > ymax:
ymax = u.max()
def approximate(f, d=1, N_e=4, numint="Gauss-Legendre2", filename="tmp.eps"):
nodes, elements = mesh(N_e, d, [0, 1])
A, b = assemble(nodes, elements, d, f, numint)
c = A.LUsolve(b)
print "nodes:", nodes
print "elements:", elements
print "A:\n", A
print "b:\n", b
print sp.latex(A, mode="plain")
# print sp.latex(b, mode='plain')
c = A.LUsolve(b)
print "c:\n", c
print "Plain interpolation:"
x = sp.Symbol("x")
f = sp.lambdify([x], f, modules="numpy")
f_at_nodes = [f(xc) for xc in nodes]
print f_at_nodes
if not symbolic:
xf = np.linspace(0, 1, 10001)
U = np.asarray(c)
xu, u = u_glob(U, elements, nodes)
from scitools.std import plot
plot(xu, u, "r-", xf, f(xf), "b-", legend=("u", "f"), savefig=filename)
def main():
v0 = 0 # start velocity
T = 23 # here this time is chosen for falling 5000 m when only gravity is applied
dt = 0.01
C_D = 1.2 # Coeffi
rho = 0.79 # density of air (assumed constant)
rho_b = 1003.0 # kgm^3;
m = 80 # Mass of the jumper and parachute in kg
R = 0.5
A = pi*R**2 # Area of the sphere pi*r^2
V = m/rho_b # Volume of the falling object
d = 2.08 # Diameter of the sphere 2*r
mu = 1.8*10**-5 # dynamic viscosity
g = 9.81 # Gravitational constant
a_s = 3*pi*rho*d*mu/(rho_b*V) # Constant to be used for the Stokes model
a_q = 0.5*C_D*rho*A/(rho_b*V) # Constant to be used for the quadratic model
b = g*(-1 + rho/rho_b) # A common constant for both Stokes and quad. model
rdm = rho*d/mu # A constant related to the Reynolds number
t,v = vm.solver(T,dt,a_s,a_q,b,v0,rdm)
plot(t,v,xlabel = 't',ylabel = 'v',title = 'Parachute Jumper')
raw_input('Press any key.')
def visualize_front(u, t, I, w, savefig=False):
"""
Visualize u and the exact solution vs t, using a
moving plot window and continuous drawing of the
curves as they evolve in time.
Makes it easy to plot very long time series.
"""
import scitools.std as st
from scitools.MovingPlotWindow import MovingPlotWindow
P = 2 * pi / w # one period
umin = 1.2 * u.min()
umax = -umin
plot_manager = MovingPlotWindow(window_width=8 * P,
dt=t[1] - t[0],
yaxis=[umin, umax],
mode='continuous drawing')
for n in range(1, len(u)):
if plot_manager.plot(n):
s = plot_manager.first_index_in_plot
st.plot(t[s:n + 1],
u[s:n + 1],
'r-1',
t[s:n + 1],
I * cos(w * t)[s:n + 1],
'b-1',
title='t=%6.3f' % t[n],
axis=plot_manager.axis(),
show=not savefig) # drop window if savefig
if savefig:
filename = 'tmp_vib%04d.png' % n
st.savefig(filename)
print 'making plot file', filename, 'at t=%g' % t[n]
plot_manager.update(n)
def approximate(f, symbolic=False, d=1, N_e=4, Omega=[0, 1], filename='tmp'):
phi = basis(d)
print 'phi basis (reference element):\n', phi
nodes, elements = mesh_uniform(N_e, d, Omega, symbolic)
A, b = assemble(nodes, elements, phi, f, symbolic=symbolic)
print 'nodes:', nodes
print 'elements:', elements
print 'A:\n', A
print 'b:\n', b
print sp.latex(A, mode='plain')
#print sp.latex(b, mode='plain')
if symbolic:
c = A.LUsolve(b)
else:
c = np.linalg.solve(A, b)
print 'c:\n', c
print 'Plain interpolation/collocation:'
x = sp.Symbol('x')
f = sp.lambdify([x], f, modules='numpy')
try:
f_at_nodes = [f(xc) for xc in nodes]
except NameError as e:
raise NameError('numpy does not support special function:\n%s' % e)
print f_at_nodes
if not symbolic and filename is not None:
xf = np.linspace(Omega[0], Omega[1], 10001)
U = np.asarray(c)
xu, u = u_glob(U, elements, nodes)
plt.plot(xu, u, '-',
xf, f(xf), '--')
plt.legend(['u', 'f'])
plt.savefig(filename + '.pdf')
plt.savefig(filename + '.png')
def run_solvers_and_check_amplitudes(solvers, timesteps_per_period=20,
num_periods=1, I=1, w=2*np.pi):
P = 2*np.pi/w # duration of one period
dt = P/timesteps_per_period
Nt = num_periods*timesteps_per_period
T = Nt*dt
t_mesh = np.linspace(0, T, Nt+1)
file_name = 'Amplitudes' # initialize filename for plot
for solver in solvers:
solver.set(f_kwargs={'w': w})
solver.set_initial_condition([0, I])
u, t = solver.solve(t_mesh)
solver_name = \
'CrankNicolson' if solver.__class__.__name__ == \
'MidpointImplicit' else solver.__class__.__name__
file_name = file_name + '_' + solver_name
minima, maxima = minmax(t, u[:,0])
a = amplitudes(minima, maxima)
plt.plot(range(len(a)), a, '-', label=solver_name)
plt.hold('on')
plt.xlabel('Number of periods')
plt.ylabel('Amplitude (absolute value)')
plt.legend(loc='upper left')
plt.savefig(file_name + '.png')
plt.savefig(file_name + '.pdf')
plt.show()
def visualize_front(u, t, window_width, savefig=False):
"""
Visualize u and the exact solution vs t, using a
moving plot window and continuous drawing of the
curves as they evolve in time.
Makes it easy to plot very long time series.
P is the approximate duration of one period.
"""
import scitools.std as st
from scitools.MovingPlotWindow import MovingPlotWindow
umin = 1.2*u.min(); umax = -umin
plot_manager = MovingPlotWindow(
window_width=window_width,
tau=t[1]-t[0],
yaxis=[umin, umax],
mode='continuous drawing')
for n in range(1,len(u)):
if plot_manager.plot(n):
s = plot_manager.first_index_in_plot
st.plot(t[s:n+1], u[s:n+1], 'r-1',
title='t=%6.3f' % t[n],
axis=plot_manager.axis(),
show=not savefig) # drop window if savefig
if savefig:
print 't=%g' % t[n]
st.savefig('tmp_vib%04d.png' % n)
plot_manager.update(n)
def random_walk_2D(np, ns, plot_step):
xpositions = numpy.zeros(np)
ypositions = numpy.zeros(np)
moves = numpy.random.random_integers(1, 4, size=ns*np)
moves.shape = (ns, np)
# Estimate max and min positions
xymax = 3*numpy.sqrt(ns); xymin = -xymax
NORTH = 1; SOUTH = 2; WEST = 3; EAST = 4 # constants
for step in range(ns):
this_move = moves[step,:]
ypositions += numpy.where(this_move == NORTH, 1, 0)
ypositions -= numpy.where(this_move == SOUTH, 1, 0)
xpositions += numpy.where(this_move == EAST, 1, 0)
xpositions -= numpy.where(this_move == WEST, 1, 0)
# Just plot every plot_step steps
if (step+1) % plot_step == 0:
plot(xpositions, ypositions, 'ko',
axis=[xymin, xymax, xymin, xymax],
title='%d particles after %d steps' %
(np, step+1),
savefig='tmp_%03d.pdf' % (step+1))
return xpositions, ypositions
def random_walk_2D(np, ns, plot_step):
xpositions = numpy.zeros(np)
ypositions = numpy.zeros(np)
# extent of the axis in the plot:
xymax = 3 * numpy.sqrt(ns)
xymin = -xymax
NORTH = 1
SOUTH = 2
WEST = 3
EAST = 4 # constants
for step in range(ns):
for i in range(np):
direction = random.randint(1, 4)
if direction == NORTH:
ypositions[i] += 1
elif direction == SOUTH:
ypositions[i] -= 1
elif direction == EAST:
xpositions[i] += 1
elif direction == WEST:
xpositions[i] -= 1
# Plot just every plot_step steps
if (step + 1) % plot_step == 0:
plot(xpositions,
ypositions,
'ko',
axis=[xymin, xymax, xymin, xymax],
title='%d particles after %d steps' % (np, step + 1),
savefig='tmp_%03d.pdf' % (step + 1))
return xpositions, ypositions
def plot(self, include_exact=True, plt=None):
"""
Add solver.u curve to scitools plotting object plt,
and include the exact solution if include_exact is True.
This plot function can be called several times (if
the solver object has computed new solutions).
"""
if plt is None:
import scitools.std as plt
plt.plot(self.solver.t, self.solver.u, '--o')
plt.hold('on')
theta = self.solver.get('theta')
theta2name = {0: 'FE', 1: 'BE', 0.5: 'CN'}
name = theta2name.get(theta, '')
legends = ['numerical %s' % name]
if include_exact:
t_e = np.linspace(0, self.problem.get('T'), 1001)
u_e = self.problem.exact_solution(t_e)
plt.plot(t_e, u_e, 'b-')
legends.append('exact')
plt.legend(legends)
plt.xlabel('t')
plt.ylabel('u')
dt = self.solver.get('dt')
plt.title('theta=%g, dt=%g' % (theta, dt))
plt.savefig('%s_%g.png' % (name, dt))
return plt
def random_walk_2D(np, ns, plot_step):
xpositions = zeros(np)
ypositions = zeros(np)
# extent of the axis in the plot:
ymin = -4
ymax = 8 * sqrt(ns)
xmax = 3 * sqrt(ns)
xmin = -xmax
for step in range(ns):
for i in range(np):
r = random.randint(0, 10)
if 0 <= r <= 5:
ypositions[i] += 1
elif r == 6:
ypositions[i] -= 1
elif 7 <= r <= 8:
xpositions[i] += 1
elif 9 <= r <= 10:
xpositions[i] -= 1
# plot just every plot_step steps:
if (step + 1) % plot_step == 0:
plot(xpositions, ypositions, 'ko',
axis=[xmin, xmax, ymin, ymax],
title='%d researchers after %d months, with ownership of strategy' % \
(np, step+1),
hardcopy='tmp_%03d.eps' % (step+1))
return xpositions, ypositions
def demo():
"""
Demonstrate difference between Euler-Cromer and the
scheme for the corresponding 2nd-order ODE.
"""
I = 1.2
V = 0.2
m = 4
b = 0.2
s = lambda u: 2 * u
F = lambda t: 0
w = np.sqrt(2. / 4) # approx freq
dt = 0.9 * 2 / w # longest possible time step
w = 0.5
P = 2 * pi / w
T = 4 * P
from vib import solver as solver2
import scitools.std as plt
for k in range(4):
u2, t2 = solver2(I, V, m, b, s, F, dt, T, 'quadratic')
u, v, t = solver(I, V, m, b, s, F, dt, T, 'quadratic')
plt.figure()
plt.plot(t, u, 'r-', t2, u2, 'b-')
plt.legend(['Euler-Cromer', 'centered scheme'])
plt.title('dt=%.3g' % dt)
raw_input()
plt.savefig('tmp_%d' % k + '.png')
plt.savefig('tmp_%d' % k + '.pdf')
dt /= 2
def random_walk_2D(np, ns, plot_step):
xpositions = numpy.zeros(np)
ypositions = numpy.zeros(np)
moves = numpy.random.random_integers(1, 4, size=ns * np)
moves.shape = (ns, np)
# Estimate max and min positions
xymax = 3 * numpy.sqrt(ns)
xymin = -xymax
NORTH = 1
SOUTH = 2
WEST = 3
EAST = 4 # constants
for step in range(ns):
this_move = moves[step, :]
ypositions += numpy.where(this_move == NORTH, 1, 0)
ypositions -= numpy.where(this_move == SOUTH, 1, 0)
xpositions += numpy.where(this_move == EAST, 1, 0)
xpositions -= numpy.where(this_move == WEST, 1, 0)
# Just plot every plot_step steps
if (step + 1) % plot_step == 0:
plot(xpositions,
ypositions,
'ko',
axis=[xymin, xymax, xymin, xymax],
title='%d particles after %d steps' % (np, step + 1),
savefig='tmp_%03d.pdf' % (step + 1))
return xpositions, ypositions
def demo_two_stars(animate=True):
# Initial condition
ic = [0.6, 0, 0, 1, # star A: velocity (0,1)
0, 0, 0, -0.5] # star B: velocity (0,-0.5)
# Solve ODEs
x_A, x_B, y_A, y_B, t = solver(
alpha=0.5, ic=ic, T=4*np.pi, dt=0.05)
if animate:
# Animate motion and draw the objects' paths in time
for i in range(len(x_A)):
plt.plot(x_A[:i+1], x_B[:i+1], 'r-',
y_A[:i+1], y_B[:i+1], 'b-',
[x_A[0], x_A[i]], [x_B[0], x_B[i]], 'r2o',
[y_A[0], y_A[i]], [y_B[0], y_B[i]], 'b4o',
daspectmode='equal', # axes aspect
legend=['A', 'B', 'A', 'B'],
axis=[-1, 1, -1, 1],
savefig='tmp_%04d.png' % i,
title='t=%.2f' % t[i])
else:
# Make a simple static plot of the solution
plt.plot(x_A, x_B, 'r-', y_A, y_B, 'b-',
daspectmode='equal', legend=['A', 'B'],
axis=[-1, 1, -1, 1], savefig='tmp_two_stars.png')
#plt.axes().set_aspect('equal') # mpl
plt.show()
def demo():
"""
Demonstrate difference between Euler-Cromer and the
scheme for the corresponding 2nd-order ODE.
"""
I = 1.2
V = 0.2
m = 4
b = 0.2
s = lambda u: 2 * u
F = lambda t: 0
w = np.sqrt(2.0 / 4) # approx freq
dt = 0.9 * 2 / w # longest possible time step
w = 0.5
P = 2 * pi / w
T = 4 * P
from vib import solver as solver2
import scitools.std as plt
for k in range(4):
u2, t2 = solver2(I, V, m, b, s, F, dt, T, "quadratic")
u, v, t = solver(I, V, m, b, s, F, dt, T, "quadratic")
plt.figure()
plt.plot(t, u, "r-", t2, u2, "b-")
plt.legend(["Euler-Cromer", "centered scheme"])
plt.title("dt=%.3g" % dt)
raw_input()
plt.savefig("tmp_%d" % k + ".png")
plt.savefig("tmp_%d" % k + ".pdf")
dt /= 2
def test_vertical_motion():
beta = 0.9
omega = np.sqrt(beta/(1-beta))
# Find num_periods. Recall that P=2*pi for scaled pendulum
# oscillations, while here we don't have gravity driven
# oscillations, but elastic oscillations with frequency omega.
period = 2*np.pi/omega
# We want T = N*period
N = 5
# simulate function has T = 2*pi*num_periods
num_periods = 5/omega
n = 600
time_steps_per_period = omega*n
y_exact = lambda t: -0.1*np.cos(omega*t)
x, y, theta, t = simulate_drag(
beta=beta, gamma=0, Theta=0, epsilon=0.1,
num_periods=num_periods,
time_steps_per_period=time_steps_per_period,
plot=False)
tol = 0.00055 # ok tolerance for the above resolution
# No motion in x direction is epxected
assert np.abs(x.max()) < tol
# Check motion in y direction
y_e = y_exact(t)
diff = np.abs(y_e - y).max()
if diff > tol: # plot
plt.plot(t, y, t, y_e, legend=['y', 'exact'])
raw_input('Type CR:')
assert diff < tol, 'diff=%g' % diff
def test_solver_exact_discrete_solution():
def tilde_w(w, dt):
return (2.0 / dt) * asin(w * dt / 2.0)
def u_numerical_exact(t):
return I * cos(tilde_w(w, dt) * t)
w = 2.5
I = 1.5
# Estimate period and time step
P = 2 * pi / w
num_periods = 4
T = num_periods * P
N = 5 # time steps per period
dt = P / N
u, t = solver(I, w, dt, T)
u_e = u_numerical_exact(t)
error = abs(u_e - u).max()
# Make a plot in a file, but not on the screen
from scitools.std import plot
plot(t, u, "bo", t, u_e, "r-", legend=("numerical", "exact"), show=False, savefig="tmp.png")
assert error < 1e-14
def approximate(f, d=1, N_e=4, numint='Gauss-Legendre2',
filename='tmp.eps'):
nodes, elements = mesh(N_e, d, [0, 1])
A, b = assemble(nodes, elements, d, f, numint)
c = A.LUsolve(b)
print 'nodes:', nodes
print 'elements:', elements
print 'A:\n', A
print 'b:\n', b
print sym.latex(A, mode='plain')
#print sym.latex(b, mode='plain')
c = A.LUsolve(b)
print 'c:\n', c
print 'Plain interpolation:'
x = sym.Symbol('x')
f = sym.lambdify([x], f, modules='numpy')
f_at_nodes = [f(xc) for xc in nodes]
print f_at_nodes
if not symbolic:
xf = np.linspace(0, 1, 10001)
U = np.asarray(c)
xu, u = u_glob(U, elements, nodes)
from scitools.std import plot
plot(xu, u, 'r-',
xf, f(xf), 'b-',
legend=('u', 'f'),
savefig=filename)
def demo():
"""
Demonstrate difference between Euler-Cromer and the
scheme for the corresponding 2nd-order ODE.
"""
I = 1.2
w = 2.0
T = 5
dt = 2 / w # longest possible time step
from vib_undamped import solver as solver2 # 2nd-order ODE
import scitools.std as plt
for k in range(4):
dt /= 4
u2, t2 = solver2(I, w, dt, T)
u, v, t = solver(I, w, dt, T)
plt.figure()
plt.plot(t,
u,
t2,
u2,
legend=('Euler-Cromer', 'center scheme for $u'
'+u=0$'),
title='dt=%.3g' % dt)
raw_input()
plt.savefig('ECvs2nd_%d' % k + '.png')
plt.savefig('ECvs2nd_%d' % k + '.pdf')
def __call__(self, u, x, t, n):
"""user_action function for solver."""
if animate == 'u and u_exact':
plt.plot(x,
u,
'r-',
x,
u_exact(x, t[n]),
'b--',
xlabel='x',
ylabel='u',
axis=[0, L, -self.ymax, self.ymax],
title='t=%f' % t[n],
show=True)
else:
error = u_exact(x, t[n]) - u
local_max_error = np.abs(error).max()
# self.max_error holds the increasing amplitude error
if self.max_error == [] or \
local_max_error > max(self.max_error):
self.max_error.append(local_max_error)
self.max_error_t.append(t[n])
# Use user's ymax until the error exceeds that value.
# This gives a growing max value of the yaxis (but
# never shrinking)
self.ymax = max(self.ymax, max(self.max_error))
plt.plot(x,
error,
'r-',
xlabel='x',
ylabel='error',
axis=[0, L, -self.ymax, self.ymax],
title='t=%f' % t[n],
show=True)
plt.savefig('%s_%04d.png' % (self.frame_name, n))
def random_walk_2D(np, ns, plot_step):
xpositions = numpy.zeros(np)
ypositions = numpy.zeros(np)
# extent of the axis in the plot:
xymax = 3 * numpy.sqrt(ns)
xymin = -xymax
NORTH = 1
SOUTH = 2
WEST = 3
EAST = 4 # constants
for step in range(ns):
for i in range(np):
direction = random.randint(1, 4)
if direction == NORTH:
ypositions[i] += 1
elif direction == SOUTH:
ypositions[i] -= 1
elif direction == EAST:
xpositions[i] += 1
elif direction == WEST:
xpositions[i] -= 1
# Plot just every plot_step steps
if (step + 1) % plot_step == 0:
plot(
xpositions,
ypositions,
"ko",
axis=[xymin, xymax, xymin, xymax],
title="%d particles after %d steps" % (np, step + 1),
savefig="tmp_%03d.eps" % (step + 1),
)
return xpositions, ypositions
def approximate(f, d=1, N_e=4, numint='Gauss-Legendre2',
filename='tmp.eps'):
nodes, elements = mesh(N_e, d, [0, 1])
A, b = assemble(nodes, elements, d, f, numint)
c = A.LUsolve(b)
print 'nodes:', nodes
print 'elements:', elements
print 'A:\n', A
print 'b:\n', b
print sm.latex(A, mode='plain')
#print sm.latex(b, mode='plain')
c = A.LUsolve(b)
print 'c:\n', c
print 'Plain interpolation:'
x = sm.Symbol('x')
f = sm.lambdify([x], f, modules='numpy')
f_at_nodes = [f(xc) for xc in nodes]
print f_at_nodes
if not symbolic:
xf = np.linspace(0, 1, 10001)
U = np.asarray(c)
xu, u = u_glob(U, elements, nodes)
from scitools.std import plot
plot(xu, u, 'r-',
xf, f(xf), 'b-',
legend=('u', 'f'),
savefig=filename)
def plot_u(u, x, t, n):
"""user_action function for solver."""
plt.plot(x, u, "r-", xlabel="x", ylabel="u", axis=[0, L, umin, umax], title="t=%f" % t[n], show=True)
# Let the initial condition stay on the screen for 2
# seconds, else insert a pause of 0.2 s between each plot
time.sleep(2) if t[n] == 0 else time.sleep(0.2)
plt.savefig("frame_%04d.png" % n) # for movie making
def main():
v0 = 0 # start velocity
T = 23 # here this time is chosen for falling 5000 m when only gravity is applied
dt = 0.01
C_D = 1.2 # Coeffi
rho = 0.79 # density of air (assumed constant)
rho_b = 1003.0 # kgm^3;
m = 80 # Mass of the jumper and parachute in kg
R = 0.5
A = pi * R**2 # Area of the sphere pi*r^2
V = m / rho_b # Volume of the falling object
d = 2.08 # Diameter of the sphere 2*r
mu = 1.8 * 10**-5 # dynamic viscosity
g = 9.81 # Gravitational constant
a_s = 3 * pi * rho * d * mu / (
rho_b * V) # Constant to be used for the Stokes model
a_q = 0.5 * C_D * rho * A / (
rho_b * V) # Constant to be used for the quadratic model
b = g * (-1 + rho / rho_b
) # A common constant for both Stokes and quad. model
rdm = rho * d / mu # A constant related to the Reynolds number
t, v = vm.solver(T, dt, a_s, a_q, b, v0, rdm)
plot(t, v, xlabel='t', ylabel='v', title='Parachute Jumper')
raw_input('Press any key.')
def test_solver_exact_discrete_solution():
def tilde_w(w, dt):
return (2. / dt) * asin(w * dt / 2.)
def u_numerical_exact(t):
return I * cos(tilde_w(w, dt) * t)
w = 2.5
I = 1.5
# Estimate period and time step
P = 2 * pi / w
num_periods = 4
T = num_periods * P
N = 5 # time steps per period
dt = P / N
u, t = solver(I, w, dt, T)
u_e = u_numerical_exact(t)
error = abs(u_e - u).max()
# Make a plot in a file, but not on the screen
from scitools.std import plot
plot(t,
u,
'bo',
t,
u_e,
'r-',
legend=('numerical', 'exact'),
show=False,
savefig='tmp.png')
assert error < 1E-14
def __call__(self, u, x, t, n):
"""user_action function for solver."""
if animate == 'u and u_exact':
plt.plot(x, u, 'r-',
x, u_exact(x, t[n]), 'b--',
xlabel='x', ylabel='u',
axis=[0, L, -self.ymax, self.ymax],
title='t=%f' % t[n], show=True)
else:
error = u_exact(x, t[n]) - u
local_max_error = np.abs(error).max()
# self.max_error holds the increasing amplitude error
if self.max_error == [] or \
local_max_error > max(self.max_error):
self.max_error.append(local_max_error)
self.max_error_t.append(t[n])
# Use user's ymax until the error exceeds that value.
# This gives a growing max value of the yaxis (but
# never shrinking)
self.ymax = max(self.ymax, max(self.max_error))
plt.plot(x, error, 'r-',
xlabel='x', ylabel='error',
axis=[0, L, -self.ymax, self.ymax],
title='t=%f' % t[n], show=True)
plt.savefig('%s_%04d.png' % (self.frame_name, n))
def plot_u(u, x, xv, y, yv, t, n):
if t[n] == 0:
time.sleep(1)
st.plot(x, u[:,1],'--',x, u_exact(x, t[n], a, Lx),'o')
filename = 'tmp_%04d.png' % n
st.savefig(filename)
time.sleep(0.1)
def test(b):
""" Runs test on SIR model, with variying beta """
beta = b #0.0005 or 0.0001 # Infection rate
v = 0.1 # Prob. of recovery per dt
S0 = 1500 # Init. No. of suspectibles
I0 = 1 # Init. No. of infected
R0 = 0. # Init. No. of resistant
U0 = [S0, I0, R0] # Initial conditions
T = 60 # Duration, days
dt = 0.5 # Time step length in days
n = T/dt # No. of solve steps
f = RHS(v, dt, T, beta) # Get right hand side of equation
solver = ForwardEuler(f) # Select ODE solver method
solver.set_initial_condition(U0)
time_points = np.linspace(0, 60, n+1)
u, t = solver.solve(time_points)
S = u[:,0] # S is all data in array no 0
I = u[:,1] # I is all data in array no 1
R = u[:,2] # R is all data in array no 2
plot(t, S, t, I, t, R,
xlabel='Days', ylabel='Persons',
legend=['Suspectibles', 'Infected', 'Resistant'],
hold=('on'))
def random_walk_2D(np, ns, plot_step):
xpositions = zeros(np)
ypositions = zeros(np)
# extent of the axis in the plot:
xymax = 3*sqrt(ns); xymin = -xymax
NORTH = 1; SOUTH = 2; WEST = 3; EAST = 4 # constants
for step in range(ns):
for i in range(np):
direction = random_number.randint(1, 4)
if direction == NORTH:
ypositions[i] += 1
elif direction == SOUTH:
ypositions[i] -= 1
elif direction == EAST:
xpositions[i] += 1
elif direction == WEST:
xpositions[i] -= 1
# plot just every plot_step steps:
if (step+1) % plot_step == 0:
plot(xpositions, ypositions, 'ko',
axis=[xymin, xymax, xymin, xymax],
title='%d particles after %d steps' % \
(np, step+1),
hardcopy='tmp_%03d.eps' % (step+1))
return xpositions, ypositions
def plot(self, include_exact=True, plt=None):
"""
Add solver.u curve to the plotting object plt,
and include the exact solution if include_exact is True.
This plot function can be called several times (if
the solver object has computed new solutions).
"""
if plt is None:
import scitools.std as plt # can use matplotlib as well
plt.plot(self.solver.t, self.solver.u, '--o')
plt.hold('on')
theta2name = {0: 'FE', 1: 'BE', 0.5: 'CN'}
name = theta2name.get(self.solver.theta, '')
legends = ['numerical %s' % name]
if include_exact:
t_e = linspace(0, self.problem.T, 1001)
u_e = self.problem.exact_solution(t_e)
plt.plot(t_e, u_e, 'b-')
legends.append('exact')
plt.legend(legends)
plt.xlabel('t')
plt.ylabel('u')
plt.title('theta=%g, dt=%g' %
(self.solver.theta, self.solver.dt))
plt.savefig('%s_%g.png' % (name, self.solver.dt))
return plt
def test_vertical_motion():
beta = 0.9
omega = np.sqrt(beta/(1-beta))
#Find num_periods. Recall that P=2*pi for scaled pendulum
#oscillations, while here we don't have gravity driven
#oscillations, but elastic oscillations with frequence omega.
period = 2*np.pi/omega
#we want T=N*period
N = 5
# simulate function has T = 2*pi*num_periods
num_periods = 5/omega
n = 600
time_steps_per_period = omega*n
y_exact = lambda t: -0.1*np.cos(omega*t)
x, y, theta, t = simulate(
beta=beta, Theta=0, epsilon=0.1,
num_periods=num_periods,
time_steps_per_period=time_steps_per_period,
plot=False)
tol = 0.00055
assert np.abs(x.max()) < tol
y_e = y_exact(t)
diff = np.abs(y_e - y).max()
if diff > tol:
plt.plot(t, y, t, y_e, legend=['y', 'exact'])
raw_input('Error in test_vertical_motion; type CR: ')
assert diff < tol, 'diff=%g' % diff
def comparison_plot(u, Omega, u_e=None, filename='tmp.eps',
plot_title='', ymin=None, ymax=None):
x = sp.Symbol('x')
u = sp.lambdify([x], u, modules="numpy")
if len(Omega) != 2:
raise ValueError('Omega=%s must be an interval (2-list)' % str(Omega))
# When doing symbolics, Omega can easily contain symbolic expressions,
# assume .evalf() will work in that case to obtain numerical
# expressions, which then must be converted to float before calling
# linspace below
if not isinstance(Omega[0], (int,float)):
Omega[0] = float(Omega[0].evalf())
if not isinstance(Omega[1], (int,float)):
Omega[1] = float(Omega[1].evalf())
resolution = 401 # no of points in plot
xcoor = linspace(Omega[0], Omega[1], resolution)
# Vectorized functions expressions does not work with
# lambdify'ed functions without the modules="numpy"
approx = u(xcoor)
plot(xcoor, approx)
legends = ['approximation']
if u_e is not None:
exact = u_e(xcoor)
hold('on')
plot(xcoor, exact)
legends = ['exact']
legend(legends)
title(plot_title)
xlabel('x')
if ymin is not None and ymax is not None:
axis([xcoor[0], xcoor[-1], ymin, ymax])
savefig(filename)
def a(N):
def psi(x, i):
#return sin((i+1)*x)
return sin((2 * i + 1) * x)
#return sin((2*(i+1))*x)
u, c = least_squares_numerical(
f,
psi,
N,
x,
#integration_method='trapezoidal',
integration_method='scipy',
orthogonal_basis=True)
os.system('rm -f *.png')
u_sum = 0
print 'XXX c', c
for i in range(N + 1):
u_sum = u_sum + c[i] * psi(x, i)
plt.plot(x,
f(x),
'-',
x,
u_sum,
'-',
legend=['exact', 'approx'],
title='Highest frequency component: sin(%d*x)' % (2 * i + 1),
axis=[x[0], x[-1], -1.5, 1.5])
plt.savefig('tmp_frame%04d.png' % i)
time.sleep(0.3)
cmd = 'avconv -r 2 -i tmp_frame%04d.png -vcodec libtheora movie.ogg'
def slides_waves():
L = 10.
a = 0.5
B0 = 1.
A = B0/5
t0 = 3.
v = 1.
x0 = lambda t: L/4 + v*t if t < t0 else L/4 + v*t0
sigma = 1.0
def bottom(x, t):
return (B(x, a, L, B0) + S(x, A, x0(t), sigma))
def depth(x, t):
return -bottom(x, t)
import scitools.std as plt
x = np.linspace(0, L, 101)
plt.plot(x, bottom(x, 0), x, depth(x, 0), legend=['bottom', 'depth'])
plt.show()
raw_input()
plt.figure()
dt_b = 0.01
solver(I=0, V=0,
f=lambda x, t: -(depth(x, t-dt_b) - 2*depth(x, t) + depth(x, t+dt_b))/dt_b**2,
c=lambda x, t: np.sqrt(np.abs(depth(x, t))),
U_0=None,
U_L=None,
L=L,
dt=0.1,
C=0.9,
T=20,
user_action=PlotSurfaceAndBottom(B, S, a, L, B0, A, x0, sigma, umin=-2, umax=2),
stability_safety_factor=0.9)
def test_vertical_motion():
beta = 0.9
omega = np.sqrt(beta/(1-beta))
# Find num_periods. Recall that P=2*pi for scaled pendulum
# oscillations, while here we don't have gravity driven
# oscillations, but elastic oscillations with frequency omega.
period = 2*np.pi/omega
# We want T = N*period
N = 5
# simulate function has T = 2*pi*num_periods
num_periods = 5/omega
n = 600
time_steps_per_period = omega*n
y_exact = lambda t: -0.1*np.cos(omega*t)
x, y, theta, t = simulate(
beta=beta, Theta=0, epsilon=0.1,
num_periods=num_periods,
time_steps_per_period=time_steps_per_period,
plot=False)
tol = 0.00055 # ok tolerance for the above resolution
# No motion in x direction is epxected
assert np.abs(x.max()) < tol
# Check motion in y direction
y_e = y_exact(t)
diff = np.abs(y_e - y).max()
if diff > tol: # plot
plt.plot(t, y, t, y_e, legend=['y', 'exact'])
raw_input('Error in test_vertical_motion; type CR:')
assert diff < tol, 'diff=%g' % diff
def visualize_front(u, t, window_width, savefig=False):
"""
Visualize u and the exact solution vs t, using a
moving plot window and continuous drawing of the
curves as they evolve in time.
Makes it easy to plot very long time series.
P is the approximate duration of one period.
"""
import scitools.std as st
from scitools.MovingPlotWindow import MovingPlotWindow
umin = 1.2 * u.min()
umax = -umin
plot_manager = MovingPlotWindow(window_width=window_width,
dt=t[1] - t[0],
yaxis=[umin, umax],
mode='continuous drawing')
for n in range(1, len(u)):
if plot_manager.plot(n):
s = plot_manager.first_index_in_plot
st.plot(t[s:n + 1],
u[s:n + 1],
'r-1',
title='t=%6.3f' % t[n],
axis=plot_manager.axis(),
show=not savefig) # drop window if savefig
if savefig:
print 't=%g' % t[n]
st.savefig('tmp_vib%04d.png' % n)
plot_manager.update(n)
def simple_visualize(t,u ):
st.plot(t, u, 'r--o',
legend=['numerical', 'exact'],
xlabel='t',
ylabel='u',
title='theta=%g, dt=%g' % (theta, dt),
savefig='%s_%g.png' % (theta2name[theta], dt),
show=True)
def plot_fe_mesh(nodes, elements, element_marker=[0, 0.1]):
"""Illustrate elements and nodes in a finite element mesh."""
plt.hold('on')
all_x_L = [nodes[elements[e][0]] for e in range(len(elements))]
element_boundaries = all_x_L + [nodes[-1]]
for x in element_boundaries:
plt.plot([x, x], element_marker, 'm--') # m gives dotted lines
plt.plot(nodes, [0]*len(nodes), 'ro2')
def plot_fe_mesh(nodes, elements, element_marker=[0, 0.1]):
"""Illustrate elements and nodes in a finite element mesh."""
plt.hold('on')
all_x_L = [nodes[elements[e][0]] for e in range(len(elements))]
element_boundaries = all_x_L + [nodes[-1]]
for x in element_boundaries:
plt.plot([x, x], element_marker, 'm--') # m gives dotted eps/pdf lines
plt.plot(nodes, [0]*len(nodes), 'ro2')
def plot_u(u, x, xv, y, yv, t, n):
if t[n] == 0:
time.sleep(1)
#mesh(x, y, u, title='t=%g' % t[n])
st.plot(y, u[1,:], title='t=%g' % t[n], ylim=[-2.4, 2.4])
filename = 'tmp_%04d.png' % n
st.savefig(filename)
time.sleep(0.4)
def _test():
# from InverseFunction import InverseFunction as I
from scitools.std import log, linspace, plot
def f(x):
return log(x)
x = linspace(1, 5, 101)
f_inv = InverseFunction(f, x)
plot(x,f(x),x, f_inv.values,legend=['log(x)', 'inv log(x)'])
raw_input('Press Enter to quit: ')
def case(N, problem=ode.problem2, showplot=False):
I = 0
t0 = time.clock()
t = np.linspace(0, 5, N)
u, t = ode.solver(problem, I, t, ode.RK2)
t1 = time.clock()
if showplot:
from scitools.std import plot
plot(t, u)
return t1 - t0
def plot_u(u, x, t, n):
"""user_action function for solver."""
st.plot(x, u, 'r-',
xlabel='x', ylabel='u',
axis=[0, L, umin, umax],
title='t=%f' % t[n])
# Let the initial condition stay on the screen for 2
# seconds, else insert a pause of 0.2 s between each plot
time.sleep(2) if t[n] == 0 else time.sleep(0.2)
st.savefig('frame_%04d.png' % n) # for movie making
def _test():
x0 = float(sys.argv[1])
x, info = Newton(_g, x0, _dg, store=True)
print('root: %.16g' % x)
for i in range(len(info)):
print('Iteration %2d: f(%g)=%g' % (i, info[i][0], info[i][1]))
x = linspace(-7, 7, 401)
y = _g(x)
plot(x, y)
def simple_visualize(t, u):
st.plot(t,
u,
'r--o',
legend=['numerical', 'exact'],
xlabel='t',
ylabel='u',
title='theta=%g, dt=%g' % (theta, dt),
savefig='%s_%g.png' % (theta2name[theta], dt),
show=True)
def plot_u(u, x, t, n):
plt.plot(x, u, "r-", axis=[0, L, umin, umax], title="t=%f" % t[n])
if framefiles:
plt.savefig("tmp_frame%04d.png" % n)
if t[n] == 0:
time.sleep(2)
elif not framefiles:
# It takes time to write files so pause is needed
# for screen only animation
time.sleep(0.2)
def plot_u(u, x, t, n):
from scitools.std import plot
umin = -0.1; umax = 1.1 # axis limits for plotting
plot(x, u, 'r-', axis=[0, L, umin, umax], title='t=%f' % t[n])
# Pause the animation initially, otherwise 0.2 s between frames
if t[n] == 0:
time.sleep(2)
else:
time.sleep(0.2)
def graph(self, resolution=1001):
'''
graphs the Lagrange polynomial over self.points
'''
#from scitools.std import *
from scitools.std import zeros, linspace, plot, array
points = self.points
xlist = linspace(points[0, 0], points[-1, 0], resolution)
ylist = self.__call__(xlist)
plot(xlist, ylist)
def plot_u(u, x, t, n):
plt.plot(x, u, 'r-', axis=[0, L, umin, umax], title='t=%f' % t[n])
if framefiles:
plt.savefig('tmp_frame%04d.png' % n)
if t[n] == 0:
time.sleep(2)
elif not framefiles:
# It takes time to write files so pause is needed
# for screen only animation
time.sleep(0.2)
def plot_u(u, x, t, n):
"""user_action function for solver."""
plt.plot(x, u, 'r-',
xlabel='x', ylabel='u',
axis=[0, L, umin, umax],
title='t=%f' % t[n], show=True)
# Let the initial condition stay on the screen for 2
# seconds, else insert a pause of 0.2 s between each plot
time.sleep(2) if t[n] == 0 else time.sleep(0.2)
plt.savefig('frame_%04d.png' % n) # for movie making