148 lines
5.5 KiB
Python
148 lines
5.5 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from iterative_solvers import *
|
|
|
|
def mat_vec_prod_p2(x_flat: np.ndarray, n: int, h_sq: float) -> np.ndarray:
|
|
"""
|
|
Computes the matrix-vector product y = Ax for the 2D PDE problem.
|
|
"""
|
|
x_2d = x_flat.reshape((n, n))
|
|
x_padded = np.zeros((n + 2, n + 2))
|
|
x_padded[1:-1, 1:-1] = x_2d
|
|
|
|
y_2d = np.zeros((n, n))
|
|
|
|
# Apply the finite difference stencil: (4+h^2)u_ij - u_{i-1,j} - u_{i+1,j} - u_{i,j-1} - u_{i,j+1}
|
|
for i in range(n):
|
|
for j in range(n):
|
|
y_2d[i, j] = (4.0 + h_sq) * x_padded[i + 1, j + 1] \
|
|
- x_padded[i, j + 1] \
|
|
- x_padded[i + 2, j + 1] \
|
|
- x_padded[i + 1, j] \
|
|
- x_padded[i + 1, j + 2]
|
|
|
|
return y_2d.flatten()
|
|
|
|
def get_diag_inv_p2(i: int, n: int, h_sq: float) -> float:
|
|
"""Returns 1/A_ii for Problem 2. Diagonal is constant: 4 + h^2."""
|
|
return 1.0 / (4.0 + h_sq)
|
|
|
|
def get_off_diag_sum_p2(x_flat: np.ndarray, i_flat: int, n: int, **kwargs) -> float:
|
|
"""Computes the off-diagonal sum for a given row `i_flat` in Problem 2."""
|
|
row = i_flat // n
|
|
col = i_flat % n
|
|
sum_val = 0.0
|
|
|
|
if row > 0: sum_val += -x_flat[i_flat - n]
|
|
if row < n - 1: sum_val += -x_flat[i_flat + n]
|
|
if col > 0: sum_val += -x_flat[i_flat - 1]
|
|
if col < n - 1: sum_val += -x_flat[i_flat + 1]
|
|
|
|
return sum_val
|
|
|
|
def solve_and_report_p2(n: int):
|
|
"""Solves the 2D PDE linear system for a given n and reports results."""
|
|
print(f"\n--- Solving for n = {n} (System size: {n*n} x {n*n}) ---")
|
|
h = 1.0 / (n + 1)
|
|
h_sq = h**2
|
|
N = n * n
|
|
|
|
b = np.full(N, h_sq)
|
|
omega = 2.0 / (1.0 + np.sin(np.pi * h))
|
|
|
|
iteration_counts = []
|
|
all_residuals = {}
|
|
all_solutions = {}
|
|
|
|
methods_to_run = {
|
|
'Jacobi': ('jacobi', {}),
|
|
'Gauss-Seidel': ('gauss_seidel', {}),
|
|
'SOR': ('sor', {'omega': omega}),
|
|
'Steepest Descent': (steepest_descent, {}),
|
|
'Conjugate Gradient': (conjugate_gradient, {})
|
|
}
|
|
|
|
print("\nStarting iterative solvers...")
|
|
for name, (method_func, params) in methods_to_run.items():
|
|
print(f"Running {name}...")
|
|
|
|
solver_kwargs = {'n': n, 'h_sq': h_sq}
|
|
|
|
if isinstance(method_func, str):
|
|
solver_args = {'b': b, 'method': method_func, 'get_diag_inv': get_diag_inv_p2,
|
|
'get_off_diag_sum': get_off_diag_sum_p2, 'mat_vec_prod': mat_vec_prod_p2,
|
|
**solver_kwargs, **params}
|
|
solution, iters, res_hist = stationary_method(**solver_args)
|
|
else:
|
|
solver_args = {'b': b, 'mat_vec_prod': mat_vec_prod_p2, **solver_kwargs}
|
|
solution, iters, res_hist = method_func(**solver_args)
|
|
|
|
iteration_counts.append((name, iters))
|
|
all_residuals[name] = res_hist
|
|
all_solutions[name] = solution if iters < MAX_ITER else None
|
|
|
|
# Report iteration counts
|
|
print("\n--- Iteration Counts ---")
|
|
for name, iters in iteration_counts:
|
|
status = "converged" if iters < MAX_ITER else "did NOT converge"
|
|
print(f"{name:<20}: {iters} iterations ({status})")
|
|
|
|
# Plot residue history
|
|
plt.figure(figsize=(12, 8))
|
|
for name, res_hist in all_residuals.items():
|
|
res_log = np.log10(np.array(res_hist) + 1e-20)
|
|
plt.plot(res_log, label=name, linestyle='-')
|
|
|
|
plt.title(f'Convergence History for PDE Problem, n = {n}')
|
|
plt.xlabel('Iteration Step (m)')
|
|
plt.ylabel('log10(||r^(m)||_2)')
|
|
plt.legend()
|
|
plt.grid(True)
|
|
|
|
current_ylim = plt.ylim()
|
|
plt.ylim(bottom=max(current_ylim[0], -13), top=current_ylim[1])
|
|
|
|
plt.savefig(f'problem2_convergence_n{n}.png')
|
|
print(f"\nConvergence plot saved to 'problem2_convergence_n{n}.png'")
|
|
|
|
if all_solutions:
|
|
print("Generating 3D surface plots of the solutions...")
|
|
|
|
# Create the X, Y grid for the surface plot, including boundaries
|
|
plot_grid = np.linspace(0, 1, n + 2)
|
|
X, Y = np.meshgrid(plot_grid, plot_grid)
|
|
|
|
# Create a figure for the subplots
|
|
fig = plt.figure(figsize=(18, 10))
|
|
fig.suptitle(f'3D Surface Plots of Numerical Solutions for n = {n}', fontsize=16)
|
|
# Loop through the stored solutions and create a subplot for each
|
|
for i, (name, solution_flat) in enumerate(all_solutions.items()):
|
|
# Reshape the 1D solution vector into a 2D grid
|
|
solution_2d = solution_flat.reshape((n, n))
|
|
|
|
# Pad the solution with zero boundaries for plotting
|
|
u_padded = np.zeros((n + 2, n + 2))
|
|
u_padded[1:-1, 1:-1] = solution_2d
|
|
|
|
# Add a new 3D subplot
|
|
ax = fig.add_subplot(2, 3, i + 1, projection='3d')
|
|
|
|
# Plot the surface
|
|
ax.plot_surface(X, Y, u_padded, cmap='viridis', edgecolor='none')
|
|
|
|
# Set titles and labels for each subplot
|
|
ax.set_title(name)
|
|
ax.set_xlabel('x')
|
|
ax.set_ylabel('y')
|
|
ax.set_zlabel('u(x,y)')
|
|
ax.view_init(elev=30, azim=-60) # Set a nice viewing angle
|
|
plt.tight_layout(rect=[0, 0.03, 1, 0.95]) # Adjust layout to make room for suptitle
|
|
plt.savefig(f'problem2_3d_solutions_n{n}.png')
|
|
print(f"3D surface comparison plot saved to 'problem2_3d_solutions_n{n}.png'")
|
|
|
|
if __name__ == "__main__":
|
|
# Solve Problem 2 for n=16 and n=32
|
|
solve_and_report_p2(n=16)
|
|
solve_and_report_p2(n=32)
|
|
|