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")