Fix Black-Scholes test

sklam · Jun 25, 2021 · May 24, 2021 · May 25, 2021 · May 25, 2021 · May 25, 2021
commit 81ed617cf6bb24b73de11e81cb2048f0f5532609
diff --git a/.flake8 b/.flake8
@@ -186,7 +186,6 @@ exclude =
     numba/tests/test_buffer_protocol.py
     numba/tests/test_casting.py
     numba/tests/test_lists.py
-    numba/tests/test_blackscholes.py
     numba/tests/test_array_analysis.py
     numba/tests/test_serialize.py
     numba/tests/test_iteration.py

diff --git a/numba/tests/test_blackscholes.py b/numba/tests/test_blackscholes.py
@@ -3,8 +3,8 @@
 import numpy as np
 
 import unittest
-from numba.core.compiler import compile_isolated, compile_extra, Flags
-from numba.core import types, typing
+from numba import njit
+from numba.extending import register_jitable
 from numba.tests.support import TestCase
 
 
@@ -20,13 +20,15 @@
 RSQRT2PI = 0.39894228040143267793994605993438
 
 
+@register_jitable
 def cnd_array(d):
     K = 1.0 / (1.0 + 0.2316419 * np.abs(d))
     ret_val = (RSQRT2PI * np.exp(-0.5 * d * d) *
                (K * (A1 + K * (A2 + K * (A3 + K * (A4 + K * A5))))))
     return np.where(d > 0, 1.0 - ret_val, ret_val)
 
 
+@register_jitable
 def cnd(d):
     K = 1.0 / (1.0 + 0.2316419 * math.fabs(d))
     ret_val = (RSQRT2PI * math.exp(-0.5 * d * d) *
@@ -36,6 +38,7 @@ def cnd(d):
     return ret_val
 
 
+@njit
 def blackscholes_arrayexpr(stockPrice, optionStrike, optionYears, Riskfree,
                            Volatility):
     S = stockPrice
@@ -56,28 +59,9 @@ def blackscholes_arrayexpr(stockPrice, optionStrike, optionYears, Riskfree,
     return callResult, putResult
 
 
-def blackscholes_arrayexpr_jitted(stockPrice, optionStrike, optionYears,
-                                  Riskfree, Volatility):
-    S = stockPrice
-    X = optionStrike
-    T = optionYears
-    R = Riskfree
-    V = Volatility
-    sqrtT = np.sqrt(T)
-    d1 = (np.log(S / X) + (R + 0.5 * V * V) * T) / (V * sqrtT)
-    d2 = d1 - V * sqrtT
-    cndd1 = cnd_array_jitted(d1)
-    cndd2 = cnd_array_jitted(d2)
-
-    expRT = np.exp(- R * T)
-
-    callResult = (S * cndd1 - X * expRT * cndd2)
-    putResult = (X * expRT * (1.0 - cndd2) - S * (1.0 - cndd1))
-    return callResult, putResult
-
-
+@njit
 def blackscholes_scalar(callResult, putResult, stockPrice, optionStrike,
-                   optionYears, Riskfree, Volatility):
+                        optionYears, Riskfree, Volatility):
     S = stockPrice
     X = optionStrike
     T = optionYears
@@ -95,87 +79,28 @@ def blackscholes_scalar(callResult, putResult, stockPrice, optionStrike,
         putResult[i] = (X[i] * expRT * (1.0 - cndd2) - S[i] * (1.0 - cndd1))
 
 
-def blackscholes_scalar_jitted(callResult, putResult, stockPrice, optionStrike,
-                               optionYears, Riskfree, Volatility):
-    S = stockPrice
-    X = optionStrike
-    T = optionYears
-    R = Riskfree
-    V = Volatility
-    for i in range(len(S)):
-        sqrtT = math.sqrt(T[i])
-        d1 = (math.log(S[i] / X[i]) + (R + 0.5 * V * V) * T[i]) / (V * sqrtT)
-        d2 = d1 - V * sqrtT
-        cndd1 = cnd_jitted(d1)
-        cndd2 = cnd_jitted(d2)
-
-        expRT = math.exp((-1. * R) * T[i])
-        callResult[i] = (S[i] * cndd1 - X[i] * expRT * cndd2)
-        putResult[i] = (X[i] * expRT * (1.0 - cndd2) - S[i] * (1.0 - cndd1))
-
-
 def randfloat(rand_var, low, high):
     return (1.0 - rand_var) * low + rand_var * high
 
 
 class TestBlackScholes(TestCase):
     def test_array_expr(self):
-        flags = Flags()
-        flags.enable_pyobject = True
-
-        global cnd_array_jitted
-        scalty = types.float64
-        arrty = types.Array(scalty, 1, 'C')
-        cr1 = compile_isolated(cnd_array, args=(arrty,), flags=flags)
-        cnd_array_jitted = cr1.entry_point
-        cr2 = compile_isolated(blackscholes_arrayexpr_jitted,
-                               args=(arrty, arrty, arrty, scalty, scalty),
-                               flags=flags)
-        jitted_bs = cr2.entry_point
-
         OPT_N = 400
-        iterations = 10
-
 
         stockPrice = randfloat(self.random.random_sample(OPT_N), 5.0, 30.0)
         optionStrike = randfloat(self.random.random_sample(OPT_N), 1.0, 100.0)
         optionYears = randfloat(self.random.random_sample(OPT_N), 0.25, 10.0)
 
         args = stockPrice, optionStrike, optionYears, RISKFREE, VOLATILITY
 
-        callResultGold, putResultGold = blackscholes_arrayexpr(*args)
-        callResultNumba, putResultNumba = jitted_bs(*args)
+        callResultGold, putResultGold = blackscholes_arrayexpr.py_func(*args)
+        callResultNumba, putResultNumba = blackscholes_arrayexpr(*args)
 
         delta = np.abs(callResultGold - callResultNumba)
-        L1norm = delta.sum() / np.abs(callResultGold).sum()
-        print("L1 norm: %E" % L1norm)
-        print("Max absolute error: %E" % delta.max())
-        self.assertEqual(delta.max(), 0)
+        self.assertAlmostEqual(delta.max(), 0)
 
     def test_scalar(self):
-        flags = Flags()
-
-        # Compile the inner function
-        global cnd_jitted
-        cr1 = compile_isolated(cnd, (types.float64,))
-        cnd_jitted = cr1.entry_point
-        # Manually type the compiled function for calling into
-        tyctx = cr1.typing_context
-        ctx = cr1.target_context
-        signature = typing.make_concrete_template("cnd_jitted", cnd_jitted,
-                                                  [cr1.signature])
-        tyctx.insert_user_function(cnd_jitted, signature)
-
-        # Compile the outer function
-        array = types.Array(types.float64, 1, 'C')
-        argtys = (array,) * 5 + (types.float64, types.float64)
-        cr2 = compile_extra(tyctx, ctx, blackscholes_scalar_jitted,
-                            args=argtys, return_type=None, flags=flags,
-                            locals={})
-        jitted_bs = cr2.entry_point
-
         OPT_N = 400
-        iterations = 10
 
         callResultGold = np.zeros(OPT_N)
         putResultGold = np.zeros(OPT_N)
@@ -189,13 +114,10 @@ def test_scalar(self):
 
         args = stockPrice, optionStrike, optionYears, RISKFREE, VOLATILITY
 
-        blackscholes_scalar(callResultGold, putResultGold, *args)
-        jitted_bs(callResultNumba, putResultNumba, *args)
+        blackscholes_scalar.py_func(callResultGold, putResultGold, *args)
+        blackscholes_scalar(callResultNumba, putResultNumba, *args)
 
         delta = np.abs(callResultGold - callResultNumba)
-        L1norm = delta.sum() / np.abs(callResultGold).sum()
-        print("L1 norm: %E" % L1norm)
-        print("Max absolute error: %E" % delta.max())
         self.assertAlmostEqual(delta.max(), 0)