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

# In[92]:


from __future__ import division
from numpy import cos, pi
import numpy as np

def eady_qy_series(N,z):
    ''' Galerkin approximatin to Eady's 
            basic state U = z + 1 
        N is the number of baroclinic modes '''

    n = np.arange(N)
    Z = np.array( np.matrix(2*n+1).T*np.matrix(z) )

    return  4.*cos(pi*Z).sum(axis=0)


# In[99]:


import numpy as np
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
plt.rcParams.update({'font.size': 25, 'legend.handlelength'  : 1.25})
import seaborn as sns
sns.set(style="darkgrid")
sns.set_context("notebook", font_scale=2, rc={"lines.linewidth": 4.})


z = np.linspace(-1.,0,300)    
U_exact = z + 1.

N = np.arange(200)
Qy_g = np.zeros((z.size,N.size))

for i in range(N.size):
    Qy_g[:,i] = eady_qy_series(N[i],z)


# In[100]:


from JSAnimation.IPython_display import display_animation
from matplotlib import animation


# In[101]:


fig = plt.figure(figsize=(7,7))
ax = plt.axes(xlim=(Qy_g[:,-1].min(),Qy_g[:,-1].max()), ylim=(-1.025,0.025))
line = ax.plot([], [], color='#003366', ls='-', lw=2.5,alpha=.5)[0]
number = ax.text(.9*Qy_g[:,-1].min(),-.075,'n = 0')


ax.spines['left'].set_position('center')
ax.spines['bottom'].set_position('center')

def animate(n):
    line.set_data(Qy_g[:,n],z)
    number.set_text('n = '+str(2*n+1))


animation.FuncAnimation(fig, animate,
                        frames=N.size, interval=100)


# In[89]:





# In[75]:





# In[ ]: