FEATURE: adding enums

Author: chemardes

pdesolvers/__init__.py

@@ -1,4 +1,5 @@
 from .pdes import *
 from .solution import *
 from .solvers import *
-from .utils import *
+from .utils import *
+from .enums import *

pdesolvers/enums/__init__.py

@@ -0,0 +1 @@
+from .option_type import OptionType

pdesolvers/enums/option_type.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+class OptionType(Enum):
+    EUROPEAN_CALL = 1
+    EUROPEAN_PUT = 2

pdesolvers/main.py

@@ -1,7 +1,6 @@
 import numpy as np
 import pdesolvers as pde
 
-
 def main():
 
     # testing for heat equation
@@ -16,7 +15,7 @@ def main():
     # solver2 = pde.Heat1DExplicitSolver(equation1)
 
     # testing for bse
-    equation2 = pde.BlackScholesEquation('call', 300, 1, 0.2, 0.05, 100, 100, 20000)
+    equation2 = pde.BlackScholesEquation(pde.OptionType.EUROPEAN_CALL, 300, 1, 0.2, 0.05, 100, 100, 20000)
 
     solver1 = pde.BlackScholesCNSolver(equation2)
     # solver2 = pde.BlackScholesExplicitSolver(equation2)

pdesolvers/pdes/black_scholes.py

@@ -1,8 +1,9 @@
 import numpy as np
+from pdesolvers.enums.option_type import OptionType
 
 class BlackScholesEquation:
 
-    def __init__(self, option_type, S_max, expiry, sigma, r, K, s_nodes=1, t_nodes=None):
+    def __init__(self, option_type: OptionType, S_max, expiry, sigma, r, K, s_nodes=1, t_nodes=None):
         """
         Initialises the solver with the necessary parameters
 
@@ -16,6 +17,8 @@ def __init__(self, option_type, S_max, expiry, sigma, r, K, s_nodes=1, t_nodes=N
         :param t_nodes: number of time nodes
         """
 
+        if not isinstance(option_type, OptionType):
+            raise TypeError(f"Option type must be of type OptionType enum" )
         self.__option_type = option_type
         self.__S_max = S_max
         self.__expiry = expiry

pdesolvers/solution/solution.py

@@ -1,4 +1,3 @@
-import pdesolvers.pdes.black_scholes as bse
 import numpy as np
 from matplotlib import pyplot as plt
 

pdesolvers/solvers/black_scholes_solvers.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pdesolvers.solution as sol
 import pdesolvers.pdes.black_scholes as bse
+import pdesolvers.enums.option_type as enum
 
 class BlackScholesExplicitSolver:
 
@@ -36,9 +37,9 @@ def solve(self):
         V = np.zeros((self.equation.s_nodes + 1, self.equation.t_nodes + 1))
 
         # setting terminal condition
-        if self.equation.option_type == 'call':
+        if self.equation.option_type == enum.OptionType.EUROPEAN_CALL:
             V[:,-1] = np.maximum((S - self.equation.strike_price), 0)
-        elif self.equation.option_type == 'put':
+        elif self.equation.option_type == enum.OptionType.EUROPEAN_PUT:
             V[:,-1] = np.maximum((self.equation.strike_price - S), 0)
         else:
             raise ValueError("Invalid option type - please choose between call/put")
@@ -74,10 +75,10 @@ def __set_boundary_conditions(self, T, tau):
 
         lower_boundary = None
         upper_boundary = None
-        if self.equation.option_type == 'call':
+        if self.equation.option_type == enum.OptionType.EUROPEAN_CALL:
             lower_boundary = 0
             upper_boundary = self.equation.S_max - self.equation.strike_price * np.exp(-self.equation.rate * (self.equation.expiry - T[tau]))
-        elif self.equation.option_type == 'put':
+        elif self.equation.option_type == enum.OptionType.EUROPEAN_PUT:
             lower_boundary = self.equation.strike_price * np.exp(-self.equation.rate * (self.equation.expiry - T[tau]))
             upper_boundary = 0
 
@@ -130,14 +131,14 @@ def solve(self):
 
 
 # setting terminal condition (for all values of S at time T)
-        if self.equation.option_type == 'call':
+        if self.equation.option_type == enum.OptionType.EUROPEAN_CALL:
             V[:,-1] = np.maximum((S - self.equation.strike_price), 0)
 
             # setting boundary conditions (for all values of t at asset prices S=0 and S=Smax)
             V[0, :] = 0
             V[-1, :] = S[-1] - self.equation.strike_price * np.exp(-self.equation.rate * (self.equation.expiry - T))
 
-        elif self.equation.option_type == 'put':
+        elif self.equation.option_type == enum.OptionType.EUROPEAN_PUT:
             V[:,-1] = np.maximum((self.equation.strike_price - S), 0)
             V[0, :] = self.equation.strike_price * np.exp(-self.equation.rate * (self.equation.expiry - T))
             V[-1, :] = 0

pdesolvers/tests/test_black_scholes.py

@@ -4,11 +4,18 @@
 import pdesolvers.pdes.black_scholes as bse
 import pdesolvers.solvers.black_scholes_solvers as solver
 import pdesolvers.utils.utility as utility
+import pdesolvers.enums.option_type as enum
+
+class TestBlackScholesEquation:
+
+    def test_check_invalid_option_type_input(self):
+        with pytest.raises(TypeError, match="Option type must be of type OptionType enum"):
+            self.equation = bse.BlackScholesEquation('woo', 300, 1, 0.2, 0.05, 100, 100, 2000)
 
 class TestBlackScholesSolvers:
 
     def setup_method(self):
-        self.equation = bse.BlackScholesEquation('call', 300, 1, 0.2, 0.05, 100, 100, 2000)
+        self.equation = bse.BlackScholesEquation(enum.OptionType.EUROPEAN_CALL, 300, 1, 0.2, 0.05, 100, 100, 2000)
 
     # explicit method tests
 
@@ -26,7 +33,7 @@ def test_check_terminal_condition_for_call_explicit(self):
         assert np.array_equal(result[:, -1], expected_payoff)
 
     def test_check_terminal_condition_for_put_explicit(self):
-        self.equation.option_type = 'put'
+        self.equation.option_type = enum.OptionType.EUROPEAN_PUT
         result = solver.BlackScholesExplicitSolver(self.equation).solve().get_result()
 
         test_asset_grid = self.equation.generate_grid(self.equation.S_max, self.equation.s_nodes)
@@ -36,7 +43,7 @@ def test_check_terminal_condition_for_put_explicit(self):
         assert np.array_equal(result[:,-1], expected_payoff)
 
     def test_check_valid_option_type(self):
-        self.equation.option_type = 'woo'
+        self.equation.option_type = "INVALID"
 
         with pytest.raises(ValueError, match="Invalid option type - please choose between call/put"):
             solver.BlackScholesExplicitSolver(self.equation).solve().get_result()
@@ -57,7 +64,7 @@ def test_check_terminal_condition_for_call_cn(self):
         assert np.array_equal(result[:, -1], expected_payoff)
 
     def test_check_terminal_condition_for_put_cn(self):
-        self.equation.option_type = 'put'
+        self.equation.option_type = enum.OptionType.EUROPEAN_PUT
         result = solver.BlackScholesCNSolver(self.equation).solve().get_result()
 
         test_asset_grid = self.equation.generate_grid(self.equation.S_max, self.equation.s_nodes)