Black source

Author: rgaveiga

CHANGELOG.md

@@ -10,6 +10,7 @@
 - Moved `create_price_samples` from support.py to a new module price_array.py and renamed it to `create_price_array`
 - Commented a code snippet in engine.py where terminal stock prices are created using `create_price_samples`, to be removed in a next version
 - Allowed a dictionary as input for `create_price_array` in price_array.py
+- Allowed a dictionary as input for `get_pop` in support.py
 
 ## 1.3.5 (2024-12-28)
 

optionlab/price_array.py

@@ -8,7 +8,6 @@
 from optionlab.models import BlackScholesModelInputs, LaplaceInputs
 
 
-@lru_cache
 def create_price_array(
     inputs_data: BlackScholesModelInputs | LaplaceInputs | dict,
     n: int = 100_000,
@@ -19,7 +18,7 @@ def create_price_array(
 
     Parameters
     ----------
-    inputs_data : BlackScholesModelInputs | LaplaceInputs
+    inputs_data : BlackScholesModelInputs | LaplaceInputs | dict
         Input data used to generate the terminal stock prices. See the documentation
         for `BlackScholesModelInputs` and `LaplaceInputs` for more details.
     n : int, optional
@@ -32,37 +31,40 @@ def create_price_array(
     numpy.ndarray
         Array of terminal prices.
     """
-       
+
     if isinstance(inputs_data, dict):
         input_type = inputs_data["model"]
-    elif isinstance(inputs_data, BlackScholesModelInputs):
-        input_type = "black-scholes"
-    elif isinstance(inputs_data, LaplaceInputs):
-        input_type = "laplace"
+
+        if input_type in ("black-scholes", "normal"):
+            inputs = BlackScholesModelInputs.model_validate(inputs_data)
+        elif input_type == "laplace":
+            inputs = LaplaceInputs.model_validate(inputs_data)
+        else:
+            raise ValueError("Inputs are not valid!")
     else:
-        raise ValueError("Inputs are not valid!")
-        
+        inputs = inputs_data
+
+        if isinstance(inputs, BlackScholesModelInputs):
+            input_type = "black-scholes"
+        elif isinstance(inputs, LaplaceInputs):
+            input_type = "laplace"
+        else:
+            raise ValueError("Inputs are not valid!")
+
     np_seed_number(seed)
 
     if input_type in ("black-scholes", "normal"):
-        arr = _get_array_price_from_BS(inputs_data, n)
+        arr = _get_array_price_from_BS(inputs, n)
     elif input_type == "laplace":
-        arr = _get_array_price_from_laplace(inputs_data, n)
+        arr = _get_array_price_from_laplace(inputs, n)
 
     np_seed_number(None)
 
     return arr
 
 
-def _get_array_price_from_BS(
-    inputs_data: BlackScholesModelInputs | dict, n: int
-) -> np.ndarray:
-    inputs = (
-        BlackScholesModelInputs.model_validate(inputs_data)
-        if isinstance(inputs_data, dict)
-        else inputs_data
-    )
-
+@lru_cache
+def _get_array_price_from_BS(inputs: BlackScholesModelInputs, n: int) -> np.ndarray:
     return exp(
         normal(
             (
@@ -80,15 +82,8 @@ def _get_array_price_from_BS(
     )
 
 
-def _get_array_price_from_laplace(
-    inputs_data: LaplaceInputs | dict, n: int
-) -> np.ndarray:
-    inputs = (
-        LaplaceInputs.model_validate(inputs_data)
-        if isinstance(inputs_data, dict)
-        else inputs_data
-    )
-
+@lru_cache
+def _get_array_price_from_laplace(inputs: LaplaceInputs, n: int) -> np.ndarray:
     return exp(
         laplace(
             (log(inputs.stock_price) + inputs.mu * inputs.years_to_target_date),

optionlab/support.py

@@ -215,7 +215,7 @@ def get_profit_range(
 
     Returns
     -------
-    list
+    list[Range]
         List of stock price pairs.
     """
 
@@ -279,10 +279,10 @@ def get_pop(
 
     if len(profit_ranges) == 0:
         return pop
-    
+
     if isinstance(inputs_data, dict):
         input_type = inputs_data["model"]
-        
+
         if input_type in ("black-scholes", "normal"):
             inputs = BlackScholesModelInputs.model_validate(inputs_data)
         elif input_type == "laplace":
@@ -293,7 +293,7 @@ def get_pop(
             raise ValueError("Inputs are not valid!")
     else:
         inputs = inputs_data
-        
+
         if isinstance(inputs, BlackScholesModelInputs):
             input_type = "black-scholes"
         elif isinstance(inputs, LaplaceInputs):

tests/test_misc.py

@@ -3,8 +3,8 @@
 
 import pytest
 
-from optionlab.models import BlackScholesModelInputs, LaplaceInputs
-from optionlab.price_array import create_price_array
+from optionlab.models import BlackScholesModelInputs
+from optionlab.price_array import create_price_array, _get_array_price_from_BS
 from optionlab.utils import get_nonbusiness_days
 
 
@@ -55,7 +55,7 @@ def test_benchmark_holidays(benchmark):
 
 
 def test_cache_price_samples():
-    create_price_array.cache_clear()
+    _get_array_price_from_BS.cache_clear()
 
     stock_price = 168.99
     volatility = 0.483
@@ -97,12 +97,13 @@ def test_cache_price_samples():
     stock_price = 167.0
 
     sample3 = create_price_array(
-        inputs_data=BlackScholesModelInputs(
-            stock_price=stock_price,
-            volatility=volatility,
-            interest_rate=interest_rate,
-            years_to_target_date=years_to_target,
-        ),
+        inputs_data={
+            "model": "normal",
+            "stock_price": stock_price,
+            "volatility": volatility,
+            "interest_rate": interest_rate,
+            "years_to_target_date": years_to_target,
+        },
         seed=0,
     )
 
@@ -113,9 +114,13 @@ def test_cache_price_samples():
     assert sample3.sum() == pytest.approx(16752035.781465728, rel=0.01)
 
     sample4 = create_price_array(
-        inputs_data=LaplaceInputs(
-            stock_price=168.99, volatility=0.483, mu=0.05, years_to_target_date=24 / 365
-        ),
+        inputs_data={
+            "model": "laplace",
+            "stock_price": 168.99,
+            "volatility": 0.483,
+            "mu": 0.05,
+            "years_to_target_date": 24 / 365,
+        },
         seed=0,
     )