poissonSoluce.py

from numpy import *

#
# -1- Importer le solveur creux de scipy :-)
#

from scipy.sparse.linalg import spsolve 
import scipy.sparse as sparse

#
# -2- Construire et résoudre le système linéaire
#     L'unique astuce est d'écrire la ligne qui permet
#     d'omettre le bon morceau de domaine 
#
#        if (i > nCut or j < nCut):  :-)
#

def poissonSolve(nCut):
  n = 2*nCut + 1; m = n*n; h = 2/(n-1)  
  A = sparse.dok_matrix(sparse.eye(m),dtype=float64)
  B = zeros(m)  
  for i in range(1,n-1):
      for j in range(1,n-1):
         if (i > nCut or j < nCut): 
            index = i + j*n
            A[index,index] = 4
            A[index,index-1] = -1
            A[index,index+1] = -1
            A[index,index+n] = -1
            A[index,index-n] = -1
            B[index] = 1            
  A = A / (h*h)  
  return spsolve(A.tocsr(),B).reshape(n,n)