BUG: Compute drawdown in a vectorized manner

Fixes https://github.com/kernc/backtesting.py/issues/46

backtesting/_plotting.py

@@ -224,9 +224,7 @@ return this.labels[index] || "";
             dd_start = dd_end = equity.index[0]
             timedelta = 0
         else:
-            dd_end = (equity[argmax:] > equity[dd_start]).idxmax()
+            dd_end = argmax
-            if dd_end == argmax:
-                dd_end = index[-1]
             if is_datetime_index and omit_missing:
                 # "Calendar" duration
                 timedelta = df.datetime.iloc[dd_end] - df.datetime.iloc[dd_start]

backtesting/backtesting.py

@@ -850,26 +850,23 @@ class Backtest:
                                                       for params in param_combos)
                 if stats['# Trades']]
 
+    @staticmethod
+    def _compute_drawdown_duration_peaks(dd: pd.Series):
+        iloc = np.unique(np.r_[(dd == 0).values.nonzero()[0], len(dd) - 1])
+        iloc = pd.Series(iloc, index=dd.index[iloc])
+        df = iloc.to_frame('iloc').assign(prev=iloc.shift())
+        df = df[df['iloc'] > df['prev'] + 1].astype(int)
+        # If no drawdown since no trade, avoid below for pandas sake and return nan series
+        if not len(df):
+            return (dd.replace(0, np.nan),) * 2
+        df['duration'] = df['iloc'].map(dd.index.__getitem__) - df['prev'].map(dd.index.__getitem__)
+        df['peak_dd'] = df.apply(lambda row: dd.iloc[row['prev']:row['iloc'] + 1].max(), axis=1)
+        df = df.reindex(dd.index)
+        return df['duration'], df['peak_dd']
+
     def _compute_stats(self, broker: _Broker, strategy: Strategy) -> pd.Series:
         data = self._data
 
-        def _drawdown_duration_peaks(dd, index):
-            # XXX: possible to vectorize any of this?
-            durations = [np.nan] * len(dd)
-            peaks = [np.nan] * len(dd)
-            i = 0
-            for j in range(1, len(dd)):
-                if dd[j] == 0:
-                    if dd[j - 1] != 0:
-                        durations[j - 1] = index[j] - index[i]
-                        peaks[j - 1] = dd[i:j].max()
-                    i = j
-            j = len(dd) - 1
-            if dd[j - 1] != 0:
-                durations[j] = index[j] - index[i]
-                peaks[j] = dd[i:j].max()
-            return pd.Series(durations), pd.Series(peaks)
-
         df = pd.DataFrame()
         df['Equity'] = pd.Series(broker.log.equity).bfill().fillna(broker._cash)
         equity = df.Equity.values
@@ -882,8 +879,8 @@ class Backtest:
         pl = df['P/L']
         df['Returns'] = returns = pl.dropna() / equity[exits.dropna().values.astype(int)]
         df['Drawdown'] = dd = 1 - equity / np.maximum.accumulate(equity)
-        dd_dur, dd_peaks = _drawdown_duration_peaks(dd, data.index)
+        dd_dur, dd_peaks = self._compute_drawdown_duration_peaks(pd.Series(dd, index=data.index))
-        df['Drawdown Duration'] = dd_dur
+        df['Drawdown Duration'] = dd_dur.values
         dd_dur = df['Drawdown Duration']
 
         df.index = data.index

backtesting/test/_test.py

@@ -219,6 +219,12 @@ class TestBacktest(TestCase):
         self.assertEqual(str(bt.run()._strategy), SmaCross.__name__)
         self.assertEqual(str(bt.run(fast=11)._strategy), SmaCross.__name__ + '(fast=11)')
 
+    def test_compute_drawdown(self):
+        dd = pd.Series([0, 1, 7, 0, 4, 0, 0])
+        durations, peaks = Backtest._compute_drawdown_duration_peaks(dd)
+        np.testing.assert_array_equal(durations, pd.Series([3, 2], index=[3, 5]).reindex(dd.index))
+        np.testing.assert_array_equal(peaks, pd.Series([7, 4], index=[3, 5]).reindex(dd.index))
+
     def test_compute_stats(self):
         stats = Backtest(GOOG, SmaCross).run()
         # Pandas compares in 'almost equal' manner