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

kernc · kernc · commit 7031f5ef8a64 · 2020-03-23T23:48:57.000+01:00
Fixes #35
diff --git a/backtesting/_plotting.py b/backtesting/_plotting.py
@@ -85,14 +85,72 @@ def lightness(color, lightness=.94):
     return color.to_rgb()
 
 
+_MAX_CANDLES = 10000
+
+
+def _maybe_resample_data(resample_rule, df, indicators, equity_data, trades):
+    if isinstance(resample_rule, str):
+        freq = resample_rule
+    else:
+        if len(df) < _MAX_CANDLES:
+            return df, indicators, equity_data, trades
+
+        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()
+
+    indicators = [_Indicator(i.df.resample(freq, label='right').mean()
+                             .dropna().reindex(df.index).values.T,
+                             **dict(i._opts, name=i.name,
+                                    # HACK: override `data` for its index
+                                    data=pd.Series(np.nan, index=df.index)))
+                  for i in indicators]
+    assert not indicators or indicators[0].df.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()
+
+    def _group_trades(column):
+        def f(s, new_index=df.index.astype(np.int64), bars=trades[column]):
+            if s.size:
+                # Via int64 because on pandas recently broken datetime
+                mean_time = int(bars.loc[s.index].view('i8').mean())
+                new_bar_idx = new_index.get_loc(mean_time, method='nearest')
+                return new_bar_idx
+        return f
+
+    if len(trades):  # Avoid pandas "resampling on Int64 index" error
+        trades = trades.assign(count=1).resample(freq, on='ExitTime', label='right').agg(dict(
+            TRADES_AGG,
+            ReturnPct=_weighted_returns,
+            count='sum',
+            EntryBar=_group_trades('EntryTime'),
+            ExitBar=_group_trades('ExitTime'),
+        )).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 +169,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 +286,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 +387,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 +416,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',
diff --git a/backtesting/autoscale_cb.js b/backtesting/autoscale_cb.js
@@ -1,40 +1,13 @@
-if (!window._bt_extremes)
-    window._bt_extremes = function (arr, initial, agg_func) {
-        const CHUNK = 32768;
-        let extreme = initial;
-        for (let i = 0, len = arr.length; i < len; i += CHUNK) {
-            const subarr = CHUNK >= len ? arr : arr.slice(i, i + CHUNK);
-            extreme = agg_func(extreme, agg_func.apply(null, subarr));
-        }
-        return extreme;
-    };
-
-if (!window._bt_bin_search)
-    window._bt_bin_search = function (index, value) {
-        let mid,
-            min = 0,
-            max = index.length - 1;
-
-        while (min < max) {
-            mid = (min + max) / 2 | 0;
-            if (index[mid] < value)
-                min = mid + 1;
-            else
-                max = mid - 1;
-        }
-        return min;
-    };
-
-if (!window._bt_scale_range)
-    window._bt_scale_range = function (range, highs, lows) {
-        const max = _bt_extremes(highs, -Infinity, Math.max),
-              min = lows && _bt_extremes(lows, Infinity, Math.min);
+if (!window._bt_scale_range) {
+    window._bt_scale_range = function (range, min, max, pad) {
+        "use strict";
         if (min !== Infinity && max !== -Infinity) {
-            const pad = (max - min) * .03;
+            pad = pad ? (max - min) * .03 : 0;
             range.start = min - pad;
             range.end = max + pad;
-        }
+        } else console.error('backtesting: scale range error:', min, max, range);
     };
+}
 
 clearTimeout(window._bt_autoscale_timeout);
 
@@ -45,20 +18,18 @@ window._bt_autoscale_timeout = setTimeout(function () {
      * @variable ohlc_range `fig_ohlc.y_range`.
      * @variable volume_range `fig_volume.y_range`.
      */
+    "use strict";
+
+    let i = Math.max(Math.floor(cb_obj.start), 0),
+        j = Math.min(Math.ceil(cb_obj.end), source.data['ohlc_high'].length);
 
-    let index = source.data['index'],
-        i = Math.max(_bt_bin_search(index, cb_obj.start) - 1, 0),
-        j = Math.min(_bt_bin_search(index, cb_obj.end) + 1, index.length);
+    let max = Math.max.apply(null, source.data['ohlc_high'].slice(i, j)),
+        min = Math.min.apply(null, source.data['ohlc_low'].slice(i, j));
+    _bt_scale_range(ohlc_range, min, max, true);
 
-    _bt_scale_range(
-        ohlc_range,
-        source.data['ohlc_high'].slice(i, j),
-        source.data['ohlc_low'].slice(i, j));
-    try {
-        _bt_scale_range(
-            volume_range,
-            source.data['Volume'].slice(i, j),
-            0);
-    } catch (e) {}  // volume_range may be undefined
+    if (volume_range) {
+        max = Math.max.apply(null, source.data['Volume'].slice(i, j));
+        _bt_scale_range(volume_range, 0, max * 1.03, false);
+    }
 
 }, 50);
diff --git a/backtesting/backtesting.py b/backtesting/backtesting.py
@@ -149,7 +149,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
@@ -977,7 +977,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
@@ -1031,6 +1032,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
 
@@ -1275,6 +1280,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.
@@ -1312,16 +1318,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.
@@ -1347,5 +1368,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)
diff --git a/backtesting/lib.py b/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:
     """
diff --git a/backtesting/test/_test.py b/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
@@ -417,6 +418,7 @@ def test_params(self):
                           plot_pl=False,
                           plot_drawdown=True,
                           superimpose=False,
+                          resample='1W',
                           smooth_equity=False,
                           relative_equity=False,
                           show_legend=False).items():
@@ -499,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):
