An important tool in the computational tool box is to find roots of equations for which no closed form solutions exist:
We want to find the roots $x_0$ of
$$ f(x_0) = 0 $$The equations of motion for the projectile with linear air resistance (see 11 ODE applications) can be solved exactly.
As a reminder: the linear drag force is \begin{align} \mathbf{F}_1 &= -b_1 \mathbf{v}\\ b &:= \frac{b_1}{m} \end{align}
Equations of motion with force due to gravity $\mathbf{g} = -g \hat{\mathbf{e}}_y$
\begin{align} \frac{d\mathbf{r}}{dt} &= \mathbf{v}\\ \frac{d\mathbf{v}}{dt} &= - g \hat{\mathbf{e}}_y -b \mathbf{v} \end{align}(Following Wang Ch 3.3.2)
Solve $x$ component of the velocity
$$ \frac{dv_x}{dt} = -b v_x $$by integration:
$$ v_x(t) = v_{0x} \exp(-bt) $$The drag force reduces the forward velocity to 0.
Integrate again to get the $x(t)$ component
$$ x(t) = x_0 + \frac{v_{0x}}{b} \left[1 - \exp(-bt)\right] $$Integrating the $y$ component of the velocity
$$ \frac{dv_y}{dt} = -g - b v_y $$gives
$$ v_y = \left(v_{0y} + \frac{g}{b}\right) \exp(-bt) - \frac{g}{b} $$and integrating again
$$ y(t) = y_0 + \frac{v_{0y} + \frac{g}{b}}{b} \left[1 - \exp(-bt)\right] - \frac{g}{b} t $$(Note: This shows immediately that the terminal velocity is
$$ \lim_{t\rightarrow\infty} v_y(t) = - \frac{g}{b}, $$i.e., the force of gravity is balanced by the drag force.)
To obtain the trajectory $y(x)$ eliminate time (and for convenience, using the origin as the initial starting point, $x_0 = 0$ and $y_0 = 0$. Solve $x(t)$ for $t$
$$ t = -\frac{1}{b} \ln \left(1 - \frac{bx}{v_{0x}}\right) $$and insert into $y(t)$:
$$ y(x) = \frac{x}{v_{0x}} \left( v_{0y} + \frac{g}{b} \right) + \frac{g}{b^2} \ln \left(1 - \frac{bx}{v_{0x}}\right) $$Plot the analytical solution $y(x)$ for $\theta = 30^\circ$ and $v_0 = 100$ m/s.
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
plt.style.use('ggplot')
The function y_lindrag()
should compute $y(x)$.
def y_lindrag(x, v0, b1=0.2, g=9.81, m=0.5):
b = b1/m
v0x, v0y = v0
return x/v0x * (v0y + g/b) + g/(b*b) * np.log(1 - b*x/v0x)
def initial_v(v, theta):
x = np.deg2rad(theta)
return v * np.array([np.cos(x), np.sin(x)])
The analytical function drops very rapidly towards the end ($> 42$ m – found by manual trial-and-error plotting...) so in order to nicely plot the function we use a fairly coarse sampling of points along $x$ for the range $0 \le x < 42$ and very fine sampling for the last 2 m ($42 \le x < 45$):
X = np.concatenate([np.linspace(0, 42, 100), np.linspace(42, 45, 10000)])
Evaluate the function for all $x$ values:
Y = y_lindrag(X, initial_v(100, 30), b1=1)
/var/folders/sm/37rm_wm16tq9qsf4n5md98ph0000gp/T/ipykernel_61984/482518341.py:4: RuntimeWarning: invalid value encountered in log return x/v0x * (v0y + g/b) + g/(b*b) * np.log(1 - b*x/v0x)
(The warning can be ignored, it just means that some of our X
values were not approriate and outside the range of validity – when the argument of the logarithm becomes ≤0).
To indicate the ground we also plot a dashed black line: note that the analytical solution goes below the dashed line.
plt.plot(X, Y)
plt.xlabel("$x$ (m)")
plt.ylabel("$y$ (m)")
plt.hlines([0], X[0], X[-1], colors="k", linestyles="--");
Compare to the numerical solution (from 11 ODE Applications):
import ode
def simulate(v0, h=0.01, b1=0.2, g=9.81, m=0.5):
def f(t, y):
# y = [x, y, vx, vy]
return np.array([y[2], y[3], -b1/m * y[2], -g - b1/m * y[3]])
vx, vy = v0
t = 0
positions = []
y = np.array([0, 0, vx, vy], dtype=np.float64)
while y[1] >= 0:
positions.append([t, y[0], y[1]]) # record t, x and y
t += h
y[:] = ode.rk4(y, f, t, h)
return np.array(positions)
r = simulate(initial_v(100, 30), h=0.01, b1=1)
plt.plot(X, Y, lw=2, label="analytical")
plt.plot(r[:, 1], r[:, 2], '--', label="RK4")
plt.legend(loc="best")
plt.xlabel("$x$ (m)"); plt.ylabel("$y$ (m)")
plt.hlines([0], X[0], X[-1], colors="k", linestyles="--");
The RK4 solution tracks the analytical solution perfectly (and we also programmed it to not go below ground...)
OPTIONAL: Show the residual
$$ r = y_\text{numerical}(x) - y_\text{analytical} $$residual = r[:, 2] - y_lindrag(r[:, 1], initial_v(100, 30), b1=1)
plt.plot(r[:, 1], residual)
plt.xlabel("$x$ (m)"); plt.ylabel("residual $r$ (m)");
How far does the ball or projectile fly, i.e., that value $x=R$ where $y(R) = 0$:
$$ \frac{R}{v_{0x}} \left( v_{0y} + \frac{g}{b} \right) + \frac{g}{b^2} \ln \left(1 - \frac{bR}{v_{0x}}\right) = 0 $$This transcendental equation can not be solved in terms of elementary functions.
Use a root finding algorithm.
Bisection is the simplest (but very robust) root finding algorithm that uses trial-and-error:
More specifically:
Test that the initial bracket contains a root; if not, raise a ValueError
.
If either of the bracket points is a root then return the bracket point.
Allow Nmax
iterations or until the convergence criterion eps
is reached.
Raise a RuntimeWarning
if no root was found after Nmax
iterations, but print the best guess and the error.
NOTE: It's better to fail with an exception than to return incorrect values. Calling code can always use try... except
to catch your exception.
def bisection(f, a, b, Nmax=100, eps=1e-14):
"""Find root for function f(x), bracketed by a <= x <= b."""
fa, fb = f(a), f(b)
if np.abs(fa) < eps:
return a
if np.abs(fb) < eps:
return b
if (fa*fb) > 0:
raise ValueError(f"bisect: Initial bracket [{a}, {b}] "
f"does not contain a single root")
for iteration in range(Nmax):
x = (a + b)/2
fx = f(x)
if np.abs(fx) < eps:
break
if f(a) * fx < 0:
# root between a and x: need to change right boundary b <- x
b = x
else:
# root between x and b: need to change left boundary a <- x
a = x
else:
raise RuntimeError(f"bisect: no root found after {Nmax} iterations (eps={eps}); "
f"best guess is {x} with error {fx}")
return x
bisection
¶It's always a good idea to test new code on input where you know the answer.
We need a function with known roots, e.g.
$$ p(x) = (x + 1)(x - 2)x $$with roots -1, 0, 2.
def p(x):
return (x+1)*(x-2)*x
X = np.linspace(-1.5, 2.5, 50)
X0 = np.array([-1, 0, 2])
plt.plot(X, p(X), 'r-')
plt.plot(X0, p(X0), 'r*');
Now find some roots with bisection
:
x0 = bisection(p, 0.3, 3)
x0, p(x0)
(1.9999999999999998, -1.3322676295501877e-15)
x0 = bisection(p, -100, -0.5)
x0, p(x0)
(-1.0000000000000022, -6.661338147750959e-15)
x0 = bisection(p, -0.5, 0.5)
x0, p(x0)
(0.0, -0.0)
All the roots are accurate to the given precision.
Also check that we correctly raise a ValueError
if the bracket is incorrect:
x0 = bisection(p, -1.5, 1)
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Input In [38], in <module>
----> 1 x0 = bisection(p, -1.5, 1)
Input In [12], in bisection(f, a, b, Nmax, eps)
8 return b
10 if (fa*fb) > 0:
---> 11 raise ValueError(f"bisect: Initial bracket [{a}, {b}] "
12 f"does not contain a single root")
14 for iteration in range(Nmax):
15 x = (a + b)/2
ValueError: bisect: Initial bracket [-1.5, 1] does not contain a single root
Define the trial function f
.
Note that our y_lindrag()
function depends on x
and v
but bisect()
only accepts functions f
that depend on a single variable, $f(x)$. We therefore have to wrap y_lindrag(x, v)
into a function f(x)
that sets v
already to a value outside the function: Python's scoping rules say that inside the function f(x)
, the variable x
has the value assigned to the argument of f(x)
but any other variables such as v
or b1
, which were not defined inside f
, will get the value that they had outside f
in the enclosing code.
def f(x):
b1 = 1.
v0 = initial_v(100, 30)
return y_lindrag(x, v0, b1=b1)
However, it is a bit messy to rely on using variables from outside the scope. It is cleaner to only use variables defined inside the function. We can achieve this effect by using keyword arguments where we set the default value to the value that we want to use inside our function: In this case, b1
and v0
are set as keyword arguments, with the understanding that when f(x)
is called later, the keyword arguments are never changed and keep their defaults:
v0 = initial_v(100, 30)
b1 = 1.0
def f(x, v0=v0, b1=b1):
return y_lindrag(x, v0, b1=b1)
The initial bracket $[a_\text{initial}, b_\text{initial}]$ is a little bit difficult for this function: choose the right bracket near the point where the argument of the logarithm becomes 0 (which is actually the maximum $x$ value $\lim_{t\rightarrow +\infty} x(t) = \frac{v_{0x}}{b}$):
$$ b_\text{initial} = \frac{v_{0x}}{b} - \epsilon' $$where $\epsilon'$ is a small number.
# use v0 and b1 from above!
m = 0.5
b = b1/m
bisection(f, 0.1, v0[0]/b - 1e-12, eps=1e-6)
43.30067423347077
Note that this solution is not the maximum value $\lim_{t\rightarrow +\infty} x(t) = \frac{v_{0x}}{b}$:
v0[0]/b
43.30127018922194
b1 = 1.
m = 0.5
b = b1/m
speed = 100
u = []
for theta in np.arange(1, 90):
v0 = initial_v(speed, theta)
def f(x, v0=v0, b1=b1):
return y_lindrag(x, v0, b1=b1)
R = bisection(f, 0.1, v0[0]/b - 1e-16, eps=1e-5)
u.append((theta, R))
u = np.array(u)
/var/folders/sm/37rm_wm16tq9qsf4n5md98ph0000gp/T/ipykernel_61984/482518341.py:4: RuntimeWarning: divide by zero encountered in log return x/v0x * (v0y + g/b) + g/(b*b) * np.log(1 - b*x/v0x)
plt.plot(u[:, 0], u[:, 1])
plt.xlabel(r"launch angle $\theta$")
plt.ylabel(r"range $R$ (m)");
# make labels with degree signs
ax = plt.gca()
ax.xaxis.set_major_formatter(plt.matplotlib.ticker.StrMethodFormatter(r"{x:g}$\!^\circ$"))
Write a function find_range()
to calculate the range for a given initial velocity $v_0$ and plot $R(\theta)$ for $10\,\text{m/s} ≤ v_0 ≤ 100\,\text{m/s}$.
For some angles, our algorithm might not convert with the default Nmax
. In order to avoid our code stopping with the RuntimeError
exception, we catch the exception and just print a message, using the try/except/else
block (note that the else
is code that is only executed if the try
succeeded).
def find_range(speed, b1=1, m=0.5):
b = b1/m
u = []
for theta in np.linspace(1, 90, 90):
v0 = initial_v(speed, theta)
def f(x, v0=v0, b1=b1):
return y_lindrag(x, v0, b1=b1)
try:
R = bisection(f, 0.01, v0[0]/b - 1e-12, eps=1e-5)
except RuntimeError:
print(f"NOTE: bisect did not converge for {theta}")
else:
u.append((theta, R))
return np.array(u)
for speed in (10, 25, 50, 75, 100):
u = find_range(speed)
plt.plot(u[:, 0], u[:, 1], label="{} m/s".format(speed))
plt.legend(loc="best")
plt.xlabel(r"launch angle $\theta$")
plt.ylabel(r"range $R$ (m)");
plt.gca().xaxis.set_major_formatter(plt.matplotlib.ticker.StrMethodFormatter(r"{x:g}$\!^\circ$"))
/var/folders/sm/37rm_wm16tq9qsf4n5md98ph0000gp/T/ipykernel_61984/482518341.py:4: RuntimeWarning: invalid value encountered in log return x/v0x * (v0y + g/b) + g/(b*b) * np.log(1 - b*x/v0x)
NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0
As a bonus, find the dependence of the optimum launch angle on the initial velocity, i.e., that angle that leads to the largest range:
np.argmax(u[:, 1])
10
velocities = np.linspace(5, 100, 100)
results = [] # (speed, theta_opt)
for speed in velocities:
u = find_range(speed)
thetas, ranges = u.transpose()
# find index for the largest range and pull corresponding theta
theta_opt = thetas[np.argmax(ranges)]
results.append((speed, theta_opt))
results = np.array(results)
plt.plot(results[:, 0], results[:, 1])
plt.xlabel(r"initial speed $v_0$ (m/s)")
plt.ylabel(r"optimal launch angle $\theta_\mathrm{best}$");
plt.gca().yaxis.set_major_formatter(plt.matplotlib.ticker.StrMethodFormatter(r"{x:g}$\!^\circ$"))
/var/folders/sm/37rm_wm16tq9qsf4n5md98ph0000gp/T/ipykernel_61984/482518341.py:4: RuntimeWarning: invalid value encountered in log return x/v0x * (v0y + g/b) + g/(b*b) * np.log(1 - b*x/v0x)
NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0 NOTE: bisect did not converge for 90.0
The launch angle decreases with the velocity.
The steps in the graph are an artifact of choosing to only calculate the trajectories for integer angles (see for theta in np.linspace(1, 90. 90)
in find_range()
).
(see derivation in class and in the second part of 12_rootfinding.pdf (PDF) or Newton's Method on MathWorld)
Implement the Newton-Raphson algorithm
while
$|f(x)| > \epsilon$
$\Delta x = -\frac{f(x))}{f'(x)}$
$x \leftarrow x + \Delta x$
Use the central difference algorithm with step size $h$ to compute $f'(x)$. (In other cases you might want to use the analytical derivative if it is available.)
Test with $g(x)$.
$$ g(x) = 2 \cos x - x $$Bonus: test performance of newton_raphson()
and bisection()
(e.g. for our function $g(x)$).
def g(x):
return 2*np.cos(x) - x
xvals = np.linspace(0, 7, 30)
plt.plot(xvals, np.zeros_like(xvals), 'k--')
plt.plot(xvals, g(xvals))
[<matplotlib.lines.Line2D at 0x113a171d0>]
def newton_raphson(f, x, h=1e-3, Nmax=100, eps=1e-14):
"""Find root x0 so that f(x0)=0 with the Newton-Raphson algorithm"""
for iteration in range(Nmax):
fx = f(x)
if np.abs(fx) < eps:
break
df = (f(x + h/2) - f(x - h/2))/h
Delta_x = -fx/df
x += Delta_x
else:
raise RuntimeError(f"Newton-Raphson: no root found after {Nmax} iterations (eps={eps}); "
f"best guess is {x} with error {fx}")
return x
x0 = newton_raphson(g, 2)
print(x0)
1.0298665293222586
g(x0)
6.661338147750939e-16
But note that the algorithm only converges well near the root. With other values it might not converge at all:
newton_raphson(g, -1000)
--------------------------------------------------------------------------- RuntimeError Traceback (most recent call last) Cell In[48], line 1 ----> 1 newton_raphson(g, -1000) Cell In[29], line 11, in newton_raphson(f, x, h, Nmax, eps) 9 x += Delta_x 10 else: ---> 11 raise RuntimeError(f"Newton-Raphson: no root found after {Nmax} iterations (eps={eps}); " 12 f"best guess is {x} with error {fx}") 13 return x RuntimeError: Newton-Raphson: no root found after 100 iterations (eps=1e-14); best guess is 66991.66522480128 with error -106154.40543292962
newton_raphson(g, 10)
1.0298665293222589
newton_raphson(g, 15)
1.0298665293222589
Let's look how Newton-Raphson iterates: also return all intermediate $x$ values:
def newton_raphson_with_history(f, x, h=1e-3, Nmax=100, eps=1e-14):
xvals = []
for iteration in range(Nmax):
fx = f(x)
if np.abs(fx) < eps:
break
df = (f(x + h/2) - f(x - h/2))/h
Delta_x = -fx/df
x += Delta_x
xvals.append(x)
else:
raise RuntimeError(f"Newton-Raphson: no root found after {Nmax} iterations (eps={eps}); "
f"best guess is {x} with error {fx}")
return x, np.array(xvals)
x = {}
x0, xvals = newton_raphson_with_history(g, 1.5)
x[1.5] = xvals
print("root x0 = {} after {} iterations".format(x0, len(xvals)))
x0, xvals = newton_raphson_with_history(g, 5)
x[5] = xvals
print("root x0 = {} after {} iterations".format(x0, len(xvals)))
x0, xvals = newton_raphson_with_history(g, 10)
x[10] = xvals
print("root x0 = {} after {} iterations".format(x0, len(xvals)))
x0, xvals = newton_raphson_with_history(g, 1500)
x[1500] = xvals
print("root x0 = {} after {} iterations".format(x0, len(xvals)))
root x0 = 1.0298665293222589 after 4 iterations root x0 = 1.0298665293222589 after 58 iterations root x0 = 1.0298665293222589 after 21 iterations root x0 = 1.0298665293222589 after 49 iterations
for xstart in sorted(x.keys()):
plt.semilogx(x[xstart], label=str(xstart))
plt.legend(loc="best")
plt.xlabel("iteration")
plt.ylabel("current guess for root $x_0$");
def g(x):
return 2*np.cos(x) - x
Starting away from the root:
%timeit newton_raphson(g, 3.5, eps=1e-14)
293 µs ± 11.3 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
%timeit bisection(g, 0, 3.5, eps=1e-14)
120 µs ± 1.42 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
Newton-Raphson can take a while to find the root.
Let's start closer to the root:
%timeit newton_raphson(g, 1.5, eps=1e-14)
16.8 µs ± 581 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
%timeit bisection(g, 0, 1.5, eps=1e-14)
123 µs ± 1.68 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
When close to a root, N-R is fast.
bisection
and newton_raphson
¶Best of both worlds:
%timeit x0 = bisection(g, 0, 3.5, eps=0.1)
18.3 µs ± 276 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
# find initial guess
x0 = bisection(g, 0, 3.5, eps=0.1)
print(f"Initial guess {x0}")
Initial guess 1.0390625
%timeit newton_raphson(g, x0, eps=1e-14)
12.8 µs ± 463 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
When starting from the initial bisection guess x0
, N-R is very fast to converge.
Let's measure the complete combined algorithm:
%%timeit
x0 = bisection(g, 0, 3.5, eps=0.1)
newton_raphson(g, x0, eps=1e-14)
31.2 µs ± 544 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
By combining bisection with Newton-Raphson we found the root in 31 µs to machine precision instead of 120 µs (bisection alone) or 293 µs (Newton-Raphson alone).