poissonSuperSoluce.py

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


def numgrid(n):
  X = ones((n,1))*linspace(-1,1,n)
  Y = flipud(X.T)
  map = (X > -1)&(X < 1)&(Y > -1)&(Y < 1)&((X > 0)|(Y > 0))
  map   = where(map,1,0) 
  index = where(map)
  map[index] = arange(len(index[0]))+1
  return map

def delsq(map):
  n = map.shape[0]; map = map.flatten()
  
  index = where(map)[0]
  m = len(index)  
  i = arange(m)
  j = arange(m)
  coeff = 4*ones(m)

  for iNeigh in [-1,1,n,-n]:
     mapNeigh = map[index+iNeigh]
     indexNeigh = where(mapNeigh)[0]
     mNeigh = len(indexNeigh)
     i = concatenate([i,map[index[indexNeigh]]-1])
     j = concatenate([j,mapNeigh[indexNeigh]-1])
     coeff = concatenate([coeff,-ones(mNeigh)])
  return sparse.csr_matrix((coeff,(i,j)),(m,m))
  

def poissonSolve(nCut) :
  n = 2*nCut + 1; m = n*n; h = 2/(n-1)
  map = numgrid(2*nCut+1)
  index = map[map>0]
  
  A = delsq(map)
  B = ones(size(index))
  U = zeros(shape(map))
  U[map>0] = (spsolve(A,B)[index-1])*(h*h)  
  return U