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