Coverage for src / monte_neo / core / evolution.py: 82%

1"""Evolutionary algorithm module.""" 

2 

3from __future__ import annotations 

4 

5import re 

6from collections.abc import Callable 

7from typing import TYPE_CHECKING 

8 

9import numpy as np 

10import pandas as pd 

11 

12from monte_neo.indicators.base import BaseIndicator 

13from monte_neo.indicators.code_gen import CodeGenerator 

14from monte_neo.indicators.dynamic import DynamicIndicator 

15from monte_neo.metrics.calculator import MetricsCalculator 

16from monte_neo.utils.ast_utils import crossover_trees 

17 

18if TYPE_CHECKING: 

19 from monte_neo.core.config import GeneratorConfig 

20 from monte_neo.utils.parallel import ParallelExecutor 

21 

22 

23class EvolutionEngine: 

24 """Evolutionary algorithm engine.""" 

25 

26 def __init__( 

27 self, 

28 config: GeneratorConfig, 

29 metrics_calc: MetricsCalculator | None = None, 

30 progress_callback: Callable[[int, int, str], None] | None = None 

31 ): 

32 """Initialize engine. 

33 

34 Args: 

35 config: Generator configuration. 

36 metrics_calc: Metrics calculator. 

37 progress_callback: Progress callback. 

38 """ 

39 self.config = config 

40 self.metrics_calc = metrics_calc or MetricsCalculator() 

41 self.rng = np.random.default_rng() 

42 self.code_gen = CodeGenerator(self.rng) 

43 self.progress_callback = progress_callback 

44 

45 def run( 

46 self, 

47 data: pd.DataFrame, 

48 initial_population: list[BaseIndicator], 

49 executor: ParallelExecutor | None = None 

50 ) -> BaseIndicator | None: 

51 """Run evolutionary optimization. 

52 

53 Args: 

54 data: OHLCV data. 

55 initial_population: Initial population. 

56 executor: Optional parallel executor for signal generation. 

57 

58 Returns: 

59 Best indicator found. 

60 """ 

61 population: list[BaseIndicator] = list(initial_population) 

62 

63 # Pad population if needed 

64 while len(population) < self.config.population_size: 

65 new_indicator = DynamicIndicator() 

66 new_indicator.set_parameter("source_code", self.code_gen.generate_code()) 

67 population.append(new_indicator) 

68 

69 for gen in range(self.config.generations): 

70 # Evaluate fitness 

71 fitness_scores: list[tuple[BaseIndicator, float]] = [] 

72 

73 # Use parallel execution for fitness evaluation to reach >2000 ops/s 

74 # Note: Evolution handles batches of DynamicIndicators 

75 if executor: 

76 from monte_neo.monte_carlo.workers import run_indicator_batch 

77 n_workers = executor.n_workers 

78 chunk_size = max(1, len(population) // n_workers) 

79 chunks = [population[i : i + chunk_size] for i in range(0, len(population), chunk_size)] 

80 # Using None for data since initializer already set it in SHARED_DATA 

81 tasks = [(chunk, None) for chunk in chunks] 

82 batch_results = executor.map(run_indicator_batch, tasks) 

83 raw_signals = [] 

84 for batch in batch_results: 

85 raw_signals.extend(batch) 

86 else: 

87 raw_signals = [ind.generate_signals_fast(data) for ind in population] 

88 

89 # Prepare for Numba batch calculation 

90 signal_matrix = np.zeros((len(population), len(data)), dtype=np.int32) 

91 from monte_neo.core.gpu_scenarios import normalize_signal_array 

92 for i, sig in enumerate(raw_signals): 

93 signal_matrix[i] = normalize_signal_array(sig, len(data)).astype(np.int32) 

94 

95 batch_metrics_arr = self.metrics_calc.calculate_batch_fast( 

96 data["close"].values, 

97 data["high"].values, 

98 data["low"].values, 

99 signal_matrix, 

100 use_sl_tp=self.config.use_sl_tp, 

101 sl_pct=self.config.stop_loss_pct, 

102 tp_pct=self.config.take_profit_pct, 

103 ) 

104 

105 for i, ind in enumerate(population): 

106 # Fitness function: Profit Factor * (1 - Max Drawdown) 

107 pf = batch_metrics_arr[i, 2] 

108 dd = batch_metrics_arr[i, 1] 

109 trades = int(batch_metrics_arr[i, 3]) 

110 

111 if trades < self.config.min_trades: 

112 score = 0.0 

113 else: 

114 score = pf * (1.0 - dd) 

115 

116 fitness_scores.append((ind, score)) 

117 

118 # Sort 

119 fitness_scores.sort(key=lambda x: x[1], reverse=True) 

120 best_gen_score = fitness_scores[0][1] 

121 

122 if self.progress_callback: 

123 self.progress_callback( 

124 gen + 1, 

125 self.config.generations, 

126 f"Evolution Gen {gen + 1}: Best Score {best_gen_score:.2f}", 

127 ) 

128 

129 # Selection (Elite + Tournament) 

130 elite_count = max(2, int(self.config.population_size * 0.1)) 

131 new_pop: list[BaseIndicator] = [x[0] for x in fitness_scores[:elite_count]] 

132 

133 while len(new_pop) < self.config.population_size: 

134 parent1 = self._tournament_select(fitness_scores) 

135 

136 if self.rng.random() < self.config.crossover_rate: 

137 parent2 = self._tournament_select(fitness_scores) 

138 child = self._crossover_indicators(parent1, parent2) 

139 else: 

140 child = self._mutate_indicator(parent1) 

141 

142 new_pop.append(child) 

143 

144 population = new_pop 

145 

146 # Return the best found 

147 if not population: 

148 return None 

149 

150 # Final evaluation 

151 final_scores: list[tuple[BaseIndicator, float]] = [] 

152 for ind in population: 

153 try: 

154 signals = ind.generate_signals(data) 

155 metrics = self.metrics_calc.calculate_all(data, signals) 

156 pf = float(metrics.get("profit_factor", 0.0)) 

157 dd = float(metrics.get("max_drawdown", 1.0)) 

158 trades = int(metrics.get("trade_count", 0)) 

159 score = pf * (1.0 - dd) if trades >= self.config.min_trades else 0.0 

160 final_scores.append((ind, score)) 

161 except Exception: 

162 final_scores.append((ind, 0.0)) 

163 

164 final_scores.sort(key=lambda x: x[1], reverse=True) 

165 return final_scores[0][0] if final_scores[0][1] > 0 else None 

166 

167 def _crossover_indicators(self, p1: BaseIndicator, p2: BaseIndicator) -> BaseIndicator: 

168 """Perform crossover.""" 

169 if not isinstance(p1, DynamicIndicator) or not isinstance(p2, DynamicIndicator): 

170 return self._mutate_indicator(p1) 

171 

172 code1 = p1.get_parameters().get("source_code", "data['close']") 

173 code2 = p2.get_parameters().get("source_code", "data['close']") 

174 

175 new_code = crossover_trees(code1, code2) 

176 

177 new_ind = DynamicIndicator() 

178 new_ind.set_parameter("source_code", new_code) 

179 return new_ind 

180 

181 def _tournament_select(self, fitness: list, k: int = 3) -> BaseIndicator: 

182 indices = self.rng.integers(0, len(fitness), size=k) 

183 best_idx = max(indices, key=lambda i: fitness[i][1]) 

184 return fitness[best_idx][0] 

185 

186 def _mutate_indicator(self, indicator: BaseIndicator) -> BaseIndicator: 

187 """Mutate an indicator.""" 

188 if not isinstance(indicator, DynamicIndicator): 

189 return indicator 

190 

191 code = indicator.get_parameters().get("source_code", "") 

192 if not code: 

193 return indicator 

194 

195 new_code = code 

196 

197 # Replace numbers 

198 def replace_num(match): 

199 val = int(match.group()) 

200 change = self.rng.choice([-1, 1]) * max(1, int(val * 0.2)) 

201 return str(max(1, val + change)) 

202 

203 if self.rng.random() < 0.5: 

204 new_code = re.sub(r"\b\d+\b", replace_num, code) 

205 else: 

206 op = self.rng.choice(["+", "-", "*"]) 

207 operand = self.rng.choice(["data['close']", "data['volume']"]) 

208 new_code = f"({code} {op} {operand})" 

209 

210 new_ind = DynamicIndicator() 

211 new_ind.set_parameter("source_code", new_code) 

212 return new_ind