also vectorize when providing forward date

1 files changed, 32 insertions, 15 deletions

diff --git a/pyisda/curve.pyx b/pyisda/curve.pyx
index 312c7d6..c88a9bf 100644
--- a/pyisda/curve.pyx
+++ b/pyisda/curve.pyx
@@ -1,7 +1,7 @@
 
 from cython.operator import dereference as deref, preincrement as preinc
 from cpython.bytes cimport PyBytes_FromStringAndSize, PyBytes_AS_STRING, PyBytes_GET_SIZE
-from libc.math cimport log1p, log, exp, isnan
+from libc.math cimport log1p, log, exp, isnan, NAN
 from libc.string cimport strcpy, strncpy, strlen
 from .date cimport (JpmcdsStringToDateInterval, pydate_to_TDate, dcc,
                    JpmcdsDateIntervalToFreq, JpmcdsDateFwdThenAdjust, TDate_to_pydate,
@@ -528,30 +528,47 @@ cdef class YieldCurve(Curve):
             np.ndarray[np.float64_t, ndim=1] r
             c_datetime.date d
             size_t i
+            np.npy_intp n
             long[::1] np_dates
             c_datetime.date[::1] np_datetimes
             TDate date_c
             long scale = 1000000000 * 86400
+            double df = NAN
 
+        if isinstance(d2, np.ndarray):
+            np_datetimes = d2
+            n = np_datetimes.shape[0]
+            r = np.PyArray_EMPTY(1, &n, np.NPY_DOUBLE, 0)
+            for i in range(n):
+                r[i] = JpmcdsZeroPrice(curve, pydate_to_TDate(np_datetimes[i]))
+        elif isinstance(d2, pd.Series):
+            np_dates = d2.values.view("int")
+            r = np.PyArray_EMPTY(1, &np_dates.shape[0], np.NPY_DOUBLE, 0)
+            for i in range(r.shape[0]):
+                date_c = np_dates[i] / scale + 134774
+                r[i] = JpmcdsZeroPrice(curve, date_c)
+        elif isinstance(d2, list):
+            n = len(d2)
+            r = np.PyArray_EMPTY(1, &n, np.NPY_DOUBLE, 0)
+            i = 0
+            for d in d2:
+                r[i] = JpmcdsZeroPrice(curve, pydate_to_TDate(d))
+                i += 1
+        else:
+            df = JpmcdsZeroPrice(curve, pydate_to_TDate(d2))
         if d1 is None:
-            if isinstance(d2, np.ndarray):
-                r = np.zeros(d2.shape[0], dtype=np.float64)
-                np_datetimes = d2
-                for i in range(np_datetimes.shape[0]):
-                    r[i] = JpmcdsZeroPrice(curve, pydate_to_TDate(np_datetimes[i]))
+            if isnan(df):
                 return r
-            elif isinstance(d2, pd.Series):
-                np_dates = d2.values.view("int")
-                r = np.zeros(np_dates.shape[0], dtype=np.float64)
+            else:
+                return df
+        else:
+            if isnan(df):
+                df = JpmcdsZeroPrice(curve, pydate_to_TDate(d1))
                 for i in range(r.shape[0]):
-                    date_c = np_dates[i] / scale + 134774
-                    r[i] = JpmcdsZeroPrice(curve, date_c)
+                    r[i] /= df
                 return r
             else:
-                return JpmcdsZeroPrice(curve, pydate_to_TDate(d2))
-        else:
-            return JpmcdsForwardZeroPrice(curve, pydate_to_TDate(d1),
-                                          pydate_to_TDate(d2))
+                return df / JpmcdsZeroPrice(curve, pydate_to_TDate(d1))
 
     @property
     def dates(self):