add project2
This commit is contained in:
147
project2/problem2.py
Normal file
147
project2/problem2.py
Normal file
@@ -0,0 +1,147 @@
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user