A..2 poisson_v2.py

テストのために同じことを実現できそうなことを二つ書いて， 結果を比較してみたり。

#!/opt/local/bin/python2
# -*- coding: utf-8 -*-
# poisson2_v2.py

import numpy as np
import scipy as sci
import scipy.sparse
import scipy.sparse.linalg
import matplotlib.pyplot as plot

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

#from pylab import *

n=50
h=1.0/n

################################################################
# in MATLAB: J=sparse(diag(ones(n-2,1),1)+diag(ones(n-2,1),-1));

# Solution 1
data=np.array([np.ones(n-1),np.ones(n-1)])
diags=np.array([1,-1])
J=sci.sparse.spdiags(data,diags,n-1,n-1)

# Solution 2
J2=sci.sparse.diags(np.ones(n-2),1)+sci.sparse.diags(np.ones(n-2),-1)

(J-J2).todense()

################################################################
# in MATLAB: I=speye(n-1,n-1);
I=sci.sparse.eye(n-1,n-1)

################################################################
# in MATLAB: A=-4*kron(I,I)+kron(J,I)+kron(I,J);

# Solution
A=-4*sci.sparse.kron(I,I)+sci.sparse.kron(J,I)+sci.sparse.kron(I,J)

################################################################
# in MATLAB: b=-h*h*ones((n-1)*(n-1),1);

# Solution 1
b=-h*h*np.ones((n-1)*(n-1))
b=sci.mat(b).T

# Solution 2
b2=-h*h*np.ones(((n-1)*(n-1),1))

x=sci.sparse.linalg.spsolve(A,b)
x2=sci.sparse.linalg.spsolve(A,b2)

x-x2

################################################################
# in MATLAB: 
#    U=zeros(n-1,n-1);
#    U(:)=A\b;
#    u=zeros(n+1,n+1);
#    u(2:n,2:n)=U;

u=np.zeros((n+1,n+1))
u[1:n,1:n]=x.reshape((n-1,n-1))

###############################################################

x=np.linspace(0.0,1.0,n+1)
y=np.linspace(0.0,1.0,n+1)
x,y=np.meshgrid(x,y)

fig=plot.figure('Poission eq')

ax1=fig.add_subplot(121, aspect='equal')
ax1.contour(x,y,u)

ax2=fig.add_subplot(122, aspect='equal', projection='3d')
surf=ax2.plot_surface(x,y,u,rstride=1,cstride=1,
                      cmap=cm.coolwarm,linewidth=0,antialiased=False)

plot.show()