Plotting pendulum motion in python using trapezoid rule

Question

I'm trying to plot how the angle θ and angular velocity ω vary with respect to time t for a linear and non-linear pendulum using the trapezoid rule to solve for the differential equations, but I'm having trouble generating the actual plot.

This is the code I'm trying to implement for the linear pendulum with no frictional dampening or driving force where θ is initialised to 0.2 and ω to 0.0

import matplotlib.pylab as plt
import math

theta = 0.2 #angle
omega = 0.0 #angular velocity
t = 0.0 #time
dt = 0.01
nsteps = 0
k = 0.0 #dampening coefficient
phi = 0.66667 #angular frequency of driving force
A = 0.0 #amplitude of driving force

#2nd order ODE for linear pendulum
def f(theta, omega, t):
    return -theta - k*omega + A*math.cos(phi*t) 

#trapezoid rule
for nsteps in range(0,1000):
    k1a= dt*omega 
    k1b = dt*f(theta, omega, t) 
    k2a = dt*(omega + k1b)
    k2b = dt*f(theta + k1a, omega + k1b, t + dt)

    theta = theta + (k1a + k2a)/2 
    omega = omega + (k1b + k2b)/2 
    t = t + dt 
    nsteps = nsteps + 1

    plt.plot(t, theta)
    plt.plot(t, omega)
    plt.axis([0, 500, -math.pi, math.pi])
    plt.title('theta = 0.2, omega = 0.0')

plt.show()

I've worked out the values for the first few iterations through the for loop by hand and it seems to behave as it should do, but the plot just gives me nothing:

I know it should get a sinusoid, so I think it might just be an issue with how I'm trying to generate the plot.

Any help would be much appreciated.

Answer 1

Welcome to MatPlotLib (and Python in general it seems). This answer will be a list of pointers more than anything else. Hopefully it will will guide you into researching the right topics next time you decide to do something like this.

The main problem with your code is that you call plt.plot for every single point in your time evolution. This creates a new plot with a pair of points in it every time. What you probably want to do is accumulate lists or arrays containing 1000 values of t, omega and theta. You can then plot those all at once outside the loop.

Before I show how to do that, here are some smaller issues I see:

1. pylab is pretty much deprecated at this point. Use pyplot instead. The imports from matplotlib import pyplot as plt and import matplotlib.pyplot as plt are equivalent.
2. You increment nsteps in the loop. This is completely unnecessary. The way for loops work is to assign a new value of nsteps in each iteration, taken from the next value in the iterable you are looping over, in this case range(1000). When the loop assigns 100 to nsteps, and you assign 999 to it in the body of the loop, the next iteration will see nsteps set to 101 whether you like it or not.
3. You are not using the right limits in your call to axis: with a dt of 0.01, 1000 steps will get you a time range of 10 seconds, not 500. Matplotlib is pretty good at setting the axes itself, so I would prefer omitting that call entirely, or at the very least fixing it to something like plt.axis([0, dt * nsteps, -math.pi, math.pi]).

Given all that, here is how I would rewrite the loop portion of your code:

from matplotlib import pyplot as plt
...

t_list = [t]
omega_list = [omega]
theta_list = [theta]

#trapezoid rule
for nsteps in range(0,1000):
    k1a = dt * omega 
    k1b = dt * f(theta, omega, t) 
    k2a = dt * (omega + k1b)
    k2b = dt * f(theta + k1a, omega + k1b, t + dt)

    theta = theta + (k1a + k2a) / 2 
    omega = omega + (k1b + k2b) / 2 
    t = t + dt 
    t_list.append(t)
    theta_list.append(theta)
    omega_list.append(omega)

plt.plot(t_list, theta_list)
plt.plot(t_list, omega_list)
plt.title('theta = 0.2, omega = 0.0')

plt.show()

The code above works, but it is not very efficient. I would look into the numpy library as a start to numerical computing in Python. It provides an array type and lots of math functions. Numpy arrays are the basis of what MatPlotLib is built on as well. All the lists you pass in are converted to arrays internally. When you have reached the limits of what numpy can do for you, look at scipy and other libraries. Python has lots of great math libraries waiting to be explored.