Coverage for src / monte_neo / indicators / custom.py: 83%

1"""Custom indicator builder. 

2 

3Build custom indicators from combinations of technical indicators. 

4""" 

5 

6from __future__ import annotations 

7 

8from collections.abc import Callable 

9from dataclasses import dataclass 

10from typing import Any 

11 

12import pandas as pd 

13 

14from monte_neo.indicators.base import BaseIndicator, IndicatorConfig 

15from monte_neo.indicators.technical import TechnicalIndicators 

16from monte_neo.utils.logger import get_logger 

17 

18logger = get_logger(__name__) 

19 

20 

21@dataclass 

22class ConditionRule: 

23 """Single condition rule.""" 

24 

25 indicator1: str 

26 operator: str # >, <, ==, crosses_above, crosses_below 

27 indicator2: str | float 

28 

29 def evaluate(self, data: pd.DataFrame) -> pd.Series: 

30 """Evaluate the condition.""" 

31 val1 = self._get_value(data, self.indicator1) 

32 val2 = self._get_value(data, self.indicator2) 

33 

34 if self.operator == ">": 

35 return val1 > val2 

36 elif self.operator == "<": 

37 return val1 < val2 

38 elif self.operator == ">=": 

39 return val1 >= val2 

40 elif self.operator == "<=": 

41 return val1 <= val2 

42 elif self.operator == "==": 

43 return val1 == val2 

44 elif self.operator == "crosses_above": 

45 return (val1 > val2) & (val1.shift(1) <= val2.shift(1)) 

46 elif self.operator == "crosses_below": 

47 return (val1 < val2) & (val1.shift(1) >= val2.shift(1)) 

48 else: 

49 raise ValueError(f"Unknown operator: {self.operator}") 

50 

51 def _get_value(self, data: pd.DataFrame, indicator: str | float) -> pd.Series: 

52 if isinstance(indicator, (int, float)): 

53 return pd.Series(indicator, index=data.index) 

54 return data[indicator] 

55 

56 

57class CustomIndicatorBuilder: 

58 """Builder for creating custom indicators.""" 

59 

60 def __init__(self) -> None: 

61 self._components: list[tuple[str, Callable, dict]] = [] 

62 self._entry_rules: list[ConditionRule] = [] 

63 self._exit_rules: list[ConditionRule] = [] 

64 self._parameters: dict[str, Any] = {} 

65 

66 def add_sma(self, name: str, period: int) -> CustomIndicatorBuilder: 

67 """Add SMA component.""" 

68 self._components.append((name, TechnicalIndicators.sma, {"period": period})) 

69 self._parameters[f"{name}_period"] = period 

70 return self 

71 

72 def add_ema(self, name: str, period: int) -> CustomIndicatorBuilder: 

73 """Add EMA component.""" 

74 self._components.append((name, TechnicalIndicators.ema, {"period": period})) 

75 self._parameters[f"{name}_period"] = period 

76 return self 

77 

78 def add_rsi(self, name: str, period: int = 14) -> CustomIndicatorBuilder: 

79 """Add RSI component.""" 

80 self._components.append((name, TechnicalIndicators.rsi, {"period": period})) 

81 self._parameters[f"{name}_period"] = period 

82 return self 

83 

84 def add_entry_rule( 

85 self, 

86 indicator1: str, 

87 operator: str, 

88 indicator2: str | float, 

89 ) -> CustomIndicatorBuilder: 

90 """Add entry condition.""" 

91 self._entry_rules.append(ConditionRule(indicator1, operator, indicator2)) 

92 return self 

93 

94 def add_exit_rule( 

95 self, 

96 indicator1: str, 

97 operator: str, 

98 indicator2: str | float, 

99 ) -> CustomIndicatorBuilder: 

100 """Add exit condition.""" 

101 self._exit_rules.append(ConditionRule(indicator1, operator, indicator2)) 

102 return self 

103 

104 def build(self, name: str = "CustomIndicator") -> CustomIndicator: 

105 """Build the custom indicator.""" 

106 config = IndicatorConfig(name=name, parameters=self._parameters) 

107 indicator = CustomIndicator(config) 

108 indicator._components = list(self._components) 

109 indicator._entry_rules = list(self._entry_rules) 

110 indicator._exit_rules = list(self._exit_rules) 

111 return indicator 

112 

113 

114class CustomIndicator(BaseIndicator): 

115 """Custom indicator built from components.""" 

116 

117 def __init__(self, config: IndicatorConfig | None = None) -> None: 

118 super().__init__(config) 

119 self._components: list[tuple[str, Callable, dict]] = [] 

120 self._entry_rules: list[ConditionRule] = [] 

121 self._exit_rules: list[ConditionRule] = [] 

122 

123 def calculate(self, data: pd.DataFrame) -> pd.DataFrame: 

124 result = data.copy() 

125 

126 for name, func, params in self._components: 

127 # Update params from current parameters 

128 updated_params = {} 

129 for key, default in params.items(): 

130 param_key = f"{name}_{key}" 

131 updated_params[key] = self._parameters.get(param_key, default) 

132 

133 if func in ( 

134 TechnicalIndicators.sma, 

135 TechnicalIndicators.ema, 

136 TechnicalIndicators.rsi, 

137 ): 

138 result[name] = func(result["close"], **updated_params) 

139 

140 return result 

141 

142 def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame: 

143 calc = self.calculate(data) 

144 signals = pd.DataFrame(index=data.index) 

145 signals["signal"] = 0 

146 

147 # Evaluate entry rules (all must be true) 

148 if self._entry_rules: 

149 entry_mask = pd.Series(True, index=data.index) 

150 for rule in self._entry_rules: 

151 entry_mask &= rule.evaluate(calc) 

152 signals.loc[entry_mask, "signal"] = 1 

153 

154 # Evaluate exit rules 

155 if self._exit_rules: 

156 exit_mask = pd.Series(True, index=data.index) 

157 for rule in self._exit_rules: 

158 exit_mask &= rule.evaluate(calc) 

159 signals.loc[exit_mask, "signal"] = -1 

160 

161 return signals 

162 

163 def get_min_periods(self) -> int: 

164 max_period = 1 

165 for _, _, params in self._components: 

166 period = params.get("period", 1) 

167 max_period = max(max_period, period) 

168 return max_period