1 files changed, 18 insertions, 22 deletions
diff --git a/python/analytics/cms_spread.py b/python/analytics/cms_spread.py
index 7d34b7b1..32c16ed1 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 .tranche_functions import GHquad
 from math import exp, sqrt, log
-from .black import bachelier
+from .black import bachelier, cnd_erf
+from numba import jit, float64, boolean
 from quantlib.time.api import (
     Date, Period, Days, Months, Years, UnitedStates, Actual365Fixed, Following)
 from quantlib.termstructures.yields.api import YieldTermStructure
@@ -21,33 +22,28 @@ from quantlib.math.matrix import Matrix
 from scipy.interpolate import RectBivariateSpline
 from db import dbconn
 
+@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)
 
-def CMS_spread(T_alpha, X, beta, gamma):
-    Z, w = GHquad(100)
-    return np.inner(f(Z), w)
-
-
-def h(v, X, S_alpha_beta, mu_beta, sigma_alpha_beta, T_alpha):
-    r = (mu_beta - 0.5 * sigma_alpha_beta * sigma_alpha_beta) * T_alpha + \
-        sigma_alpha_beta * sqrt(T_alpha) * v
-    return X + S_alpha_beta * exp(r)
-
-
-def f(v, X, S_alpha_beta, S_alpha_gamma, mu_beta, mu_gamma, T_alpha, rho):
-    h = h(v, X, S_alpha_beta, mu_beta, sigma_alpha_beta, T_alpha)
-    u = rho * sigma_alpha_gamma * sqrt(T_alpha) * v
-    d = sigma_alpha_gamma * sqrt(T_alpha) * sqrt(1 - rho ** 2)
-    r = mu_gamma * T_alpha - 0.5 * rho * rho * sigma_alpha_gamma ** 2 * T_alpha + u
-    u0 = log(S_alpha_gamma / h) + u
-    u1 = u0 + (mu_gamma + (0.5 - rho ** 2) * sigma_alpha_gamma**2) * T_alpha
-    u2 = u0 + (mu_gamma - 0.5 * sigma_alpha_gamma**2) * T_alpha
-    return 0.5 * (S_alpha_gamma * exp(r) * cnd_erf(u1 / d) - h * cnd_erf(u2 / d))
+    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))
 
 
 def get_fixings(conn, tenor1, tenor2, fixing_date=None):
     if fixing_date:
         sql_str = f'SELECT "{tenor1}y" ,"{tenor2}y" FROM USD_swap_fixings ' \
-                  'WHERE fixing_date=%s'
+            'WHERE fixing_date=%s'
         with conn.cursor() as c:
             c.execute(sql_str, (fixing_date,))
             try: