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
1from __future__ import annotations
3import functools
4from typing import TYPE_CHECKING
6import numpy as np
8from pandas.compat._optional import import_optional_dependency
10if TYPE_CHECKING:
11 from collections.abc import Callable
14@functools.cache
15def generate_online_numba_ewma_func(
16 nogil: bool,
17 parallel: bool,
18) -> Callable[
19 [np.ndarray, np.ndarray, int, float, float, np.ndarray, bool, bool],
20 tuple[np.ndarray, np.ndarray],
21]:
22 """
23 Generate a numba jitted groupby ewma function specified by values
24 from engine_kwargs.
26 Parameters
27 ----------
28 nogil : bool
29 nogil to be passed into numba.jit
30 parallel : bool
31 parallel to be passed into numba.jit
33 Returns
34 -------
35 Numba function
36 """
37 if TYPE_CHECKING:
38 import numba
39 else:
40 numba = import_optional_dependency("numba")
42 @numba.jit(nogil=nogil, parallel=parallel)
43 def online_ewma(
44 values: np.ndarray,
45 deltas: np.ndarray,
46 minimum_periods: int,
47 old_wt_factor: float,
48 new_wt: float,
49 old_wt: np.ndarray,
50 adjust: bool,
51 ignore_na: bool,
52 ) -> tuple[np.ndarray, np.ndarray]:
53 """
54 Compute online exponentially weighted mean per column over 2D values.
56 Takes the first observation as is, then computes the subsequent
57 exponentially weighted mean accounting minimum periods.
58 """
59 result = np.empty(values.shape)
60 weighted_avg = values[0].copy()
61 nobs = (~np.isnan(weighted_avg)).astype(np.int64)
62 result[0] = np.where(nobs >= minimum_periods, weighted_avg, np.nan)
64 for i in range(1, len(values)):
65 cur = values[i]
66 is_observations = ~np.isnan(cur)
67 nobs += is_observations.astype(np.int64)
68 for j in numba.prange(len(cur)):
69 if not np.isnan(weighted_avg[j]):
70 if is_observations[j] or not ignore_na:
71 # note that len(deltas) = len(vals) - 1 and deltas[i] is to be
72 # used in conjunction with vals[i+1]
73 old_wt[j] *= old_wt_factor ** deltas[i - 1]
74 if is_observations[j]:
75 # avoid numerical errors on constant series
76 if weighted_avg[j] != cur[j]:
77 weighted_avg[j] = (
78 (old_wt[j] * weighted_avg[j]) + (new_wt * cur[j])
79 ) / (old_wt[j] + new_wt)
80 if adjust:
81 old_wt[j] += new_wt
82 else:
83 old_wt[j] = 1.0
84 elif is_observations[j]:
85 weighted_avg[j] = cur[j]
87 result[i] = np.where(nobs >= minimum_periods, weighted_avg, np.nan)
89 return result, old_wt
91 return online_ewma
94class EWMMeanState:
95 def __init__(
96 self,
97 com: float,
98 adjust: bool,
99 ignore_na: bool,
100 shape: tuple[int, ...],
101 ) -> None:
102 alpha = 1.0 / (1.0 + com)
103 self.shape = shape
104 self.adjust = adjust
105 self.ignore_na = ignore_na
106 self.new_wt = 1.0 if adjust else alpha
107 self.old_wt_factor = 1.0 - alpha
108 self.old_wt = np.ones(self.shape[-1])
109 self.last_ewm: np.ndarray | None = None
111 def run_ewm(
112 self,
113 weighted_avg: np.ndarray,
114 deltas: np.ndarray,
115 min_periods: int,
116 ewm_func: Callable[
117 [np.ndarray, np.ndarray, int, float, float, np.ndarray, bool, bool],
118 tuple[np.ndarray, np.ndarray],
119 ],
120 ) -> np.ndarray:
121 result, old_wt = ewm_func(
122 weighted_avg,
123 deltas,
124 min_periods,
125 self.old_wt_factor,
126 self.new_wt,
127 self.old_wt,
128 self.adjust,
129 self.ignore_na,
130 )
131 self.old_wt = old_wt
132 self.last_ewm = result[-1]
133 return result
135 def reset(self) -> None:
136 self.old_wt = np.ones(self.shape[-1])
137 self.last_ewm = None