Show time step stability constraint for explicit Euler and leapfrog methods for ODEs

stability.py

"""Euler and leapfrog stability.

Stability for both is guaranteed by

    dt < - 2 / lambda,

where lambda is from the ODE

    dv/dt = lambda * v.
"""

import numpy as np
import matplotlib.pyplot as plt


def euler(position, velocity, acceleration, dt, func, **kwargs):
    """Euler timestep.

    N.B. func is the acceleration function.
    """
    position += dt * velocity
    acceleration = func(velocity, **kwargs)
    velocity += dt * acceleration
    return position, velocity, acceleration


def leapfrog(position, velocity, acceleration, dt, func, **kwargs):
    """Leapfrog timestep.

    N.B. func is the acceleration function.
    """
    v_half = velocity + 0.5 * dt * acceleration
    position += dt * v_half
    acceleration = func(v_half, **kwargs)
    velocity = v_half + 0.5 * dt * acceleration
    return position, velocity, acceleration


# Dictionary of methods: key is name, value is function
METHODS = {'Euler': euler, 'leapfrog': leapfrog}


def integrate(x0, v0, dt, tfinal, method, func, **kwargs):
    """Integrate."""
    num_points = int(tfinal/dt)
    position = np.zeros(num_points)
    velocity = np.zeros(num_points)
    acceleration = np.zeros(num_points)
    position[0], velocity[0] = x0, v0
    acceleration[0] = func(v0, **kwargs)
    for i in range(num_points - 1):
        position[i + 1], velocity[i + 1], acceleration[i + 1] = METHODS[method](
            position[i], velocity[i], acceleration[i], dt, func, **kwargs
        )
    time = np.linspace(0, tfinal, num_points)
    return time, position, velocity


def solve_and_plot(x0, v0, dts, tfinal, func, stopping_time):

    def _solve_and_plot(x0, v0, dt, tfinal, func, stopping_time, method, ax):
        time, position, velocity = integrate(
            x0=x0,
            v0=v0,
            dt=dt,
            tfinal=tfinal,
            method=method,
            func=func,
            stopping_time=stopping_time,
        )
        [line] = ax.plot(time, velocity, alpha=0.5)
        ax.plot(
            time,
            velocity,
            'x',
            color=line.get_color(),
            label=rf'$dt/t_s={dt/stopping_time:.2f}$',
        )

    fig, axs = plt.subplots(ncols=2, figsize=(10, 5))
    for method, ax in zip(METHODS, axs):
        for dt in dts:
            _solve_and_plot(
                x0=x0,
                v0=v0,
                dt=dt,
                tfinal=tfinal,
                func=func,
                stopping_time=stopping_time,
                method=method,
                ax=ax,
            )

        ax.legend()
        ax.set_title(rf'{method} with $t_s={stopping_time}$')
        ax.set_xlabel('time')
        ax.set_ylabel('velocity')

    plt.show()


if __name__ == '__main__':

    x0 = 0.0
    v0 = 1.0
    stopping_time = 2

    tfinals = [10, 500]

    def drag(velocity, stopping_time=1):
        return - velocity / stopping_time

    dts = [dt * stopping_time for dt in [0.01, 0.1, 1.0, 1.5, 1.9, 2.1]]

    for tfinal in tfinals:
        solve_and_plot(
            x0=x0,
            v0=v0,
            dts=dts,
            tfinal=tfinal,
            func=drag,
            stopping_time=stopping_time,
        )

Comparing Euler and leapfrog with stopping time of 2 up to t=10 for the velocity equation:

dv/dt = - v / t_s

Showing blow-up for Euler and leapfrog with dt > 2 t_s.