Solvers¶
Background¶
This page presents the Python implementation for Solvers.
Code¶
```python """ Solvers
Educational script demonstrating solvers concepts. """
============================================================================¶
heat_equation_2d/solvers.py¶
============================================================================¶
import numpy as np import scipy.sparse.linalg as spsolve from .matrices import ( construct_backward_matrix_2d, construct_crank_nicolson_matrices_2d,
if name == "main": )
def step_forward_euler_2d(u: np.ndarray, D: float, dt: float,
dx2: float, dy2: float) -> np.ndarray:
"""
Perform one Forward Euler time step in 2D.
Args:
u: Current solution (2D array)
D: Thermal diffusivity
dt: Time step size
dx2: dx^2
dy2: dy^2
Returns:
Updated solution array
"""
u_new = u.copy()
# Compute second derivatives using finite differences
u_xx = (u[2:, 1:-1] - 2 * u[1:-1, 1:-1] + u[:-2, 1:-1]) / dx2
u_yy = (u[1:-1, 2:] - 2 * u[1:-1, 1:-1] + u[1:-1, :-2]) / dy2
# Update interior points
u_new[1:-1, 1:-1] += D * dt * (u_xx + u_yy)
return u_new
def solve_forward_euler_2d(u_initial, params):
"""Solve using Forward Euler method."""
# Check stability
if params.rx + params.ry > 0.5:
raise ValueError(f"Unstable: rx+ry = {params.rx+params.ry:.3f} > 0.5")
u = u_initial.copy()
dx2, dy2 = params.dx**2, params.dy**2
for _ in range(params.Nt):
u_new = u.copy()
# Compute finite differences for interior points
u_xx = (u[2:, 1:-1] - 2*u[1:-1, 1:-1] + u[:-2, 1:-1]) / dx2
u_yy = (u[1:-1, 2:] - 2*u[1:-1, 1:-1] + u[1:-1, :-2]) / dy2
u_new[1:-1, 1:-1] += params.D * params.dt * (u_xx + u_yy)
u = u_new
return u
def solve_backward_euler_2d(u_initial, params):
"""Solve using Backward Euler method."""
A = construct_backward_matrix_2d(params)
u = u_initial.copy()
for _ in range(params.Nt):
u_flat = u.reshape(-1)
u_new_flat = spsolve.spsolve(A, u_flat)
u = u_new_flat.reshape(params.Nx, params.Ny)
return u
def solve_crank_nicolson_2d(u_initial, params):
"""Solve using Crank-Nicolson method."""
A, B = construct_crank_nicolson_matrices_2d(params)
u = u_initial.copy()
for _ in range(params.Nt):
u_flat = u.reshape(-1)
rhs = B @ u_flat
u_new_flat = spsolve.spsolve(A, rhs)
u = u_new_flat.reshape(params.Nx, params.Ny)
return u
def solve_crank_nicolson_adi_2d(u_initial, params):
"""
Solve 2D heat equation using Crank-Nicolson ADI method.
Args:
u_initial: Initial condition (2D array)
params: Grid parameters
Returns:
Final solution array (2D)
"""
Nx, Ny = params.Nx, params.Ny
rx, ry = params.rx / 2, params.ry / 2 # Half time steps for ADI
u = u_initial.copy()
for _ in range(params.Nt):
# First half-step: solve in x-direction
u_half = solve_adi_x_step(u, rx, ry)
# Second half-step: solve in y-direction
u = solve_adi_y_step(u_half, rx, ry)
return u
def solve_adi_x_step(u, rx, ry):
"""
Solve ADI x-direction step using tridiagonal solver.
"""
Nx, Ny = u.shape
u_half = np.zeros_like(u)
# Apply boundary conditions (Dirichlet: u = 0 on boundaries)
u_half[0, :] = 0
u_half[-1, :] = 0
u_half[:, 0] = 0
u_half[:, -1] = 0
# Solve for each row (y-constant lines)
for j in range(1, Ny-1):
# Set up tridiagonal system for x-direction
a = -rx * np.ones(Nx-3) # sub-diagonal
b = (1 + 2*rx) * np.ones(Nx-2) # main diagonal
c = -rx * np.ones(Nx-3) # super-diagonal
# Right-hand side with y-direction explicit terms
rhs = u[1:-1, j].copy()
if j > 1:
rhs += ry * (u[1:-1, j-1] - 2*u[1:-1, j] + u[1:-1, j+1])
# Solve tridiagonal system
u_half[1:-1, j] = solve_tridiagonal(a, b, c, rhs)
return u_half
def solve_adi_y_step(u_half, rx, ry):
"""
Solve ADI y-direction step using tridiagonal solver.
"""
Nx, Ny = u_half.shape
u_new = np.zeros_like(u_half)
# Apply boundary conditions
u_new[0, :] = 0
u_new[-1, :] = 0
u_new[:, 0] = 0
u_new[:, -1] = 0
# Solve for each column (x-constant lines)
for i in range(1, Nx-1):
# Set up tridiagonal system for y-direction
a = -ry * np.ones(Ny-3) # sub-diagonal
b = (1 + 2*ry) * np.ones(Ny-2) # main diagonal
c = -ry * np.ones(Ny-3) # super-diagonal
# Right-hand side with x-direction explicit terms
rhs = u_half[i, 1:-1].copy()
if i > 1:
rhs += rx * (u_half[i-1, 1:-1] - 2*u_half[i, 1:-1] + u_half[i+1, 1:-1])
# Solve tridiagonal system
u_new[i, 1:-1] = solve_tridiagonal(a, b, c, rhs)
return u_new
def solve_tridiagonal(a, b, c, d):
"""
Solve tridiagonal system using Thomas algorithm.
Args:
a: sub-diagonal (length n-1)
b: main diagonal (length n)
c: super-diagonal (length n-1)
d: right-hand side (length n)
Returns:
Solution vector
"""
n = len(d)
# Forward elimination
for i in range(1, n):
w = a[i-1] / b[i-1]
b[i] -= w * c[i-1]
d[i] -= w * d[i-1]
# Back substitution
x = np.zeros(n)
x[n-1] = d[n-1] / b[n-1]
for i in range(n-2, -1, -1):
x[i] = (d[i] - c[i] * x[i+1]) / b[i]
return x
def compare_2d_methods(u_initial, params):
"""Compare different 2D methods."""
results = {}
# Forward Euler
try:
results["Forward Euler"] = solve_forward_euler_2d(u_initial, params)
print("✓ Forward Euler completed")
except ValueError as e:
print(f"✗ Forward Euler failed: {e}")
results["Forward Euler"] = None
except Exception as e:
print(f"✗ Forward Euler failed: {e}")
results["Forward Euler"] = None
# Backward Euler
try:
results["Backward Euler"] = solve_backward_euler_2d(u_initial, params)
print("✓ Backward Euler completed")
except Exception as e:
print(f"✗ Backward Euler failed: {e}")
results["Backward Euler"] = None
# Crank-Nicolson
try:
results["Crank-Nicolson"] = solve_crank_nicolson_2d(u_initial, params)
print("✓ Crank-Nicolson completed")
except Exception as e:
print(f"✗ Crank-Nicolson failed: {e}")
results["Crank-Nicolson"] = None
# ADI (if you want to include it)
try:
results["ADI"] = solve_crank_nicolson_adi_2d(u_initial, params)
print("✓ ADI completed")
except Exception as e:
print(f"✗ ADI failed: {e}")
results["ADI"] = None
return results
```
Exercises¶
Exercise 1. Explain the Thomas algorithm for solving tridiagonal systems. What is its computational complexity compared to general Gaussian elimination?
Solution to Exercise 1
The Thomas algorithm solves \(A\mathbf{x} = \mathbf{d}\) where \(A\) is tridiagonal with sub-diagonal \(a\), main diagonal \(b\), and super-diagonal \(c\).
Forward elimination: For \(i = 2, \ldots, n\): \(w = a_{i-1}/b_{i-1}\), \(b_i \leftarrow b_i - w c_{i-1}\), \(d_i \leftarrow d_i - w d_{i-1}\).
Back substitution: \(x_n = d_n/b_n\); for \(i = n-1, \ldots, 1\): \(x_i = (d_i - c_i x_{i+1})/b_i\).
Total cost: \(O(n)\) operations (specifically \(5n - 4\) multiplications/divisions and \(3n - 3\) additions). General Gaussian elimination costs \(O(n^3)\), so the Thomas algorithm is dramatically faster for tridiagonal systems.
Exercise 2. Compare the Forward Euler, Backward Euler, and Crank-Nicolson methods for the 2D heat equation in terms of stability and per-step cost.
Solution to Exercise 2
| Method | Stability | Per-step cost |
|---|---|---|
| Forward Euler | Conditional: \(r_x + r_y \le 1/2\) | \(O(N_x N_y)\) explicit update |
| Backward Euler | Unconditional | \(O(N_x N_y)\) sparse solve |
| Crank-Nicolson | Unconditional | \(O(N_x N_y)\) sparse multiply + solve |
The sparse solve for 2D problems costs more than the explicit update, but Backward Euler and Crank-Nicolson allow much larger time steps. Crank-Nicolson has the best accuracy (\(O(\Delta t^2)\)) among the three.
Exercise 3. The ADI method splits the 2D problem into a sequence of 1D problems. Explain why this is advantageous over directly solving the full 2D implicit system.
Solution to Exercise 3
Direct 2D implicit methods require solving a linear system of size \(N_x N_y\). Even with sparse solvers, the cost is \(O((N_x N_y)^{3/2})\) for direct methods or requires iterative solvers.
ADI reduces this to solving \(N_y\) tridiagonal systems of size \(N_x\) (x-sweep) plus \(N_x\) tridiagonal systems of size \(N_y\) (y-sweep). Each tridiagonal solve costs \(O(N)\), so the total is \(O(N_x N_y)\) -- the same as the explicit method but with unconditional stability. ADI also has better cache locality since it operates on rows/columns sequentially.
Exercise 4.
The compare_2d_methods function tries each method and catches exceptions. Why might Forward Euler fail while the other methods succeed?
Solution to Exercise 4
Forward Euler fails when the stability condition \(r_x + r_y \le 1/2\) is violated. This happens when the time step \(\Delta t\) is too large relative to the spatial discretization. The other methods (Backward Euler, Crank-Nicolson, ADI) are unconditionally stable, meaning they produce bounded solutions for any \(\Delta t\).
The compare_2d_methods function explicitly checks the stability condition before attempting Forward Euler and raises a ValueError if violated. The try-except block catches this and reports the failure while allowing the comparison to continue with the stable methods.