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