# -*- coding: utf-8 -*-
"""
Created on Wed Jan 28 15:17:47 2015

Listing 13.3 : Résolution de l'équation de la chaleur
par Euler explicite, Euler Implicite ou Crank-Nicolson

@author: ribot/grivet
"""
import numpy as np

#  Discrétisation en espace

xmin=0; xmax=1; npt=21; nx=npt-2
h = (xmax-xmin)/(npt -1.0)
xx = np.linspace(xmin,xmax,npt)
xx = np.transpose(xx)
xxint = xx[1:nx+1]

# Matrice de discrétisation en espace 1D

K = 2*np.eye(nx,nx) - np.diag(np.ones(nx-1),1) - np.diag(np.ones(nx-1),-1)
K1D = K/h**2

#  Discrétisation en temps

meth = input("choix de la methode (une chaine parmi ee, ei ou cn): ")

#  Cas explicite
if meth == "ee":
    Tfin = 1; dt = 0.9*h**2/2; ntps = int(Tfin/dt)

#  Cas Implicite ou Crank-Nicolson
if meth == "ei" or meth == "cn":
    	Tfin = 1; dt = 0.1; ntps = int(Tfin/dt)

#u1 = xxint[xxint<=0.5]
#u2 =  1-xxint[xxint>0.5]
u0 = np.hstack((xxint[xxint <= 0.5],1.0 - xxint[xxint > 0.5]))
#u0 = array([0.05,0.1,0.15,0.2,0.25,0.30,0.35,0.4,0.45,0.5,0.45,0.4,0.35,\
#0.3,0.25,0.2,0.15,0.1,0.05])
u = np.copy(u0)
for t in range(ntps):
    #  Euler Explicite
    if meth == "ee":
        u = u-dt*np.dot(K1D,u)
    #  Euler Implicite
    elif meth == "ei":
        u = np.linalg.solve(np.eye(nx,nx) + dt*K1D,u)
    #  Crank-Nicolson
    elif meth == "cn":
        u = np.linalg.solve(np.eye(nx ,nx)+dt*K1D/2, u-dt*np.dot(K1D,u/2))
    else:
        break
uu = hstack([0,u,0]) 
plot(xx,uu);