FEATURE: 2d heateq to pdesolver pkg

.gitignore

@@ -24,3 +24,5 @@ pdesolvers.egg-info/
 # Other files
 *.log
 /.coverage
+*.gif
+*.csv

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()

pdesolvers/pdes/__init__.py

@@ -1,2 +1,3 @@
 from .heat_1d import HeatEquation
-from .black_scholes import BlackScholesEquation
+from .black_scholes import BlackScholesEquation
+from .heat_2d import HeatEquation2D

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

pdesolvers/solution/__init__.py

@@ -1 +1 @@
-from .solution import Solution1D, SolutionBlackScholes
+from .solution import Solution1D, SolutionBlackScholes, Heat2DSolution

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')
+        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
+

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