Week 3 - Solving second order ODEs

import math
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import odeint
# defining function
def function(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

# defining the value of input arguements
b=0.5
g=9.81
l=1
m=1
# inital condition
theta_0 = [0,5]
# creating an array for time
t = np.linspace(0,10,20)
# defining Ode function
theta = odeint(function,theta_0,t,args = (b,g,l,m))
# to create animation
ct = 1
x0 = 0
y0 = 0
for i in theta[:,0]:
    x1 = (x0+(l*math.sin(i)))
    y1 = (y0+(-l*math.cos(i)))
    filename = str(ct)+'.png'
    ct = ct+1

    plt.figure()
    plt.plot([-1,1],[0,0],"black")
    plt.plot([x0,x1],[y0,y1],"red")
    plt.plot(x1,y1,'o',markersize = '20')
    plt.xlim([-1.5,1.5])
    plt.ylim([-1.5,1.0])
    plt.savefig(filename)