scipy odeint

1import numpy as np
2import matplotlib.pyplot as plt
3from scipy.integrate import odeint
4
5def f(y, t, params):
6    theta, omega = y      # unpack current values of y
7    Q, d, Omega = params  # unpack parameters
8    derivs = [omega,      # list of dy/dt=f functions
9             -omega/Q + np.sin(theta) + d*np.cos(Omega*t)]
10    return derivs
11
12# Parameters
13Q = 2.0          # quality factor (inverse damping)
14d = 1.5          # forcing amplitude
15Omega = 0.65     # drive frequency
16
17# Initial values
18theta0 = 0.0     # initial angular displacement
19omega0 = 0.0     # initial angular velocity
20
21# Bundle parameters for ODE solver
22params = [Q, d, Omega]
23
24# Bundle initial conditions for ODE solver
25y0 = [theta0, omega0]
26
27# Make time array for solution
28tStop = 200.
29tInc = 0.05
30t = np.arange(0., tStop, tInc)
31
32# Call the ODE solver
33psoln = odeint(f, y0, t, args=(params,))
34
35# Plot results
36fig = plt.figure(1, figsize=(8,8))
37
38# Plot theta as a function of time
39ax1 = fig.add_subplot(311)
40ax1.plot(t, psoln[:,0])
41ax1.set_xlabel('time')
42ax1.set_ylabel('theta')
43
44# Plot omega as a function of time
45ax2 = fig.add_subplot(312)
46ax2.plot(t, psoln[:,1])
47ax2.set_xlabel('time')
48ax2.set_ylabel('omega')
49
50# Plot omega vs theta
51ax3 = fig.add_subplot(313)
52twopi = 2.0*np.pi
53ax3.plot(psoln[:,0]%twopi, psoln[:,1], '.', ms=1)
54ax3.set_xlabel('theta')
55ax3.set_ylabel('omega')
56ax3.set_xlim(0., twopi)
57
58plt.tight_layout()
59plt.show()
60