Coverage for src / monte_neo / core / acceleration / indicators.py: 52%

1""" 

2MLX-accelerated indicators. 

3""" 

4 

5from typing import Any 

6 

7import mlx.core as mx 

8import numpy as np 

9 

10 

11class MLXIndicator: 

12 """Base class for MLX indicators.""" 

13 def compute(self, close: mx.array) -> mx.array: 

14 """ 

15 Compute indicator. 

16 Args: 

17 close: (N, T) matrix of close prices. 

18 Returns: 

19 (N, T) matrix of indicator values. 

20 """ 

21 raise NotImplementedError 

22 

23class MLXSMA(MLXIndicator): 

24 def __init__(self, period: int): 

25 self.period = period 

26 # Create kernel: (Out=1, Kernel=P, In=1) 

27 # Simple average = sum / P 

28 self.weight = mx.full((1, period, 1), 1.0 / period) 

29 

30 def compute(self, close: mx.array) -> mx.array: 

31 """Compute SMA using 1D convolution.""" 

32 if close.dtype != mx.float32 and close.dtype != mx.float16: 

33 close = close.astype(mx.float32) 

34 # Reshape for conv1d: [batch, length, channels] 

35 # x is [scenarios, time] -> [scenarios, time, 1] 

36 x = close[..., None] 

37 

38 # Pad beginning with first value to keep same length 

39 pad_size = self.period - 1 

40 padding = [(0, 0), (pad_size, 0), (0, 0)] 

41 x_padded = mx.pad(x, padding, constant_values=0) 

42 # Fix: instead of 0, use first value to avoid spikes 

43 # Actually, MLX pad constant_values can be an array/scalar. 

44 # But we want to pad with the first value of EACH scenario. 

45 # For now, let's use a simpler approach: 

46 x_padded = mx.concatenate([mx.repeat(x[:, :1, :], pad_size, axis=1), x], axis=1) 

47 

48 out = mx.conv1d(x_padded, self.weight, stride=1, padding=0) 

49 return out.squeeze(-1) 

50 

51class MLXRSI(MLXIndicator): 

52 """RSI indicator on MLX.""" 

53 def __init__(self, period: int = 14): 

54 self.period = period 

55 

56 def compute(self, close: mx.array) -> mx.array: 

57 # Diff 

58 diff = close[:, 1:] - close[:, :-1] 

59 # First diff is 0 

60 diff = mx.concatenate([mx.zeros((close.shape[0], 1)), diff], axis=1) 

61 

62 gain = mx.where(diff > 0, diff, 0.0) 

63 loss = mx.where(diff < 0, -diff, 0.0) 

64 

65 # SMA of gains/losses (Simplified version of Wilders) 

66 # Real RSI uses SMMA/EMA, but SMA is often used as approx or in some variants. 

67 # Let's use SMA for simplicity in MLX for now. 

68 sma_gain = MLXSMA(self.period).compute(gain) 

69 sma_loss = MLXSMA(self.period).compute(loss) 

70 

71 rs = mx.where(sma_loss > 0, sma_gain / sma_loss, 100.0) 

72 rsi = 100.0 - (100.0 / (1.0 + rs)) 

73 return rsi 

74 

75class MLXRollingMax(MLXIndicator): 

76 def __init__(self, period: int): 

77 self.period = period 

78 

79 def compute(self, close: mx.array) -> mx.array: 

80 # MLX doesn't have a direct rolling_max with window, 

81 # but we can use a trick with reshape or just use a loop for small windows. 

82 # For large windows, we might need a more efficient implementation. 

83 # For now, let's use a simple implementation. 

84 n, t = close.shape 

85 res = mx.zeros_like(close) 

86 for i in range(t): 

87 start = max(0, i - self.period + 1) 

88 res[:, i] = mx.max(close[:, start:i+1], axis=1) 

89 return res 

90 

91class MLXDynamicStrategy: 

92 """Fallback strategy that runs Python/Numba logic on MLX data.""" 

93 def __init__(self, indicator: Any): 

94 self.indicator = indicator 

95 

96 def generate_signals(self, close: mx.array) -> mx.array: 

97 # Convert to numpy 

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

99 

100 if close_np.ndim == 1: 

101 # Single scenario 

102 signals_np = self.indicator.generate_signals_fast(close_np) 

103 return mx.array(signals_np) 

104 

105 # Multiple scenarios - run in loop for now 

106 n_scenarios = close_np.shape[0] 

107 n_steps = close_np.shape[1] 

108 all_signals = np.zeros((n_scenarios, n_steps), dtype=np.float32) 

109 

110 for i in range(n_scenarios): 

111 all_signals[i] = self.indicator.generate_signals_fast(close_np[i]) 

112 

113 return mx.array(all_signals) 

114 

115class MLXCrossStrategy: 

116 """Simple Crossover Strategy on GPU.""" 

117 

118 def __init__(self, indicator: MLXIndicator, mode: str = "greater"): 

119 self.indicator = indicator 

120 self.mode = mode 

121 

122 def generate_signals(self, close: mx.array) -> mx.array: 

123 ind_vals = self.indicator.compute(close) 

124 if self.mode == "greater": 

125 state = mx.where(close > ind_vals, 1, -1) 

126 else: 

127 state = mx.where(close < ind_vals, 1, -1) 

128 

129 prev_state = mx.concatenate([state[:, :1], state[:, :-1]], axis=1) 

130 return mx.where(state != prev_state, state, 0) 

131 

132class MLXSMACrossStrategy: 

133 """SMA Crossover Strategy on GPU (Fast vs Slow).""" 

134 

135 def __init__(self, fast_period: int, slow_period: int): 

136 self.fast_sma = MLXSMA(fast_period) 

137 self.slow_sma = MLXSMA(slow_period) 

138 

139 def generate_signals(self, close: mx.array) -> mx.array: 

140 fast = self.fast_sma.compute(close) 

141 slow = self.slow_sma.compute(close) 

142 state = mx.where(fast > slow, 1, -1) 

143 prev_state = mx.concatenate([state[:, :1], state[:, :-1]], axis=1) 

144 return mx.where(state != prev_state, state, 0)