plot.generate_candlestick_graph Added short equivelent, separating plotting scatter creation to a function

This commit is contained in:
Sam Germain 2022-02-02 16:02:54 -06:00
parent a31cf236e4
commit 1f74cfe841
2 changed files with 52 additions and 48 deletions

View File

@ -1,6 +1,6 @@
import logging import logging
from pathlib import Path from pathlib import Path
from typing import Any, Dict, List from typing import Any, Dict, List, Optional
import pandas as pd import pandas as pd
@ -385,6 +385,35 @@ def add_areas(fig, row: int, data: pd.DataFrame, indicators) -> make_subplots:
return fig return fig
def create_scatter(
data,
column_name,
color,
direction
) -> Optional[go.Scatter]:
if column_name in data.columns:
df_short = data[data[column_name] == 1]
if len(df_short) > 0:
shorts = go.Scatter(
x=df_short.date,
y=df_short.close,
mode='markers',
name=column_name,
marker=dict(
symbol=f"triangle-{direction}-dot",
size=9,
line=dict(width=1),
color=color,
)
)
return shorts
else:
logger.warning(f"No {column_name}-signals found.")
return None
def generate_candlestick_graph(pair: str, data: pd.DataFrame, trades: pd.DataFrame = None, *, def generate_candlestick_graph(pair: str, data: pd.DataFrame, trades: pd.DataFrame = None, *,
indicators1: List[str] = [], indicators1: List[str] = [],
indicators2: List[str] = [], indicators2: List[str] = [],
@ -431,44 +460,15 @@ def generate_candlestick_graph(pair: str, data: pd.DataFrame, trades: pd.DataFra
) )
fig.add_trace(candles, 1, 1) fig.add_trace(candles, 1, 1)
# TODO-lev: Needs short equivalent longs = create_scatter(data, 'enter_long', 'green', 'up')
if 'enter_long' in data.columns: exit_longs = create_scatter(data, 'exit_long', 'red', 'down')
df_buy = data[data['enter_long'] == 1] shorts = create_scatter(data, 'enter_short', 'blue', 'down')
if len(df_buy) > 0: exit_shorts = create_scatter(data, 'exit_short', 'violet', 'up')
buys = go.Scatter(
x=df_buy.date, for scatter in [longs, exit_longs, shorts, exit_shorts]:
y=df_buy.close, if scatter:
mode='markers', fig.add_trace(scatter, 1, 1)
name='buy',
marker=dict(
symbol='triangle-up-dot',
size=9,
line=dict(width=1),
color='green',
)
)
fig.add_trace(buys, 1, 1)
else:
logger.warning("No buy-signals found.")
if 'exit_long' in data.columns:
df_sell = data[data['exit_long'] == 1]
if len(df_sell) > 0:
sells = go.Scatter(
x=df_sell.date,
y=df_sell.close,
mode='markers',
name='sell',
marker=dict(
symbol='triangle-down-dot',
size=9,
line=dict(width=1),
color='red',
)
)
fig.add_trace(sells, 1, 1)
else:
logger.warning("No sell-signals found.")
# Add Bollinger Bands # Add Bollinger Bands
fig = plot_area(fig, 1, data, 'bb_lowerband', 'bb_upperband', fig = plot_area(fig, 1, data, 'bb_lowerband', 'bb_upperband',
label="Bollinger Band") label="Bollinger Band")

View File

@ -202,6 +202,8 @@ def test_generate_candlestick_graph_no_signals_no_trades(default_conf, mocker, t
datadir=testdatadir, timerange=timerange) datadir=testdatadir, timerange=timerange)
data['enter_long'] = 0 data['enter_long'] = 0
data['exit_long'] = 0 data['exit_long'] = 0
data['enter_short'] = 0
data['exit_short'] = 0
indicators1 = [] indicators1 = []
indicators2 = [] indicators2 = []
@ -222,8 +224,10 @@ def test_generate_candlestick_graph_no_signals_no_trades(default_conf, mocker, t
assert row_mock.call_count == 2 assert row_mock.call_count == 2
assert trades_mock.call_count == 1 assert trades_mock.call_count == 1
assert log_has("No buy-signals found.", caplog) assert log_has("No enter_long-signals found.", caplog)
assert log_has("No sell-signals found.", caplog) assert log_has("No exit_long-signals found.", caplog)
assert log_has("No enter_short-signals found.", caplog)
assert log_has("No exit_short-signals found.", caplog)
def test_generate_candlestick_graph_no_trades(default_conf, mocker, testdatadir): def test_generate_candlestick_graph_no_trades(default_conf, mocker, testdatadir):
@ -249,7 +253,7 @@ def test_generate_candlestick_graph_no_trades(default_conf, mocker, testdatadir)
assert fig.layout.title.text == pair assert fig.layout.title.text == pair
figure = fig.layout.figure figure = fig.layout.figure
assert len(figure.data) == 6 assert len(figure.data) == 8
# Candlesticks are plotted first # Candlesticks are plotted first
candles = find_trace_in_fig_data(figure.data, "Price") candles = find_trace_in_fig_data(figure.data, "Price")
assert isinstance(candles, go.Candlestick) assert isinstance(candles, go.Candlestick)
@ -257,15 +261,15 @@ def test_generate_candlestick_graph_no_trades(default_conf, mocker, testdatadir)
volume = find_trace_in_fig_data(figure.data, "Volume") volume = find_trace_in_fig_data(figure.data, "Volume")
assert isinstance(volume, go.Bar) assert isinstance(volume, go.Bar)
buy = find_trace_in_fig_data(figure.data, "buy") enter_long = find_trace_in_fig_data(figure.data, "enter_long")
assert isinstance(buy, go.Scatter) assert isinstance(enter_long, go.Scatter)
# All buy-signals should be plotted # All buy-signals should be plotted
assert int(data['enter_long'].sum()) == len(buy.x) assert int(data['enter_long'].sum()) == len(enter_long.x)
sell = find_trace_in_fig_data(figure.data, "sell") exit_long = find_trace_in_fig_data(figure.data, "exit_long")
assert isinstance(sell, go.Scatter) assert isinstance(exit_long, go.Scatter)
# All buy-signals should be plotted # All buy-signals should be plotted
assert int(data['exit_long'].sum()) == len(sell.x) assert int(data['exit_long'].sum()) == len(exit_long.x)
assert find_trace_in_fig_data(figure.data, "Bollinger Band") assert find_trace_in_fig_data(figure.data, "Bollinger Band")