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

# In[ ]:


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


# # Coupled Oscillators
# ## Lecture 10

# ## Mass-spring system
# 
# Equation of motion:
# 
# $$ x'' + \omega_0^2 x = 0 $$
# 
# where $\omega_0 = \sqrt{k/m}$ is the natural angular frequency. 
# 
# The solution is
# 
# $$ x(t) = A \cos(\omega_0 t + \phi) $$
# 
# where the coefficients $A$ and $\phi$ are determined from the 
#  initial conditions for $x(0)$ and $x'(0)$.
#  
# *See handwritten notes for lecture.*

# ## Coupled mass-spring systems
# 
# 

# For a system of coupled harmonic oscillators, the motion can be understood as a superposition of independent normal modes.
# 
# The trick is finding the normal modes coordinates
# 
# $$ q_1, q_2, \ldots $$
# 
# that decouple the equations.

# ### Matrix Approach

# By rewriting the equations of motion as matrix equation
# 
# $$ x'' = A x $$
# 
# the normal modes are the *eigenvectors* of $A$ and the *eigenvalues* are the frequencies of the normal modes.

# ### Finding Eigenvectors and Eigenvalues with numpy

# We can create array using the np.array( ) function:

# In[ ]:


A = np.array([[-2, 1],[1,-2]])
print(A)


# NumPy comes with many mathematical tools that are useful, including a sub-package for methods from linear algebra.

# In[ ]:


import numpy.linalg as LA
help(LA)


# We need to calculate the eigenvectors and eigenvalues of the matrix $A$.  This can be done with the `LA.eig()` function.

# In[ ]:


help(LA.eig)


# In[ ]:


LA.eig(A)


# In[ ]:


w, v = LA.eig(A)
print(w[0], v[:,0])
print(w[1], v[:,1]) 


# We can recognize these at the eigenvectors we found earlier
# 
# $$ q_1 = \frac{1}{\sqrt{2}} \Big( \begin{matrix}
#       1\\
#       1\\
#    \end{matrix}\Big) $$
# 
# $$ q_2 = \frac{1}{\sqrt{2}} \Big( \begin{matrix}
#       -1\\
#       1\\
#    \end{matrix}\Big) $$
#       
# with the corresponding eigenvalues -1 and -3.

# Any general solution for the coupled oscillator is a linear combination of these two normal modes.

# In[ ]:


t = np.arange(0, 30, 0.01)
q1 = v[:,0].reshape(2,1) * np.cos(np.sqrt(-w[0])* t)
q2 = v[:,1].reshape(2,1) * np.cos(np.sqrt(-w[1])* t)


# Note: the `.reshape(2,1)` is needed because we need to the eigenvector to be written as a column vector and not row vector.

# With these two normal modes, we can generate arbitrary solutions.

# In[ ]:


fig, axes = plt.subplots(2, 1)

C1 = 1
C2 = 1

x = C1*q1 + C2*q2

plt.sca(axes[0])
plt.plot(t, x[0])
plt.xlabel('$t$')
plt.ylabel('$x_1$')

plt.sca(axes[1])
plt.plot(t, x[1])
plt.xlabel('$t$')
plt.ylabel('$x_2$')


plt.show()


# To visualize the different modes of the coupled mass-spring system, try this:
# 
# https://phet.colorado.edu/sims/normal-modes/normal-modes_en.html
# 
# (two masses, horizontal polarization)

# ## Four mass, five spring system

# We can extend this analysis to many more masses and many more springs:

# In[ ]:


A = np.array([[-2, 1, 0, 0],
              [ 1,-2, 1, 0],
              [ 0, 1,-2, 1],
              [ 0, 0, 1,-2]])
print(A)


# In[ ]:


w, v = LA.eig(A)
print(w[0], v[:,0])
print(w[1], v[:,1]) 
print(w[2], v[:,2])
print(w[3], v[:,3])


# Here, the eigenvalues are not in order.  Let's sort them.

# In[ ]:


idx = w.argsort()[::-1]   
w = w[idx]
v = v[:,idx]


# In[ ]:


for i in range(4):
    print(w[i], v[:,i])


# Now we can define our eigenmodes.

# In[ ]:


t = np.arange(0, 30, 0.01)
q1 = v[:,0].reshape(4,1) * np.cos(np.sqrt(-w[0])* t)
q2 = v[:,1].reshape(4,1) * np.cos(np.sqrt(-w[1])* t)
q3 = v[:,2].reshape(4,1) * np.cos(np.sqrt(-w[2])* t)
q4 = v[:,3].reshape(4,1) * np.cos(np.sqrt(-w[3])* t)


# And create arbitrary solutions!

# In[ ]:


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

C1 = 1
C2 = 1
C3 = 1
C4 = 1

x = C1*q1 + C2*q2 + C3*q3 + C4*q4

for i in range(4):
    plt.sca(axes[i])
    plt.plot(t, x[i])  
    plt.ylabel('$x_{}$'.format(i))
plt.xlabel('$t$')

plt.show()


# Again, compare with the animation
# 
# https://phet.colorado.edu/sims/normal-modes/normal-modes_en.html