diff --git a/.gitignore b/.gitignore index 31b5570..43104ef 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,5 @@ pdesolvers.egg-info/ # Other files *.log /.coverage +*.gif +*.csv diff --git a/null b/null new file mode 100644 index 0000000..78b522a Binary files /dev/null and b/null differ diff --git a/pdesolvers/main.py b/pdesolvers/main.py index 7b2aafe..97df7c9 100644 --- a/pdesolvers/main.py +++ b/pdesolvers/main.py @@ -15,35 +15,58 @@ def main(): # solver1 = pde.Heat1DCNSolver(equation1) # solver2 = pde.Heat1DExplicitSolver(equation1) - # testing for monte carlo pricing + # testing 2d heat equation + + xLength = 10 # Lx + yLength = 10 # Ly + maxTime = 0.5 # tmax + diffusivityConstant = 4 # kappa + numPointsSpace = 50 # x_points = y_points + numPointsTime = 2000 # t_points + + equation = (pde.HeatEquation2D(maxTime, numPointsTime, diffusivityConstant, xLength, numPointsSpace)) + equation.set_initial_temp(lambda x, y: 10 * np.exp(-((x - xLength/2)**2 + (y - yLength/2)**2) / 2)) + equation.set_right_boundary_temp(lambda t, y: 20 + 100 * y * (yLength - y)**3 * (t - 1)**2 * (t > 1)) + equation.set_left_boundary_temp(lambda t, y: 20 + 10 * y * (yLength - y) * t**2) + equation.set_top_boundary_temp(lambda t, x: 20 + 5 * x * (xLength - x) * t**4) + equation.set_bottom_boundary_temp(lambda t, x: 20) - ticker = 'AAPL' + solver1 = pde.Heat2DExplicitSolver(equation) + solver1 = pde.Heat2DCNSolver(equation) + solution1 = solver1.solve() + solution1.animate(filename="Explicit") + solver2 = pde.Heat2DCNSolver(equation) + solution2 = solver2.solve() + solution2.animate(filename="Crank-Nicolson") + + # testing for monte carlo pricing + # ticker = 'AAPL' - # STOCK - historical_data = pde.HistoricalStockData(ticker) - historical_data.fetch_stock_data( "2024-03-21","2025-03-21") - sigma, r = historical_data.estimate_metrics() - current_price = historical_data.get_latest_stock_price() + # # STOCK + # historical_data = pde.HistoricalStockData(ticker) + # historical_data.fetch_stock_data( "2024-03-21","2025-03-21") + # sigma, r = historical_data.estimate_metrics() + # current_price = historical_data.get_latest_stock_price() - equation2 = pde.BlackScholesEquation(pde.OptionType.EUROPEAN_CALL, current_price, 100, r, sigma, 1, 100, 20000) + # equation2 = pde.BlackScholesEquation(pde.OptionType.EUROPEAN_CALL, current_price, 100, r, sigma, 1, 100, 20000) - solver1 = pde.BlackScholesCNSolver(equation2) - solver2 = pde.BlackScholesExplicitSolver(equation2) - sol1 = solver1.solve() - sol1.plot() + # solver1 = pde.BlackScholesCNSolver(equation2) + # solver2 = pde.BlackScholesExplicitSolver(equation2) + # sol1 = solver1.solve() + # sol1.plot() - # COMPARISON - # look to see the corresponding option price for the expiration date and strike price - pricing_1 = pde.BlackScholesFormula(pde.OptionType.EUROPEAN_CALL, current_price, 100, r, sigma, 1) - pricing_2 = pde.MonteCarloPricing(pde.OptionType.EUROPEAN_CALL, current_price, 100, r, sigma, 1, 365, 1000) + # # COMPARISON + # # look to see the corresponding option price for the expiration date and strike price + # pricing_1 = pde.BlackScholesFormula(pde.OptionType.EUROPEAN_CALL, current_price, 100, r, sigma, 1) + # pricing_2 = pde.MonteCarloPricing(pde.OptionType.EUROPEAN_CALL, current_price, 100, r, sigma, 1, 365, 1000) - bs_price = pricing_1.get_black_scholes_merton_price() - monte_carlo_price = pricing_2.get_monte_carlo_option_price() + # bs_price = pricing_1.get_black_scholes_merton_price() + # monte_carlo_price = pricing_2.get_monte_carlo_option_price() - pde_price = sol1.get_result()[-1, 0] - print(f"PDE Price: {pde_price}") - print(f"Black-Scholes Price: {bs_price}") - print(f"Monte-Carlo Price: {monte_carlo_price}") + # pde_price = sol1.get_result()[-1, 0] + # print(f"PDE Price: {pde_price}") + # print(f"Black-Scholes Price: {bs_price}") + # print(f"Monte-Carlo Price: {monte_carlo_price}") if __name__ == "__main__": main() \ No newline at end of file diff --git a/pdesolvers/pdes/__init__.py b/pdesolvers/pdes/__init__.py index 07d6670..2ccc5c3 100644 --- a/pdesolvers/pdes/__init__.py +++ b/pdesolvers/pdes/__init__.py @@ -1,2 +1,3 @@ from .heat_1d import HeatEquation -from .black_scholes import BlackScholesEquation \ No newline at end of file +from .black_scholes import BlackScholesEquation +from .heat_2d import HeatEquation2D \ No newline at end of file diff --git a/pdesolvers/pdes/heat_2d.py b/pdesolvers/pdes/heat_2d.py new file mode 100644 index 0000000..6cba046 --- /dev/null +++ b/pdesolvers/pdes/heat_2d.py @@ -0,0 +1,91 @@ +import numpy as np + +class HeatEquation2D: + def __init__(self, time, t_nodes, k, xlength, x_nodes, ylength=None, y_nodes=None): + + assert time > 0, "Time must be positive" + assert t_nodes > 1, "Number of time nodes must be greater than 1" + assert k > 0, "Diffusivity constant k must be positive" + assert xlength > 0, "X-length must be positive" + assert x_nodes > 2, "Number of x nodes must be greater than 2" + if ylength is not None: + assert ylength > 0, "Y-length must be positive" + if y_nodes is not None: + assert y_nodes > 2, "Number of y nodes must be greater than 2" + + self.__time = time + self.__t_nodes = t_nodes + self.__k = k + self.__xlength = xlength + self.__x_nodes = x_nodes + self.__ylength = ylength if ylength is not None else xlength + self.__y_nodes = y_nodes if y_nodes is not None else x_nodes + # Initialize boundary conditions to None + self.__initial_temp = None + self.__left_boundary = None + self.__right_boundary = None + self.__top_boundary = None + self.__bottom_boundary = None + + def set_initial_temp(self, u0): + self.__initial_temp = u0 + + def set_left_boundary_temp(self, left): + self.__left_boundary = left + + def set_right_boundary_temp(self, right): + self.__right_boundary = right + + def set_top_boundary_temp(self, top): + self.__top_boundary = top + + def set_bottom_boundary_temp(self, bottom): + self.__bottom_boundary = bottom + + @property + def xlength(self): + return self.__xlength + + @property + def ylength(self): + return self.__ylength + + @property + def x_nodes(self): + return self.__x_nodes + + @property + def y_nodes(self): + return self.__y_nodes + + @property + def time(self): + return self.__time + + @property + def t_nodes(self): + return self.__t_nodes + + @property + def k(self): + return self.__k + + @property + def initial_temp(self): + return self.__initial_temp + + @property + def left_boundary(self): + return self.__left_boundary + + @property + def right_boundary(self): + return self.__right_boundary + + @property + def top_boundary(self): + return self.__top_boundary + + @property + def bottom_boundary(self): + return self.__bottom_boundary \ No newline at end of file diff --git a/pdesolvers/solution/__init__.py b/pdesolvers/solution/__init__.py index febf7aa..b5ca3d0 100644 --- a/pdesolvers/solution/__init__.py +++ b/pdesolvers/solution/__init__.py @@ -1 +1 @@ -from .solution import Solution1D, SolutionBlackScholes \ No newline at end of file +from .solution import Solution1D, SolutionBlackScholes, Heat2DSolution \ No newline at end of file diff --git a/pdesolvers/solution/solution.py b/pdesolvers/solution/solution.py index a518e22..cbce941 100644 --- a/pdesolvers/solution/solution.py +++ b/pdesolvers/solution/solution.py @@ -1,3 +1,4 @@ +from matplotlib.animation import FuncAnimation import numpy as np import pdesolvers.utils.utility as utility import pdesolvers.enums.enums as enum @@ -165,4 +166,64 @@ def _set_plot_labels(self, ax): ax.set_xlabel('Time') ax.set_ylabel('Asset Price') ax.set_zlabel(f'{self.option_type.value} Option Value') - ax.set_title(f'{self.option_type.value} Option Value Surface Plot') \ No newline at end of file + ax.set_title(f'{self.option_type.value} Option Value Surface Plot') + +class Heat2DSolution: + def __init__(self, result, x_grid, y_grid, t_grid, dx, dy, dt, duration): + self.result = result + self.x_grid = x_grid + self.y_grid = y_grid + self.t_grid = t_grid + self.dx = dx + self.dy = dy + self.dt = dt + self.duration = duration + + @staticmethod + def __plot_surface(u_k, k, ax, xDomain, yDomain, dt): + ax.clear() + X, Y = np.meshgrid(xDomain, yDomain) + # Transpose u_k to match meshgrid orientation + surf = ax.plot_surface(X, Y, u_k.T, + cmap='hot', + alpha=0.9) + ax.set_xlabel('X Position') + ax.set_ylabel('Y Position') + ax.set_zlabel('Temperature') + ax.set_title(f'2D Heat Equation: t = {k*dt:.4f}') + ax.view_init(elev=30, azim=45) + ax.set_zlim(0, 100) + + return surf + + def animate(self, export=False, filename="heat_equation_2d_plot.gif"): + print("Creating animation...") + self + fig = plt.figure(figsize=(12, 8)) + ax = fig.add_subplot(111, projection='3d') + + def animateFrame(k): + return Heat2DSolution.__plot_surface(self.result[k], k, ax, self.x_grid, self.y_grid, self.dt) + + anim = FuncAnimation(fig, animateFrame, interval=100, frames=len(self.t_grid), repeat=True) + + if export: + anim.save(filename+'.gif', writer='pillow', fps=5) + + plt.show() + + def get_result(self): + """ + Gets the grid of computed temperature values + + :return: grid result + """ + return self.result + + def get_execution_time(self): + """ + Gets the time taken for the solver to solve the equation + :return: duration + """ + return self.duration + \ No newline at end of file diff --git a/pdesolvers/solvers/__init__.py b/pdesolvers/solvers/__init__.py index 3f2ed36..2dafc44 100644 --- a/pdesolvers/solvers/__init__.py +++ b/pdesolvers/solvers/__init__.py @@ -1,2 +1,3 @@ from .heat_solvers import Heat1DExplicitSolver, Heat1DCNSolver +from .heat2d_solvers import Heat2DExplicitSolver, Heat2DCNSolver from .black_scholes_solvers import BlackScholesExplicitSolver, BlackScholesCNSolver diff --git a/pdesolvers/solvers/heat2d_solvers.py b/pdesolvers/solvers/heat2d_solvers.py new file mode 100644 index 0000000..6ee67f3 --- /dev/null +++ b/pdesolvers/solvers/heat2d_solvers.py @@ -0,0 +1,144 @@ +import time +import logging + +import numpy as np +import pdesolvers.solution as sol +import pdesolvers.pdes.heat_2d as heat +import pdesolvers.utils.utility as utility + +from scipy.sparse.linalg import spsolve +from pdesolvers.solvers.solver import Solver + +logging.basicConfig( + level = logging.INFO, + format = "{asctime} - {levelname} - {message}", + style="{", + datefmt="%Y-%m-%d %H:%M", +) + +class Heat2DExplicitSolver (Solver): + def __init__(self, equation: heat.HeatEquation2D): + self.equation = equation + + def solve(self): + logging.info(f"Starting {self.__class__.__name__} with {self.equation.x_nodes+1} spatial nodes and {self.equation.t_nodes+1} time nodes.") + start = time.perf_counter() + + if (self.equation.left_boundary is None or + self.equation.right_boundary is None or + self.equation.top_boundary is None or + self.equation.bottom_boundary is None): + raise ValueError("All boundary conditions must be set before solving") + if self.equation.initial_temp is None: + raise ValueError("Initial condition must be set before solving") + + x = np.linspace(0, self.equation.xlength, self.equation.x_nodes) + y = np.linspace(0, self.equation.ylength, self.equation.y_nodes) + t = np.linspace(0, self.equation.time, self.equation.t_nodes) + + dx = x[1]-x[0] + dy = y[1]-y[0] + dt = t[1]-t[0] + assert dt < (dx*dy)/(4*self.equation.k), "Time-step size too large!! violates CFL condtion for Forward Euler method." + lambdaConstant = (self.equation.k * dt) + + print("Initializing matrix...") + U = utility.Heat2DHelper.initMatrix(self.equation.t_nodes, + self.equation.x_nodes, + self.equation.y_nodes, + self.equation.left_boundary, + self.equation.right_boundary, + self.equation.bottom_boundary, + self.equation.top_boundary, + self.equation.initial_temp, + x, y, t) + + print(f"Calculating temperature evolution with {self.equation.t_nodes-1} iterations...", flush=True) + for tau in range(self.equation.t_nodes-1): + for i in range(1, self.equation.x_nodes-1): + for j in range(1, self.equation.y_nodes-1): + # 5-point stencil for 2D heat equation + U[tau+1, i, j] = U[tau, i, j] + lambdaConstant * ( + # Second derivative in x-direction + (1/dx**2)*(U[tau, i-1, j] - 2*U[tau, i, j] + U[tau, i+1, j]) + + # Second derivative in y-direction + (1/dy**2)*(U[tau, i, j-1] - 2*U[tau, i, j] + U[tau, i, j+1]) + ) + + end = time.perf_counter() + duration = end - start + logging.info(f"Solver completed in {duration} seconds.") + return sol.Heat2DSolution(U, x, y, t, dx, dy, dt, duration) + +class Heat2DCNSolver (Solver): + def __init__(self, equation: heat.HeatEquation2D): + self.equation = equation + + def solve(self): + logging.info(f"Starting {self.__class__.__name__} with {self.equation.x_nodes+1} spatial nodes and {self.equation.t_nodes+1} time nodes.") + start = time.perf_counter() + + if (self.equation.left_boundary is None or + self.equation.right_boundary is None or + self.equation.top_boundary is None or + self.equation.bottom_boundary is None): + raise ValueError("All boundary conditions must be set before solving") + if self.equation.initial_temp is None: + raise ValueError("Initial condition must be set before solving") + + x = np.linspace(0, self.equation.xlength, self.equation.x_nodes) + y = np.linspace(0, self.equation.ylength, self.equation.y_nodes) + t = np.linspace(0, self.equation.time, self.equation.t_nodes) + + dx = x[1]-x[0] + dy = y[1]-y[0] + dt = t[1]-t[0] + c = self.equation.k * dt / 2 + cx = c / (dx**2) + cy = c / (dy**2) + alpha = 1 + 2*cx + 2*cy + beta = 1 - 2*cx - 2*cy + + print("Initializing matrix...") + U = utility.Heat2DHelper.initMatrix(self.equation.t_nodes, + self.equation.x_nodes, + self.equation.y_nodes, + self.equation.left_boundary, + self.equation.right_boundary, + self.equation.bottom_boundary, + self.equation.top_boundary, + self.equation.initial_temp, + x, y, t) + + # create sparse matrix + G, n_interior_x, n_interior_y = utility.Heat2DHelper.innitTriDiagMatrix(self.equation.x_nodes, self.equation.y_nodes, cx, cy, alpha) + + # time-stepping loop + print(f"Calculating temperature evolution with {self.equation.t_nodes-1} iterations...", flush=True) + for tau in range(self.equation.t_nodes - 1): + rhs = np.zeros(n_interior_x * n_interior_y) + idx = 0 + for j in range(1, self.equation.y_nodes-1): + for i in range(1, self.equation.x_nodes-1): + # RHS = β*U_τ + cx*(neighbors_x) + cy*(neighbors_y) + boundary_terms + rhs[idx] = beta * U[tau, i, j] + rhs[idx] += cx * (U[tau, i-1, j] + U[tau, i+1, j]) + rhs[idx] += cy * (U[tau, i, j-1] + U[tau, i, j+1]) + # Boundary contributions + if i == 1: + rhs[idx] += cx * U[tau+1, 0, j] + if i == self.equation.x_nodes-2: + rhs[idx] += cx * U[tau+1, -1, j] + if j == 1: + rhs[idx] += cy * U[tau+1, i, 0] + if j == self.equation.y_nodes-2: + rhs[idx] += cy * U[tau+1, i, -1] + idx += 1 + # Solve G*u_{τ+1} = rhs + u_next_interior = spsolve(G, rhs) + U[tau+1, 1:-1, 1:-1] = u_next_interior.reshape((n_interior_x, n_interior_y)) + + end = time.perf_counter() + duration = end - start + logging.info(f"Solver completed in {duration} seconds.") + return sol.Heat2DSolution(U, x, y, t, dx, dy, dt, duration) \ No newline at end of file diff --git a/pdesolvers/tests/test_heat2d_solvers.py b/pdesolvers/tests/test_heat2d_solvers.py new file mode 100644 index 0000000..107d673 --- /dev/null +++ b/pdesolvers/tests/test_heat2d_solvers.py @@ -0,0 +1,457 @@ +import pytest +import numpy as np + +import pdesolvers.pdes.heat_2d as heat2d +import pdesolvers.solvers.heat2d_solvers as solver2d +import pdesolvers.utils.utility as utility + +class TestHeat2DSolvers: + + def setup_method(self): + """Setup basic 2D heat equation for each test""" + self.equation = heat2d.HeatEquation2D( + time=0.5, t_nodes=800, + k=1.0, + xlength=1.0, x_nodes=20, + ylength=1.0, y_nodes=20 + ) + + # Set up boundary and initial conditions + self.equation.set_initial_temp(lambda x, y: np.sin(np.pi * x) * np.sin(np.pi * y)) + self.equation.set_left_boundary_temp(lambda t, y: 0) + self.equation.set_right_boundary_temp(lambda t, y: 0) + self.equation.set_bottom_boundary_temp(lambda t, x: 0) + self.equation.set_top_boundary_temp(lambda t, x: 0) + + def test_check_boundary_conditions_consistency_at_time_zero(self): + """Test that boundary conditions match initial conditions at corners""" + # Test corner consistency + tolerance = 1e-12 + + # Bottom-left corner (0,0) + assert abs(self.equation.left_boundary(0, 0) - self.equation.initial_temp(0, 0)) < tolerance + assert abs(self.equation.bottom_boundary(0, 0) - self.equation.initial_temp(0, 0)) < tolerance + + # Bottom-right corner (1,0) + assert abs(self.equation.right_boundary(0, 0) - self.equation.initial_temp(1.0, 0)) < tolerance + assert abs(self.equation.bottom_boundary(0, 1.0) - self.equation.initial_temp(1.0, 0)) < tolerance + + # Top-left corner (0,1) + assert abs(self.equation.left_boundary(0, 1.0) - self.equation.initial_temp(0, 1.0)) < tolerance + assert abs(self.equation.top_boundary(0, 0) - self.equation.initial_temp(0, 1.0)) < tolerance + + # Top-right corner (1,1) + assert abs(self.equation.right_boundary(0, 1.0) - self.equation.initial_temp(1.0, 1.0)) < tolerance + assert abs(self.equation.top_boundary(0, 1.0) - self.equation.initial_temp(1.0, 1.0)) < tolerance + + def test_solution_dimensions_explicit(self): + """Test that explicit solver solution has correct dimensions""" + result = solver2d.Heat2DExplicitSolver(self.equation).solve() + solution = result.get_result() + + expected_shape = (self.equation.t_nodes, self.equation.x_nodes, self.equation.y_nodes) + assert solution.shape == expected_shape, f"Expected shape {expected_shape}, got {solution.shape}" + + def test_solution_dimensions_crank_nicolson(self): + """Test that Crank-Nicolson solver solution has correct dimensions""" + result = solver2d.Heat2DCNSolver(self.equation).solve() + solution = result.get_result() + + expected_shape = (self.equation.t_nodes, self.equation.x_nodes, self.equation.y_nodes) + assert solution.shape == expected_shape, f"Expected shape {expected_shape}, got {solution.shape}" + + def test_convergence_between_explicit_and_crank_nicolson(self): + """Test that explicit and Crank-Nicolson methods converge to similar solutions""" + result_explicit = solver2d.Heat2DExplicitSolver(self.equation).solve().get_result() + result_cn = solver2d.Heat2DCNSolver(self.equation).solve().get_result() + + # Compare solutions at final time + diff = np.abs(result_explicit[-1] - result_cn[-1]) + max_diff = np.max(diff) + + assert max_diff < 0.1, f"Methods differ by {max_diff}, expected < 0.1" + + def test_maximum_principle_explicit(self): + """Test that maximum principle holds for explicit solver""" + result = solver2d.Heat2DExplicitSolver(self.equation).solve().get_result() + + initial_max = np.max(result[0]) + solution_max = np.max(result) + + tolerance = 1e-10 + assert solution_max <= initial_max + tolerance, f"Maximum principle violated: max = {solution_max}, initial_max = {initial_max}" + + def test_maximum_principle_crank_nicolson(self): + """Test that maximum principle holds for Crank-Nicolson solver""" + result = solver2d.Heat2DCNSolver(self.equation).solve().get_result() + + initial_max = np.max(result[0]) + solution_max = np.max(result) + + tolerance = 1e-10 + assert solution_max <= initial_max + tolerance, f"Maximum principle violated: max = {solution_max}, initial_max = {initial_max}" + + def test_steady_state_convergence_with_constant_boundaries(self): + """Test convergence to steady state with constant boundary conditions""" + steady_eq = heat2d.HeatEquation2D( + time=5.0, t_nodes=1000, k=1.0, + xlength=1.0, x_nodes=20, + ylength=1.0, y_nodes=20 + ) + + # Constant boundary conditions + steady_eq.set_initial_temp(lambda x, y: 0) + steady_eq.set_left_boundary_temp(lambda t, y: 1.0) + steady_eq.set_right_boundary_temp(lambda t, y: 0.0) + steady_eq.set_bottom_boundary_temp(lambda t, x: 0.0) + steady_eq.set_top_boundary_temp(lambda t, x: 0.0) + + result = solver2d.Heat2DCNSolver(steady_eq).solve().get_result() + + # Check that solution has reached steady state (small time derivative) + final_change = np.abs(result[-1] - result[-10]) + max_change = np.max(final_change) + + assert max_change < 1e-3, f"Solution should reach steady state: max_change = {max_change}" + + def test_symmetry_preservation(self): + """Test that symmetric initial conditions preserve symmetry""" + sym_eq = heat2d.HeatEquation2D( + time=0.1, t_nodes=50, k=1.0, + xlength=2.0, x_nodes=20, + ylength=2.0, y_nodes=20 + ) + + # Symmetric initial condition (Gaussian centered at (1,1)) + sym_eq.set_initial_temp(lambda x, y: np.exp(-((x-1)**2 + (y-1)**2))) + sym_eq.set_left_boundary_temp(lambda t, y: 0) + sym_eq.set_right_boundary_temp(lambda t, y: 0) + sym_eq.set_bottom_boundary_temp(lambda t, x: 0) + sym_eq.set_top_boundary_temp(lambda t, x: 0) + + result = solver2d.Heat2DCNSolver(sym_eq).solve().get_result() + final_solution = result[-1] + + # Test x-symmetry (solution should be symmetric about x=1) + left_half = final_solution[:10, :] + right_half = np.flip(final_solution[10:, :], axis=0) + + symmetry_diff = np.max(np.abs(left_half - right_half)) + assert symmetry_diff < 1e-2, f"X-symmetry not preserved: diff = {symmetry_diff}" + + def test_grid_refinement_convergence(self): + """Test that finer grids give more accurate solutions""" + # Coarse grid + coarse_eq = heat2d.HeatEquation2D( + time=0.1, t_nodes=50, k=1.0, + xlength=1.0, x_nodes=10, + ylength=1.0, y_nodes=10 + ) + coarse_eq.set_initial_temp(lambda x, y: np.sin(np.pi * x) * np.sin(np.pi * y)) + coarse_eq.set_left_boundary_temp(lambda t, y: 0) + coarse_eq.set_right_boundary_temp(lambda t, y: 0) + coarse_eq.set_bottom_boundary_temp(lambda t, x: 0) + coarse_eq.set_top_boundary_temp(lambda t, x: 0) + + # Fine grid + fine_eq = heat2d.HeatEquation2D( + time=0.1, t_nodes=100, k=1.0, + xlength=1.0, x_nodes=20, + ylength=1.0, y_nodes=20 + ) + fine_eq.set_initial_temp(lambda x, y: np.sin(np.pi * x) * np.sin(np.pi * y)) + fine_eq.set_left_boundary_temp(lambda t, y: 0) + fine_eq.set_right_boundary_temp(lambda t, y: 0) + fine_eq.set_bottom_boundary_temp(lambda t, x: 0) + fine_eq.set_top_boundary_temp(lambda t, x: 0) + + coarse_result = solver2d.Heat2DCNSolver(coarse_eq).solve().get_result() + fine_result = solver2d.Heat2DCNSolver(fine_eq).solve().get_result() + + # Sample at center points for comparison + coarse_center = coarse_result[-1, 5, 5] + fine_center = fine_result[-1, 10, 10] + + # Fine grid should be closer to analytical solution + analytical_center = np.exp(-2 * np.pi**2 * 0.1) * np.sin(np.pi * 0.5) * np.sin(np.pi * 0.5) + + coarse_error = abs(coarse_center - analytical_center) + fine_error = abs(fine_center - analytical_center) + + assert fine_error < coarse_error, f"Fine grid should be more accurate: coarse_error={coarse_error}, fine_error={fine_error}" + + @pytest.mark.parametrize("solver_class", [solver2d.Heat2DExplicitSolver, solver2d.Heat2DCNSolver]) + def test_interpolation_convergence_between_methods(self, solver_class): + """Test spatial interpolation at final time""" + result = solver_class(self.equation).solve() + u = result.get_result() # Shape: (t_nodes, x_nodes, y_nodes) + + # Calculate correct spatial grid spacing + hx = self.equation.xlength / (self.equation.x_nodes - 1) + hy = self.equation.ylength / (self.equation.y_nodes - 1) + + try: + # Interpolate in space at final time + final_field = u[-1] # 2D spatial field at final time + interpolator = utility.RBFInterpolator(final_field, hx, hy) + interp_value = interpolator.interpolate(0.5, 0.5) + + assert np.isfinite(interp_value), "Interpolated value should be finite" + + # For sin(πx)sin(πy) initial condition with zero boundaries, + # value should be reasonable (can be negative) + assert abs(interp_value) < 10, f"Interpolated value seems unreasonable: {interp_value}" + + except Exception as e: + pytest.skip(f"RBF spatial interpolation failed: {e}") + def test_energy_dissipation_homogeneous_boundaries(self): + """Test that energy decreases monotonically with homogeneous boundaries""" + result = solver2d.Heat2DCNSolver(self.equation).solve().get_result() + + # Calculate total energy (L2 norm) at each time step + energies = [np.sum(result[t]**2) for t in range(result.shape[0])] + + # Energy should decrease monotonically (or stay constant) + tolerance = 1e-10 + for i in range(1, len(energies)): + assert energies[i] <= energies[i-1] + tolerance, f"Energy increased at step {i}: {energies[i]} > {energies[i-1]}" + + def test_explicit_solver_stability_violation_detection(self): + """Test that explicit solver detects CFL condition violations""" + unstable_eq = heat2d.HeatEquation2D( + time=1.0, t_nodes=10, # Large time step + k=1.0, xlength=1.0, x_nodes=50, # Fine spatial grid + ylength=1.0, y_nodes=50 + ) + unstable_eq.set_initial_temp(lambda x, y: 1.0) + unstable_eq.set_left_boundary_temp(lambda t, y: 0) + unstable_eq.set_right_boundary_temp(lambda t, y: 0) + unstable_eq.set_bottom_boundary_temp(lambda t, x: 0) + unstable_eq.set_top_boundary_temp(lambda t, x: 0) + + with pytest.raises(AssertionError, match="CFL"): + solver2d.Heat2DExplicitSolver(unstable_eq).solve() + + def test_analytical_solution_comparison(self): + """Compare with known analytical solution for simple case""" + # Use equation with known analytical solution + # u(x,y,t) = exp(-2π²kt) * sin(πx) * sin(πy) + + result = solver2d.Heat2DCNSolver(self.equation).solve() + numerical_solution = result.get_result() + + # Create analytical solution + x = np.linspace(0, 1, self.equation.x_nodes) + y = np.linspace(0, 1, self.equation.y_nodes) + t_final = 0.5 + + X, Y = np.meshgrid(x, y) + analytical = np.exp(-2 * np.pi**2 * t_final) * np.sin(np.pi * X) * np.sin(np.pi * Y) + + # Compare at final time + numerical_final = numerical_solution[-1] + + error = np.max(np.abs(numerical_final.T - analytical)) + assert error < 0.1, f"Error vs analytical solution: {error}" + + @pytest.mark.parametrize("solver_class", [solver2d.Heat2DExplicitSolver, solver2d.Heat2DCNSolver]) + def test_execution_time_tracking(self, solver_class): + """Test that solvers track execution time""" + result = solver_class(self.equation).solve() + + execution_time = result.get_execution_time() + assert execution_time > 0, f"Execution time should be positive: {execution_time}" + assert isinstance(execution_time, float), f"Execution time should be float: {type(execution_time)}" + + def test_non_negative_temperatures_with_non_negative_conditions(self): + """Test that non-negative initial and boundary conditions maintain non-negative temperatures""" + non_neg_eq = heat2d.HeatEquation2D( + time=0.1, t_nodes=50, k=1.0, + xlength=1.0, x_nodes=15, + ylength=1.0, y_nodes=15 + ) + + # Non-negative conditions + non_neg_eq.set_initial_temp(lambda x, y: 10 * np.exp(-((x-0.5)**2 + (y-0.5)**2) / 0.1)) + non_neg_eq.set_left_boundary_temp(lambda t, y: 5.0) + non_neg_eq.set_right_boundary_temp(lambda t, y: 3.0) + non_neg_eq.set_bottom_boundary_temp(lambda t, x: 2.0) + non_neg_eq.set_top_boundary_temp(lambda t, x: 4.0) + + result = solver2d.Heat2DCNSolver(non_neg_eq).solve().get_result() + + # All temperatures should remain non-negative + min_temp = np.min(result) + assert min_temp >= -1e-10, f"Negative temperature found: {min_temp}" + + def test_boundary_conditions_preserved_throughout_simulation(self): + """Test that boundary conditions are maintained at all time steps""" + result = solver2d.Heat2DCNSolver(self.equation).solve() + solution = result.get_result() + + # Check boundary conditions at a few time steps + times_to_check = [0, solution.shape[0]//2, -1] + tolerance = 1e-12 + + for t_idx in times_to_check: + t_val = t_idx * (0.5 / (solution.shape[0] - 1)) if t_idx >= 0 else 0.5 + + # Check left boundary (x=0) + for j in range(solution.shape[2]): + y_val = j / (solution.shape[2] - 1) + expected = self.equation.left_boundary(t_val, y_val) + actual = solution[t_idx, 0, j] + assert abs(actual - expected) < tolerance, f"Left boundary violated at t={t_val}, y={y_val}" + + # Check right boundary (x=1) + for j in range(solution.shape[2]): + y_val = j / (solution.shape[2] - 1) + expected = self.equation.right_boundary(t_val, y_val) + actual = solution[t_idx, -1, j] + assert abs(actual - expected) < tolerance, f"Right boundary violated at t={t_val}, y={y_val}" + + def test_invalid_equation_parameters(self): + """Test that invalid equation parameters raise appropriate errors""" + + # Negative time + with pytest.raises((ValueError, AssertionError)): + heat2d.HeatEquation2D(time=-1.0, t_nodes=100, k=1.0, xlength=1.0, x_nodes=20, ylength=1.0, y_nodes=20) + + # Zero or negative spatial dimensions + with pytest.raises((ValueError, AssertionError)): + heat2d.HeatEquation2D(time=1.0, t_nodes=100, k=1.0, xlength=0, x_nodes=20, ylength=1.0, y_nodes=20) + + # Negative diffusivity constant + with pytest.raises((ValueError, AssertionError)): + heat2d.HeatEquation2D(time=1.0, t_nodes=100, k=-1.0, xlength=1.0, x_nodes=20, ylength=1.0, y_nodes=20) + + # Too few nodes + with pytest.raises((ValueError, AssertionError)): + heat2d.HeatEquation2D(time=1.0, t_nodes=1, k=1.0, xlength=1.0, x_nodes=2, ylength=1.0, y_nodes=2) + + def test_missing_boundary_conditions(self): + """Test that missing boundary conditions are detected""" + incomplete_eq = heat2d.HeatEquation2D(time=0.1, t_nodes=50, k=1.0, xlength=1.0, x_nodes=10, ylength=1.0, y_nodes=10) + + # Only set some boundary conditions, leave others undefined + incomplete_eq.set_initial_temp(lambda x, y: 1.0) + incomplete_eq.set_left_boundary_temp(lambda t, y: 0) + # Missing right, top, bottom boundaries + + with pytest.raises((AttributeError, ValueError)): + solver2d.Heat2DExplicitSolver(incomplete_eq).solve() + + def test_missing_initial_condition(self): + """Test that missing initial condition is detected""" + incomplete_eq = heat2d.HeatEquation2D(time=0.1, t_nodes=50, k=1.0, xlength=1.0, x_nodes=10, ylength=1.0, y_nodes=10) + + # Set boundaries but no initial condition + incomplete_eq.set_left_boundary_temp(lambda t, y: 0) + incomplete_eq.set_right_boundary_temp(lambda t, y: 0) + incomplete_eq.set_bottom_boundary_temp(lambda t, x: 0) + incomplete_eq.set_top_boundary_temp(lambda t, x: 0) + + with pytest.raises((AttributeError, ValueError)): + solver2d.Heat2DExplicitSolver(incomplete_eq).solve() + + def test_boundary_function_exceptions(self): + """Test handling of boundary functions that raise exceptions""" + + def bad_boundary(t, coord): + if t > 0.05: + raise ValueError("Boundary function failed") + return 1.0 + + bad_eq = heat2d.HeatEquation2D(time=0.1, t_nodes=50, k=1.0, xlength=1.0, x_nodes=10, ylength=1.0, y_nodes=10) + bad_eq.set_initial_temp(lambda x, y: 1.0) + bad_eq.set_left_boundary_temp(bad_boundary) + bad_eq.set_right_boundary_temp(lambda t, y: 0) + bad_eq.set_bottom_boundary_temp(lambda t, x: 0) + bad_eq.set_top_boundary_temp(lambda t, x: 0) + + with pytest.raises(ValueError): + solver2d.Heat2DCNSolver(bad_eq).solve() + + def test_boundary_function_returns_invalid_values(self): + """Test handling of boundary functions returning NaN or inf""" + + bad_eq = heat2d.HeatEquation2D(time=0.1, t_nodes=50, k=1.0, xlength=1.0, x_nodes=10, ylength=1.0, y_nodes=10) + bad_eq.set_initial_temp(lambda x, y: 1.0) + bad_eq.set_left_boundary_temp(lambda t, y: np.nan) # Returns NaN + bad_eq.set_right_boundary_temp(lambda t, y: 0) + bad_eq.set_bottom_boundary_temp(lambda t, x: 0) + bad_eq.set_top_boundary_temp(lambda t, x: 0) + + result = solver2d.Heat2DCNSolver(bad_eq).solve().get_result() + + # Should detect NaN in solution + assert np.any(np.isnan(result)), "NaN should propagate through solution" + + def test_initial_condition_returns_invalid_values(self): + """Test handling of initial condition returning invalid values""" + + bad_eq = heat2d.HeatEquation2D(time=0.1, t_nodes=50, k=1.0, xlength=1.0, x_nodes=10, ylength=1.0, y_nodes=10) + bad_eq.set_initial_temp(lambda x, y: np.inf if x > 0.5 else 1.0) + bad_eq.set_left_boundary_temp(lambda t, y: 0) + bad_eq.set_right_boundary_temp(lambda t, y: 0) + bad_eq.set_bottom_boundary_temp(lambda t, x: 0) + bad_eq.set_top_boundary_temp(lambda t, x: 0) + + result = solver2d.Heat2DCNSolver(bad_eq).solve().get_result() + + # Should detect inf in solution + assert np.any(np.isinf(result)), "Inf should propagate through solution" + + def test_solution_object_completeness(self): + """Test that solution object contains all required information""" + + result = solver2d.Heat2DCNSolver(self.equation).solve() + + # Test all expected attributes exist + assert hasattr(result, 'get_result'), "Solution should have get_result method" + assert hasattr(result, 'get_execution_time'), "Solution should have get_execution_time method" + + # Test that methods return valid data + solution_array = result.get_result() + assert isinstance(solution_array, np.ndarray), "get_result should return numpy array" + + exec_time = result.get_execution_time() + assert isinstance(exec_time, (int, float)), "get_execution_time should return number" + assert exec_time >= 0, "Execution time should be non-negative" + + @pytest.mark.parametrize("solver_class", [solver2d.Heat2DExplicitSolver, solver2d.Heat2DCNSolver]) + def test_solver_with_None_equation(self, solver_class): + """Test solver behavior with None equation""" + + with pytest.raises((TypeError, AttributeError)): + solver_class(None).solve() + + def test_solver_reusability(self): + """Test that solvers can be reused safely""" + + solver = solver2d.Heat2DCNSolver(self.equation) + + # First solve + result1 = solver.solve() + + # Second solve should give same result + result2 = solver.solve() + + np.testing.assert_array_equal(result1.get_result(), result2.get_result()) + + def test_equation_modification_after_solver_creation(self): + """Test behavior when equation is modified after solver creation""" + + eq = heat2d.HeatEquation2D(time=0.1, t_nodes=50, k=1.0, xlength=1.0, x_nodes=10, ylength=1.0, y_nodes=10) + eq.set_initial_temp(lambda x, y: 1.0) + eq.set_left_boundary_temp(lambda t, y: 0) + eq.set_right_boundary_temp(lambda t, y: 0) + eq.set_bottom_boundary_temp(lambda t, x: 0) + eq.set_top_boundary_temp(lambda t, x: 0) + + solver = solver2d.Heat2DCNSolver(eq) + + # Modify equation after solver creation + eq.set_initial_temp(lambda x, y: 10.0) + + result = solver.solve() \ No newline at end of file diff --git a/pdesolvers/utils/utility.py b/pdesolvers/utils/utility.py index f67dcdb..edba1f3 100644 --- a/pdesolvers/utils/utility.py +++ b/pdesolvers/utils/utility.py @@ -1,6 +1,7 @@ import numpy as np import pandas as pd import matplotlib.pyplot as plt +from scipy.sparse import diags from scipy.sparse import csc_matrix import pdesolvers.enums.enums as enum @@ -28,6 +29,111 @@ def _build_tridiagonal_matrix(a, b, c, nodes): return matrix +class Heat2DHelper: + @staticmethod + def initMatrix(t_nodes, x_nodes, y_nodes, left, right, bottom, top, u0, xDomain, yDomain, tDomain): + """ + Initialize a 3D matrix with boundary and initial conditions for the 2D heat equation. + + Parameters: + ----------- + t_nodes : int + Number of time points + x_nodes : int + Number of x spatial points + y_nodes : int + Number of y spatial points + left : callable + Left boundary condition function: left(t, y) -> temperature + right : callable + Right boundary condition function: right(t, y) -> temperature + bottom : callable + Bottom boundary condition function: bottom(t, x) -> temperature + top : callable + Top boundary condition function: top(t, x) -> temperature + u0 : callable + Initial condition function: u0(x, y) -> temperature + xDomain : np.array + X coordinate array + yDomain : np.array + Y coordinate array + tDomain : np.array + Time array + + Returns: + -------- + np.array + 3D array of shape (t_nodes, x_nodes, y_nodes) with initialized conditions + """ + + matrix = np.zeros((t_nodes, x_nodes, y_nodes)) + # sp.sparse() + for tau in range(t_nodes): + t = tDomain[tau] + + # Left and right boundaries (x = 0 and x = xLength) + for i in range(y_nodes): + y = yDomain[i] + matrix[tau, 0, i] = left(t, y) # Left boundary (x=0) + matrix[tau, -1, i] = right(t, y) # Right boundary (x=xLength) + + # Bottom and top boundaries (y = 0 and y = yLength) + for j in range(x_nodes): + x = xDomain[j] + matrix[tau, j, 0] = bottom(t, x) # Bottom boundary (y=0) + matrix[tau, j, -1] = top(t, x) # Top boundary (y=yLength) + + # Set initial condition at t=0 + for i in range(x_nodes): + for j in range(y_nodes): + try: + initial_val = u0(xDomain[i], yDomain[j]) + if hasattr(initial_val, '__iter__') and not isinstance(initial_val, str): + matrix[0, i, j] = float(initial_val.flat[0]) if hasattr(initial_val, 'flat') else float(initial_val[0]) + else: + matrix[0, i, j] = float(initial_val) + except (TypeError, IndexError, AttributeError): + matrix[0, i, j] = float(u0(xDomain[i], yDomain[j])) + + for tau in range(t_nodes): + t = tDomain[tau] + matrix[tau, 0, 0] = left(t, yDomain[0]) # Bottom-left + matrix[tau, 0, -1] = left(t, yDomain[-1]) # Top-left + matrix[tau, -1, 0] = right(t, yDomain[0]) # Bottom-right + matrix[tau, -1, -1] = right(t, yDomain[-1]) # Top-right + + return matrix + + def innitTriDiagMatrix(nx, ny, cx, cy, alpha): + """ + Build sparse matrix G for: -cx*U[i-1,j] + α*U[i,j] - cx*U[i+1,j] - cy*U[i,j-1] - cy*U[i,j+1] = RHS + """ + n_interior_x = nx - 2 + n_interior_y = ny - 2 + n_total = n_interior_x * n_interior_y + + # Main diagonal: α coefficients + main_diagonal = np.full(n_total, alpha) + + # x-direction coupling: -cx coefficients (±1 in flattened index) + x_off_diagonal = np.full(n_total - 1, -cx) + # Zero out connections across y-boundaries + for i in range(n_interior_x - 1, n_total - 1, n_interior_x): + if i < len(x_off_diagonal): + x_off_diagonal[i] = 0 + + # y-direction coupling: -cy coefficients (±n_interior_x in flattened index) + y_off_diagonal = np.full(n_total - n_interior_x, -cy) + + # Assemble sparse matrix + diagonals = [y_off_diagonal, x_off_diagonal, main_diagonal, + x_off_diagonal[:n_total-1], y_off_diagonal] + offsets = [-n_interior_x, -1, 0, 1, n_interior_x] + + G = diags(diagonals, offsets=offsets, shape=(n_total, n_total), format='csr') + + return G, n_interior_x, n_interior_y + class BlackScholesHelper: @staticmethod diff --git a/playground/2D_Cranky.py b/playground/2D_Cranky.py new file mode 100644 index 0000000..d583b2e --- /dev/null +++ b/playground/2D_Cranky.py @@ -0,0 +1,186 @@ +import numpy as np +import matplotlib.pyplot as plt +from scipy.sparse import diags +from scipy.sparse.linalg import spsolve +from matplotlib.animation import FuncAnimation + +# Global Constants +xLength = 10 +yLength = 10 +maxTime = 0.5 +diffusivityConstant = 4 +numPointsSpace = 50 +numPointsTime = 2000 + +# Domain setup +xDomain = np.linspace(0, xLength, numPointsSpace) +yDomain = np.linspace(0, yLength, numPointsSpace) +timeDomain = np.linspace(0, maxTime, numPointsTime) + +# Step sizes +timeStepSize = timeDomain[1] - timeDomain[0] +spaceStepSize = xDomain[1] - xDomain[0] + +print(f"Time step: {timeStepSize:.6f}") +print(f"Space step: {spaceStepSize:.6f}") + +# Crank-Nicolson parameters (following your notes) +c = diffusivityConstant * timeStepSize / 2 # Note: factor of 2 for Crank-Nicolson +cx = c / (spaceStepSize**2) +cy = c / (spaceStepSize**2) +alpha = 1 + 2*cx + 2*cy +beta = 1 - 2*cx - 2*cy + +print(f"c = {c:.6f}") +print(f"cx = cy = {cx:.6f}") +print(f"alpha = {alpha:.6f}") +print(f"beta = {beta:.6f}") + +# Initial condition - a hot spot in the center +u0 = lambda x, y: 10 * np.exp(-((x - xLength/2)**2 + (y - yLength/2)**2) / 2) +left = lambda t, y: 20 + 10 * y * (yLength - y) * t**2 +right = lambda t, y: 20 + 100 * y * (yLength - y)**3 * (t - 1)**2 * (t > 1) +top = lambda t, x: 20 + 5 * x * (xLength - x) * t**4 +bottom = lambda t, x: 20 + +def initMatrix(t_nodes, x_nodes, y_nodes, left, right, bottom, top, u0, xDomain, yDomain, tDomain): + """Initialize matrix with boundary and initial conditions""" + matrix = np.zeros((t_nodes, x_nodes, y_nodes)) + + # Set boundary conditions for all time steps + for tau in range(t_nodes): + t = tDomain[tau] + # Left and right boundaries + for i in range(y_nodes): + y = yDomain[i] + matrix[tau, 0, i] = left(t, y) + matrix[tau, -1, i] = right(t, y) + # Bottom and top boundaries + for j in range(x_nodes): + x = xDomain[j] + matrix[tau, j, 0] = bottom(t, x) + matrix[tau, j, -1] = top(t, x) + + # Set initial condition at t=0 + for i in range(x_nodes): + for j in range(y_nodes): + matrix[0, i, j] = u0(xDomain[i], yDomain[j]) + + return matrix + +def build_2d_crank_nicolson_matrix(nx, ny, cx, cy): + """ + Build sparse matrix G for: -cx*U[i-1,j] + α*U[i,j] - cx*U[i+1,j] - cy*U[i,j-1] - cy*U[i,j+1] = RHS + """ + n_interior_x = nx - 2 + n_interior_y = ny - 2 + n_total = n_interior_x * n_interior_y + + # Main diagonal: α coefficients + main_diagonal = np.full(n_total, alpha) + + # x-direction coupling: -cx coefficients (±1 in flattened index) + x_off_diagonal = np.full(n_total - 1, -cx) + # Zero out connections across y-boundaries + for i in range(n_interior_x - 1, n_total - 1, n_interior_x): + if i < len(x_off_diagonal): + x_off_diagonal[i] = 0 + + # y-direction coupling: -cy coefficients (±n_interior_x in flattened index) + y_off_diagonal = np.full(n_total - n_interior_x, -cy) + + # Assemble sparse matrix + diagonals = [y_off_diagonal, x_off_diagonal, main_diagonal, + x_off_diagonal[:n_total-1], y_off_diagonal] + offsets = [-n_interior_x, -1, 0, 1, n_interior_x] + + G = diags(diagonals, offsets=offsets, shape=(n_total, n_total), format='csr') + + return G, n_interior_x, n_interior_y + +def calculateTemperatureCN(U): + """Crank-Nicolson time stepping""" + nx, ny = numPointsSpace, numPointsSpace + + G, n_interior_x, n_interior_y = build_2d_crank_nicolson_matrix(nx, ny, cx, cy) + + for tau in range(numPointsTime - 1): + rhs = np.zeros(n_interior_x * n_interior_y) + + idx = 0 + for j in range(1, ny-1): + for i in range(1, nx-1): + # RHS = β*U_τ + cx*(neighbors_x) + cy*(neighbors_y) + boundary_terms + rhs[idx] = beta * U[tau, i, j] + rhs[idx] += cx * (U[tau, i-1, j] + U[tau, i+1, j]) + rhs[idx] += cy * (U[tau, i, j-1] + U[tau, i, j+1]) + + # Boundary contributions + if i == 1: + rhs[idx] += cx * U[tau+1, 0, j] + if i == nx-2: + rhs[idx] += cx * U[tau+1, -1, j] + if j == 1: + rhs[idx] += cy * U[tau+1, i, 0] + if j == ny-2: + rhs[idx] += cy * U[tau+1, i, -1] + + idx += 1 + + # Solve G*u_{τ+1} = rhs + u_next_interior = spsolve(G, rhs) + U[tau+1, 1:-1, 1:-1] = u_next_interior.reshape((n_interior_x, n_interior_y)) + + return U + +def plot_surface(u_k, k, ax): + ax.clear() + X, Y = np.meshgrid(xDomain, yDomain) + + surf = ax.plot_surface(X, Y, u_k.T, + cmap='hot', + alpha=0.9) + + ax.set_xlabel('X Position') + ax.set_ylabel('Y Position') + ax.set_zlabel('Temperature') + ax.set_title(f'2D Heat Equation (Crank-Nicolson): t = {k*timeStepSize:.4f}') + ax.view_init(elev=30, azim=45) + + # Dynamic z-limits + z_max = max(100, np.max(u_k)) + ax.set_zlim(0, z_max) + + return surf + +# Main execution +if __name__ == "__main__": + print("Initializing matrix...") + emptyMatrix = initMatrix( + t_nodes=numPointsTime, + x_nodes=numPointsSpace, + y_nodes=numPointsSpace, + left=left, + right=right, + bottom=bottom, + top=top, + u0=u0, + xDomain=xDomain, + yDomain=yDomain, + tDomain=timeDomain + ) + + print("Calculating temperature evolution using Crank-Nicolson...") + tempMatrix = calculateTemperatureCN(emptyMatrix) + + # Create animation + print("Creating animation...") + fig = plt.figure(figsize=(12, 8)) + ax = fig.add_subplot(111, projection='3d') + + def animate(k): + return plot_surface(tempMatrix[k], k, ax) + + anim = FuncAnimation(fig, animate, interval=100, frames=min(numPointsTime, 200), repeat=True) + anim.save("heat_equation_2d_crank_nicolson.gif", writer='pillow', fps=10) + plt.show() \ No newline at end of file diff --git a/playground/2D_explicit.py b/playground/2D_explicit.py index 5a32277..eeace39 100644 --- a/playground/2D_explicit.py +++ b/playground/2D_explicit.py @@ -1,118 +1,143 @@ -# We will use the 5-point Laplacian aprroximation stencil explicit method -# below to simulate the 2D Heat Equation +# Corrected 2D Heat Equation simulation using explicit finite difference method import numpy as np import matplotlib.pyplot as plt - from matplotlib.animation import FuncAnimation # Global Constants -xLength = 50 -yLength = 50 -maxTime = 1 -diffusivityConstant = 1 -numPointsSpace = 200 -numPointsTime = 60 - -# Length Vector plotted on x-axis and y-axis +xLength = 10 # Lx +yLength = 10 # Ly +maxTime = 0.5 # tmax +diffusivityConstant = 4 # kappa +numPointsSpace = 50 # x_points = y_points +numPointsTime = 2000 # t_points + +# Domain setup xDomain = np.linspace(0, xLength, numPointsSpace) yDomain = np.linspace(0, yLength, numPointsSpace) -# Time Vector plotted on z-axis timeDomain = np.linspace(0, maxTime, numPointsTime) -# Step-sizes- dt, dx, dy +# Step sizes timeStepSize = timeDomain[1] - timeDomain[0] spaceStepSize = xDomain[1] - xDomain[0] -# Stability condition: dt <= dx^2/(2*k) -err_dt= np.abs((spaceStepSize**2)/(2*diffusivityConstant)) -print(err_dt) -assert(timeStepSize <= err_dt) +# Stability condition: dt <= dx^2/(4*k) for 2D +# stability_limit = (spaceStepSize**2)/(4*diffusivityConstant) +# print(f"Stability limit: {stability_limit:.6f}") +# print(f"Time step: {timeStepSize:.6f}") +# assert(timeStepSize <= stability_limit), "Time step too large for stability!" -# Calculate lambda constant for x and y (they're the same if dx = dy) +# Lambda constant lambdaConstant = (diffusivityConstant * timeStepSize) / (spaceStepSize**2) -# Stability condition: kdt/dx^2 <= 0.5 -assert(lambdaConstant <= 0.5) -print(lambdaConstant) +# print(f"Lambda: {lambdaConstant:.6f}") +# assert(lambdaConstant <= 0.25), "Lambda too large for 2D stability!" # 0.25 for 2D + +# Initial condition - a hot spot in the center +u0 = lambda x, y: 10 * np.exp(-((x - xLength/2)**2 + (y - yLength/2)**2) / 2) + +# Boundary conditions based on your MATLAB code -# lambda functions for initial condition and boundary condtions -# 2 sin x sin 2y + 3 sin 4x sin 5y + t -u0 = lambda x, y: 2 * np.sin(x) * np.sin(2*y) + 3 * np.sin(4*x) * np.sin(5*y) -left = lambda t, y: 2 * np.sin(0) * np.sin(2*y) + 3 * np.sin(4*0) * np.sin(5*y) + t -right = lambda t, y: 2 * np.sin(xLength) * np.sin(2*y) + 3 * np.sin(4*xLength) * np.sin(5*y) + t -bottom = lambda t, x: 2 * np.sin(x) * np.sin(2*0) + 3 * np.sin(4*x) * np.sin(5*0) + t -top = lambda t, x: 2 * np.sin(x) * np.sin(2*yLength) + 3 * np.sin(4*x) * np.sin(5*yLength) + t +# left = @(t,y) 20 + 10*y*(Ly-y)*t^2; +left = lambda t, y: 20 + 10 * y * (yLength - y) * t**2 +# right = @(t, y) 20 + 100*y*(Ly-y)^3*(t-1)^2*(t>1); +right = lambda t, y: 20 + 100 * y * (yLength - y)**3 * (t - 1)**2 * (t > 1) -# error assertion for intial & boundary conditions -err1_u0 = np.abs(u0(0, 0) - left(0, 0)) -assert(err1_u0 < 1e-12) -err2_u0 = np.abs(u0(0, yLength) - right(0, 0)) -assert(err2_u0 < 1e-12) -err3_u0 = np.abs(u0(0, 0) - bottom(0, 0)) -assert(err3_u0 < 1e-12) -err4_u0 = np.abs(u0(0, xLength) - top(0, 0)) -assert(err4_u0 < 1e-12) +# up = @(t, x) 20 + 5*x*(Lx-x)*t^4; +top = lambda t, x: 20 + 5 * x * (xLength - x) * t**4 + +# down = @(t, x) 20; +bottom = lambda t, x: 20 # Warm top boundary def initMatrix(): - # initialize an empty matrix - matrix = np.empty((numPointsTime, numPointsSpace, numPointsSpace)) - # set the boundary conditions for the entire 2D space + matrix = np.zeros((numPointsTime, numPointsSpace, numPointsSpace)) + + # Set boundary conditions for all time steps for tau in range(numPointsTime): + t = timeDomain[tau] + # Left and right boundaries (x = 0 and x = xLength) for i in range(numPointsSpace): - matrix[tau, i, 0] = left(tau, xDomain[i]) - matrix[tau, i, -1] = right(tau, xDomain[i]) - for tau in range(numPointsTime): + y = yDomain[i] + matrix[tau, 0, i] = left(t, y) # Left boundary (x=0) + matrix[tau, -1, i] = right(t, y) # Right boundary (x=xLength) + + # Bottom and top boundaries (y = 0 and y = yLength) for j in range(numPointsSpace): - matrix[tau, 0, j] = bottom(tau, yDomain[j]) - matrix[tau, -1, j] = top(tau, yDomain[j]) - # set the initial condition for the entire 2D space at t=0 - for i in range(len(xDomain)): - for j in range(len(yDomain)): - matrix[0,i,j] = u0(xDomain[i], yDomain[j]) + x = xDomain[j] + matrix[tau, j, 0] = bottom(t, x) # Bottom boundary (y=0) + matrix[tau, j, -1] = top(t, x) # Top boundary (y=yLength) + + # Set initial condition at t=0 + for i in range(numPointsSpace): + for j in range(numPointsSpace): + matrix[0, i, j] = u0(xDomain[i], yDomain[j]) + + # Ensure corners are consistent (use boundary values) + for tau in range(numPointsTime): + t = timeDomain[tau] + matrix[tau, 0, 0] = left(t, yDomain[0]) # Bottom-left + matrix[tau, 0, -1] = left(t, yDomain[-1]) # Top-left + matrix[tau, -1, 0] = right(t, yDomain[0]) # Bottom-right + matrix[tau, -1, -1] = right(t, yDomain[-1]) # Top-right return matrix - -# Time-stepping loop def calculateTemperature(U): - for tau in range(0, numPointsTime-1, 1): - for i in range(1, numPointsSpace-1, 1): - for j in range(1, numPointsSpace-1, 1): - # 5-point stencil implementation - U[tau+1,i,j] = U[tau,i,j] + lambdaConstant * ( - # x-direction terms - (U[tau,i-1,j] - 2*U[tau,i,j] + U[tau,i+1,j]) + - # y-direction terms - (U[tau,i,j-1] - 2*U[tau,i,j] + U[tau,i,j+1]) + for tau in range(numPointsTime-1): + for i in range(1, numPointsSpace-1): + for j in range(1, numPointsSpace-1): + # 5-point stencil for 2D heat equation + U[tau+1, i, j] = U[tau, i, j] + lambdaConstant * ( + # Second derivative in x-direction + (U[tau, i-1, j] - 2*U[tau, i, j] + U[tau, i+1, j]) + + # Second derivative in y-direction + (U[tau, i, j-1] - 2*U[tau, i, j] + U[tau, i, j+1]) ) return U - def plot_surface(u_k, k, ax): ax.clear() X, Y = np.meshgrid(xDomain, yDomain) - surf = ax.plot_surface(X, Y, u_k, - cmap='jet', - vmin=0, vmax=100) + # Transpose u_k to match meshgrid orientation + surf = ax.plot_surface(X, Y, u_k.T, + cmap='hot', + alpha=0.9) ax.set_xlabel('X Position') - ax.set_ylabel('Y Position') + ax.set_ylabel('Y Position') ax.set_zlabel('Temperature') - ax.set_title(f'Temperature at t = {k*timeStepSize:.3f}') + ax.set_title(f'2D Heat Equation: t = {k*timeStepSize:.4f}') ax.view_init(elev=30, azim=45) + + # Set consistent z-axis limits for better visualization + ax.set_zlim(0, 100) + return surf -# Create figure and 3D axes outside animation loop -fig = plt.figure(figsize=(10, 7)) -ax = fig.add_subplot(111, projection='3d') - -def animate(k): - plot_surface(tempMatrix[k], k, ax) - -emptyMatrix = initMatrix() -tempMatrix = calculateTemperature(emptyMatrix) - - -anim = FuncAnimation(fig, animate, interval=10, frames=numPointsTime, repeat=False) -anim.save("heat_equation_solution_surf_explicit.gif", writer='pillow') \ No newline at end of file +# Main execution +if __name__ == "__main__": + # Initialize and solve + print("Initializing matrix...") + emptyMatrix = initMatrix() + print("Calculating temperature evolution...") + tempMatrix = calculateTemperature(emptyMatrix) + + # # Save temperature matrix to CSV file + # print("Saving temperature matrix to file...") + # np.savetxt("temperature_data.csv", tempMatrix.reshape(numPointsTime, -1), delimiter=",") + + # Create animation + print("Creating animation...") + fig = plt.figure(figsize=(12, 8)) + ax = fig.add_subplot(111, projection='3d') + + def animate(k): + return plot_surface(tempMatrix[k], k, ax) + + anim = FuncAnimation(fig, animate, interval=100, frames=numPointsTime, repeat=True) + + # Show the plot + # plt.show() + + # Optionally save animation + anim.save("heat_equation_2d_explicit.gif", writer='pillow', fps=10) \ No newline at end of file