# -*- coding: utf-8 -*-
"""
Created on Mon Feb 02 12:04:27 2015

Listing 13.5 : Equation de la chaleur, résolution par 
diverses méthodes de factorisation (splitting)

@author: Ribot/grivet
"""
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
#  Discrétisation en espace

xmin = 0.0; xmax = 1.0; nptx = 21; nx = nptx-2  
hx = (xmax-xmin)/(nptx -1)
xx = np.linspace(xmin,xmax,nptx) 
xx = xx.transpose()
xxint = xx[1:nx+1]
ymin = 0.0; ymax = 4.0; npty = 41; ny = npty-2 
hy = (ymax-ymin)/(npty -1)
yy = np.linspace(ymin,ymax,npty)
yy=yy.transpose() 
yyint = yy[1:ny+1]

#  Matrice de discrétisation en espace 2D

Kx = 2*np.eye(nx,nx)-np.diag(np.ones(nx-1),1)-np.diag(np.ones(nx-1),-1)
Ky = 2*np.eye(ny,ny)-np.diag(np.ones(ny-1),1)-np.diag(np.ones(ny-1),-1)
K2Dx = np.kron(np.eye(ny,ny),Kx/hx**2)
K2Dy = np.kron(Ky,np.eye(nx,nx)/hy**2);

#   Discrétisation en temps
meth = input('classe de methode (une des chaines exp ou imp): ' );
if meth == 'exp':
    #  Cas explicite
    Tfin = 1; dt = 0.9*hx**2/2; ntps = int(Tfin/dt)
elif meth == 'imp':
    #  Cas Implicite
    Tfin = 1; dt = 0.1; ntps = int(Tfin/dt)
u0 = np.kron(np.hstack([yyint[yyint<=2.0],4-yyint[yyint >2.0]]),\
     np.hstack([xxint[xxint<=0.5],1-xxint[xxint>0.5]]))
u = np.copy(u0)
meth = input('choix de la methode (une des chaines spe ou spi ou stre ou stri): ')
for t in range(ntps):
    #  splitting explicite
	if meth == 'spe':
         u = u - dt*np.dot(K2Dx,u) 
         u = u - dt*np.dot(K2Dy,u)
    #   splitting implicite
	elif meth == 'spi':
		u = np.linalg.solve(np.eye(nx*ny,nx*ny) + dt*K2Dx,u)
		u = np.linalg.solve(np.eye(nx*ny,nx*ny) + dt*K2Dy,u)
	#  Strang explicite 
	elif meth == 'stre':
		u = u-dt*np.dot(K2Dx,u)/2
		u = u-dt*np.dot(K2Dy,u)
		u = u-dt*np.dot(K2Dx,u/2)
    # Strang implicit 
	elif meth == 'stri':
		u = np.linalg.solve(np.eye(nx*ny,nx*ny) + dt*K2Dx/2,u)
		u = np.linalg.solve(np.eye(nx*ny,nx*ny) + dt*K2Dy,u);
		u = np.linalg.solve(np.eye(nx*ny,nx*ny) + dt*K2Dx/2,u)
	else:
		break
uuu = np.reshape(u,(nx ,ny),order = 'F');
uu = np.vstack((np.zeros((1,ny+2)),np.hstack((np.zeros((nx,1)),uuu,np.zeros((nx,1)))),np.zeros((1,ny+2))))
#uu = uu/abs(max(uu));
X,Y = np.meshgrid(xx,yy)
()
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X,Y,uu.T,rstride = 1, cstride = 1);