# -*- coding: utf-8 -*-
"""
Created on Fri Jan 23 17:46:45 2015

Listing 13.1 : résolution par différences finies de 
l'équation laplacien(u) = 100x(1-x)y(4-y)

@author: ribot/grivet
"""

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

#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 = np.transpose(xx)
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 = np.transpose(yy)
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))
K2D = np.kron(np.eye(ny,ny),Kx)/hx**2\
             + np.kron(Ky,np.eye(nx ,nx))/hy**2

# Résolution du système 2D

F = 100*np.kron(yyint*(4-yyint),xxint*(1-xxint));
u = np.linalg.solve(K2D,F) 
uu = np.reshape(u,(nx,ny),order = 'F');
ub = np.vstack((np.zeros((1,ny+2)),np.hstack((np.zeros((nx ,1)),uu,np.zeros((nx ,1)))),np.zeros((1 ,ny+2))))
#uu1 = [zeros(nx ,1) ,uu , zeros(nx ,1)];
#uu2 = [zeros(1 ,ny+2);uu1;zeros(1 ,ny+2)];

X,Y = np.meshgrid(xx,yy)
ax.plot_surface(X,Y,ub.T,rstride = 1, cstride = 1);