xref: /petsc/src/binding/petsc4py/demo/legacy/poisson3d/poisson3d.py (revision 5a48edb989d3ea10d6aff6c0e26d581c18691deb)
1*55a74a43SLisandro Dalcinimport sys, petsc4py
2*55a74a43SLisandro Dalcinpetsc4py.init(sys.argv)
3*55a74a43SLisandro Dalcin
4*55a74a43SLisandro Dalcinfrom petsc4py import PETSc
5*55a74a43SLisandro Dalcinfrom del2mat import Del2Mat
6*55a74a43SLisandro Dalcin
7*55a74a43SLisandro Dalcin# this a sequential example
8*55a74a43SLisandro Dalcinassert PETSc.COMM_WORLD.getSize() == 1
9*55a74a43SLisandro Dalcin
10*55a74a43SLisandro Dalcin# number of nodes in each direction
11*55a74a43SLisandro Dalcin# excluding those at the boundary
12*55a74a43SLisandro Dalcinn = 32
13*55a74a43SLisandro Dalcinh = 1.0/(n+1) # grid spacing
14*55a74a43SLisandro Dalcin
15*55a74a43SLisandro Dalcin# setup linear system matrix
16*55a74a43SLisandro DalcinA = PETSc.Mat().create()
17*55a74a43SLisandro DalcinA.setSizes([n**3, n**3])
18*55a74a43SLisandro DalcinA.setType('python')
19*55a74a43SLisandro Dalcinshell = Del2Mat(n) # shell context
20*55a74a43SLisandro DalcinA.setPythonContext(shell)
21*55a74a43SLisandro DalcinA.setUp()
22*55a74a43SLisandro Dalcin
23*55a74a43SLisandro Dalcin# setup linear system vectors
24*55a74a43SLisandro Dalcinx, b = A.createVecs()
25*55a74a43SLisandro Dalcinx.set(0.0)
26*55a74a43SLisandro Dalcinb.set(1.0)
27*55a74a43SLisandro Dalcin
28*55a74a43SLisandro Dalcin# setup Krylov solver
29*55a74a43SLisandro Dalcinksp = PETSc.KSP().create()
30*55a74a43SLisandro Dalcinpc = ksp.getPC()
31*55a74a43SLisandro Dalcinksp.setType('cg')
32*55a74a43SLisandro Dalcinpc.setType('none')
33*55a74a43SLisandro Dalcin
34*55a74a43SLisandro Dalcin# iteratively solve linear
35*55a74a43SLisandro Dalcin# system of equations A*x=b
36*55a74a43SLisandro Dalcinksp.setOperators(A)
37*55a74a43SLisandro Dalcinksp.setFromOptions()
38*55a74a43SLisandro Dalcinksp.solve(b, x)
39*55a74a43SLisandro Dalcin
40*55a74a43SLisandro Dalcin# scale solution vector to
41*55a74a43SLisandro Dalcin# account for grid spacing
42*55a74a43SLisandro Dalcinx.scale(h**2)
43*55a74a43SLisandro Dalcin
44*55a74a43SLisandro DalcinOptDB = PETSc.Options()
45*55a74a43SLisandro Dalcinif OptDB.getBool('plot_mpl', False):
46*55a74a43SLisandro Dalcin    try:
47*55a74a43SLisandro Dalcin        from matplotlib import pylab
48*55a74a43SLisandro Dalcin    except ImportError:
49*55a74a43SLisandro Dalcin        PETSc.Sys.Print("matplotlib not available")
50*55a74a43SLisandro Dalcin    else:
51*55a74a43SLisandro Dalcin        from numpy import mgrid
52*55a74a43SLisandro Dalcin        X, Y =  mgrid[0:1:1j*n,0:1:1j*n]
53*55a74a43SLisandro Dalcin        Z = x[...].reshape(n,n,n)[:,:,n/2-2]
54*55a74a43SLisandro Dalcin        pylab.contourf(X, Y, Z)
55*55a74a43SLisandro Dalcin        pylab.axis('equal')
56*55a74a43SLisandro Dalcin        pylab.colorbar()
57*55a74a43SLisandro Dalcin        pylab.show()
58