ENH: Add bt.plot(resample=) and auto-downsample large data

Fixes #35

Author: kernc

backtesting/_plotting.py

@@ -85,14 +85,65 @@ def lightness(color, lightness=.94):
     return color.to_rgb()
 
 
+_MAX_CANDLES = 10000
+
+
+def _maybe_resample_data(resample_rule, df, indicators, equity_data, trades):
+    if not resample_rule or len(df) < _MAX_CANDLES:
+        return df, indicators, equity_data, trades
+
+    if isinstance(resample_rule, str):
+        freq = resample_rule
+    else:
+        from_index = dict(day=-2, hour=-6, minute=1, second=0, millisecond=0,
+                          microsecond=0, nanosecond=0)[df.index.resolution]
+        FREQS = ('1T', '5T', '10T', '15T', '30T', '1H', '2H', '4H', '8H', '1D', '1W', '1M')
+        freq = next((f for f in FREQS[from_index:]
+                     if len(df.resample(f)) <= _MAX_CANDLES), FREQS[-1])
+        warnings.warn("Data contains too many candlesticks to plot; downsampling to {!r}. "
+                      "See `Backtest.plot(resample=...)`".format(freq))
+
+    from .lib import OHLCV_AGG, TRADES_AGG, _EQUITY_AGG
+    df = df.resample(freq, label='right').agg(OHLCV_AGG).dropna()
+
+    # XXX: copy(True) pandas bug https://github.com/pandas-dev/pandas/issues/31710
+    indicators = [_Indicator(i.s.copy(True).resample(freq).mean().dropna().reindex(df.index),
+                             **dict(i._opts,
+                                    # HACK: override `data` for its index
+                                    data=pd.Series(np.nan, index=df.index)))
+                  for i in indicators]
+    assert not indicators or indicators[0].s.index.equals(df.index)
+
+    equity_data = equity_data.resample(freq, label='right').agg(_EQUITY_AGG).dropna(how='all')
+    assert equity_data.index.equals(df.index)
+
+    def _weighted_returns(s, trades=trades):
+        df = trades.loc[s.index]
+        return ((df['Size'].abs() * df['ReturnPct']) / df['Size'].abs().sum()).sum()
+
+    trades = trades.assign(count=1).resample(freq, on='ExitTime', label='right').agg(dict(
+        TRADES_AGG,
+        ReturnPct=_weighted_returns,
+        count='sum',
+        # XXX: Can this prettier?
+        EntryBar=(lambda s, trades=trades, index=df.index:
+                  index.get_loc(trades.loc[s.index]['EntryTime'].mean(), method='nearest')),
+        ExitBar=(lambda s, trades=trades, index=df.index:
+                 index.get_loc(trades.loc[s.index]['ExitTime'].mean(), method='nearest')),
+    )).dropna()
+
+    return df, indicators, equity_data, trades
+
+
 def plot(*, results: pd.Series,
          df: pd.DataFrame,
          indicators: List[_Indicator],
          filename='', plot_width=None,
          plot_equity=True, plot_pl=True,
          plot_volume=True, plot_drawdown=False,
          smooth_equity=False, relative_equity=True,
-         superimpose=True, show_legend=True, open_browser=True):
+         superimpose=True, resample=True,
+         show_legend=True, open_browser=True):
     """
     Like much of GUI code everywhere, this is a mess.
     """
@@ -111,15 +162,19 @@ def plot(*, results: pd.Series,
     trades = results['_trades']
 
     plot_volume = plot_volume and not df.Volume.isnull().all()
-    time_resolution = getattr(df.index, 'resolution', None)
     is_datetime_index = df.index.is_all_dates
 
     from .lib import OHLCV_AGG
     # ohlc df may contain many columns. We're only interested in, and pass on to Bokeh, these
     df = df[list(OHLCV_AGG.keys())].copy(deep=False)
+
+    # Limit data to max_candles
+    if is_datetime_index:
+        df, indicators, equity_data, trades = _maybe_resample_data(
+            resample, df, indicators, equity_data, trades)
+
     df.index.name = None  # Provides source name @index
     df['datetime'] = df.index  # Save original, maybe datetime index
-
     df = df.reset_index(drop=True)
     equity_data = equity_data.reset_index(drop=True)
     index = df.index
@@ -224,6 +279,10 @@ def _plot_equity_section():
                 x1, x2 = dd_end - 1, dd_end
                 y, y1, y2 = equity[dd_start], equity[x1], equity[x2]
                 dd_end -= (1 - (y - y1) / (y2 - y1)) * (dd_end - x1)  # y = a x + b
+                # If _plot_resample_data() was applied,
+                # the agg'd equity might have "stretched" the calculation
+                # XXX: test this?
+                dd_end = min(dd_end, equity.index[-1])
 
         if smooth_equity:
             interest_points = pd.Index([
@@ -321,11 +380,15 @@ def _plot_pl_section():
         trade_source.add(returns_long, 'returns_long')
         trade_source.add(returns_short, 'returns_short')
         trade_source.add(size, 'marker_size')
+        if 'count' in trades:
+            trade_source.add(trades['count'], 'count')
         r1 = fig.scatter('index', 'returns_long', source=trade_source, fill_color=cmap,
                          marker='triangle', line_color='black', size='marker_size')
         r2 = fig.scatter('index', 'returns_short', source=trade_source, fill_color=cmap,
                          marker='inverted_triangle', line_color='black', size='marker_size')
         tooltips = [("Size", "@size{0,0}")]
+        if 'count' in trades:
+            tooltips.append(("Count", "@count{0,0}"))
         set_tooltips(fig, tooltips + [("P/L", "@returns_long{+0.[000]%}")],
                      vline=False, renderers=[r1])
         set_tooltips(fig, tooltips + [("P/L", "@returns_short{+0.[000]%}")],
@@ -346,8 +409,9 @@ def _plot_volume_section():
 
     def _plot_superimposed_ohlc():
         """Superimposed, downsampled vbars"""
+        time_resolution = pd.DatetimeIndex(df['datetime']).resolution
         resample_rule = (superimpose if isinstance(superimpose, str) else
-                         dict(day='W',
+                         dict(day='M',
                               hour='D',
                               minute='H',
                               second='T',

backtesting/backtesting.py

@@ -141,7 +141,7 @@ def init():
 
         value = _Indicator(value, name=name, plot=plot, overlay=overlay,
                            color=color, scatter=scatter,
-                           # lib.resample_apply() uses this:
+                           # _Indicator.s Series accessor uses this:
                            data=self.data)
         self._indicators.append(value)
         return value
@@ -952,7 +952,8 @@ def __init__(self,
                           stacklevel=2)
             data = data.sort_index()
         if not data.index.is_all_dates:
-            warnings.warn('Data index is not datetime. Assuming simple periods.',
+            warnings.warn('Data index is not datetime. Assuming simple periods, '
+                          'but `pd.DateTimeIndex` is advised.',
                           stacklevel=2)
 
         self._data = data   # type: pd.DataFrame
@@ -1006,6 +1007,10 @@ def run(self, **kwargs) -> pd.Series:
                 # Next tick, a moment before bar close
                 strategy.next()
 
+        # Set data back to full length
+        # for future `indicator._opts['data'].index` calls to work
+        data._set_length(len(self._data))
+
         self._results = self._compute_stats(broker, strategy)
         return self._results
 
@@ -1250,6 +1255,7 @@ def plot(self, *, results: pd.Series = None, filename=None, plot_width=None,
              plot_volume=True, plot_drawdown=False,
              smooth_equity=False, relative_equity=True,
              superimpose: Union[bool, str] = True,
+             resample=True,
              show_legend=True, open_browser=True):
         """
         Plot the progression of the last backtest run.
@@ -1287,16 +1293,31 @@ def plot(self, *, results: pd.Series = None, filename=None, plot_width=None,
         If `relative_equity` is `True`, scale and label equity graph axis
         with return percent, not absolute cash-equivalent values.
 
-        If `superimpose` is `True`, superimpose downsampled candlesticks
-        over the original candlestick chart. Default downsampling is:
-        weekly for daily data, daily for hourly data, hourly for minute data,
-        and minute for second and sub-second data.
-        `superimpose` can also be a string, in which case it is a valid
-        [Pandas offset string], such as `'5T'` or `'5min'`.
+        If `superimpose` is `True`, superimpose larger-timeframe candlesticks
+        over the original candlestick chart. Default downsampling rule is:
+        monthly for daily data, daily for hourly data, hourly for minute data,
+        and minute for (sub-)second data.
+        `superimpose` can also be a valid [Pandas offset string],
+        such as `'5T'` or `'5min'`, in which case this frequency will be
+        used to superimpose.
         Note, this only works for data with a datetime index.
 
+        If `resample` is `True`, the OHLC data is resampled in a way that
+        makes the upper number of candles for Bokeh to plot limited to 10_000.
+        This may, in situations of overabundant data,
+        improve plot's interactive performance and avoid browser's
+        `Javascript Error: Maximum call stack size exceeded` or similar.
+        Equity & dropdown curves and individual trades data is,
+        likewise, [reasonably _aggregated_][TRADES_AGG].
+        `resample` can also be a [Pandas offset string],
+        such as `'5T'` or `'5min'`, in which case this frequency will be
+        used to resample, overriding above numeric limitation.
+        Note, all this only works for data with a datetime index.
+
         [Pandas offset string]: \
-            http://pandas.pydata.org/pandas-docs/stable/timeseries.html#offset-aliases
+            https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#dateoffset-objects
+
+        [TRADES_AGG]: lib.html#backtesting.lib.TRADES_AGG
 
         If `show_legend` is `True`, the resulting plot graphs will contain
         labeled legends.
@@ -1322,5 +1343,6 @@ def plot(self, *, results: pd.Series = None, filename=None, plot_width=None,
             smooth_equity=smooth_equity,
             relative_equity=relative_equity,
             superimpose=superimpose,
+            resample=resample,
             show_legend=show_legend,
             open_browser=open_browser)

backtesting/lib.py

@@ -40,6 +40,31 @@
     df.resample('4H', label='right').agg(OHLCV_AGG)
 """
 
+TRADES_AGG = OrderedDict((
+    ('Size', 'sum'),
+    ('EntryBar', 'first'),
+    ('ExitBar', 'last'),
+    ('EntryPrice', 'mean'),
+    ('ExitPrice', 'mean'),
+    ('PnL', 'sum'),
+    ('ReturnPct', 'mean'),
+    ('EntryTime', 'first'),
+    ('ExitTime', 'last'),
+    ('Duration', 'sum'),
+))
+"""Dictionary of rules for aggregating resampled trades data,
+e.g.
+
+    stats['_trades'].resample('1D', on='ExitTime',
+                              label='right').agg(TRADES_AGG)
+"""
+
+_EQUITY_AGG = {
+    'Equity': 'mean',
+    'DrawdownPct': 'max',
+    'DrawdownDuration': 'max',
+}
+
 
 def barssince(condition: Sequence[bool], default=np.inf) -> int:
     """

backtesting/test/_test.py

@@ -8,6 +8,7 @@
 from runpy import run_path
 from tempfile import NamedTemporaryFile, gettempdir
 from unittest import TestCase
+from unittest.mock import patch
 
 import numpy as np
 import pandas as pd
@@ -500,6 +501,17 @@ def next(self):
             # Give browser time to open before tempfile is removed
             time.sleep(1)
 
+    def test_resample(self):
+        bt = Backtest(GOOG, SmaCross)
+        bt.run()
+        import backtesting._plotting
+        with _tempfile() as f,\
+                patch.object(backtesting._plotting, '_MAX_CANDLES', 10),\
+                self.assertWarns(UserWarning):
+            bt.plot(filename=f, resample=True)
+            # Give browser time to open before tempfile is removed
+            time.sleep(1)
+
     def test_indicator_color(self):
         class S(Strategy):
             def init(self):