Coverage for src / monte_neo / visualization / charts.py: 86%

1"""Chart visualization module. 

2 

3Generate charts for indicator visualization. 

4""" 

5 

6from __future__ import annotations 

7 

8from typing import TYPE_CHECKING 

9 

10import pandas as pd 

11 

12from monte_neo.utils.logger import get_logger 

13 

14if TYPE_CHECKING: 

15 pass 

16 

17logger = get_logger(__name__) 

18 

19 

20class ChartGenerator: 

21 """Generate trading charts.""" 

22 

23 def __init__(self) -> None: 

24 """Initialize chart generator.""" 

25 self._has_mplfinance = False 

26 try: 

27 import mplfinance # noqa 

28 

29 self._has_mplfinance = True 

30 except ImportError: 

31 logger.warning("mplfinance not available, using plotext") 

32 

33 def plot_candlestick( 

34 self, 

35 data: pd.DataFrame, 

36 title: str = "Price Chart", 

37 save_path: str | None = None, 

38 ) -> None: 

39 """Plot candlestick chart. 

40 

41 Args: 

42 data: OHLCV DataFrame. 

43 title: Chart title. 

44 save_path: Path to save image. 

45 """ 

46 if self._has_mplfinance: 

47 self._plot_mpl(data, title, save_path) 

48 else: 

49 self._plot_terminal(data, title) 

50 

51 def _plot_mpl( 

52 self, 

53 data: pd.DataFrame, 

54 title: str, 

55 save_path: str | None, 

56 ) -> None: 

57 """Plot with mplfinance.""" 

58 import mplfinance as mpf 

59 

60 # Prepare data 

61 df = data.copy() 

62 if not isinstance(df.index, pd.DatetimeIndex): 

63 df.index = pd.to_datetime(df.index) 

64 

65 kwargs = { 

66 "type": "candle", 

67 "volume": True, 

68 "title": title, 

69 "style": "charles", 

70 "figsize": (12, 8), 

71 } 

72 

73 if save_path: 

74 kwargs["savefig"] = save_path 

75 mpf.plot(df, **kwargs) 

76 else: 

77 mpf.plot(df, **kwargs) 

78 

79 def _plot_terminal( 

80 self, 

81 data: pd.DataFrame, 

82 title: str, 

83 ) -> None: 

84 """Plot in terminal with plotext.""" 

85 import plotext as plt 

86 

87 plt.clear_figure() 

88 plt.title(title) 

89 

90 # Use candlestick if OHLC data is available 

91 if all(col in data.columns for col in ["open", "high", "low", "close"]): 

92 if hasattr(data.index, "astype"): 

93 dates = data.index.astype(str).tolist() 

94 else: 

95 dates = list(range(len(data))) 

96 # Plotext expects lists 

97 plt.candlestick( 

98 dates, 

99 data["open"].tolist(), 

100 data["high"].tolist(), 

101 data["low"].tolist(), 

102 data["close"].tolist(), 

103 label="OHLC", 

104 orientation="vertical" 

105 ) 

106 else: 

107 plt.plot(data["close"].values, label="Close") 

108 

109 plt.show() 

110 

111 def plot_with_signals( 

112 self, 

113 data: pd.DataFrame, 

114 signals: pd.DataFrame, 

115 title: str = "Chart with Signals", 

116 ) -> None: 

117 """Plot chart with entry/exit signals. 

118 

119 Args: 

120 data: OHLCV DataFrame. 

121 signals: Signals DataFrame. 

122 title: Chart title. 

123 """ 

124 import plotext as plt 

125 

126 plt.clear_figure() 

127 plt.title(title) 

128 

129 # Price 

130 close = data["close"].values 

131 plt.plot(close, label="Price") 

132 

133 # Entry signals 

134 if signals is not None: 

135 # Normalize signals to Series if it's a dict or DataFrame 

136 if isinstance(signals, dict): 

137 sig_series = pd.Series(signals.get("signal", 0), index=data.index) 

138 elif isinstance(signals, pd.DataFrame): 

139 sig_series = signals["signal"] if "signal" in signals.columns else signals.iloc[:, 0] 

140 else: 

141 sig_series = pd.Series(signals, index=data.index) 

142 

143 entries = sig_series == 1 

144 exits = sig_series == -1 

145 

146 # Mark entries and exits 

147 for i in range(len(sig_series)): 

148 if entries.iloc[i]: 

149 plt.scatter([i], [close[i]], marker="▲", color="green") 

150 elif exits.iloc[i]: 

151 plt.scatter([i], [close[i]], marker="▼", color="red") 

152 

153 plt.show()