xref: /petsc/src/binding/petsc4py/demo/legacy/kspsolve/test_mat_cg.py (revision 5a48edb989d3ea10d6aff6c0e26d581c18691deb)
1try:
2    execfile
3except NameError:
4    def execfile(file, globals=globals(), locals=locals()):
5        fh = open(file, "r")
6        try: exec(fh.read()+"\n", globals, locals)
7        finally: fh.close()
8
9import petsc4py, sys
10petsc4py.init(sys.argv)
11
12from petsc4py import PETSc
13
14execfile('petsc-mat.py')
15execfile('petsc-cg.py')
16
17x, b = A.createVecs()
18
19ksp = PETSc.KSP().create()
20ksp.setType('cg')
21ksp.getPC().setType('none')
22ksp.setOperators(A)
23ksp.setFromOptions()
24
25ksp.max_it = 100
26ksp.rtol = 1e-5
27ksp.atol = 0
28x.set(0)
29b.set(1)
30ksp.solve(b,x)
31print("iterations: %d residual norm: %g" % (ksp.its, ksp.norm))
32
33x.set(0)
34b.set(1)
35its, norm = cg(A,b,x,100,1e-5)
36print("iterations: %d residual norm: %g" % (its, norm))
37
38OptDB = PETSc.Options()
39
40if OptDB.getBool('plot', True):
41    da = PETSc.DMDA().create([m,n])
42    u = da.createGlobalVec()
43    x.copy(u)
44    draw = PETSc.Viewer.DRAW()
45    OptDB['draw_pause'] = 1
46    draw(u)
47
48if OptDB.getBool('plot_mpl', False):
49    try:
50        from matplotlib import pylab
51    except ImportError:
52        print("matplotlib not available")
53    else:
54        from numpy import mgrid
55        X, Y =  mgrid[0:1:1j*m,0:1:1j*n]
56        Z = x[...].reshape(m,n)
57        pylab.figure()
58        pylab.contourf(X,Y,Z)
59        pylab.plot(X.ravel(),Y.ravel(),'.k')
60        pylab.axis('equal')
61        pylab.colorbar()
62        pylab.show()
63