Function cg_solver

python/test/conftest.py:150–200

Conjugate Gradient solver for SPD problems, which can work in serial or parallel for testing use. Not suitable for large problems.

()

@pytest.fixture(scope="function", autouse=True)
def cg_solver():
    """Conjugate Gradient solver for SPD problems, which can work in
    serial or parallel for testing use. Not suitable for large
    problems.
    """

    def _cg(comm, A, b, x, maxit=500, rtol=None):
        rtol2 = 10 * np.finfo(x.array.dtype).eps if rtol is None else rtol**2

        def _global_dot(comm, v0, v1):
            return comm.allreduce(np.vdot(v0, v1), MPI.SUM)

        A_op = A.to_scipy()
        nr = A_op.shape[0]
        assert nr == A.index_map(0).size_local * A.block_size[0]

        # Create larger ghosted vector based on matrix column space
        # and get initial y = A.x
        p = dolfinx_vector(A.index_map(1), bs=A.block_size[1], dtype=x.array.dtype)
        p.array[:nr] = x.array[:nr]
        p.scatter_forward()
        y = A_op @ p.array

        # Copy residual to p
        r = b.array[:nr] - y
        p.array[:nr] = r

        # Iterations of CG
        rnorm0 = _global_dot(comm, r, r)
        rnorm = rnorm0
        k = 0
        while k < maxit:
            k += 1
            p.scatter_forward()
            y = A_op @ p.array
            alpha = rnorm / _global_dot(comm, p.array[:nr], y)
            x.array[:nr] += alpha * p.array[:nr]
            r -= alpha * y
            rnorm_new = _global_dot(comm, r, r)
            beta = rnorm_new / rnorm
            rnorm = rnorm_new
            if rnorm / rnorm0 < rtol2:
                x.scatter_forward()
                return
            p.array[:nr] = beta * p.array[:nr] + r

        raise RuntimeError(
            f"Solver exceeded max iterations ({maxit}). Relative residual={rnorm / rnorm0}."
        )

    return _cg