1 files changed, 27 insertions, 17 deletions
diff --git a/python/analytics/cms_spread.py b/python/analytics/cms_spread.py
index e518c498..402d0413 100644
--- a/python/analytics/cms_spread.py
+++ b/python/analytics/cms_spread.py
@@ -2,7 +2,8 @@ import numpy as np
 import matplotlib.pyplot as plt
 from math import exp, sqrt, log
 from .black import bachelier, cnd_erf
-from numba import cfunc, int64, float64, boolean, types
+from numba import (cfunc, types, jit, float64, boolean,
+                   optional, vectorize)
 from quantlib.time.api import (
     Date, Period, Days, Months, Years, UnitedStates, Actual365Fixed, Following)
 from quantlib.termstructures.yields.api import YieldTermStructure
@@ -20,26 +21,35 @@ from quantlib.quotes import SimpleQuote
 from quantlib.math.matrix import Matrix
 from scipy.interpolate import RectBivariateSpline
 from db import dbconn
-from numba import cfunc
 
-# @jit(float64(float64, float64, float64, float64, float64, float64, float64,
-#              float64, float64, boolean), cache=True, nopython=True)
-# def h(z, K, S1, S2, mu_x, mu_y, sigma_x, sigma_y, rho, call=True):
-#     # z = (y - mu_y) / sigma_y
-#     u1 = mu_x + rho * sigma_x * z
-#     Ktilde = K + S2 * exp(mu_y + sigma_y * z)
-#     u2 = log(Ktilde / S1)
+@vectorize([float64(float64, float64, float64, float64, float64, float64, float64,
+                    float64, float64)], cache=True, nopython=True)
+def h_call(z, K, S1, S2, mu_x, mu_y, sigma_x, sigma_y, rho):
+    # z = (y - mu_y) / sigma_y
+    u1 = mu_x + rho * sigma_x * z
+    Ktilde = K + S2 * exp(mu_y + sigma_y * z)
+    u2 = log(Ktilde / S1)
 
-#     v = sigma_x * sqrt(1 - rho * rho)
-#     v2 = sigma_x * sigma_x * (1 - rho * rho)
-#     if call:
-#         x = (u1 - u2) / v
-#         return 0.5 * (S1 * exp(u1 + 0.5 * v2) * cnd_erf(x + v) - Ktilde * cnd_erf(x))
-#     else:
-#         x = (u2 - u1) / v
-#         return 0.5 * (Ktilde * cnd_erf(x) - S1 * exp(u1 + 0.5 * v2) * cnd_erf(x - v))
+    v = sigma_x * sqrt(1 - rho * rho)
+    v2 = sigma_x * sigma_x * (1 - rho * rho)
+    x = (u1 - u2) / v
+    return 0.5 * (S1 * exp(u1 + 0.5 * v2) * cnd_erf(x + v) - Ktilde * cnd_erf(x))
+
+@vectorize([float64(float64, float64, float64, float64, float64, float64, float64,
+                    float64, float64)], cache=True, nopython=True)
+def h_put(z, K, S1, S2, mu_x, mu_y, sigma_x, sigma_y, rho,):
+    # z = (y - mu_y) / sigma_y
+    u1 = mu_x + rho * sigma_x * z
+    Ktilde = K + S2 * exp(mu_y + sigma_y * z)
+    u2 = log(Ktilde / S1)
+
+    v = sigma_x * sqrt(1 - rho * rho)
+    v2 = sigma_x * sigma_x * (1 - rho * rho)
+    x = (u2 - u1) / v
+    return 0.5 * (Ktilde * cnd_erf(x) - S1 * exp(u1 + 0.5 * v2) * cnd_erf(x - v))
 
 sig = types.double(types.intc, types.CPointer(types.double))
+
 @cfunc(sig, cache=True, nopython=True)
 def h1(n, args):
     z = args[0]