Coverage for src / monte_neo / core / acceleration / engine.py: 100%

93 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-28 16:27 +0200

1""" 

2GPU Acceleration Engine (MLX). 

3""" 

4 

5import time 

6from typing import Any 

7 

8import mlx.core as mx 

9import numpy as np 

10import pandas as pd 

11 

12from monte_neo.core.acceleration.indicators import MLXSMA, MLXCrossStrategy 

13from monte_neo.core.acceleration.tensor_ops import generate_noise_scenarios, generate_shuffle_scenarios, to_tensor 

14 

15 

16class GpuAccelerationEngine: 

17 """High-performance GPU engine.""" 

18 

19 def __init__(self, batch_size: int = 50000, precision: str = "float32", metal_driver: str = "cpp", initial_capital: float = 100000.0, leverage: float = 1.0): 

20 """ 

21 Initialize GPU engine. 

22  

23 Args: 

24 batch_size: Number of scenarios to process in each batch 

25 precision: 'float32', 'float16', 'float8_e4m3', or 'float8_e5m2' 

26 metal_driver: Metal driver to use ('cpp', 'objc', 'swift') 

27 initial_capital: Initial account balance. 

28 leverage: Trading leverage. 

29 """ 

30 self.batch_size = batch_size 

31 self.precision = precision 

32 self.metal_driver = metal_driver 

33 self.initial_capital = initial_capital 

34 self.leverage = leverage 

35 

36 # Initialize float8 encoder if needed 

37 self.float8_encoder = None 

38 self.metal_engine = None 

39 

40 if precision.startswith("float8"): 

41 # For float8 we still use the specialized MetalFloat8Engine for now 

42 # but we can pass the driver if it supports it in the future 

43 try: 

44 from monte_neo.core.native.metal_engine import MetalFloat8Engine 

45 self.metal_engine = MetalFloat8Engine() 

46 except (ImportError, RuntimeError) as e: 

47 print(f"Warning: Could not initialize Metal engine: {e}. Falling back to Python encoder.") 

48 

49 if not self.metal_engine: 

50 from monte_neo.core.acceleration.float8 import Float8Encoder 

51 format_type = "e4m3" if "e4m3" in precision else "e5m2" 

52 self.float8_encoder = Float8Encoder(format_type) 

53 

54 def _reconstruct_strategy(self, mlx_strategy: Any) -> Any: 

55 """Reconstruct MLX strategy from dictionary if needed.""" 

56 if not isinstance(mlx_strategy, dict): 

57 return mlx_strategy 

58 

59 strat_type = mlx_strategy.get("type") 

60 if strat_type == "sma_crossover": 

61 period = mlx_strategy.get("period", 20) 

62 return MLXCrossStrategy(MLXSMA(period)) 

63 # Add more types as needed 

64 return mlx_strategy 

65 

66 def run_simulation( 

67 self, 

68 data: pd.DataFrame, 

69 mlx_strategy: Any, 

70 n_scenarios: int, 

71 method: str = "shuffling", 

72 seed: int = 42 

73 ) -> list[dict[str, Any]]: 

74 """ 

75 Run generic simulation on GPU. 

76  

77 Args: 

78 data: OHLCV DataFrame. 

79 mlx_strategy: Strategy object or dict representation. 

80 n_scenarios: Number of scenarios. 

81 method: 'shuffling' or 'noise'. 

82 """ 

83 # 0. Reconstruct strategy if needed 

84 strategy = self._reconstruct_strategy(mlx_strategy) 

85 

86 # 1. To Tensor (Done once) 

87 tensors = to_tensor(data) 

88 close = tensors["close"] 

89 

90 processed = 0 

91 all_results = [] 

92 

93 while processed < n_scenarios: 

94 current_batch = min(self.batch_size, n_scenarios - processed) 

95 

96 # 2. Scenarios 

97 if self.metal_engine and self.precision.startswith("float8"): 

98 # Use Metal engine for float8 scenarios 

99 # First convert close to float8 using Metal 

100 close_np = np.array(close).astype(np.float32) 

101 if self.precision == "float8_e4m3": 

102 close_f8 = self.metal_engine.encode_float32_to_e4m3(close_np) 

103 scenarios_f8 = self.metal_engine.generate_scenarios_e4m3(close_f8, current_batch, seed=seed+processed) 

104 scenarios_np = self.metal_engine.decode_e4m3_to_float32(scenarios_f8) 

105 scenarios = mx.array(scenarios_np) 

106 else: 

107 # Fallback to MLX for e5m2 if not fully implemented in Metal yet 

108 scenarios = generate_shuffle_scenarios(close, current_batch, seed=seed+processed) 

109 elif method == "shuffling": 

110 scenarios = generate_shuffle_scenarios(close, current_batch, seed=seed+processed) 

111 elif method == "noise": 

112 scenarios = generate_noise_scenarios(close, current_batch, seed=seed+processed) 

113 else: 

114 # Default or error 

115 scenarios = generate_shuffle_scenarios(close, current_batch, seed=seed+processed) 

116 

117 # 3. Signals 

118 signals = strategy.generate_signals(scenarios) 

119 

120 # 4. Backtest 

121 returns = (scenarios[:, 1:] / scenarios[:, :-1]) - 1.0 

122 strat_returns = (signals[:, :-1] * returns) * self.leverage 

123 

124 # Metrics 

125 equity = self.initial_capital * mx.exp(mx.cumsum(mx.log1p(strat_returns), axis=1)) 

126 final_balance = equity[:, -1] 

127 

128 # Max DD 

129 running_max = mx.cummax(equity, axis=1) 

130 max_dds = mx.max((running_max - equity) / running_max, axis=1) 

131 

132 # Profit Factor 

133 wins = mx.where(strat_returns > 0, strat_returns, 0) 

134 losses = mx.where(strat_returns < 0, strat_returns, 0) 

135 gross_profit = mx.sum(wins, axis=1) 

136 gross_loss = mx.abs(mx.sum(losses, axis=1)) 

137 profit_factor = mx.where(gross_loss > 0, gross_profit / gross_loss, 100.0) 

138 

139 # Evaluate batch 

140 mx.eval(final_balance, max_dds, profit_factor) 

141 

142 # Convert to numpy/list for result 

143 # We can't keep all results in GPU memory if N is huge? 

144 # Actually we just keep scalars. 

145 

146 fb_np = np.array(final_balance) 

147 mdd_np = np.array(max_dds) 

148 pf_np = np.array(profit_factor) 

149 

150 for i in range(current_batch): 

151 all_results.append({ 

152 "total_return": (float(fb_np[i]) / self.initial_capital) - 1.0, 

153 "final_balance": float(fb_np[i]), 

154 "total_profit_abs": float(fb_np[i]) - self.initial_capital, 

155 "max_drawdown": float(mdd_np[i]), 

156 "profit_factor": float(pf_np[i]), 

157 "passed": bool(fb_np[i] > self.initial_capital and mdd_np[i] < 0.2), # Default criteria 

158 "metrics": { 

159 "total_return": (float(fb_np[i]) / self.initial_capital) - 1.0, 

160 "final_balance": float(fb_np[i]), 

161 "total_profit_abs": float(fb_np[i]) - self.initial_capital, 

162 "max_drawdown": float(mdd_np[i]), 

163 "profit_factor": float(pf_np[i]), 

164 } 

165 }) 

166 

167 processed += current_batch 

168 

169 return all_results 

170 

171 def run_benchmark_simulation(self, data: pd.DataFrame, n_scenarios: int) -> dict: 

172 """ 

173 Run a full simulation pipeline on GPU to benchmark performance. 

174 Pipeline: 

175 1. Data -> Tensor 

176 2. Shuffle Scenarios (N x T) 

177 3. Indicator Calc (SMA) -> Signals (N x T) 

178 4. Backtest (Vectorized) 

179 """ 

180 start_time = time.time() 

181 

182 # 1. To Tensor (Done once) 

183 tensors = to_tensor(data) 

184 close = tensors["close"] 

185 

186 processed = 0 

187 

188 while processed < n_scenarios: 

189 current_batch = min(self.batch_size, n_scenarios - processed) 

190 

191 # 2. Scenarios 

192 scenarios = generate_shuffle_scenarios(close, current_batch) 

193 

194 # 3. Signals 

195 sma = MLXSMA(10) 

196 strat = MLXCrossStrategy(sma) 

197 signals = strat.generate_signals(scenarios) 

198 

199 # 4. Backtest 

200 returns = (scenarios[:, 1:] / scenarios[:, :-1]) - 1.0 

201 strat_returns = signals[:, :-1] * returns 

202 equity = mx.exp(mx.cumsum(mx.log1p(strat_returns), axis=1)) 

203 final_return = equity[:, -1] 

204 

205 # Force computation 

206 mx.eval(final_return) 

207 

208 processed += current_batch 

209 

210 elapsed = time.time() - start_time 

211 

212 return { 

213 "elapsed": elapsed, 

214 "ops_per_sec": n_scenarios / elapsed, 

215 "scenarios_processed": n_scenarios 

216 }