Welcome to Auro's World!

""" 
General Numerical Solver for the 1D Time-Dependent Schrodinger's equation.

adapted from code at 
    https://jakevdp.github.io/downloads/code/double_pendulum.py
    
    adapted from code at 
    http://matplotlib.sourceforge.net/examples/animation/double_pendulum_animated.py

    Double pendulum formula translated from the C code at
    http://www.physics.usyd.edu.au/~wheat/dpend_html/solve_dpend.c

author: Auro Prasad Mohanty
email: aurogreenindia@gmail.com
website: mycrofd.gihub.io
license: MIT
Please feel free to use and modify this, but keep the above information. Thanks!
""" 

from numpy import sin, cos
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as integrate
import matplotlib.animation as animation

thr = np.pi/180.
print "Give the initial angle in degrees."
theta = thr*float(raw_input())
print "Give the initial angular velocity in degrees per sec."
omega = thr*float(raw_input())
class Pendulum:
    """Pendulum Class
    init_state is [theta, omega] in radians,
    where theta, omega are the angular position and velocity of the
    pendulum arm
    """
    def __init__(self,
                 init_state = [theta, omega ],
                 L=1.0,  # length of pendulum in m
                 M=1.0,  # mass of pendulum in kg
                 G=9.8,  # acceleration due to gravity, in m/s^2
                 origin=(0, 0)):
        self.init_state = np.asarray(init_state, dtype='float')
        self.params = (L, M, G)
        self.origin = origin
        self.time_elapsed = 0
        self.state = self.init_state    
    def position(self):
        """compute the current x,y positions of the pendulum arms"""
        (L, M, G) = self.params
        x = np.cumsum([self.origin[0],
                       L * sin(self.state[0]),
                       ])
        y = np.cumsum([self.origin[1],
                       -L * cos(self.state[0]),
                       ])
        return (x, y)
    def energy(self):
        """compute the energy of the current state"""
        (L, M, G) = self.params

        x = np.cumsum([L * sin(self.state[0])])
        y = np.cumsum([-L * cos(self.state[0])])
        vx = np.cumsum([L * self.state[1] * cos(self.state[0])])
        vy = np.cumsum([L * self.state[1] * sin(self.state[0])])
        U = M * G * y[0]
        K = 0.5 * M * (np.dot(vx, vx) + np.dot(vy, vy))

        return U + K
    def dstate_dt(self, state, t):
        """compute the derivative of the given state"""
        (L, M, G) = self.params
        dydx = np.zeros_like(state)
        dydx[0] = state[1]
        dydx[1] = -(G/L)*np.sin(state[0])
        return dydx
    def step(self, dt):
        """execute one time step of length dt and update state"""
        self.state = integrate.odeint(self.dstate_dt, self.state, [0, dt])[1]
        self.time_elapsed += dt

#------------------------------------------------------------
# set up initial state and global variables
pendulum = Pendulum([theta, omega])
dt = 1./30 # 30 fps
#------------------------------------------------------------
# set up figure and animation
fig = plt.figure()
ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,
                     xlim=(-2, 2), ylim=(-2, 2))
ax.grid()

line, = ax.plot([], [], '-o',color='black', lw=2)
time_text = ax.text(0.02, 0.95,'' , transform=ax.transAxes)
energy_text = ax.text(0.02, 0.90,'' , transform=ax.transAxes)

def init():
    """initialize animation"""
    line.set_data([], [])
    time_text.set_text('')
    energy_text.set_text('')
    return line, time_text, energy_text

def animate(i):
    """perform animation step"""
    global pendulum, dt
    pendulum.step(dt)

    line.set_data(*pendulum.position())
    time_text.set_text('time = %.1f' % pendulum.time_elapsed)
    energy_text.set_text('energy = %.3f J' % pendulum.energy())
    return line, time_text, energy_text

# choose the interval based on dt and the time to animate one step
from time import time
t0 = time()
animate(0)
t1 = time()
interval = 1000 * dt - (t1 - t0)

ani = animation.FuncAnimation(fig, animate, frames=3,
                              interval=interval, blit=True, init_func=init)
for i in xrange(10):
    plt.plot([-0.5+i/10.,-0.4+i/10.],[0,0.1],'-',color='black')
plt.plot([-0.5,0.5],[0,0.],'-',color='black')
# save the animation as an mp4.  This requires ffmpeg or mencoder to be
# installed.  The extra_args ensure that the x264 codec is used, so that
# the video can be embedded in html5.  You may need to adjust this for
# your system: for more information, see
# http://matplotlib.sourceforge.net/api/animation_api.html
#ani.save('pendulum.mp4', fps=30, extra_args=['-vcodec, 'libx264'])

plt.show()