added test suite and validation

Author: debdal6

pdesolvers/main.py

@@ -15,7 +15,8 @@ def main():
     # solver1 = pde.Heat1DCNSolver(equation1)
     # solver2 = pde.Heat1DExplicitSolver(equation1)
 
-    #testing Heat 2d
+    # testing 2d heat equation
+    
     xLength = 10  # Lx
     yLength = 10  # Ly  
     maxTime = 0.5  # tmax
@@ -30,12 +31,13 @@ def main():
     equation.set_top_boundary_temp(lambda t, x: 20 + 5 * x * (xLength - x) * t**4)
     equation.set_bottom_boundary_temp(lambda t, x: 20)
 
-    solver = pde.Heat2DExplicitSolver(equation)
-    solution = solver.solve()
-    # Save temperature matrix to CSV file
-    # print("Saving temperature matrix to file...")
-    # np.savetxt("temperature_data_pkg.csv", solution.result.reshape(numPointsTime, -1), delimiter=",")
-    solution.animate(export=True)
+    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'

pdesolvers/pdes/heat_2d.py

@@ -2,6 +2,17 @@
 
 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

pdesolvers/solution/solution.py

@@ -196,7 +196,7 @@ def __plot_surface(u_k, k, ax, xDomain, yDomain, dt):
 
         return surf
 
-    def animate(self, export=False):
+    def animate(self, export=False, filename="heat_equation_2d_plot.gif"):
         print("Creating animation...")
         self
         fig = plt.figure(figsize=(12, 8))
@@ -208,7 +208,7 @@ def animateFrame(k):
         anim = FuncAnimation(fig, animateFrame, interval=100, frames=len(self.t_grid), repeat=True)
 
         if export:
-            anim.save("heat_equation_2d_plot.gif", writer='pillow', fps=5)
+            anim.save(filename+'.gif', writer='pillow', fps=5)
 
         plt.show()
 

pdesolvers/solvers/__init__.py

@@ -1,3 +1,3 @@
 from .heat_solvers import Heat1DExplicitSolver, Heat1DCNSolver
-from .heat2d_solvers import Heat2DExplicitSolver
+from .heat2d_solvers import Heat2DExplicitSolver, Heat2DCNSolver
 from .black_scholes_solvers import BlackScholesExplicitSolver, BlackScholesCNSolver

pdesolvers/solvers/heat2d_solvers.py

@@ -24,6 +24,14 @@ 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)
@@ -45,7 +53,7 @@ def solve(self):
                                                  self.equation.initial_temp,
                                                  x, y, t)
 
-        print("Calculating temperature evolution...")
+        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):
@@ -69,13 +77,27 @@ def __init__(self, equation: heat.HeatEquation2D):
     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,
@@ -89,9 +111,33 @@ def solve(self):
                                                  x, y, t)
 
         # create sparse matrix
-
-        # time-stepping loop
+        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.")