Coverage for src / monte_neo / core / acceleration / tensor_ops.py: 94%

1""" 

2Tensor operations for MLX-based acceleration. 

3""" 

4 

5import mlx.core as mx 

6import numpy as np 

7import pandas as pd 

8 

9 

10class TensorOps: 

11 """Collection of MLX-based tensor operations.""" 

12 

13 @staticmethod 

14 def to_tensor(data: pd.DataFrame) -> dict[str, mx.array]: 

15 """Convert DataFrame to dictionary of MLX arrays.""" 

16 return { 

17 "open": mx.array(data["open"].values.astype(np.float32)), 

18 "high": mx.array(data["high"].values.astype(np.float32)), 

19 "low": mx.array(data["low"].values.astype(np.float32)), 

20 "close": mx.array(data["close"].values.astype(np.float32)), 

21 "volume": mx.array(data["volume"].values.astype(np.float32)), 

22 } 

23 

24 @staticmethod 

25 def moving_average(data: mx.array, period: int) -> mx.array: 

26 """ 

27 Calculate Simple Moving Average on MLX. 

28 Uses a sliding window approach with convolution. 

29 """ 

30 if period <= 1: 

31 return data 

32 

33 # MLX conv1d default: input (N, L, C), weight (O, K, I) 

34 # N: batch, L: length, C: channels 

35 # O: out_channels, K: kernel_width, I: in_channels 

36 

37 is_1d = len(data.shape) == 1 

38 if is_1d: 

39 # (1, T, 1) 

40 x = data[None, :, None] 

41 else: 

42 # (N, T, 1) 

43 x = data[:, :, None] 

44 

45 # Pad to keep same length 

46 # Padding for 'same' with causal behavior (pad left) 

47 # padding is list of (before, after) for each dimension 

48 padding = [(0, 0), (period - 1, 0), (0, 0)] 

49 x_padded = mx.pad(x, padding, constant_values=x[:, 0:1, :]) 

50 

51 # Kernel: (1, period, 1) 

52 kernel = mx.ones((1, period, 1)) / period 

53 

54 out = mx.conv1d(x_padded, kernel) 

55 

56 if is_1d: 

57 return out.reshape(-1) 

58 else: 

59 return out.reshape(data.shape) 

60 

61 @staticmethod 

62 def generate_shuffle_scenarios( 

63 close: mx.array, n_scenarios: int, seed: int = 42 

64 ) -> mx.array: 

65 """ 

66 Generate shuffled return scenarios on GPU. 

67 """ 

68 time_steps = close.shape[0] 

69 returns = (close[1:] / close[:-1]) - 1.0 

70 

71 key = mx.random.key(seed) 

72 indices = mx.random.randint(0, time_steps-1, (n_scenarios, time_steps-1), key=key) 

73 

74 shuffled_returns = returns[indices] 

75 cum_returns = mx.cumprod(1 + shuffled_returns, axis=1) 

76 

77 ones = mx.ones((n_scenarios, 1)) 

78 factors = mx.concatenate([ones, cum_returns], axis=1) 

79 

80 start_price = close[0] 

81 scenarios = start_price * factors 

82 

83 return scenarios 

84 

85 @staticmethod 

86 def generate_noise_scenarios( 

87 close: mx.array, n_scenarios: int, std_dev: float = 0.01, seed: int = 42 

88 ) -> mx.array: 

89 """ 

90 Generate scenarios with Gaussian noise injected into returns. 

91 """ 

92 time_steps = close.shape[0] 

93 returns = (close[1:] / close[:-1]) - 1.0 

94 

95 base_returns = mx.broadcast_to(returns, (n_scenarios, time_steps-1)) 

96 

97 key = mx.random.key(seed) 

98 noise = mx.random.normal((n_scenarios, time_steps-1), scale=std_dev, key=key) 

99 

100 noisy_returns = base_returns + noise 

101 

102 cum_returns = mx.cumprod(1 + noisy_returns, axis=1) 

103 ones = mx.ones((n_scenarios, 1)) 

104 factors = mx.concatenate([ones, cum_returns], axis=1) 

105 

106 start_price = close[0] 

107 scenarios = start_price * factors 

108 

109 return scenarios 

110 

111 

112# Maintain backward compatibility for functional imports if needed 

113to_tensor = TensorOps.to_tensor 

114generate_shuffle_scenarios = TensorOps.generate_shuffle_scenarios 

115generate_noise_scenarios = TensorOps.generate_noise_scenarios