In [1]:
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import solve_ivp
%matplotlib widget
In [2]:
b = 1      # the damping coefficient (you can change this value)
m = 1      # the particle mass (you can change this, too)
v0 = 5     # the initial velocity
dt = 0.1   # the time step (you can also change this, if you like)
v = [v0,]  # initialize an array of velocities
t = [0, ]  # and an array of times
while t[-1] < 1:
    dv = -b / m * v[-1] * dt  # compute the change in velocity, assumed constant over time interval
    v.append(v[-1] + dv)      # append the new velocity to the growing array of velocities
    t.append(t[-1] + dt)      # and append the new time

# Solving exactly gives the equation v(t) = v0 * exp(-b/m t)
# Let's plot the comparison between the Euler’s-method solution and the true function

fig, ax = plt.subplots()   # create a plot
ax.plot(t, v, 'ro', label='Euler')
ax.plot(t, v0 * np.exp(-b/m * np.array(t)), 'b-', label='true')
ax.set_xlabel('$t$')
ax.set_ylabel('$v$')
ax.legend();
Figure
No description has been provided for this image
In [3]:
def Euler(tvals, v0, b=1, m=1):
    """Starting at v0 at the first entry in tvals, use Euler's
    method to integrate the first-order differential equation to
    find v at each of the remaining tvals.
    """
    v = [v0,] # initialize the output array
    for n in range(1, len(tvals)):
        dt = tvals[n] - tvals[n-1] # in case the time steps aren't even
        dv = -b * v[-1] * dt / m
        v.append(v[-1] + dv)
    return np.array(v)
In [4]:
def EulerError(nsteps:list, b=1, m=1, v0=5):
    """Prepare a graph showing errors accumulated by using Euler's
    method to integrate a first-order differential equation.
    """
    assert isinstance(nsteps, (tuple, list, np.ndarray)), "Pass in a list of positive integers"
    fig, axes = plt.subplots(1, 2, figsize=(8, 4))  # create an empty plot
    ax1, ax2 = axes
    ax1.set_xlabel('$$t$$')
    ax1.set_ylabel('$$v$$')
    ax2.set_xlabel('$$t$$')
    ax2.set_ylabel('Error')
    for n in nsteps:
        tvals = np.linspace(0, 1.0, n+1) # linspace divides the span between the arguments into n equal steps
        vvals = Euler(tvals, v0, b, m)
        ax1.plot(tvals, vvals, '.-', label=str(n))
        ax2.plot(tvals, np.abs(vvals - v0 * np.exp(-tvals * b/m)), '.-', label=str(n))
    ax1.legend()

EulerError([2, 4, 8, 16]) # Cover the [0,1] range in 2, 4, 8, and 16 steps
Figure
No description has been provided for this image

Now using solve_ivp¶

In [5]:
def myderiv(t, v, b, m):
    """Given the current value of v and t, and the parameters
    b (damping) and m (mass), compute and return the derivative."""
    return -v * b / m
In [6]:
b, m = 1, 1
res = solve_ivp(myderiv, [0, 1], [5.0], t_eval=np.linspace(0, 1, 11), args=(b, m))
res
Out[6]:
  message: The solver successfully reached the end of the integration interval.
  success: True
   status: 0
        t: [ 0.000e+00  1.000e-01  2.000e-01  3.000e-01  4.000e-01
             5.000e-01  6.000e-01  7.000e-01  8.000e-01  9.000e-01
             1.000e+00]
        y: [[ 5.000e+00  4.524e+00  4.093e+00  3.703e+00  3.350e+00
              3.031e+00  2.742e+00  2.482e+00  2.247e+00  2.034e+00
              1.840e+00]]
      sol: None
 t_events: None
 y_events: None
     nfev: 14
     njev: 0
      nlu: 0
In [7]:
errors = res.y[0,:] - 5.0 * np.exp(-res.t * b / m)
fig, ax = plt.subplots()
ax.plot(res.t, errors, 'ro')
ax.set_xlabel("$t$ (s)")
ax.set_ylabel("Error");
Figure
No description has been provided for this image

Exploring how tolerances influence the error at the end of the interval¶

In [8]:
logtol = np.arange(3, 10.1, 0.25)

def verr(rtol, atol):
    res = solve_ivp(myderiv, [0,1], [5.0], t_eval=np.linspace(0,1,11), args=(b,m),
                    rtol=rtol, atol=atol)
    err = np.abs(5 * np.exp( - 1 * b / m) - res.y[0,-1])
    return err

rt, at = np.meshgrid(logtol, logtol)
logerrors = np.zeros((len(logtol), len(logtol)))
for i in range(len(logtol)):
    for j in range(len(logtol)):
        logerrors[i,j] = np.log10(verr(10**(-rt[i,j]), 10**(-at[i,j])))
In [10]:
plt.close('all')
fig, ax = plt.subplots()
im = ax.contourf(rt, at, logerrors)
ax.set_xlabel(r"$-\log_{10}(\mathrm{atol})$")
ax.set_ylabel(r"$-\log_{10}(\mathrm{rtol})$");
plt.colorbar(im);
Figure
No description has been provided for this image

Handling second-order equations¶

In [11]:
def pendulum_derivs(t, Y, g, l):
    θ, ω = Y      # split apart the dependent variable vector
    return [ω, - g * np.sin(θ) / l]
In [12]:
def pend_bottom(t, Y, g, l):
    return Y[0]
pend_bottom.terminal = True # stop when we hit bottom
In [13]:
g = 9.8 # m/s^2
l = g / (2 * np.pi)**2  # m, chosen so small amplitudes give period 1
t_range = (0, 10)
res = solve_ivp(pendulum_derivs,
                t_range,
                [3.14,0], # (θ, ω)
                t_eval = np.arange(*t_range,0.025),
                args=(g, l),
                atol=1e-8,
                rtol=1e-8,
                events=pend_bottom
               )
res
Out[13]:
  message: A termination event occurred.
  success: True
   status: 1
        t: [ 0.000e+00  2.500e-02 ...  1.325e+00  1.350e+00]
        y: [[ 3.140e+00  3.140e+00 ...  3.906e-01  7.897e-02]
            [ 0.000e+00 -1.578e-03 ... -1.233e+01 -1.256e+01]]
      sol: None
 t_events: [array([ 1.356e+00])]
 y_events: [array([[ 6.661e-16, -1.257e+01]])]
     nfev: 362
     njev: 0
      nlu: 0
In [14]:
def period(θ0):
    res = solve_ivp(pendulum_derivs,
                t_range,
                [θ0,0], # (θ, ω)
                args=(g, l),
                atol=1e-8,
                rtol=1e-8,
                events=pend_bottom
               )
    return 4 * res.t_events[0][0]
In [15]:
period(0.01)
Out[15]:
np.float64(1.0000062172405908)
In [16]:
θ = np.radians(np.arange(0.5, 180, 0.5))
T = [period(q) for q in θ]
plt.close('all')
fig, ax = plt.subplots()
ax.plot(np.degrees(θ), T)
ax.set_xlabel(r"$\theta$ (°)")
ax.set_ylabel(r"$T$ (s)")
ax.grid();
Figure
No description has been provided for this image
In [ ]: