#!/usr/bin/env python
# coding: utf-8

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'inline')
import numpy as np
import matplotlib.pyplot as plt


# # Forced Pendulums
# ## Lecture 8

# ## Solution for a damped, driven pendulum
# 
# In the damped pendulum, the motion decay from an initial perturbation until it comes to rest.
# 
# To keep the pendulum swinging, we need to provide a source of energy to drive the motion.
# 
# One way of modeling such a system is to drive it with a driving amplitude $F_D$ and driving angular frequency $\Omega_D$
# 
# \begin{align}
# \frac{d^2\theta}{dt^2} &= - \frac{g}{L} \theta -q \frac{d \theta}{dt}  + F_D \sin(\Omega_D t)\;, \\
# \end{align}
# 
# This driving force will pump energy into (or out of) the system.  The external forcing frequency, $\Omega_D$, "competes" with the natural frequency of the pendulum, $\Omega = \sqrt{g/L}$.

# In[ ]:


# constant parameters

#pendulum length in metres
length = 9.8

#acceleration due to gravity
g = 9.8


# In[ ]:


def Pendulum(N=1500, dt=0.04, F_D=0.0, Ω_D=0.0, q=0.5, θ_0=0.2):
    """
    Euler Cromer Solution for damped, driven pendulum
    
    Keyword arguments: 
        N   number of timesteps
        dt  time step
        F_D driving force amplitude
        Ω_D driving force frequency
        q   damping strength
        θ_0 initial displacement
        
    Returns
        t time
        ω angular velocity
        θ angular displacement
    """
    
    #initializes omega, a vector of dimension N,to being all zeros
    ω = np.zeros(N)
    #initializes theta, a vector of dimensionN,to being all zeros
    θ = np.zeros(N)
    #this initializes the vector time to being all zeros
    t = np.zeros(N);
    
    #you need to have some initial displacement, otherwise the pendulum will not swing
    θ[0] = θ_0
    ω[0] = 0

    #loop over the timesteps.
    for i in range(N-1):
        ω[i+1] = ω[i] + (-(g/length)*θ[i] - q*ω[i] + F_D*np.sin(Ω_D*t[i]))*dt
        θ[i+1] = θ[i] + ω[i+1]*dt  
        t[i+1] = t[i] + dt
        
    return t, ω, θ


# ### Plots the numerical solution for F_Drive = 0.5

# In[ ]:


t, ω, θ = Pendulum(N=15000, F_D=0.5, Ω_D=2/3)

fig, axes = plt.subplots(figsize=(8,3))

plt.plot(t, θ, 'b-')
# label axes
plt.xlabel('$t$ (s) ')
plt.ylabel('$\\theta$ (rad)')
# axes scales
plt.xlim(0, 60);
plt.ylim(-2.0, 2.0);

plt.title('Results from Physical pendulum,  F_D = 0.5')
plt.show()


# ### Plots the numerical solution for F_Drive = 1.2

# In[ ]:


t, ω, θ = Pendulum(N=15000, F_D=1.2, Ω_D=2/3)

fig, axes = plt.subplots(figsize=(8,3))

plt.plot(t, θ, 'b-')
# label axes
plt.xlabel('$t$ (s) ')
plt.ylabel('$\\theta$ (rad)')
# axes scales
plt.xlim(0, 60)
plt.ylim(-1.5, 1.5)

plt.title('Results from Physical pendulum,  F_D = 1.2')
plt.show()


# ### Exercise
# > 1. What is the natural frequency of a pendulum?
# > 2. Write code using a loop to investigate how the steady-state amplitude of driven, damped pendulum depends on
# > the forcing frequency (while keeping $q$ and $F_D$ unchanged, see below!)
# > 3. What happens when the forcing frequency is near to the natural frequency?
# > 4. Explore the effect of changing the damping parameter, $q$.

# In[ ]:


fig, axes = plt.subplots(2,1, figsize=(8,6))

for Ω_D in np.arange(0.2, 0.2, 0.1):  ## <-- FIX THIS
    dt = 0.04
    tmax = 100
    N = int(tmax // dt)
    t, ω, θ = Pendulum(N=N, dt=dt, F_D=1.2, 
                       Ω_D=Ω_D, 
                       q=0.5,   
                       θ_0=0.2)
    
    # plot the solution on the upper axes
    axes[0].plot(t, θ,'-')
    
    ## EXERCISE: 
    # Consider the long term (n>N//2) behaviour only
    max_θ = 0 ## REPLACE THIS with an expression involving θ
    
    # plot the maximum value on the lower axes
    axes[1].plot(Ω_D, max_θ, 'ko')
    
# label axes
plt.sca(axes[0])
plt.xlabel('$t $(s) ')
plt.ylabel('$\\theta$ (rad)')
plt.ylim(-4, 4)

plt.sca(axes[1])
plt.xlabel('$\Omega_D$ (s$^{-1}$) ')
plt.ylabel('Amplitude (maximum angle)')
plt.axvline(1, color='k', linestyle=':')
plt.ylim(0, 4)

plt.suptitle('Amplitude for different values of $\Omega_D$')
plt.show()


# ### Discussion: Resonance

# ## Non-linear Pendulums

# The final improvement to model a more realistic pendulum is to no longer assume the amplitude is *small*. Recall back to the model we made when deriving the equation of motion of a simple harmonic oscillator for the pendulum.  We assumed
# 
# $$ \sin(\theta) \approx \theta $$
# 
# when $\theta << 1$.  The assumption of small oscillations is what leads to the *simple* harmonic oscillator model.
# 
# For forced pendulums, we observed that the oscillations can exceed small values.  Indeed, for large enough values of the driving force, especially near resonance, the values of $\theta$ are much larger than $\pi$!  That is the pendulum is actually swinging completely around!
# 
# More realistically we should consider the equation of motion:
# 
# \begin{align}
# \frac{d^2\theta}{dt^2} &= - \frac{g}{L} \sin(\theta) \;, \\
# \end{align}
# 
# 
# And if we include both damping and forcing we have the following equation of model:
# \begin{align}
# \frac{d^2\theta}{dt^2} &= - \frac{g}{L} \sin(\theta) -q \frac{d \theta}{dt}  + F_D \sin(\Omega_D t)\;, \\
# \end{align}
# 
# 

# ## Solution for a non-linear, damped, driven pendulum:
# 

# In[ ]:


#pendulum length in metres
length = 9.8
#acceleration due to gravity
g = 9.8

def Nonlinear_Pendulum(N=1500, dt=0.04, F_D=0.0, Ω_D=0.0, q=0.5, θ_0=0.2):
    """
    Euler Cromer Solution for non-linear, damped, driven pendulum
    
    Keyword arguments: 
        N   number of timesteps
        dt  time step
        F_D driving force amplitude
        Ω_D driving force frequency
        q   damping strength
        θ_0 initial displacement
        
    Returns
        t time
        ω angular velocity
        θ angular displacement
    """

    #initialize vectors of dimension N to being all zeros
    ω = np.zeros(N)
    θ = np.zeros(N)
    t = np.zeros(N)
    
    #you need to have some initial displacement, otherwise the pendulum will not swing
    θ[0] = θ_0
    ω[0] = 0

    #loop over the timesteps.
    for i in range(N-1):
        ω[i+1] = ω[i] + (-(g/length)*np.sin(θ[i]) - q*ω[i] + F_D*np.sin(Ω_D*t[i]))*dt
        temp_theta = θ[i]+ω[i+1]*dt
    
        # We need to adjust theta after each iteration so as to keep it between +/-π
        # The pendulum can now swing right around the pivot, corresponding to theta>2π.
        # Theta is an angular variable so values of theta that differ by 2π
        # correspond to the same position.
        # For plotting purposes it is nice to keep (-π < θ < π).
        # So, if theta is <-π, add 2π.If theta is > π, subtract 2π
        #********************************************************************************************
        if (temp_theta < -np.pi):
            temp_theta = temp_theta+2*np.pi 
        elif (temp_theta > np.pi):
            temp_theta = temp_theta-2*np.pi
        #********************************************************************************************
        
        # Update theta array
        θ[i+1]=temp_theta;
        t[i+1] = t[i] + dt
        
    return t, ω, θ


# ### Numerical solution for F_D = 0.5

# In[ ]:


t, ω, θ = Nonlinear_Pendulum(F_D=0.5, q=0.5, Ω_D=2/3)

fig, axes = plt.subplots(figsize=(8,3))

plt.plot(t, θ, 'b-')
# label axes
plt.xlabel('t (s) ')
plt.ylabel('$\\theta$ (rad)')
# axes scales 
plt.xlim(0, 60)
plt.ylim(-1, 1)

plt.title('Results from Physical pendulum,  F_D = 0.5')
plt.show()


# The pendulum oscillates with a uniform period.  Let's estimate this period.
# 
# 
# Choose two peaks in the plot and estimate the period

# In[ ]:


t1 = 30.0 # CHANGE THIS SO THAT
t2 = 40.0 # WE PICK OUT PEAKS IN THE PLOT

fig, axes = plt.subplots(figsize=(8,3))
plt.plot(t, θ, 'b-')
plt.xlabel('t (s) ')
plt.ylabel('$\\theta$ (rad)') 
plt.xlim(0, 60)
plt.ylim(-1, 1)

plt.axvline(t1)
plt.axvline(t2)

plt.show()


# In[ ]:


T = t2 - t1
print(T)


# So the frequency is 

# In[ ]:


print(2*np.pi/T)


# Observe that this is the driven frequency $\Omega_D$.
# 
# This is the same as we saw with the *linear* pendulum.

# ### Numerical solution for F_Drive = 1.2

# In[ ]:


t, ω, θ = Nonlinear_Pendulum(N=15000, dt=0.04, F_D=1.2, Ω_D=2/3)

fig, axes = plt.subplots(figsize=(8,3))

plt.plot(t, ω, 'b-');
# label axes
plt.xlabel('t (s) ');
plt.ylabel('$\\theta$ (rad)');
# axes scales 
plt.xlim(0, 60);
#plt.ylim(-4, 4);

plt.title('Results from Physical pendulum,  F_D = 1.2')
plt.show()


# Now the motion for the non-linear pendulum is completely different compared to the linear pendulum.  In fact, the motion appears random and unpredictable. However the motion is completely deterministic -- run the code again and you'll get exactly the same plot.
# 
# Here we have an example of a **chaotic** system.

# ### Exercise
# > Explore the behaviour of the non-linear pendulum by adjusting $F_D$ and $\Omega_D$. 
# 
# > Changing the parameters in the above plot from the original values of `F_D=1.2` and `Ω_D=2/3`.
# 
# > What kinds of motion can you discover?

# ## What is chaos?
# 
# 

# 
# One way of understanding chaotic systems is to study how small changes in the initial conditions lead to potentially large difference in the solution.
# 
# First consider two solutions that differ only slightly in their initial conditions for a non-chaotic solution. 

# In[ ]:


θ_0 = 0.2    # intitial condition 1
Δθ = 0.001   # initial condition 2: θ_0 + Δθ

F_D = 0.5 # driving force

t, ω1, θ1 = Nonlinear_Pendulum(N=15000, dt=0.04, F_D=F_D, Ω_D=2/3, θ_0 = θ_0)
t, ω2, θ2 = Nonlinear_Pendulum(N=15000, dt=0.04, F_D=F_D, Ω_D=2/3, θ_0 = θ_0 + Δθ)

fig, axes = plt.subplots()

plt.plot(t, abs(θ2-θ1), 'b-')

# label axes
plt.xlabel('t (s) ')
plt.ylabel('$\Delta\\theta$ (rad)')

# axes scales 
plt.yscale('log')
plt.xlim(0, 100)

plt.title('Comparison between two solution,  F_D = {:.1f}'.format(F_D))
plt.show()


# The error decreases exponentially. This solution is *stable*.
# 
# Now repeat the exercise with a larger driving force that leads to a chaotic solution.

# In[ ]:


θ_0 = 0.2    # intitial condition 1
Δθ = 0.001   # initial condition 2: θ_0 + Δθ

F_D = 1.2 # driving force

t, ω1, θ1 = Nonlinear_Pendulum(N=15000, dt=0.04, F_D=F_D, Ω_D=2/3, θ_0 = θ_0)
t, ω2, θ2 = Nonlinear_Pendulum(N=15000, dt=0.04, F_D=F_D, Ω_D=2/3, θ_0 = θ_0 + Δθ)

fig, axes = plt.subplots()

plt.plot(t, abs(θ1-θ2), 'b-')

# label axes
plt.xlabel('t (s) ')
plt.ylabel('$\Delta\\theta$ (rad)')

# axes scales 
plt.yscale('log')
plt.xlim(0, 120)

plt.title('Comparison between two solution,  F_D = {:.1f}'.format(F_D))
plt.show()


# Now, the difference between the two solution diverges exponentially.  The divergent behaviour is key attribute of chaotic system.

# ## Phase plots
# 
# There is another way of visualizing the motion of a pendulum.  We can plot the solution in *phase-space*.

# In[ ]:


t, ω, θ = Nonlinear_Pendulum(N=1500, dt=0.04, F_D=0.5, Ω_D=2/3, θ_0=0.75)

fig, axes = plt.subplots()

plt.plot(θ, ω, 'r.', markersize=0.5)
# label axes
plt.xlabel('$\\theta$ (rad) ')
plt.ylabel('$\omega$ (rad/s)')
# axes scales
plt.xlim(-1, 1)
plt.ylim(-0.8, 0.8)

plt.title('Results from Physical pendulum, F_D = 0.5')
plt.show()


# ### Plots the numerical solution for F_Drive = 1.2

# In[ ]:


t, ω, θ = Nonlinear_Pendulum(N=15000, dt=0.04, F_D=1.2, Ω_D=2/3, θ_0=0.2)

fig, axes = plt.subplots()

plt.plot(θ, ω, 'r.', markersize=0.5)
# label axes
plt.xlabel('$\\theta$ (rad) ')
plt.ylabel('$\omega$ (rad/s)')

# axes scales
plt.xlim(-np.pi, np.pi)
plt.ylim(-2.5, 2.5)

plt.title('Results from Physical pendulum, F_D = 1.2')
plt.show()


# If you want higher resolution, simply increase the resolution by changing `N`.