Implement the Conjugate Gradient Method for Solving Linear Systems

Problem

Implement the Conjugate Gradient method for solving a system of linear equations Ax = b, where A is a symmetric positive-definite matrix. Return the solution vector x.

Solution

import numpy as np

def conjugate_gradient(A, b, x=None, tol=1e-10, max_iter=1000):
    n = len(b)
    if x is None:
        x = np.zeros(n, dtype=float)
    b = b.astype(float)
    A = A.astype(float)

    r = b - A @ x
    p = r.copy()
    rs_old = np.dot(r, r)

    for i in range(max_iter):
        if np.sqrt(rs_old) < tol:
            break
        Ap = A @ p
        alpha = rs_old / np.dot(p, Ap)
        x = x + alpha * p
        r = r - alpha * Ap
        rs_new = np.dot(r, r)
        beta = rs_new / rs_old
        p = r + beta * p
        rs_old = rs_new

    return x

Explanation

1. Start with an initial guess x (default zeros) and compute the initial residual r = b - Ax.
2. Set the initial search direction p = r.
3. At each iteration:
4. - Compute step size alpha = r^T r / (p^T A p).
5. - Update the solution: x = x + alpha * p.
6. - Update the residual: r = r - alpha * A p.
7. - Compute the new direction: p = r + beta * p where beta = r_new^T r_new / r_old^T r_old.
8. Converge when the residual norm falls below tolerance.
9. For an n x n SPD matrix, CG converges in at most n iterations (in exact arithmetic).

Complexity

• Time: O(k * n^2) where k is the number of iterations (at most n) and n is matrix dimension
• Space: O(n) for storing vectors r, p, and Ap