Merge pull request #22043 from anntzer/dps

QuLogic · web-flow · commit 14a94c33b07c · 2022-01-05T15:13:55.000-05:00
Cleanup differential equations examples.
diff --git 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
@@ -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")
