From daac83c73b631b7278062286037624271e20a73f Mon Sep 17 00:00:00 2001 From: Kevin Taylor <2325494+tkdtaylor@users.noreply.github.com> Date: Mon, 7 Apr 2025 21:39:58 -0400 Subject: [PATCH] Add shape and size for scatter plot --- backtesting/_plotting.py | 38 ++++++++++++++---- backtesting/backtesting.py | 9 +++++ backtesting/test/_test.py | 80 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 120 insertions(+), 7 deletions(-) diff --git a/backtesting/_plotting.py b/backtesting/_plotting.py index faccc61d..e348c0ac 100644 --- a/backtesting/_plotting.py +++ b/backtesting/_plotting.py @@ -522,6 +522,16 @@ def _plot_ohlc_trades(): legend_label=f'Trades ({len(trades)})', line_width=8, line_alpha=1, line_dash='dotted') + MARKER_FUNCTIONS = { + 'circle': lambda fig, **kwargs: fig.scatter(marker='circle', **kwargs), + 'square': lambda fig, **kwargs: fig.scatter(marker='square', **kwargs), + 'triangle': lambda fig, **kwargs: fig.scatter(marker='triangle', **kwargs), + 'diamond': lambda fig, **kwargs: fig.scatter(marker='diamond', **kwargs), + 'cross': lambda fig, **kwargs: fig.scatter(marker='cross', **kwargs), + 'x': lambda fig, **kwargs: fig.scatter(marker='x', **kwargs), + 'star': lambda fig, **kwargs: fig.scatter(marker='star', **kwargs), + } + def _plot_indicators(): """Strategy indicators""" @@ -563,7 +573,19 @@ def __eq__(self, other): tooltips = [] colors = value._opts['color'] colors = colors and cycle(_as_list(colors)) or ( - cycle([next(ohlc_colors)]) if is_overlay else colorgen()) + cycle([next(ohlc_colors)]) if is_overlay else colorgen() + ) + + marker = value._opts.get('marker', 'circle') + if marker not in MARKER_FUNCTIONS: + warnings.warn(f"Unknown marker type '{marker}', falling back to 'circle'") + marker = 'circle' + value._opts['marker'] = marker + marker_func = MARKER_FUNCTIONS[marker] + + marker_size = value._opts.get('marker_size') + if marker_size is None: + marker_size = BAR_WIDTH / 2 * (.9 if is_overlay else .6) if isinstance(value.name, str): tooltip_label = value.name @@ -582,11 +604,12 @@ def __eq__(self, other): if is_overlay: ohlc_extreme_values[source_name] = arr if is_scatter: - fig.circle( - 'index', source_name, source=source, + marker_func( + fig, + x='index', y=source_name, source=source, legend_label=legend_labels[j], color=color, line_color='black', fill_alpha=.8, - radius=BAR_WIDTH / 2 * .9) + size=marker_size) else: fig.line( 'index', source_name, source=source, @@ -594,10 +617,11 @@ def __eq__(self, other): line_width=1.3) else: if is_scatter: - r = fig.circle( - 'index', source_name, source=source, + r = marker_func( + fig, + x='index', y=source_name, source=source, legend_label=legend_labels[j], color=color, - radius=BAR_WIDTH / 2 * .6) + size=marker_size) else: r = fig.line( 'index', source_name, source=source, diff --git a/backtesting/backtesting.py b/backtesting/backtesting.py index a85f5c9d..a938b641 100644 --- a/backtesting/backtesting.py +++ b/backtesting/backtesting.py @@ -74,6 +74,7 @@ def _check_params(self, params): def I(self, # noqa: E743 func: Callable, *args, name=None, plot=True, overlay=None, color=None, scatter=False, + marker='circle', marker_size=None, **kwargs) -> np.ndarray: """ Declare an indicator. An indicator is just an array of values @@ -106,6 +107,13 @@ def I(self, # noqa: E743 If `scatter` is `True`, the plotted indicator marker will be a circle instead of a connected line segment (default). + `marker` sets the marker shape for scatter plots. Available options: + 'circle', 'square', 'triangle', 'diamond', 'cross', 'x', 'star'. + Default is 'circle'. + + `marker_size` sets the size of scatter plot markers. If None, + defaults to a size relative to the bar width. + Additional `*args` and `**kwargs` are passed to `func` and can be used for parameters. @@ -173,6 +181,7 @@ def _format_name(name: str) -> str: value = _Indicator(value, name=name, plot=plot, overlay=overlay, color=color, scatter=scatter, + marker=marker, marker_size=marker_size, # _Indicator.s Series accessor uses this: index=self.data.index) self._indicators.append(value) diff --git a/backtesting/test/_test.py b/backtesting/test/_test.py index 366d54c4..26c5477a 100644 --- a/backtesting/test/_test.py +++ b/backtesting/test/_test.py @@ -32,6 +32,7 @@ resample_apply, ) from backtesting.test import BTCUSD, EURUSD, GOOG, SMA +import itertools SHORT_DATA = GOOG.iloc[:20] # Short data for fast tests with no indicator lag @@ -1138,3 +1139,82 @@ def test_optimize_datetime_index_with_timezone(self): data.index = data.index.tz_localize('Asia/Kolkata') res = Backtest(data, SmaCross).optimize(fast=range(2, 3), slow=range(4, 5)) self.assertGreater(res['# Trades'], 0) + + +class TestPlotting(unittest.TestCase): + def setUp(self): + self.data = GOOG.copy() + + def test_marker_shapes_and_sizes(self): + class MarkerStrategy(Strategy): + shapes = [] + invalid_checks = [] + default_checks = [] + + def init(self): + # Test all marker shapes with different sizes + markers = ['circle', 'square', 'triangle', 'diamond', 'cross', 'x', 'star'] + sizes = [6, 12, 18] # Test different sizes + + for i, (marker, size) in enumerate(itertools.product(markers, sizes)): + # Copy data.Close and add a value to it so the shapes don't overlap + close = self.data.Close.copy() + close += i * 10 + size * 2 + # Clear values except when i*7 divides into it + close[close % (i*3) != 0] = np.nan + + ind = self.I(lambda x: x, close, + name=f'{marker}_{size}', + scatter=True, + marker=marker, + marker_size=size, + overlay=i % 2 == 0) # Alternate between overlay and separate + self.shapes.append(ind) + + # Test invalid marker fallback + invalid_close = self.data.Close.copy() + invalid_close[invalid_close % 10 != 0] = np.nan + self.invalid = self.I(lambda x: x, invalid_close, + scatter=True, + marker='invalid_shape', + marker_size=10) + self.invalid_checks.append(self.invalid) + + # Test default size + default_close = self.data.Close.copy() + default_close[default_close % 15 != 0] = np.nan + self.default_size = self.I(lambda x: x, default_close, + scatter=True, + marker='circle') + self.default_checks.append(self.default_size) + + def next(self): + pass + + def get_default_size(self): + return self.default_size + + bt = Backtest(self.data, MarkerStrategy) + stats = bt.run() + fig = bt.plot() + + # Verify all indicators were created + strategy = bt._strategy + self.assertEqual(len(strategy.shapes), 7 * 3) + + # Verify each indicator has correct options + markers = ['circle', 'square', 'triangle', 'diamond', 'cross', 'x', 'star'] + sizes = [6, 12, 18] + + for ind, (marker, size) in zip(strategy.shapes, itertools.product(markers, sizes)): + self.assertTrue(ind._opts['scatter']) + self.assertEqual(ind._opts['marker'], marker) + self.assertEqual(ind._opts['marker_size'], size) + + # Verify invalid marker falls back to circle + self.assertEqual(strategy.invalid_checks[0]._opts['marker'], 'circle') + self.assertEqual(strategy.invalid_checks[0]._opts['marker_size'], 10) + + # Verify default size is None (will be set relative to bar width) + self.assertEqual(strategy.default_checks[0]._opts['marker'], 'circle') + self.assertIsNone(strategy.default_checks[0]._opts['marker_size'])