Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/indexers/objects.py: 24%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

144 statements  

1"""Indexer objects for computing start/end window bounds for rolling operations""" 

2 

3from __future__ import annotations 

4 

5from datetime import timedelta 

6from typing import TYPE_CHECKING 

7 

8import numpy as np 

9 

10from pandas._libs.tslibs import BaseOffset 

11from pandas._libs.window.indexers import calculate_variable_window_bounds 

12from pandas.util._decorators import set_module 

13 

14from pandas.core.dtypes.common import ensure_platform_int 

15 

16from pandas.core.indexes.datetimes import DatetimeIndex 

17 

18from pandas.tseries.offsets import Nano 

19 

20if TYPE_CHECKING: 

21 from collections.abc import Hashable 

22 from typing import Any 

23 

24 

25@set_module("pandas.api.indexers") 

26class BaseIndexer: 

27 """ 

28 Base class for window bounds calculations. 

29 

30 Subclasses should implement ``get_window_bounds`` to define custom 

31 windowing logic for use with :meth:`DataFrame.rolling` or 

32 :meth:`Series.rolling`. 

33 

34 Parameters 

35 ---------- 

36 index_array : np.ndarray, default None 

37 Array-like structure representing the indices for the data points. 

38 If None, the default indices are assumed. This can be useful for 

39 handling non-uniform indices in data, such as in time series 

40 with irregular timestamps. 

41 window_size : int, default 0 

42 Size of the moving window. This is the number of observations used 

43 for calculating the statistic. The default is to consider all 

44 observations within the window. 

45 **kwargs 

46 Additional keyword arguments passed to the subclass's methods. 

47 

48 See Also 

49 -------- 

50 DataFrame.rolling : Provides rolling window calculations on dataframe. 

51 Series.rolling : Provides rolling window calculations on series. 

52 

53 Examples 

54 -------- 

55 >>> from pandas.api.indexers import BaseIndexer 

56 >>> class CustomIndexer(BaseIndexer): 

57 ... def get_window_bounds(self, num_values, min_periods, center, closed, step): 

58 ... start = np.arange(num_values, dtype=np.int64) 

59 ... end = np.arange(num_values, dtype=np.int64) + self.window_size 

60 ... return start, end 

61 >>> df = pd.DataFrame({"values": range(5)}) 

62 >>> indexer = CustomIndexer(window_size=2) 

63 >>> df.rolling(indexer).sum() 

64 values 

65 0 1.0 

66 1 3.0 

67 2 5.0 

68 3 7.0 

69 4 4.0 

70 """ 

71 

72 def __init__( 

73 self, index_array: np.ndarray | None = None, window_size: int = 0, **kwargs 

74 ) -> None: 

75 self.index_array = index_array 

76 self.window_size = window_size 

77 # Set user defined kwargs as attributes that can be used in get_window_bounds 

78 for key, value in kwargs.items(): 

79 setattr(self, key, value) 

80 

81 def get_window_bounds( 

82 self, 

83 num_values: int = 0, 

84 min_periods: int | None = None, 

85 center: bool | None = None, 

86 closed: str | None = None, 

87 step: int | None = None, 

88 ) -> tuple[np.ndarray, np.ndarray]: 

89 """ 

90 Computes the bounds of a window. 

91 

92 Parameters 

93 ---------- 

94 num_values : int, default 0 

95 number of values that will be aggregated over 

96 min_periods : int, default None 

97 min_periods passed from the top level rolling API 

98 center : bool, default None 

99 center passed from the top level rolling API 

100 closed : str, default None 

101 closed passed from the top level rolling API 

102 step : int, default None 

103 step passed from the top level rolling API 

104 

105 Returns 

106 ------- 

107 A tuple of ndarray[int64]s, indicating the boundaries of each 

108 window 

109 """ 

110 raise NotImplementedError 

111 

112 

113class FixedWindowIndexer(BaseIndexer): 

114 """Creates window boundaries that are of fixed length.""" 

115 

116 def get_window_bounds( 

117 self, 

118 num_values: int = 0, 

119 min_periods: int | None = None, 

120 center: bool | None = None, 

121 closed: str | None = None, 

122 step: int | None = None, 

123 ) -> tuple[np.ndarray, np.ndarray]: 

124 """ 

125 Computes the bounds of a window. 

126 

127 Parameters 

128 ---------- 

129 num_values : int, default 0 

130 number of values that will be aggregated over 

131 min_periods : int, default None 

132 min_periods passed from the top level rolling API 

133 center : bool, default None 

134 center passed from the top level rolling API 

135 closed : str, default None 

136 closed passed from the top level rolling API 

137 step : int, default None 

138 step passed from the top level rolling API 

139 

140 Returns 

141 ------- 

142 A tuple of ndarray[int64]s, indicating the boundaries of each 

143 window 

144 """ 

145 if center or self.window_size == 0: 

146 offset = (self.window_size - 1) // 2 

147 else: 

148 offset = 0 

149 

150 end = np.arange(1 + offset, num_values + 1 + offset, step, dtype="int64") 

151 start = end - self.window_size 

152 if closed in ["left", "both"]: 

153 start -= 1 

154 if closed in ["left", "neither"]: 

155 end -= 1 

156 

157 end = np.clip(end, 0, num_values) 

158 start = np.clip(start, 0, num_values) 

159 

160 return start, end 

161 

162 

163class VariableWindowIndexer(BaseIndexer): 

164 """Creates window boundaries that are of variable length, namely for time series.""" 

165 

166 def get_window_bounds( 

167 self, 

168 num_values: int = 0, 

169 min_periods: int | None = None, 

170 center: bool | None = None, 

171 closed: str | None = None, 

172 step: int | None = None, 

173 ) -> tuple[np.ndarray, np.ndarray]: 

174 """ 

175 Computes the bounds of a window. 

176 

177 Parameters 

178 ---------- 

179 num_values : int, default 0 

180 number of values that will be aggregated over 

181 min_periods : int, default None 

182 min_periods passed from the top level rolling API 

183 center : bool, default None 

184 center passed from the top level rolling API 

185 closed : str, default None 

186 closed passed from the top level rolling API 

187 step : int, default None 

188 step passed from the top level rolling API 

189 

190 Returns 

191 ------- 

192 A tuple of ndarray[int64]s, indicating the boundaries of each 

193 window 

194 """ 

195 assert self.index_array is not None 

196 if (index_length := len(self.index_array)) < num_values: 

197 raise ValueError( 

198 "Variable rolling window requires the index to be at least as long " 

199 f"as the 'other' index. Got {index_length} < {num_values}. " 

200 "Please align 'other' to the rolling object's index using " 

201 "reindex_like() or similar method." 

202 ) 

203 # error: Argument 4 to "calculate_variable_window_bounds" has incompatible 

204 # type "Optional[bool]"; expected "bool" 

205 return calculate_variable_window_bounds( 

206 num_values, 

207 self.window_size, 

208 min_periods, 

209 center, # type: ignore[arg-type] 

210 closed, 

211 self.index_array, 

212 ) 

213 

214 

215@set_module("pandas.api.indexers") 

216class VariableOffsetWindowIndexer(BaseIndexer): 

217 """ 

218 Calculate window boundaries based on a non-fixed offset such as a BusinessDay. 

219 

220 Unlike the default fixed-size window, this indexer uses a 

221 :class:`DateOffset` to determine variable-width window boundaries 

222 based on the datetime index of the data. 

223 

224 Parameters 

225 ---------- 

226 index_array : np.ndarray, default 0 

227 Array-like structure specifying the indices for data points. 

228 This parameter is currently not used. 

229 

230 window_size : int, optional, default 0 

231 Specifies the number of data points in each window. 

232 This parameter is currently not used. 

233 

234 index : DatetimeIndex, optional 

235 ``DatetimeIndex`` of the labels of each observation. 

236 

237 offset : BaseOffset, optional 

238 ``DateOffset`` representing the size of the window. 

239 

240 **kwargs 

241 Additional keyword arguments passed to the parent class ``BaseIndexer``. 

242 

243 See Also 

244 -------- 

245 api.indexers.BaseIndexer : Base class for all indexers. 

246 DataFrame.rolling : Rolling window calculations on DataFrames. 

247 offsets : Module providing various time offset classes. 

248 

249 Examples 

250 -------- 

251 >>> from pandas.api.indexers import VariableOffsetWindowIndexer 

252 >>> df = pd.DataFrame(range(10), index=pd.date_range("2020", periods=10)) 

253 >>> offset = pd.offsets.BDay(1) 

254 >>> indexer = VariableOffsetWindowIndexer(index=df.index, offset=offset) 

255 >>> df 

256 0 

257 2020-01-01 0 

258 2020-01-02 1 

259 2020-01-03 2 

260 2020-01-04 3 

261 2020-01-05 4 

262 2020-01-06 5 

263 2020-01-07 6 

264 2020-01-08 7 

265 2020-01-09 8 

266 2020-01-10 9 

267 >>> df.rolling(indexer).sum() 

268 0 

269 2020-01-01 0.0 

270 2020-01-02 1.0 

271 2020-01-03 2.0 

272 2020-01-04 3.0 

273 2020-01-05 7.0 

274 2020-01-06 12.0 

275 2020-01-07 6.0 

276 2020-01-08 7.0 

277 2020-01-09 8.0 

278 2020-01-10 9.0 

279 """ 

280 

281 def __init__( 

282 self, 

283 index_array: np.ndarray | None = None, 

284 window_size: int = 0, 

285 index: DatetimeIndex | None = None, 

286 offset: BaseOffset | None = None, 

287 **kwargs, 

288 ) -> None: 

289 super().__init__(index_array, window_size, **kwargs) 

290 if not isinstance(index, DatetimeIndex): 

291 raise ValueError("index must be a DatetimeIndex.") 

292 self.index = index 

293 if not isinstance(offset, BaseOffset): 

294 raise ValueError("offset must be a DateOffset-like object.") 

295 self.offset = offset 

296 

297 def get_window_bounds( 

298 self, 

299 num_values: int = 0, 

300 min_periods: int | None = None, 

301 center: bool | None = None, 

302 closed: str | None = None, 

303 step: int | None = None, 

304 ) -> tuple[np.ndarray, np.ndarray]: 

305 """ 

306 Computes the bounds of a window. 

307 

308 Parameters 

309 ---------- 

310 num_values : int, default 0 

311 number of values that will be aggregated over 

312 min_periods : int, default None 

313 min_periods passed from the top level rolling API 

314 center : bool, default None 

315 center passed from the top level rolling API 

316 closed : str, default None 

317 closed passed from the top level rolling API 

318 step : int, default None 

319 step passed from the top level rolling API 

320 

321 Returns 

322 ------- 

323 A tuple of ndarray[int64]s, indicating the boundaries of each 

324 window 

325 """ 

326 if step is not None: 

327 raise NotImplementedError("step not implemented for variable offset window") 

328 if num_values <= 0: 

329 return np.empty(0, dtype="int64"), np.empty(0, dtype="int64") 

330 

331 # if windows is variable, default is 'right', otherwise default is 'both' 

332 if closed is None: 

333 closed = "right" if self.index is not None else "both" 

334 

335 right_closed = closed in ["right", "both"] 

336 left_closed = closed in ["left", "both"] 

337 

338 if self.index[num_values - 1] < self.index[0]: 

339 index_growth_sign = -1 

340 else: 

341 index_growth_sign = 1 

342 offset_diff = index_growth_sign * self.offset 

343 

344 start = np.empty(num_values, dtype="int64") 

345 start.fill(-1) 

346 end = np.empty(num_values, dtype="int64") 

347 end.fill(-1) 

348 

349 start[0] = 0 

350 

351 # right endpoint is closed 

352 if right_closed: 

353 end[0] = 1 

354 # right endpoint is open 

355 else: 

356 end[0] = 0 

357 

358 zero = timedelta(0) 

359 # start is start of slice interval (including) 

360 # end is end of slice interval (not including) 

361 for i in range(1, num_values): 

362 end_bound = self.index[i] 

363 start_bound = end_bound - offset_diff 

364 

365 # left endpoint is closed 

366 if left_closed: 

367 start_bound -= Nano(1) 

368 

369 # advance the start bound until we are 

370 # within the constraint 

371 start[i] = i 

372 for j in range(start[i - 1], i): 

373 start_diff = (self.index[j] - start_bound) * index_growth_sign 

374 if start_diff > zero: 

375 start[i] = j 

376 break 

377 

378 # end bound is previous end 

379 # or current index 

380 end_diff = (self.index[end[i - 1]] - end_bound) * index_growth_sign 

381 if end_diff == zero and not right_closed: 

382 end[i] = end[i - 1] + 1 

383 elif end_diff <= zero: 

384 end[i] = i + 1 

385 else: 

386 end[i] = end[i - 1] 

387 

388 # right endpoint is open 

389 if not right_closed: 

390 end[i] -= 1 

391 

392 return start, end 

393 

394 

395class ExpandingIndexer(BaseIndexer): 

396 """Calculate expanding window bounds, mimicking df.expanding()""" 

397 

398 def get_window_bounds( 

399 self, 

400 num_values: int = 0, 

401 min_periods: int | None = None, 

402 center: bool | None = None, 

403 closed: str | None = None, 

404 step: int | None = None, 

405 ) -> tuple[np.ndarray, np.ndarray]: 

406 """ 

407 Computes the bounds of a window. 

408 

409 Parameters 

410 ---------- 

411 num_values : int, default 0 

412 number of values that will be aggregated over 

413 min_periods : int, default None 

414 min_periods passed from the top level rolling API 

415 center : bool, default None 

416 center passed from the top level rolling API 

417 closed : str, default None 

418 closed passed from the top level rolling API 

419 step : int, default None 

420 step passed from the top level rolling API 

421 

422 Returns 

423 ------- 

424 A tuple of ndarray[int64]s, indicating the boundaries of each 

425 window 

426 """ 

427 return ( 

428 np.zeros(num_values, dtype=np.int64), 

429 np.arange(1, num_values + 1, dtype=np.int64), 

430 ) 

431 

432 

433@set_module("pandas.api.indexers") 

434class FixedForwardWindowIndexer(BaseIndexer): 

435 """ 

436 Creates window boundaries for fixed-length windows that include the current row. 

437 

438 This indexer produces forward-looking windows, where each window starts at 

439 the current row and extends ``window_size`` rows ahead, unlike the default 

440 backward-looking behavior of :meth:`DataFrame.rolling`. 

441 

442 Parameters 

443 ---------- 

444 index_array : np.ndarray, default None 

445 Array-like structure representing the indices for the data points. 

446 If None, the default indices are assumed. This can be useful for 

447 handling non-uniform indices in data, such as in time series 

448 with irregular timestamps. 

449 window_size : int, default 0 

450 Size of the moving window. This is the number of observations used 

451 for calculating the statistic. The default is to consider all 

452 observations within the window. 

453 **kwargs 

454 Additional keyword arguments passed to the subclass's methods. 

455 

456 See Also 

457 -------- 

458 DataFrame.rolling : Provides rolling window calculations. 

459 api.indexers.VariableWindowIndexer : Calculate window bounds based on 

460 variable-sized windows. 

461 

462 Examples 

463 -------- 

464 >>> df = pd.DataFrame({"B": [0, 1, 2, np.nan, 4]}) 

465 >>> df 

466 B 

467 0 0.0 

468 1 1.0 

469 2 2.0 

470 3 NaN 

471 4 4.0 

472 

473 >>> indexer = pd.api.indexers.FixedForwardWindowIndexer(window_size=2) 

474 >>> df.rolling(window=indexer, min_periods=1).sum() 

475 B 

476 0 1.0 

477 1 3.0 

478 2 2.0 

479 3 4.0 

480 4 4.0 

481 """ 

482 

483 def get_window_bounds( 

484 self, 

485 num_values: int = 0, 

486 min_periods: int | None = None, 

487 center: bool | None = None, 

488 closed: str | None = None, 

489 step: int | None = None, 

490 ) -> tuple[np.ndarray, np.ndarray]: 

491 """ 

492 Computes the bounds of a window. 

493 

494 Parameters 

495 ---------- 

496 num_values : int, default 0 

497 number of values that will be aggregated over 

498 min_periods : int, default None 

499 min_periods passed from the top level rolling API 

500 center : bool, default None 

501 center passed from the top level rolling API 

502 closed : str, default None 

503 closed passed from the top level rolling API 

504 step : int, default None 

505 step passed from the top level rolling API 

506 

507 Returns 

508 ------- 

509 A tuple of ndarray[int64]s, indicating the boundaries of each 

510 window 

511 """ 

512 if center: 

513 raise ValueError("Forward-looking windows can't have center=True") 

514 if closed is not None: 

515 raise ValueError( 

516 "Forward-looking windows don't support setting the closed argument" 

517 ) 

518 if step is None: 

519 step = 1 

520 

521 start = np.arange(0, num_values, step, dtype="int64") 

522 end = start + self.window_size 

523 if self.window_size: 

524 end = np.clip(end, 0, num_values) 

525 

526 return start, end 

527 

528 

529class GroupByIndexer(BaseIndexer): 

530 """Calculate bounds to compute groupby rolling, mimicking df.groupby().rolling()""" 

531 

532 def __init__( 

533 self, 

534 index_array: np.ndarray | None = None, 

535 window_size: int | BaseIndexer = 0, 

536 groupby_indices: dict[Hashable, np.ndarray] | None = None, 

537 window_indexer: type[BaseIndexer] = BaseIndexer, 

538 indexer_kwargs: dict[str, Any] | None = None, 

539 **kwargs, 

540 ) -> None: 

541 """ 

542 Parameters 

543 ---------- 

544 index_array : np.ndarray or None 

545 np.ndarray of the index of the original object that we are performing 

546 a chained groupby operation over. This index has been pre-sorted relative to 

547 the groups 

548 window_size : int or BaseIndexer 

549 window size during the windowing operation 

550 groupby_indices : dict or None 

551 dict of {group label: [positional index of rows belonging to the group]} 

552 window_indexer : BaseIndexer 

553 BaseIndexer class determining the start and end bounds of each group 

554 indexer_kwargs : dict or None 

555 Custom kwargs to be passed to window_indexer 

556 **kwargs : 

557 keyword arguments that will be available when get_window_bounds is called 

558 """ 

559 self.groupby_indices = groupby_indices or {} 

560 self.window_indexer = window_indexer 

561 self.indexer_kwargs = indexer_kwargs.copy() if indexer_kwargs else {} 

562 super().__init__( 

563 index_array=index_array, 

564 window_size=self.indexer_kwargs.pop("window_size", window_size), 

565 **kwargs, 

566 ) 

567 

568 def get_window_bounds( 

569 self, 

570 num_values: int = 0, 

571 min_periods: int | None = None, 

572 center: bool | None = None, 

573 closed: str | None = None, 

574 step: int | None = None, 

575 ) -> tuple[np.ndarray, np.ndarray]: 

576 """ 

577 Computes the bounds of a window. 

578 

579 Parameters 

580 ---------- 

581 num_values : int, default 0 

582 number of values that will be aggregated over 

583 min_periods : int, default None 

584 min_periods passed from the top level rolling API 

585 center : bool, default None 

586 center passed from the top level rolling API 

587 closed : str, default None 

588 closed passed from the top level rolling API 

589 step : int, default None 

590 step passed from the top level rolling API 

591 

592 Returns 

593 ------- 

594 A tuple of ndarray[int64]s, indicating the boundaries of each 

595 window 

596 """ 

597 # 1) For each group, get the indices that belong to the group 

598 # 2) Use the indices to calculate the start & end bounds of the window 

599 # 3) Append the window bounds in group order 

600 start_arrays = [] 

601 end_arrays = [] 

602 window_indices_start = 0 

603 for indices in self.groupby_indices.values(): 

604 index_array: np.ndarray | None 

605 

606 if self.index_array is not None: 

607 index_array = self.index_array.take(ensure_platform_int(indices)) 

608 else: 

609 index_array = self.index_array 

610 indexer = self.window_indexer( 

611 index_array=index_array, 

612 window_size=self.window_size, 

613 **self.indexer_kwargs, 

614 ) 

615 start, end = indexer.get_window_bounds( 

616 len(indices), min_periods, center, closed, step 

617 ) 

618 start = start.astype(np.int64) 

619 end = end.astype(np.int64) 

620 assert len(start) == len(end), ( 

621 "these should be equal in length from get_window_bounds" 

622 ) 

623 # Cannot use groupby_indices as they might not be monotonic with the object 

624 # we're rolling over 

625 window_indices = np.arange( 

626 window_indices_start, window_indices_start + len(indices) 

627 ) 

628 window_indices_start += len(indices) 

629 # Extend as we'll be slicing window like [start, end) 

630 window_indices = np.append(window_indices, [window_indices[-1] + 1]).astype( 

631 np.int64, copy=False 

632 ) 

633 start_arrays.append(window_indices.take(ensure_platform_int(start))) 

634 end_arrays.append(window_indices.take(ensure_platform_int(end))) 

635 if len(start_arrays) == 0: 

636 return np.array([], dtype=np.int64), np.array([], dtype=np.int64) 

637 start = np.concatenate(start_arrays) 

638 end = np.concatenate(end_arrays) 

639 return start, end 

640 

641 

642class ExponentialMovingWindowIndexer(BaseIndexer): 

643 """Calculate ewm window bounds (the entire window)""" 

644 

645 def get_window_bounds( 

646 self, 

647 num_values: int = 0, 

648 min_periods: int | None = None, 

649 center: bool | None = None, 

650 closed: str | None = None, 

651 step: int | None = None, 

652 ) -> tuple[np.ndarray, np.ndarray]: 

653 """ 

654 Computes the bounds of a window. 

655 

656 Parameters 

657 ---------- 

658 num_values : int, default 0 

659 number of values that will be aggregated over 

660 min_periods : int, default None 

661 min_periods passed from the top level rolling API 

662 center : bool, default None 

663 center passed from the top level rolling API 

664 closed : str, default None 

665 closed passed from the top level rolling API 

666 step : int, default None 

667 step passed from the top level rolling API 

668 

669 Returns 

670 ------- 

671 A tuple of ndarray[int64]s, indicating the boundaries of each 

672 window 

673 """ 

674 return np.array([0], dtype=np.int64), np.array([num_values], dtype=np.int64)