From e67514f9da817a778027e412dca5ef309b6f1c98 Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Wed, 22 Sep 2021 17:32:51 +0200 Subject: [PATCH] Cleanup differential equations examples. Slightly cleanup the lorenz_attractor example; also use Euler's method for the double pendulum. It is less accurate, but sufficient for illustrative purposes (also, 1. the double pendulum is chaotic anyways so even RK4 will end up quite far from the actual behavior, if one waits for long enough, as can be checked by trying various integrators in solve_ivp; 2. we're fine with using Euler's method for the also chaotic lorenz_attractor). The point is also to make fewer examples dependent on scipy. --- examples/animation/double_pendulum.py | 20 +++++++++----- examples/mplot3d/lorenz_attractor.py | 40 ++++++++++++--------------- 2 files changed, 31 insertions(+), 29 deletions(-) diff --git a/examples/animation/double_pendulum.py b/examples/animation/double_pendulum.py index 8f50f6cb195d..fba7ea4ec9e2 100644 --- a/examples/animation/double_pendulum.py +++ b/examples/animation/double_pendulum.py @@ -12,7 +12,6 @@ from numpy import sin, cos import numpy as np import matplotlib.pyplot as plt -import scipy.integrate as integrate import matplotlib.animation as animation from collections import deque @@ -22,13 +21,13 @@ L = L1 + L2 # maximal length of the combined pendulum M1 = 1.0 # mass of pendulum 1 in kg M2 = 1.0 # mass of pendulum 2 in kg -t_stop = 5 # how many seconds to simulate +t_stop = 2.5 # how many seconds to simulate history_len = 500 # how many trajectory points to display -def derivs(state, t): - +def derivs(t, state): dydx = np.zeros_like(state) + dydx[0] = state[1] delta = state[2] - state[0] @@ -51,7 +50,7 @@ def derivs(state, t): return dydx # create a time array from 0..t_stop sampled at 0.02 second steps -dt = 0.02 +dt = 0.01 t = np.arange(0, t_stop, dt) # th1 and th2 are the initial angles (degrees) @@ -64,8 +63,15 @@ def derivs(state, t): # initial state state = np.radians([th1, w1, th2, w2]) -# integrate your ODE using scipy.integrate. -y = integrate.odeint(derivs, state, t) +# integrate the ODE using Euler's method +y = np.empty((len(t), 4)) +y[0] = state +for i in range(1, len(t)): + y[i] = y[i - 1] + derivs(t[i - 1], y[i - 1]) * dt + +# A more accurate estimate could be obtained e.g. using scipy: +# +# y = scipy.integrate.solve_ivp(derivs, t[[0, -1]], state, t_eval=t).y.T x1 = L1*sin(y[:, 0]) y1 = -L1*cos(y[:, 0]) diff --git a/examples/mplot3d/lorenz_attractor.py b/examples/mplot3d/lorenz_attractor.py index 921bf5a47099..7162c12d2dce 100644 --- a/examples/mplot3d/lorenz_attractor.py +++ b/examples/mplot3d/lorenz_attractor.py @@ -18,45 +18,41 @@ import matplotlib.pyplot as plt -def lorenz(x, y, z, s=10, r=28, b=2.667): +def lorenz(xyz, *, s=10, r=28, b=2.667): """ - Given: - x, y, z: a point of interest in three dimensional space - s, r, b: parameters defining the lorenz attractor - Returns: - x_dot, y_dot, z_dot: values of the lorenz attractor's partial - derivatives at the point x, y, z + Parameters + ---------- + xyz : array-like, shape (3,) + Point of interest in three dimensional space. + s, r, b : float + Parameters defining the Lorenz attractor. + + Returns + ------- + xyz_dot : array, shape (3,) + Values of the Lorenz attractor's partial derivatives at *xyz*. """ + x, y, z = xyz x_dot = s*(y - x) y_dot = r*x - y - x*z z_dot = x*y - b*z - return x_dot, y_dot, z_dot + return np.array([x_dot, y_dot, z_dot]) dt = 0.01 num_steps = 10000 -# Need one more for the initial values -xs = np.empty(num_steps + 1) -ys = np.empty(num_steps + 1) -zs = np.empty(num_steps + 1) - -# Set initial values -xs[0], ys[0], zs[0] = (0., 1., 1.05) - +xyzs = np.empty((num_steps + 1, 3)) # Need one more for the initial values +xyzs[0] = (0., 1., 1.05) # Set initial values # Step through "time", calculating the partial derivatives at the current point # and using them to estimate the next point for i in range(num_steps): - x_dot, y_dot, z_dot = lorenz(xs[i], ys[i], zs[i]) - xs[i + 1] = xs[i] + (x_dot * dt) - ys[i + 1] = ys[i] + (y_dot * dt) - zs[i + 1] = zs[i] + (z_dot * dt) - + xyzs[i + 1] = xyzs[i] + lorenz(xyzs[i]) * dt # Plot ax = plt.figure().add_subplot(projection='3d') -ax.plot(xs, ys, zs, lw=0.5) +ax.plot(*xyzs.T, lw=0.5) ax.set_xlabel("X Axis") ax.set_ylabel("Y Axis") ax.set_zlabel("Z Axis")