xref: /petsc/src/ksp/ksp/tutorials/example100.py (revision c20d77252dee0f9c80fc6f8b1a6f948e11175edb) 
1from __future__ import print_function
2
3# --------------------------------------------------------------------
4
5from petsc4py import PETSc
6
7# --------------------------------------------------------------------
8
9OptDB = PETSc.Options()
10
11INFO = OptDB.hasName('info')
12
13def LOG(arg):
14    if INFO:
15        print(arg)
16
17# --------------------------------------------------------------------
18
19class Laplace1D(object):
20
21    def create(self, A):
22        LOG('Laplace1D.create()')
23        M, N = A.getSize()
24        assert M == N
25
26    def destroy(self, A):
27        LOG('Laplace1D.destroy()')
28
29    def view(self, A, vw):
30        LOG('Laplace1D.view()')
31
32    def setFromOptions(self, A):
33        LOG('Laplace1D.setFromOptions()')
34
35    def setUp(self, A):
36        LOG('Laplace1D.setUp()')
37
38    def assemblyBegin(self, A, flag):
39        LOG('Laplace1D.assemblyBegin()')
40
41    def assemblyEnd(self, A, flag):
42        LOG('Laplace1D.assemblyEnd()')
43
44    def getDiagonal(self, A, d):
45        LOG('Laplace1D.getDiagonal()')
46        M, N = A.getSize()
47        h = 1.0/(M-1)
48        d.set(2.0/h**2)
49
50    def mult(self, A, x, y):
51        LOG('Laplace1D.mult()')
52        M, N = A.getSize()
53        xx = x.getArray(readonly=1) # to numpy array
54        yy = y.getArray(readonly=0) # to numpy array
55        yy[0]    =  2.0*xx[0] - xx[1]
56        yy[1:-1] = - xx[:-2] + 2.0*xx[1:-1] - xx[2:]
57        yy[-1]   = - xx[-2] + 2.0*xx[-1]
58        h = 1.0/(M-1)
59        yy *= 1.0/h**2
60
61    def multTranspose(self, A, x, y):
62        LOG('Laplace1D.multTranspose()')
63        self.mult(A, x, y)
64
65
66# --------------------------------------------------------------------
67
68class Jacobi(object):
69
70    def create(self, pc):
71        LOG('Jacobi.create()')
72        self.diag = None
73
74    def destroy(self, pc):
75        LOG('Jacobi.destroy()')
76        if self.diag:
77            self.diag.destroy()
78
79    def view(self, pc, vw):
80        LOG('Jacobi.view()')
81
82    def setFromOptions(self, pc):
83        LOG('Jacobi.setFromOptions()')
84
85    def setUp(self, pc):
86        LOG('Jacobi.setUp()')
87        A, B = pc.getOperators()
88        self.diag = B.getDiagonal(self.diag)
89
90    def apply(self, pc, x, y):
91        LOG('Jacobi.apply()')
92        y.pointwiseDivide(x, self.diag)
93
94    def applyTranspose(self, pc, x, y):
95        LOG('Jacobi.applyTranspose()')
96        self.apply(pc, x, y)
97
98# --------------------------------------------------------------------
99
100class ConjGrad(object):
101
102    def create(self, ksp):
103        LOG('ConjGrad.create()')
104        self.work = []
105
106    def destroy(self, ksp):
107        LOG('ConjGrad.destroy()')
108        for vec in self.work:
109            if vec:
110                vec.destroy()
111        self.work = []
112
113    def view(self, ksp, viewer):
114        LOG('ConjGrad.view()')
115
116    def setUp(self, ksp):
117        LOG('ConjGrad.setUp()')
118        self.work[:] = ksp.getWorkVecs(right=3, left=None)
119
120    def solve(self, ksp, b, x):
121        LOG('ConjGrad.solve()')
122        A, P = get_op_pc(ksp, transpose=False)
123        pcg(ksp, A, P, b, x, *self.work)
124
125    def solveTranspose(self, ksp, b, x):
126        LOG('ConjGrad.solveTranspose()')
127        A, P = get_op_pc(ksp, transpose=True)
128        pcg(ksp, A, P, b, x, *self.work)
129
130def get_op_pc(ksp, transpose=False):
131    op, _ = ksp.getOperators()
132    pc = ksp.getPC()
133    if not transpose:
134        A = op.mult
135        P = pc.apply
136    else:
137        A = op.multTranspose
138        P = pc.applyTranspose
139    return A, P
140
141def do_loop(ksp, r):
142    its = ksp.getIterationNumber()
143    rnorm = r.norm()
144    ksp.setResidualNorm(rnorm)
145    ksp.logConvergenceHistory(rnorm)
146    ksp.monitor(its, rnorm)
147    reason = ksp.callConvergenceTest(its, rnorm)
148    if not reason:
149        ksp.setIterationNumber(its+1)
150    else:
151        ksp.setConvergedReason(reason)
152    return reason
153
154def pcg(ksp, A, P, b, x, r, z, p):
155    A(x, r)
156    r.aypx(-1, b)
157    P(r, z)
158    delta = r.dot(z)
159    z.copy(p)
160    while not do_loop(ksp, z):
161        A(p, z)
162        alpha = delta / z.dot(p)
163        x.axpy(+alpha, p)
164        r.axpy(-alpha, z)
165        P(r, z)
166        delta_old = delta
167        delta = r.dot(z)
168        beta = delta / delta_old
169        p.aypx(beta, z)
170
171def richardson(ksp, A, P, b, x, r, z):
172    A(x, r)
173    r.aypx(-1, b)
174    P(r, z)
175    x.axpy(1, z)
176    while not do_loop(ksp, z):
177        A(x, r)
178        r.aypx(-1, b)
179        P(r, z)
180        x.axpy(1, z)
181
182# --------------------------------------------------------------------
183