fix tweak_curves

1 files changed, 37 insertions, 23 deletions

diff --git a/pyisda/curve.pyx b/pyisda/curve.pyx
index dd9d66b..49e2522 100644
--- a/pyisda/curve.pyx
+++ b/pyisda/curve.pyx
@@ -314,22 +314,28 @@ cdef class YieldCurve(Curve):
                                        <double>1, self._thisptr.fDayCountConv)
         return yc
 
+@cython.cdivision(True)
 cdef inline void tweak_curve(TCurve* sc, TCurve* sc_tweaked, double epsilon,
                              double* h, double* T, bint* mask, int n):
-    cdef h1, h2, t1, t2
+    ## We want to tweak in the forward space, so we convert the hazard rates
+    ## into forward rates and then back
+    cdef double h1, h2, t1, t2
     h1 = t1 = 0
+    cdef size_t i
     for i in range(n):
         h2 = sc.fArray[i].fRate
         t2 = T[i]
         h[i] = (h2 * t2 - h1 * t1) / (t2 - t1)
-        if mask[i]:
+        if mask is NULL or (mask is not NULL and mask[i]):
             h[i] *= (1 + epsilon)
-        h1, t1 = h2, t2
+        h1 = h2
+        t1 = t2
+
     t1 = 0
     cdef double c = 0
     for i in range(n):
         c += (T[i] - t1) * h[i]
-        sc_tweaked.fArray[i].fRate = c/T[i]
+        sc_tweaked.fArray[i].fRate = c / T[i]
         t1 = T[i]
 
 cdef class SpreadCurve(Curve):
@@ -428,9 +434,9 @@ cdef class SpreadCurve(Curve):
                                        <double>basis, dcc(day_count_conv))
         return sc
 
-    @cython.cdivision(True)
+    @cython.boundscheck(False)
     def tweak_curve(self, double epsilon, bint multiplicative=True,
-                    bint[:] mask=None, bint inplace=False):
+                    bint[::1] mask=None, bint inplace=False):
         """
         Tweak the survival curve in place.
 
@@ -444,30 +450,38 @@ cdef class SpreadCurve(Curve):
            If None (default), tweak everything, otherwise only tweak values
            in the mask.
         """
-        ## We want to tweak in the forward space, so we convert the hazard rates
-        ## into forward rates and then back
-        cdef double t1, h1, t2, h2
-        cdef TCurve* curve = self._thisptr
-        cdef TCurve* curve_tweaked
-        cdef SpreadCurve sc
-        cdef int num_items = curve.fNumItems
+        cdef:
+            TCurve* curve_tweaked
+            SpreadCurve sc
+            int num_items = self._thisptr.fNumItems
+            double *h
+            double *T
+            size_t i
+            bint* mask_ptr
+
         if not inplace:
             sc = SpreadCurve.__new__(SpreadCurve)
-            sc._thisptr = curve_tweaked = JpmcdsCopyCurve(curve)
+            curve_tweaked = JpmcdsCopyCurve(self._thisptr)
+            sc._thisptr = curve_tweaked
         else:
-            curve_tweaked = curve
-        cdef double* h = <double*>malloc(num_items * sizeof(double))
-        cdef double* T = <double*>malloc(num_items * sizeof(double))
+            sc = self
+            curve_tweaked = self._thisptr
+
+        h = <double*>malloc(num_items * sizeof(double))
+        T = <double*>malloc(num_items * sizeof(double))
+        for i in range(num_items):
+            T[i] = (curve_tweaked.fArray[i].fDate - curve_tweaked.fBaseDate) / 365.
+
         if mask is not None:
-            if mask.size != curve.fNumItems:
+            if mask.shape[0] != self._thisptr.fNumItems:
                 raise ValueError("mask size need to be the same as the number of Items")
-        tweak_curve(curve, curve_tweaked, epsilon, h, T, &mask[0], num_items)
+            mask_ptr = &mask[0]
+        else:
+            mask_ptr = NULL
+        tweak_curve(self._thisptr, curve_tweaked, epsilon, h, T, mask_ptr, num_items)
         free(h)
         free(T)
-        if inplace:
-            return self
-        else:
-            return sc
+        return sc
 
 cdef class CreditIndex:
     def __init__(self, base_date, start_date, end_date, end_dates, curves):