Krylov Subspace Methods - Executable Engineering: Interactive Numerical Methods in Python

import numpy as np
from scipy.sparse.linalg import cg, gmres
import plotly.graph_objects as go

try:
    import executable_engineering as exe
except ImportError:
    %pip install -q executable_engineering
    import executable_engineering as exe

The Subspace Philosophy¶

Unlike stationary methods (which continually update a single guess vector), Krylov subspace methods construct a growing set of vector bases (directions) and search within that spanning subspace for the optimal solution.

Because we are building a basis for an $n$ -dimensional vector space, we would theoretically need at most $n$ basis vectors. However, the true power of these methods is that they often converge to an acceptable tolerance using significantly fewer than $n$ vectors!

Conjugate Gradient (CG)¶

The Conjugate Gradient (CG) method is the gold standard for solving symmetric positive definite (SPD) matrices.

Instead of directly solving $\mathbf{A}\mathbf{x}=\mathbf{b}$ , we can reframe the problem as an optimization task. Consider the quadratic surface defined by:

f(\mathbf{x}^{(k)}) = \frac{1}{2} {\mathbf{x}^{(k)}}^T \mathbf{A} \mathbf{x}^{(k)} - \mathbf{b}^T \mathbf{x}^{(k)}

(1)

This surface has an extremum where its gradient is zero:

\begin{aligned} \nabla f = \frac{d f}{ d \mathbf{x}} = \vec{0} &= \mathbf{A} \mathbf{x} - \mathbf{b} \\ \vec{0} &= \mathbf{r} \end{aligned}

(2)

Notice that the gradient is exactly the residual $\mathbf{r}$ . Therefore, finding the minimum of the surface $f$ is mathematically identical to solving the linear system!

The surface conception¶

Why can we now think of a quadratic surface? Why does the matrix $\mathbf{A}$ need to be symmetric positive definite? This geometric interpretation is only valid if we can think of $\mathbf{A}$ as the matrix of mixed second derivatives of the surface $f$ , $\begin{bmatrix} \frac{\partial^2 f}{\partial x_1^2} & \frac{\partial^2 f}{\partial x_1 \partial x_2} \\ \frac{\partial^2 f}{\partial x_2 \partial x_1} & \frac{\partial^2 f}{\partial x_2^2} \\ \end{bmatrix}$ (called the Hessian of $f$ ). This is only possible if $\mathbf{A}$ is symmetric!

Furthermore, if $\mathbf{A}$ is positive definite ( $\mathbf{x}^T \mathbf{A} \mathbf{x} > 0$ ), the quadratic surface is strictly convex (opening upwards like a bowl). This guarantees that the extremum is a global minimum, and that walking “downhill” on the surface will inevitably lead us to the exact solution!

If $\mathbf{A}$ isn’t positive definite the surface is called a saddle point because if curves upwards in some direction and downwards in others.

Step length¶

To walk down the bowl towards the minimum, we start at a guess $\mathbf{x}^{(k)}$ . We must choose a step direction $\mathbf{s}^{(k)}$ and a step length $\alpha_k$ to reach our next, better guess $\mathbf{x}^{(k+1)}$ :

\mathbf{x}^{(k+1)} = \mathbf{x}^{(k)} + \alpha_k \mathbf{s}^{(k)}

(3)

The magic of these methods lies entirely in how we choose the step direction. To understand this, we will approach this problem backwards:

Find the optimal $\alpha_k$ assuming we already have $\mathbf{s}^{(k)}$
Find the suitable direction $\mathbf{s}^{(k)}$

Given a step direction, an obvious choice for the length $\alpha_k$ is whatever minimizes $f(\mathbf{x}^{(k+1)})$ along that specific ray. By setting the derivative with respect to $\alpha_k$ to zero, we find the optimal step length:

\begin{align} f(\vec{x}+\alpha^k \vec{s}^k) &= \frac{1}{2} [\vec{x}+\alpha^k \vec{s}^k]^T A [\vec{x}+\alpha^k \vec{s}^k] -\vec{b}^T [\vec{x}+\alpha^k \vec{s}^k] \\ \frac{\partial f}{\partial \alpha^k} = 0 &= {\vec{s}^k}^T A [\vec{x}+\alpha^k \vec{s}^k] -\vec{b}^T \vec{s}^k\\ &= {\vec{s}^k}^T A \vec{x}^k+ {\vec{s}^k}^T A \alpha^k \vec{s}^k -{\vec{s}^k}^T \vec{b} \\ &= {\vec{s}^k}^T [A \vec{x}^k-\vec{b}] + \alpha^k {\vec{s}^k}^T A \vec{s}^k \end{align}

(4)

and therefore the optimal step size is:

\begin{aligned} \alpha_k &= \frac{{\mathbf{s}^{(k)}}^T \mathbf{r}^{(k)}}{{\mathbf{s}^{(k)}}^T \mathbf{A} \mathbf{s}^{(k)}} \end{aligned}

(5)

So, given any step direction $\mathbf{s}^{(k)}$ , we can always calculate the mathematically optimal distance to travel.

Step Direction: Steepest Descent¶

The most intuitive choice is to simply walk straight downhill following the steepest gradient of the surface. We already proved the gradient is the residual, so:

\mathbf{s}^{(k)} = -\nabla f = \mathbf{r}^{(k)}

(6)

Substituting this into our step length formula yields the classic Steepest Descent algorithm! Let’s code it up and visualize the path it takes down the 2D contour.

Notice the zig-zag path the Steepest Descent algorithm takes. It frantically zig-zags back and forth across the ravine! Why? Because taking the steepest path down often forces the next step to partially undo the progress of the previous step. In an $n$ -dimensional space, this inefficiency becomes catastrophic.

A = np.array([[4, 1], [1, 3]])
b = np.array([5, 6])
x0 = np.array([0., 0.])

# Calculate Steepest Descent
x_final, iterations = exe.steepest_descent(A, b, x0, track_history=True)

# Visualize the path over the purely quadratic surface
fig = exe.visualize_convergence_2d(A, b, iterations, surface_type='quadratic')
fig.show()

Step Direction: Conjugate Gradient¶

To fix the zig-zagging, we must choose step directions that strictly do not undo each other. We enforce this by requiring that every new step direction is conjugate (orthogonal with respect to $\mathbf{A}$ ) to all previous step directions:

{\mathbf{s}^{(k+1)}}^T \mathbf{A} \mathbf{s}^{(k)} = 0

(7)

This ensures the steps are linearly independent. By building the next direction from a combination of the current residual and the previous direction, we derive the Conjugate Gradient update rule:

\mathbf{s}^{(k+1)} = \mathbf{r}^{(k+1)} - \frac{{\mathbf{r}^{(k+1)}}^T \mathbf{A} \mathbf{s}^{(k)}}{{\mathbf{s}^{(k)}}^T \mathbf{A} \mathbf{s}^{(k)}} \mathbf{s}^{(k)}

(8)

Let’s use scipy.sparse.linalg.cg to see how beautifully this solves the zig-zag problem!

A = np.array([[4, 1], [1, 3]])
b = np.array([5, 6])
x0 = np.array([0., 0.])

# We use a clean tracker from our package to avoid boilerplate callbacks
tracker = exe.IterationTracker(x0, A, b, print_residual=True)

# Execute the native SciPy solver
solution, info = cg(A, b, x0=x0, callback=tracker)

# Visualize the path over the purely quadratic surface
fig = exe.visualize_convergence_2d(A, b, tracker.iterations, surface_type='quadratic')
fig.show()

iterations = np.array(tracker.iterations)
print('Converged in ', len(tracker.iterations)-1, ' steps.')

Verifing the properties¶

As shown in the plot, CG walks directly to the center of the bowl. We can mathematically verify that the steps are indeed conjugate in $\mathbf{A}$ , ( ${\mathbf{s}^{2}}^T \mathbf{A} \mathbf{s}^{1} =0$ ). It can also be shown that subsequent residuals are also (standard) orthogonal.

s1 = iterations[1]-iterations[0]
s2 = iterations[2]-iterations[1]

print('The steps are ', s1,s2)

print('The dot product of steps: ', np.dot(s1,s2)) 
print('The dot product of a step and the transformed previous step', np.dot(s2,A@s1)) 

r0 = A@iterations[0]-b
r1 = A@iterations[1]-b
r2 = A@iterations[2]-b

print('Dot product of subsequent residuals, ', r0.dot(r1), r1.dot(r2))

Is CG a direct method?¶

Note the space of the solution vector $\mathbf{x}$ is just its dimension $n$ , and if CG determins linearly independent vectors, it can only find, at most $n$ ! Therefore, CG will find the exact solution in $n$ iterations (baring roundoff error) which some may take to imply it is a direct method. However, the situation is even better, which qualifies it as an iterative technique! Let’s see what happens with a larger system:

n = 8
A_r = np.random.rand(n, n)
A_rand = A_r @ A_r.T  # Ensure A is symmetric positive definite
b_rand = np.random.rand(n)
x0_rand = np.zeros(n)

cg_tracker = exe.IterationTracker(x0_rand, A_rand, b_rand, print_residual=True)
solution, info = cg(A_rand, b_rand, x0=x0_rand, atol=1e-6, rtol=1e-10, callback=cg_tracker)

np.set_printoptions(precision=4)
print("True solution:", np.linalg.solve(A_rand, b_rand))
print("CG solution:  ", solution)
print(f"\nConverged in {len(cg_tracker.iterations)-1} iterations (n={n})")

Restart and Matrix free methods¶

Restart: In extreme-scale computations, round-off error can slowly destroy the strict conjugacy of the vectors. Many implementations include a restart parameter that periodically flushes the basis history to keep the algorithm stable.

Matrix-free methods Notice that CG only ever uses $\mathbf{A}$ to compute matrix-vector products ( $\mathbf{A} \mathbf{x}$ ). If we can calculate $\mathbf{A}\mathbf{x}$ directly on-the-fly without ever assembling or storing the massive dense matrix $\mathbf{A}$ , we can save colossal amounts of RAM. This is the foundation of matrix-free HPC solvers!

Generalized Minimal Residual Method (GMRES)¶

What if $\mathbf{A}$ is not symmetric?

The GMRES method¶

f $\mathbf{A}$ is non-symmetric, it cannot be a Hessian, and the beautiful quadratic bowl surface analogy falls apart. We can no longer enforce $\mathbf{A}$ -conjugacy.

Instead, we use the Generalized Minimal Residual (GMRES) method. GMRES abandons the bowl analogy and directly minimizes the Euclidean norm of the residual $||\mathbf{A}\mathbf{x}-\mathbf{b}||$ . Step directions, $\mathbf{s}^k$ are chosen to be orthogonal which is a more general, but less powerful condition. The method will still converge in at most $n$ iterations (barring roundoff error), but typically converges sooner (not as quickly as CG however).

Because it must enforce standard orthogonality against all previous search directions, its memory footprint grows with every iteration. Therefore, GMRES almost always requires a restart parameter to prevent memory exhaustion on large systems!

# We explicitly define a non-symmetric matrix!
A = np.array([[4, -1], [2, 3]])
b = np.array([5, 6])
x0 = np.array([0., 0.])

tracker = exe.IterationTracker(x0, A, b, print_residual=True)

# Execute the native SciPy solver
# Note: callback_type='pr_norm' provides the residual at each inner step, but GMRES only calculates the physical vector at the very end!
solution, info = gmres(A, b, x0=x0, callback=tracker, callback_type='pr_norm')

print("True solution:", np.linalg.solve(A, b))
print("GMRES solution:", solution)

# Because GMRES only calculates the final vector, we manually append it for the visualizer
tracker.iterations.append(solution)
print(f"Converged in {tracker.step} steps.")

CG vs GMRES: Residual Comparison¶

Now that we understand GMRES, let’s see what happens if we use it to solve a newly generated random $8 \times 8$ symmetric positive definite matrix and compare it directly to Conjugate Gradient.

Notice something fascinating in the plot below: The Euclidean residual for Conjugate Gradient is not strictly monotonically decreasing. This is a fundamental mathematical property of Conjugate Gradient.

Methods like GMRES are designed to strictly minimize the Euclidean norm of the residual at every step.
Conjugate Gradient minimizes the energy norm; $\sqrt{ \Delta \mathbf{x}^T \mathbf{A} \Delta \mathbf{x} }$ .

While the energy norm drops monotonically at every single step of CG, the standard Euclidean residual is free to oscillate along the way. This is a brilliant pedagogical reminder of exactly what surface the CG algorithm is navigating!

Notice also how GMRES never exceeds $n=8$ iterations, but CG often takes 9 or 10? In exact mathematical theory, both algorithms guarantee exact convergence in at most $n$ steps. However, in finite precision floating-point arithmetic, CG’s reliance on only the previous two vectors (the 3-term recurrence) causes it to slowly lose orthogonality over time, wandering back into previously searched subspaces. GMRES explicitly stores and orthogonalizes against every previous vector, remaining mathematically rigid.

# Generate a new random 8x8 symmetric positive definite matrix on every run
n = 8
A_comp = np.random.rand(n, n)
A_comp = A_comp @ A_comp.T
b_comp = np.random.rand(n)
x0_comp = np.zeros(n)

# Solve with Conjugate Gradient
cg_comp_tracker = exe.IterationTracker(x0_comp, A_comp, b_comp)
cg(A_comp, b_comp, x0=x0_comp, callback=cg_comp_tracker, atol=1e-6, rtol=1e-10)

# Solve with GMRES
gmres_comp_tracker = exe.IterationTracker(x0_comp, A_comp, b_comp)
gmres(A_comp, b_comp, x0=x0_comp, callback=gmres_comp_tracker, callback_type='pr_norm', atol=1e-6, rtol=1e-10)

# Plot the residuals
fig = go.Figure()

fig.add_trace(go.Scatter(
    y=cg_comp_tracker.residuals,
    mode='lines+markers',
    name='CG (Minimizes A-norm)',
    line=dict(color='blue')
))

fig.add_trace(go.Scatter(
    y=gmres_comp_tracker.residuals,
    mode='lines+markers',
    name='GMRES (Minimizes Euclidean norm)',
    line=dict(color='red')
))

fig.update_layout(
    title='CG vs GMRES Residual Convergence (8x8 Symmetric System)',
    xaxis_title='Iteration',
    yaxis_title='Euclidean Residual ||Ax - b||',
    yaxis_type='log',
    template='plotly_white',
    height=500
)

fig.show()

Why use CG if GMRES is smoother?¶

Looking at the plot above, GMRES seems vastly superior. Its residual drops perfectly smoothly, while CG violently spikes up and down. Since GMRES strictly minimizes the Euclidean residual at every step, why do we use CG for symmetric positive definite matrices?

The answer is Computational Cost and Memory:

GMRES remembers everything: To ensure each new search direction is perfectly orthogonal, GMRES must store every single previous step direction in RAM, and compute dot products against all of them. For a 1-million dimension system taking 10,000 iterations, GMRES would exhaust all available memory, requiring harsh restart parameters that cripple its convergence.
The CG Miracle: Because the matrix is symmetric, the math in Conjugate Gradient collapses beautifully. CG only needs to remember the last two search directions to guarantee conjugacy against all previous steps! It requires constant memory and constant compute per iteration.

So while GMRES takes a more direct path, it does massive amounts of hidden work at each step to achieve it. CG meanders in the Euclidean norm, but it computes incredibly cheaply with a microscopic memory footprint!

Performance Benchmark¶

To truly appreciate the computational difference between these two algorithms, let’s race them on a much larger system.

We will generate a $2000 \times 2000$ symmetric positive definite matrix and use Jupyter’s %timeit magic command to see how long it takes each solver to converge. Because Conjugate Gradient uses only 2 search directions to compute the next, it should reliably outperform GMRES (which must orthogonalize against a growing history of vectors) on symmetric systems.

try:
    import memory_profiler
except ImportError:
    %pip install -q memory_profiler
%load_ext memory_profiler

# Generate a massive 2000x2000 symmetric positive definite system
n_large = 2000
A_large_rand = np.random.rand(n_large, n_large)
# We add a small diagonal component to ensure it converges in a reasonable timeframe for the benchmark
A_large = A_large_rand @ A_large_rand.T + np.eye(n_large) * 5
b_large = np.random.rand(n_large)

print("Timing Conjugate Gradient (CG)...")
%memit -r 3 cg(A_large, b_large, atol=1e-5, rtol=1e-5)

print("Timing GMRES...")
# We must allow GMRES enough restarts/iterations to solve a system this large
%memit -r 3 gmres(A_large, b_large, atol=1e-5, rtol=1e-5, restart=200)