Coverage for src / monte_neo / data / sampler.py: 51%

1"""Data sampling module. 

2 

3Generate samples for Monte Carlo simulations. 

4""" 

5 

6from __future__ import annotations 

7 

8from typing import TYPE_CHECKING 

9 

10import numpy as np 

11import pandas as pd 

12 

13from monte_neo.utils.logger import get_logger 

14 

15if TYPE_CHECKING: 

16 pass 

17 

18logger = get_logger(__name__) 

19 

20 

21class DataSampler: 

22 """Generate samples for Monte Carlo simulations.""" 

23 

24 def __init__(self, random_seed: int | None = None) -> None: 

25 """Initialize sampler. 

26 

27 Args: 

28 random_seed: Random seed for reproducibility. 

29 """ 

30 self.rng = np.random.default_rng(random_seed) 

31 

32 def bootstrap( 

33 self, 

34 data: pd.DataFrame, 

35 n_samples: int = 1000, 

36 sample_size: int | None = None, 

37 ) -> list[pd.DataFrame]: 

38 """Generate bootstrap samples. 

39 

40 Args: 

41 data: Source DataFrame. 

42 n_samples: Number of samples to generate. 

43 sample_size: Size of each sample (default: same as data). 

44 

45 Returns: 

46 List of sampled DataFrames. 

47 """ 

48 if sample_size is None: 

49 sample_size = len(data) 

50 

51 samples = [] 

52 indices = np.arange(len(data)) 

53 

54 for _ in range(n_samples): 

55 sampled_idx = self.rng.choice(indices, size=sample_size, replace=True) 

56 sampled_idx.sort() # Keep time order 

57 samples.append(data.iloc[sampled_idx].copy()) 

58 

59 logger.debug(f"Generated {n_samples} bootstrap samples") 

60 return samples 

61 

62 def block_bootstrap( 

63 self, 

64 data: pd.DataFrame, 

65 n_samples: int = 1000, 

66 block_size: int | None = None, 

67 ) -> list[pd.DataFrame]: 

68 """Generate block bootstrap samples (preserves time structure). 

69 

70 Args: 

71 data: Source DataFrame. 

72 n_samples: Number of samples to generate. 

73 block_size: Size of each block (default: sqrt(len(data))). 

74 

75 Returns: 

76 List of sampled DataFrames. 

77 """ 

78 n = len(data) 

79 if block_size is None: 

80 block_size = max(1, int(np.sqrt(n))) 

81 

82 n_blocks = n // block_size 

83 samples = [] 

84 

85 # Vectorized index generation helper 

86 indices_range = np.arange(block_size) 

87 

88 for _ in range(n_samples): 

89 # Randomly select block start positions 

90 block_starts = self.rng.choice( 

91 n - block_size + 1, 

92 size=n_blocks, 

93 replace=True, 

94 ) 

95 

96 # Construct all indices at once 

97 # Broadcasting: (n_blocks, 1) + (block_size,) -> (n_blocks, block_size) 

98 full_indices = (block_starts[:, None] + indices_range).ravel() 

99 

100 # Single slice operation is much faster than loop + concat 

101 sample = data.iloc[full_indices].copy() 

102 sample.reset_index(drop=True, inplace=True) 

103 samples.append(sample) 

104 

105 logger.debug( 

106 f"Generated {n_samples} block bootstrap samples (block_size={block_size})" 

107 ) 

108 return samples 

109 

110 def circular_block_bootstrap( 

111 self, 

112 data: pd.DataFrame, 

113 n_samples: int = 1000, 

114 block_size: int | None = None, 

115 ) -> list[pd.DataFrame]: 

116 """Generate circular block bootstrap samples. 

117 

118 Treats data as circular to avoid edge effects. 

119 

120 Args: 

121 data: Source DataFrame. 

122 n_samples: Number of samples to generate. 

123 block_size: Size of each block. 

124 

125 Returns: 

126 List of sampled DataFrames. 

127 """ 

128 n = len(data) 

129 if block_size is None: 

130 block_size = max(1, int(np.sqrt(n))) 

131 

132 n_blocks = n // block_size 

133 samples = [] 

134 

135 # Create circular data 

136 circular_data = pd.concat([data, data], ignore_index=True) 

137 indices_range = np.arange(block_size) 

138 

139 for _ in range(n_samples): 

140 block_starts = self.rng.choice(n, size=n_blocks, replace=True) 

141 

142 # Vectorized index generation 

143 full_indices = (block_starts[:, None] + indices_range).ravel() 

144 

145 sample = circular_data.iloc[full_indices].copy() 

146 sample.reset_index(drop=True, inplace=True) 

147 samples.append(sample) 

148 

149 logger.debug(f"Generated {n_samples} circular block bootstrap samples") 

150 return samples 

151 

152 def stratified_sample( 

153 self, 

154 data: pd.DataFrame, 

155 n_samples: int = 1000, 

156 strata_column: str = "returns", 

157 n_strata: int = 10, 

158 ) -> list[pd.DataFrame]: 

159 """Generate stratified samples based on return distribution. 

160 

161 Args: 

162 data: Source DataFrame. 

163 n_samples: Number of samples to generate. 

164 strata_column: Column to stratify by. 

165 n_strata: Number of strata. 

166 

167 Returns: 

168 List of sampled DataFrames. 

169 """ 

170 # Calculate returns if not present 

171 df = data.copy() 

172 if strata_column not in df.columns: 

173 df["returns"] = df["close"].pct_change() 

174 strata_column = "returns" 

175 

176 # Create strata labels 

177 df["strata"] = pd.qcut( 

178 df[strata_column].fillna(0), 

179 q=n_strata, 

180 labels=False, 

181 duplicates="drop", 

182 ) 

183 

184 samples = [] 

185 strata_groups = df.groupby("strata") 

186 

187 for _ in range(n_samples): 

188 sampled_parts = [] 

189 for _, group in strata_groups: 

190 n_from_strata = max(1, len(group) // n_strata) 

191 sampled = group.sample(n=min(n_from_strata, len(group)), replace=True) 

192 sampled_parts.append(sampled) 

193 

194 sample = pd.concat(sampled_parts, ignore_index=True).sort_index() 

195 sample = sample.drop(columns=["strata"]) 

196 samples.append(sample) 

197 

198 logger.debug(f"Generated {n_samples} stratified samples") 

199 return samples 

200 

201 def inject_noise( 

202 self, 

203 data: pd.DataFrame, 

204 noise_scale: float = 0.001, 

205 ) -> pd.DataFrame: 

206 """Inject random noise into price data. 

207 

208 Args: 

209 data: Source DataFrame. 

210 noise_scale: Scale of random noise (percentage). 

211 

212 Returns: 

213 DataFrame with noisy prices. 

214 """ 

215 noisy = data.copy() 

216 

217 # Apply noise to close price 

218 noise = self.rng.normal(0, noise_scale, len(data)) 

219 noisy["close"] = noisy["close"] * (1 + noise) 

220 

221 # Adjust other prices to be consistent 

222 noisy["open"] = noisy["open"] * (1 + self.rng.normal(0, noise_scale, len(data))) 

223 noisy["high"] = noisy[["open", "close", "high"]].max(axis=1) 

224 noisy["low"] = noisy[["open", "close", "low"]].min(axis=1) 

225 

226 logger.debug(f"Injected noise with scale {noise_scale}") 

227 return noisy 

228 

229 def synthetic_data( 

230 self, 

231 data: pd.DataFrame, 

232 n_samples: int = 1000, 

233 method: str = "returns", 

234 ) -> list[pd.DataFrame]: 

235 """Generate synthetic data based on statistical properties. 

236 

237 Args: 

238 data: Source DataFrame. 

239 n_samples: Number of samples to generate. 

240 method: Generation method ('returns', 'garch'). 

241 

242 Returns: 

243 List of synthetic DataFrames. 

244 """ 

245 samples = [] 

246 returns = data["close"].pct_change().dropna() 

247 

248 mean_return = returns.mean() 

249 std_return = returns.std() 

250 

251 for _ in range(n_samples): 

252 # Generate synthetic returns 

253 synthetic_returns = self.rng.normal(mean_return, std_return, len(data)) 

254 

255 # Convert to prices 

256 initial_price = data["close"].iloc[0] 

257 synthetic_prices = initial_price * np.cumprod(1 + synthetic_returns) 

258 

259 # Create synthetic OHLCV 

260 sample = pd.DataFrame( 

261 { 

262 "open": synthetic_prices 

263 * (1 + self.rng.uniform(-0.002, 0.002, len(data))), 

264 "high": synthetic_prices 

265 * (1 + self.rng.uniform(0, 0.01, len(data))), 

266 "low": synthetic_prices 

267 * (1 - self.rng.uniform(0, 0.01, len(data))), 

268 "close": synthetic_prices, 

269 "volume": data["volume"].values 

270 * self.rng.uniform(0.5, 1.5, len(data)), 

271 } 

272 ) 

273 samples.append(sample) 

274 

275 logger.debug(f"Generated {n_samples} synthetic samples") 

276 return samples