improve survival_matrix

1 files changed, 15 insertions, 5 deletions

diff --git a/pyisda/credit_index.pyx b/pyisda/credit_index.pyx
index 27511f9..b59e19c 100644
--- a/pyisda/credit_index.pyx
+++ b/pyisda/credit_index.pyx
@@ -23,7 +23,7 @@ cimport numpy as np
 np.import_array()
 import pandas as pd
 import warnings
-
+cimport cython
 
 ctypedef TFeeLeg* TFeeLeg_ptr
 ctypedef TContingentLeg* TContingentLeg_ptr
@@ -570,23 +570,33 @@ cdef class CreditIndex(CurveList):
 
 
 
-    def survival_matrix(self, TDate[:] schedule, double epsilon=0.):
+    def survival_matrix(self, const TDate[::1] schedule=None, double epsilon=0.):
         cdef:
             shared_ptr[TCurve] sc
             pair[CurveName, size_t] p
             size_t i
             np.npy_intp[2] n
+            const TDate* schedule_ptr
+
+        if schedule is None:
+            schedule_ptr = self._maturities.const_data()
+            n[1] = self._maturities.size()
+        else:
+            schedule_ptr = &schedule[0]
+            n[1] = schedule.shape[0]
         n[0] = self._curves.size()
-        n[1] = schedule.shape[0]
+
         cdef:
-            np.ndarray[np.float64_t, ndim=2] sp = np.PyArray_EMPTY(2, n, np.NPY_DOUBLE, 1)
+            np.ndarray sp = np.PyArray_EMPTY(2, n, np.NPY_DOUBLE, 1)
             np.ndarray tickers = np.PyArray_EMPTY(1, n, np.NPY_OBJECT, 1)
+            double[::1,:] sp_view = sp
 
         for p in self.names:
             sc = self._curves[p.second]
             tickers[p.second] = p.first.ticker
             for i in range(n[1]):
-                sp[p.second, i] = survival_prob(sc.get(), self.base_date, schedule[i], epsilon)
+                sp_view[p.second, i] = survival_prob(sc.get(), self.base_date,
+                                                     schedule_ptr[i], epsilon)
         return sp, tickers
 
 cdef unsigned long fill_mask(const TDate maturity, const vector[TDate]& maturities,