xref: /petsc/src/binding/petsc4py/demo/legacy/kspsolve/petsc-cg.py (revision 5a48edb989d3ea10d6aff6c0e26d581c18691deb)
1def cg(A, b, x, imax=50, eps=1e-6):
2    """
3    A, b, x  : matrix, rhs, solution
4    imax     : maximum allowed iterations
5    eps      : tolerance for convergence
6    """
7    # allocate work vectors
8    r = b.duplicate()
9    d = b.duplicate()
10    q = b.duplicate()
11    # initialization
12    i = 0
13    A.mult(x, r)
14    r.aypx(-1, b)
15    r.copy(d)
16    delta_0 = r.dot(r)
17    delta = delta_0
18    # enter iteration loop
19    while i < imax and \
20          delta > delta_0 * eps**2:
21        A.mult(d, q)
22        alpha = delta / d.dot(q)
23        x.axpy(+alpha, d)
24        r.axpy(-alpha, q)
25        delta_old = delta
26        delta = r.dot(r)
27        beta = delta / delta_old
28        d.aypx(beta, r)
29        i = i + 1
30    return i, delta**0.5
31