%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('default')
Consider the intial value problem of solving the following ordinary differential equation:
\begin{align} \frac{dy}{dt} = f(y), \quad y(0) = y_0 \end{align}In such problems, $y(t)$ is an unknown function that needs to be determined. The right hand side of the differential equation, $f(y)$, represents a known function of $y$ that needs to be given explicitly.
We have seen a number of methods (Euler, Midpoint, RK4) that can be used to numerically solve such problems.
At its simplest, projectile motion involves solving
$$ \begin{align} \frac{d^2 x}{dt^2} &= 0 \\ \frac{d^2 y}{dt^2} &= -g \end{align} $$given the initial conditions for $x(0)$, $y(0)$, and the initial velocities.
We have learned that to solve such problems we can rewrite the equations as a system of four first order differential equations:
given the initial conditions $x(0) = x_0$, $y(0) = y_0$, $v_x(0) = v_{x0}$, $v_y(0) = v_{y0}$
Let's solve the projectile motion problem using the midpoint scheme.
1. Estimate slope $s_1$ at $t$
$$ s_1 = f(y(t), t) $$2. Use $s_1$ to estimate the midpoint between $t$ and $t + \Delta t$
\begin{align} y^* &= y ( t + \Delta t) \\ &= y(t) + \frac{\Delta t}{2} s_1 \end{align}3. Use $y^*$ to get the the slope $s_2$ at the midpoint
$$ s_2 = f( y^*, t + \frac{\Delta t}{2} ) $$4. Use $s_2$ to estimate $y(t + \Delta t)$
$$ y(t + \Delta t) = y(t) + \Delta t s_2 $$def solve_1(tmax = 2, dt = 0.01,
x0 = 0, y0 = 0, vx0 = 3, vy0 = 10):
g = 9.81
# Allocate memory for arrays and set to zero
N = round(tmax/dt)
x = np.zeros(N)
y = np.zeros(N)
vx = np.zeros(N)
vy = np.zeros(N)
t = np.zeros(N)
# Initial values
x[0] = x0
y[0] = y0
vx[0] = vx0
vy[0] = vy0
t[0] = 0
# Calculate the solution using the Midpoint Method:
for i in range(N-1):
# Estimate slope s1 at t
s1_x = vx[i]
s1_y = vy[i]
s1_vx = 0
s1_vy = -g
# Use s1 to estimate the midpoint between t and t+Δt
x_tmp = x[i] + s1_x*dt/2
y_tmp = y[i] + s1_y*dt/2
vx_tmp = vx[i] + s1_vx*dt/2
vy_tmp = vy[i] + s1_vy*dt/2
# Estimate the slope s2 at the midpoint
s2_x = vx_tmp
s2_y = vy_tmp
s2_vx = 0
s2_vy = -g
# Use s2 to estimate y(t+Δt)
x[i+1] = x[i] + s2_x*dt
y[i+1] = y[i] + s2_y*dt
vx[i+1] = vx[i] + s2_vx*dt
vy[i+1] = vy[i] + s2_vy*dt
t[i+1] = t[i] + dt
return x, y, t
x, y, t = solve_1()
plt.plot(x, y)
plt.xlabel('x (m)')
plt.ylabel('y (m)')
plt.title('Projectile Motion solved with Midpoint Scheme')
Going back to the general form of an ODE we can also think about the unknown function as being a vector-valued function such as
\begin{align} \frac{d \vec{q} }{dt} = \vec{F}( \vec{q} ), \quad \vec{q}(0) = \vec{q}_0 \end{align}where $ \vec{q}(t) = \left( q_1(t), q_2(t), \ldots, q_m(t) \right)$ is a vector in which every component is a function.
We can define a vector for the projectile motion problem like this
$$ \vec{q}(t) = \left( x(t), y(t), v_x(t), v_y(t) \right) $$This vector represents the state of the projectile motion system. Likewise, we can write down
\begin{align} \vec{F}( \vec{q} ) &= \vec{F}(x, y, v_x, v_y) = (v_x, v_y, 0, -g) \\ \vec{q}_0 &= \left( x_0, y_0, v_{x0}, v_{y0} \right) \\ \end{align}This is just a compact rewriting of our system of four first-order DEs in a vector formulation.
This allows us to generalize our midpoint scheme as follows:
def solver(F, q0, t):
"""
Solve the ODE dq/dt = F(q) with q(0) = q0
where F is a vector valued function
q0 is a vector of initial conditions
t is time array variable
"""
# count number points of time to solve
N = len(t)
# count number of indepedent variables
m = len(q0)
# Allocate memory set arrays to zero
q = np.zeros((m, N))
# Initial values
q[:, 0] = q0
# Calculate the solution using the Midpoint Method:
for i in range(N-1):
dt = t[i+1] - t[i]
# Estimate slope s1 at t
s1 = F( q[:, i], t[i] )
s1 = np.asarray(s1) # ensure s1 is a numpy array
# Use s1 to estimate the midpoint between t and t+Δt
q_tmp = q[:, i] + s1*dt/2
# Estimate the slope s2 at the midpoint
s2 = F( q_tmp, t[i]+ dt/2 )
s2 = np.asarray(s2) # ensure s2 is a numpy array
# Use s2 to estimate y(t+Δt)
q[:, i+1] = q[:, i] + s2*dt
return q
Notice that our solver is now very general. The same code could be used for many different differential equations.
We can formulate our specific projectile motion problem by defining a particular right-hand-side function.
g = 9.81
tmax = 2
dt = 0.01
# create an array for the time variables
t = np.arange(0, tmax, dt)
# define the right hand side, F(q)
def F( q, t ):
# separate out the variables
x, y, vx, vy = q
# evaluate the right hand side of the ODE
dqdt = [vx, vy, 0, -g ]
return dqdt
# define initial values (x0, y0, vx0, vy0)
q0 = [0, 0, 3, 10]
Then all we need to do is pass in this specific function for F and the initial conditions q0 into our general solver.
x, y, vx, vy = solver(F, q0, t)
plt.plot(x,y)
plt.xlabel('x (m)')
plt.ylabel('y (m)')
plt.title('Projectile Motion solved with Midpoint Scheme')
This abstraction separating out the physical model from the numerical solver, also allows us to try new solvers with almost no changes the code at all.
The Python package scipy comes with its own ODE solver. Here is the same projectile motion problem using a different solver.
# we need to import the SciPy integrate module
from scipy import integrate
sol = integrate.odeint(F, q0, t)
x, y, vx, vy = sol.T
plt.plot(x,y)
plt.xlabel('x (m)')
plt.ylabel('y (m)')
plt.title('Projectile Motion solved with integrate.odeint')
Compare the previous two cell blocks. What differences can you observe?
where
$$ v = v(x, y) $$is the velocity of a seismic waves that depends on the physical composition of the rock within the ground.
$ v(x,y)$ is assumed to be known and the goal is to determine the path of the rays.
Measurements of the travel times of the rays are known and the goal is to estimate $v(x,y)$.
Let's start with a rock structure that is getting linearly denser with depth.
were $v$ is in km/s if $\rho$ is in $10^3$ km/m$^3$.
So,
y = np.arange(-40, 0, 0.1)
x = np.arange(0, 120, 0.1)
def ρ(x, y):
return (2000 - 1500/40*y)/1000
def v(x, y):
return 0.33 + 2.20*ρ(x, y)
fig, axes = plt.subplots(1,2, figsize=(10,6))
plt.sca(axes[0])
plt.plot(ρ(0, y), y)
plt.xlabel(r'$\rho$')
plt.ylabel(r'$y$ (km)')
plt.sca(axes[1])
plt.plot(v(0, y), y)
plt.xlabel(r'$v$ km/s')
plt.ylabel(r'$y$ (km)')
plt.show()
Let's send out a ray!
sin = np.sin
cos = np.cos
π = np.pi
tmax = 22
dt = 0.1
# create an array for the time variables
t = np.arange(0, tmax, dt)
# define the right hand side, F(q)
def F( q, t ):
# separate out the variables
x, y, θ = q
# evaluate the right hand side of the ODE
dxdt = v(x,y)*sin(θ)
dydt = -v(x,y)*cos(θ)
dθdt = -cos(θ)*dvdx(x, y) - sin(θ)*dvdy(x,y)
dqdt = [dxdt, dydt, dθdt]
return dqdt
# define initial values (x0, y0, θ0)
q0 = [0, 0, π/4]
We'll also need the derivatives of $v(x,y)$. We can estimate those numerically
# Centered difference scheme
def NDc(f, x0, dx):
return (f(x0+dx) - f(x0-dx)) / (2*dx)
def dvdx(x, y):
return NDc(lambda x: v(x,y), x, 1)
def dvdy(x, y):
return NDc(lambda y: v(x,y), y, 1)
sol = integrate.odeint(F, q0, t)
x, y, θ = sol.T
plt.plot(x, y)
plt.xlabel('x (km)')
plt.ylabel('y (km)')
plt.title('Ray Trace')
plt.show()
We could do this for a number of different initial angles:
def vary_theta(x0=0, y0=0, θmin = π/4, θmax = π/2):
for θ0 in np.arange(θmin, θmax, 0.1):
q0 = [x0, y0, θ0]
sol = integrate.odeint(F, q0, t)
x, y, θ = sol.T
plt.plot(x, y, 'k-')
plt.xlabel('x (km)')
plt.ylabel('y (km)')
plt.ylim(-30, 0)
plt.xlim(0, 120)
plt.title('Ray Traces')
vary_theta()
Probably a good idea to show the density/velocity as a background to these traces.
def plot_velocity():
X, Y = np.mgrid[0:120:0.1, -30:0:0.1]
V = v(X,Y)
plt.contourf(X, Y, V, 10, cmap=plt.cm.Blues)
plt.xlabel('x (km)')
plt.ylabel('y (km)')
plt.colorbar()
plot_velocity()
Of course we could put these together:
plot_velocity()
vary_theta()
What if there were some variations of rock density.
def ρ_perturbation(x, y):
r =0
r += 1*np.exp(-(x-60)**2/ 20 - (y+15)**2/20)
r += 0.1*np.exp(-(x-20)**2/ 30 - (y+5)**2/20)
return r
def ρ(x, y):
return (2000 - 1500/40*y)/1000 + ρ_perturbation(x, y)
plot_velocity()
vary_theta()
Seismic pulse at depth
plot_velocity()
vary_theta(x0=80, y0 = -25, θmin = -3*π/4, θmax = -π/2)