'''
Solve the coupled system of ODEs

y'   = Ay
y(0) = y0

where A is the tri-diagonal circular matrix.

Such systems result from the finite difference discretizations
of the heat conduction systems.


dy/dt = d^2 y/dx^2 

considered with the periodic boundary conditions
'''


import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as sla
from matplotlib import pyplot as plt
from matplotlib import animation

N = 300
x = np.linspace(0,1,N)

dx = x[1]-x[0]
d1 = -2/dx**2 * np.ones(N)
d2 =  1/dx**2 * np.ones(N)

A  = sp.diags([d1,d2,d2,d2,d2],[0,-1,1,-N+1,N-1]).tocsc()



omega0 = 2*np.pi
omega1 = 4*np.pi
omega2 = 6*np.pi
omega3 = 8*np.pi
t0     = 0
y0     = np.sin(omega0*x) + np.sin(omega1*x) + np.sin(omega2*x) + np.sin(omega3*x)
h      = 1.0E-05

def ydot(t,y):
    return A.dot(y)

def eulerstep(t,y,h):
    k1=ydot(t,y)
    return y + h*k1

def midstep(t,y,h):
    k1=ydot(t,y)
    k2=ydot(t+h/2,y+h/2*k1)
    return y + h*k2

# First set up the figure, the axis, and the plot element we want to animate
fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(-3, 3))
line, = ax.plot([], [], lw=2)

# initialization function: plot the background of each frame
def init():
    line.set_data([], [])
    return line,

# animation function.  This is called sequentially
def animate(i):
    t = t0
    y = y0.copy()
    for j in range(i):
        y=midstep(t,y,h)
        t=t+h
    line.set_data(x, y)
    return line,

# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=250, interval=20, blit=True)
plt.show()