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.

Author: anntzer

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

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