Lagrangian double pendulum with T/U statistic plotting

from visual import *
from math import sin,cos,pi
from random import randint
import matplotlib.pyplot as plt

class pendulum():
    def __init__(self):
        self.simlength = int(input('How many seconds should the simulation run for?: '))
        self.disp = display(title='Double pendulum', x=0, y=0, width=1920, height=1080, background=(0,0,0))
        self.c = sphere(pos = (0,0,0), radius = 0.3)
        self.l1 = randint(3,7)
        self.l2 = randint(3,7)
        self.disp.range = 50
        self.m1 = 1
        self.m2 = 2
        self.theta1 = (randint(0,100)/100)*2*pi
        self.theta2 = (randint(0,100)/100)*2*pi
        self.vtheta1 = randint(-3,3)
        self.vtheta2 = randint(-3,3)
        self.atheta1 = 0
        self.atheta2 = 0
        self.g = 9.81
        self.dt = 1/30
        self.mass1 = sphere(pos = (self.l1*sin(self.theta1),-self.l1*cos(self.theta1),0), radius = ((3*self.m1)/(4*pi))**(1/3), color = color.red, make_trail = True, retain = 500)
        self.mass2 = sphere(pos = (self.l1*sin(self.theta1)+self.l2*(self.theta2),-self.l1*cos(self.theta1)-self.l2*cos(self.theta2),0), radius = ((3*self.m2)/(4*pi))**(1/3), color = color.blue, make_trail = True, retain = 500)
        self.stick1 = cylinder(pos=(0,0,0), axis = self.mass1.pos, radius = 0.1)
        self.stick2 = cylinder(pos=self.mass1.pos, axis = self.mass2.pos-self.mass1.pos, radius = 0.1)
        self.stats = [[],[],[],[],[],[]]
        self.t=0

    def acc1(self):
        self.atheta1 = (-self.g*(2*self.m1+self.m2)*sin(self.theta1) -self.m2*sin(self.theta1-2*self.theta2)-2*sin(self.theta1-self.theta2)*self.m2*((self.vtheta2**2)*self.l2+(self.vtheta1**2)*self.l1*cos(self.theta1-self.theta2)))/(self.l1*(2*self.m1+self.m2-self.m2*cos(2*self.theta1+2*self.theta2)))

    def acc2(self):
        self.atheta2 = (2*sin(self.theta1-self.theta2)*((self.vtheta1**2)*self.l1*(self.m1+self.m2)+self.g*(self.m1+self.m2)*cos(self.theta1)+(self.vtheta2**2)*self.l2*self.m2*cos(self.theta1-self.theta2))/(self.l2*(2*self.m1 +self.m2-self.m2*cos(2*self.theta1-2*self.theta2))))

    def move(self):
        self.vtheta1 += self.atheta1*self.dt
        self.vtheta2 += self.atheta2*self.dt
        self.theta1 += self.vtheta1*self.dt
        self.theta2 += self.vtheta2*self.dt

    def anglim(self):
        if self.theta1 >= 2*pi:
            self.theta1 -= 2*pi
        if self.theta1 <0:
            self.theta1 += 2*pi
        if self.theta2 >= 2*pi:
            self.theta2 -= 2*pi
        if self.theta2 <0:
            self.theta2 += 2*pi

    def convert(self):
        self.mass1.pos = (self.l1*sin(self.theta1),-self.l1*cos(self.theta1),0)
        self.mass2.pos = (self.l1*sin(self.theta1)+self.l2*sin(self.theta2),-self.l1*cos(self.theta1)-self.l2*cos(self.theta2),0)
        self.stick1.axis = self.mass1.pos
        self.stick2.pos = self.mass1.pos
        self.stick2.axis = self.mass2.pos-self.mass1.pos

    def telemetry(self):
        self.stats[0].append(self.t)
        self.stats[1].append(0.5*self.m1*(self.l1**2)*(self.vtheta1**2))
        self.stats[2].append(self.l1*self.m1*self.g*cos(self.theta1))
        self.stats[3].append(0.5*self.m2*((self.l1**2)*(self.vtheta1**2)+(self.l2**2)*(self.vtheta2**2)+(2*self.l1*self.l2*self.vtheta1*self.vtheta2*cos(self.theta1-self.theta2))))
        self.stats[4].append(-self.l1*self.m2*self.g*cos(self.theta1)-self.l2*self.m2*self.g*cos(self.theta2))

    def step(self):
        self.acc1()
        self.acc2()
        self.move()
        self.anglim()
        self.convert()
        if self.simlength > 0:
            self.telemetry()
        self.t += self.dt
p = pendulum()

while True:
    rate(1/p.dt)
    p.step()
    if p.t >= p.simlength and p.simlength != 0:
        break
plt.figure(1)
#plt.xkcd()

plt.subplot(211)
plt.xlabel('Time (s)')
plt.ylabel('Energy (J)')
plt.title('Mass 1')
plt.grid(True)
plt.plot(p.stats[0],p.stats[1],p.stats[0],p.stats[2])
plt.legend(['Kinetic Energy','Potential Energy'])

plt.subplot(212)
plt.xlabel('Time (s)')
plt.ylabel('Energy (J)')
plt.title('Mass 2')
plt.grid(True)
plt.plot(p.stats[0],p.stats[3],p.stats[0],p.stats[4])
plt.legend(['Kinetic Energy','Potential Energy'])

plt.show()