from IPython.core.display import HTML
css_file = './example.css'
HTML(open(css_file, "r").read())

%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from clawpack.visclaw.JSAnimation import IPython_display

m=2**8-1                           # Number of grid points
h=1./(m+1)                         # Mesh width
x=np.linspace(0,1,m+2); x=x[1:-1]  # grid
alpha=1.; beta=3.                  # boundary values
phi = lambda x: 20.* np.pi * x**3
f = lambda x: -20 + 0.5*120*np.pi*x * np.cos(phi(x)) - \
                0.5*(60*np.pi*x**2)**2 * np.sin(phi(x)) # RHS
plt.plot(x,f(x)); plt.xlabel('x'); plt.ylabel('f(x)')

U=np.linspace(alpha,beta,m)       # Just use a straight line as initial guess
UU=[U]
F=0.5*h**2*f(x)                   # Construct the RHS
F[0]-=alpha/2.; F[-1]-=beta/2.    # boundary conditions
e=np.ones(m-1)
G=0.5*(np.diag(e,-1)+np.diag(e,1)) # Jacobi matrix (average neighbors)

for i in range(100):
    U=np.dot(G,U) - F
    UU.append(U)

xf = np.linspace(0,1,1000); 
uf = 1.+12.*xf-10.*xf**2 + 0.5*np.sin(phi(xf))  # Exact solution on fine grid
u = 1.+12.*x-10.*x**2 + 0.5*np.sin(phi(x))  # Exact solution on computational grid

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(-3, 6))
ax.plot(xf,uf)
line1, = ax.plot([],[],'-')
line2, = ax.plot([],[],'-o')
plt.legend(['Exact','Jacobi','Error'],loc=2)

def fplot(i):
    line1.set_data(x,UU[i])
    line2.set_data(x,UU[i]-u)
    ax.set_title('# of iterations: '+str(i))
    return line1, 

animation.FuncAnimation(fig, fplot, frames=range(len(UU)))

M=30; H=1./(M+1)
omegas=np.linspace(0.,1.)
p=np.arange(1,M+1)

fig = plt.figure()
ax = plt.axes(xlim=(0, M), ylim=(0,1))
ax.set_xlabel('p', fontsize=20)
ax.set_ylabel('$\widehat{\gamma}_p$', fontsize=20)
line1, = ax.plot([],[],'o')


def fplot(i):
    omega = omegas[i]
    gamma = abs(1.+omega*(np.cos(p*np.pi*H)-1))
    line1.set_data(p,gamma)
    ax.set_title('$\omega = %02.2f $' % omega, fontsize=20)
    return line1, 

animation.FuncAnimation(fig, fplot, frames=range(50))

omega = 2./3.
U=np.linspace(alpha,beta,m)  # Just use a straight line as initial guess
UU=[U]
F=0.5*h**2*f(x)
F[0]-=alpha/2.; F[-1]-=beta/2.    # Construct the RHS, including boundary conditions
e=np.ones(m-1)
G=0.5*(np.diag(e,-1)+np.diag(e,1))
niter=100

for i in range(niter):
    U=(1.-omega)*U+omega*(np.dot(G,U)-F)
    UU.append(U)
    
    
xf = np.linspace(0,1,1000); 
uf = 1.+12.*xf-10.*xf**2 + 0.5*np.sin(phi(xf))  # Exact solution on fine grid
u = 1.+12.*x-10.*x**2 + 0.5*np.sin(phi(x))  # Exact solution on computational grid

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(-3, 6))
ax.plot(xf,uf)
line1, = ax.plot([],[],'-')
line2, = ax.plot([],[],'-o')
plt.legend(['Exact','Jacobi','Error'],loc=2)

def fplot(i):
    line1.set_data(x,UU[i])
    line2.set_data(x,UU[i]-u)
    ax.set_title('# of iterations: '+str(i))
    return line1, 

animation.FuncAnimation(fig, fplot, frames=range(niter))

from matplotlib import animation
import matplotlib.pyplot as plt
from clawpack.visclaw.JSAnimation import IPython_display

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(-1.2, 1.2))

p=5
xf=np.linspace(0,1,1000) # fine grid

line1, = ax.plot(xf, np.sin(p*np.pi*xf), '-r', lw=2)
line2, = ax.plot([],[],'o-',lw=2)

def fplot(m):
    x=np.linspace(0,1,m+2); # grid
    line2.set_data(x,np.sin(p*np.pi*x))
    ax.set_title('m='+str(m))
    return line2,

animation.FuncAnimation(fig, fplot, frames=(63,31,15,7,3,1))

def coarsen(f):
    return f[1::2]  # This slicing just takes the odd-numbered points

def interpolate(f,alpha,beta):
    m_coarse=len(f)
    m_fine  =2*m_coarse+1
    f_interp = np.zeros(m_fine)
    f_interp[1::2]=f                          #Set the values of the odd numbered points
    f_interp[2:-1:2]=0.5*(f[:-1]+f[1:])       #Set the values of the (interior) even numbered points
    f_interp[0]=0.5*(f_interp[1]+alpha)     #Set the values of the endpoints
    f_interp[-1]=0.5*(f_interp[-2]+beta)
    return f_interp

#compute the residual
A=2./h**2*(G-np.eye(m))
F=f(x); F[0]-=alpha/h**2; F[-1]-=beta/h**2
r=F-np.dot(A,U)
print max(abs(r))
omega=2./3.
m2=(m-1)/2.               # Number of points in coarse grid
h2=1./(m2+1)              # coarse mesh width
x2=coarsen(x)             # coarse grid points
r2=coarsen(r)             # residual restricted to coarse grid
err2=np.zeros(m2)         # initial guess for the error
e2=np.ones(m2-1); G2=0.5*(np.diag(e2,-1)+np.diag(e2,1))
for i in range(3):                        #Solve Ae=-r by Jacobi iteration
    err2=(1.-omega)*err2+omega*(np.dot(G2,err2)+0.5*h2**2*r2)
err=interpolate(err2,0,0)     # interpolate the error
print err[0]
U=U-err                   # and use it to correct our solution
plt.clf()
plt.plot(x,u,x,U,x,U-u)
plt.legend(['Exact','Jacobi','Error'],loc='best')

def Jacobi(U,f,alpha,beta,m,nu):
    """Perform nu Jacobi iterations on a grid with m points, with initial guess U, right hand side function f and 
       Dirichlet boundary conditions with values alpha and beta.  The function returns both the approximate
       solution and the residual."""
    omega=2./3.
    h=1./(m+1)
    F=0.5*h**2*f.copy()
    F[0]-=alpha/2.; F[-1]-=beta/2.
    e=np.ones(m-1)
    G=0.5*(np.diag(e,-1)+np.diag(e,1))
    for i in range(nu):
        U=(1.-omega)*U + omega*(np.dot(G,U)-F)
    A=2./h**2*(G-np.eye(m))
    FF=f.copy(); FF[0]-=alpha/h**2; FF[-1]-=beta/h**2
    rr=FF-np.dot(A,U)
    return U,rr

k=8;
m=2**k-1
rdep=k-1  # Recursion depth; this is how many grids down we want to go
          # rdep=k-1 gives a full V-cycle
nu=3      # Number of Jacobi iterations to take at each step
U=np.linspace(alpha,beta,m)  # Initial guess
x=np.linspace(0,1,m+2); x=x[1:-1]  # grid
phi = lambda x: 20.* np.pi * x**3
u = 1.+12.*x-10.*x**2 + 0.5*np.sin(phi(x))
f = lambda x: -20 + 0.5*120*np.pi*x * np.cos(phi(x)) - 0.5*(60*np.pi*x**2)**2 * np.sin(phi(x))
alpha=1.; beta=3.
F=f(x)
r=[None]*(rdep+1); error=[None]*(rdep+1)        # This just initializes these lists to have the right length
U,rr=Jacobi(U,F,alpha,beta,m,nu)  # Initial iteration on fine grid
for i in range(1,rdep+1): # Going down
    m=(m-1)/2 #  = 2**(k-i) - 1
    r[i]=coarsen(rr)             # residual restricted to next coarser grid
    error[i],rr=Jacobi(np.zeros(m),-r[i],0.,0.,m,nu)
for i in range(1,rdep): # Coming up
    m=2*m+1
    err=error[rdep-i]-interpolate(error[rdep+1-i],0,0)     # Interpolate and subtract the correction
    error[-i-1],rr=Jacobi(err,-r[rdep-i],0.,0.,m,nu)
m=2*m+1
U=U-interpolate(error[1],0,0)                   # final solution correction
U,rr=Jacobi(U,F,alpha,beta,m,nu)                #Final iterations on original grid
plt.clf()
plt.plot(x,u,x,U,'o-',x,U-u)
plt.legend(['Exact','Jacobi','Error'],loc='best')

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(-5,5))
line1, = ax.plot([],[],'-ok',lw=2)

def fplot(i):
    y = error[-i]
    m = len(y)
    x=np.linspace(0,1,m+2); x=x[1:-1]    
    line1.set_data(x,y)
    return line1, 

animation.FuncAnimation(fig, fplot, frames=range(1,len(error)))





from IPython.display import display, HTML
from IPython.html import widgets
from IPython.html.widgets import interact, interactive

def Jacobi(U,f,alpha,beta,m,nu,omega=2./3):
    """Perform nu Jacobi iterations on a grid with m points, with initial guess U, right hand side function f and 
       Dirichlet boundary conditions with values alpha and beta.  The function returns both the approximate
       solution and the residual."""
    h=1./(m+1)
    F=0.5*h**2*f.copy()
    F[0]-=alpha/2.; F[-1]-=beta/2.
    e=np.ones(m-1)
    G=0.5*(np.diag(e,-1)+np.diag(e,1))
    for i in range(nu):
        U=(1.-omega)*U + omega*(np.dot(G,U)-F)
    A=2./h**2*(G-np.eye(m))
    FF=f.copy(); FF[0]-=alpha/h**2; FF[-1]-=beta/h**2
    rr=FF-np.dot(A,U)
    return U,rr

def Vcycle(k=8,nu=3,omega=2./3,smooth_after_interp=True):
    m=2**k-1
    rdep=k-1  # Recursion depth; this is how many grids down we want to go
              # rdep=k-1 gives a full V-cycle
    alpha=1.; beta=3.
    U=np.linspace(alpha,beta,m)  # Initial guess
    x=np.linspace(0,1,m+2); x=x[1:-1]  # grid
    u = 1.+12.*x-10.*x**2 + 0.5*np.sin(phi(x))

    F=f(x)
    r=[None]*(rdep+1); error=[None]*(rdep+1) # This just initializes these lists to have the right length
    U,rr=Jacobi(U,F,alpha,beta,m,nu,omega)  # Initial iteration on fine grid
    for i in range(1,rdep+1): # Going down
        m=(m-1)/2 #  = 2**(k-i) - 1
        r[i]=coarsen(rr)             # residual restricted to next coarser grid
        error[i],rr=Jacobi(np.zeros(m),-r[i],0.,0.,m,nu,omega)
    for i in range(1,rdep): # Coming up
        m=2*m+1
        err=error[rdep-i]-interpolate(error[rdep+1-i],0,0)     # Interpolate and subtract the correction
        if smooth_after_interp:
            error[-i-1],rr=Jacobi(err,-r[rdep-i],0.,0.,m,nu,omega)
        else:
            error[-i-1] = err

    m=2*m+1
    U=U-interpolate(error[1],0,0)                   # final solution correction
    U,rr=Jacobi(U,F,alpha,beta,m,nu,omega)                #Final iterations on original grid
    
    fig, ax = plt.subplots(figsize=(8,6),subplot_kw={'axisbg':'#EEEEEE',
                                       'axisbelow':True})
    ax = plt.axes(xlim=(0, 1), ylim=(-5,6))
    ax.plot(x,u,x,U,'o-',x,U-u)
    plt.legend(['Exact','Jacobi','Error'],loc='best')

interact(Vcycle, k=(4,10,1), nu=(1,4,1), omega=(0,2,1./6))