Week 3 - Solving second order ODEs

import numpy as np 
from scipy.integrate import odeint 
import matplotlib.pyplot as plt
import math

# function that returns dz/dt
def model(theta,t,b,g,l,m):

    theta1 = theta[0]
    theta2 = theta[1]
    dtheta1_dt = theta2
    dtheta2_dt = -(b/m)*theta2 - (g/l)*math.sin(theta1)
    dtheta_dt = [dtheta1_dt, dtheta2_dt]
    return dtheta_dt

b=0.3
g=9.81
l=1
m=1
ct=1
#initial condition
theta_0 = [0,3]

# time points 
t = np.linspace(0,20,150)
print(t)
# solve ODE
theta = odeint(model,theta_0,t,args = (b,g,l,m))

# plot results
plt.plot(t,theta[:,0],label='angulardisplacement')
plt.plot(t,theta[:,1],label='angularvelocity')
plt.ylabel('plot')
plt.xlabel('time')
plt.legend()
plt.show()
angles = theta[:,1]
for angle in angles:
    x1=5
    y1=5
    x2=x1+l*math.sin(angle)
    y2=y1-l*math.cos(angle)

    filename = str(ct)+'.png'

    ct=ct+1

    plt.figure(2)
    plt.plot([x1,x2],[y1,y2])
    plt.plot(x2,y2,'o','marker',markersize=20,color='k')

    plt.xlim([3.5,6.5])
    plt.ylim([3.5,6.5])
    plt.title('pendulum animation')
    plt.savefig(filename)
    plt.show()