speed up survival_probability

1 files changed, 13 insertions, 7 deletions

diff --git a/pyisda/curve.pyx b/pyisda/curve.pyx
index a307b63..b0c0e29 100644
--- a/pyisda/curve.pyx
+++ b/pyisda/curve.pyx
@@ -794,6 +794,7 @@ cdef class SpreadCurve(Curve):
     cdef inline char* name(self) nogil:
         return self.buf.get() + self.offset_name
 
+    @cython.boundscheck(False)
     def survival_probability(self, d2, d1=None):
         """ computes the survival probability at a given date.
 
@@ -807,20 +808,25 @@ cdef class SpreadCurve(Curve):
         """
         cdef:
             const TCurve* curve = self.get_TCurve()
-            np.ndarray r
+            np.ndarray[np.float64_t, ndim=1] r
             c_datetime.date d
+            c_datetime.date[::1] date_vec
             size_t i
+            list r_list = []
 
         if d1 is None:
             if isinstance(d2, np.ndarray):
-                r = np.zeros_like(d2)
-                for i, e in enumerate(d2):
-                    r[i] = JpmcdsForwardZeroPrice(curve, curve.fBaseDate,
-                                                  pydate_to_TDate(d2[i]))
+                date_vec = d2
+                r = np.zeros(date_vec.shape[0])
+                for i in range(date_vec.shape[0]):
+                    r[i] = JpmcdsZeroPrice(curve, pydate_to_TDate(date_vec[i]))
                 return r
+            elif isinstance(d2, list):
+                for d in d2:
+                    r_list.append(JpmcdsZeroPrice(curve, pydate_to_TDate(d)))
+                return r_list
             else:
-                return JpmcdsForwardZeroPrice(curve, curve.fBaseDate,
-                                              pydate_to_TDate(d2))
+                return JpmcdsZeroPrice(curve, pydate_to_TDate(d2))
         else:
             return JpmcdsForwardZeroPrice(curve, pydate_to_TDate(d1),
                                           pydate_to_TDate(d2))