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 (
5 TYPE_CHECKING,
6 Any,
7)
9import numpy as np
11from pandas.compat._optional import import_optional_dependency
13from pandas.core.util.numba_ import jit_user_function
15if TYPE_CHECKING:
16 from collections.abc import Callable
18 from pandas._typing import Scalar
21@functools.cache
22def generate_numba_apply_func(
23 func: Callable[..., Scalar],
24 nopython: bool,
25 nogil: bool,
26 parallel: bool,
27):
28 """
29 Generate a numba jitted apply function specified by values from engine_kwargs.
31 1. jit the user's function
32 2. Return a rolling apply function with the jitted function inline
34 Configurations specified in engine_kwargs apply to both the user's
35 function _AND_ the rolling apply function.
37 Parameters
38 ----------
39 func : function
40 function to be applied to each window and will be JITed
41 nopython : bool
42 nopython to be passed into numba.jit
43 nogil : bool
44 nogil to be passed into numba.jit
45 parallel : bool
46 parallel to be passed into numba.jit
48 Returns
49 -------
50 Numba function
51 """
52 numba_func = jit_user_function(func)
53 if TYPE_CHECKING:
54 import numba
55 else:
56 numba = import_optional_dependency("numba")
58 @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)
59 def roll_apply(
60 values: np.ndarray,
61 begin: np.ndarray,
62 end: np.ndarray,
63 minimum_periods: int,
64 *args: Any,
65 ) -> np.ndarray:
66 result = np.empty(len(begin))
67 for i in numba.prange(len(result)):
68 start = begin[i]
69 stop = end[i]
70 window = values[start:stop]
71 count_nan = np.sum(np.isnan(window))
72 if len(window) - count_nan >= minimum_periods:
73 result[i] = numba_func(window, *args)
74 else:
75 result[i] = np.nan
76 return result
78 return roll_apply
81@functools.cache
82def generate_numba_ewm_func(
83 nopython: bool,
84 nogil: bool,
85 parallel: bool,
86 com: float,
87 adjust: bool,
88 ignore_na: bool,
89 deltas: tuple,
90 normalize: bool,
91):
92 """
93 Generate a numba jitted ewm mean or sum function specified by values
94 from engine_kwargs.
96 Parameters
97 ----------
98 nopython : bool
99 nopython to be passed into numba.jit
100 nogil : bool
101 nogil to be passed into numba.jit
102 parallel : bool
103 parallel to be passed into numba.jit
104 com : float
105 adjust : bool
106 ignore_na : bool
107 deltas : tuple
108 normalize : bool
110 Returns
111 -------
112 Numba function
113 """
114 if TYPE_CHECKING:
115 import numba
116 else:
117 numba = import_optional_dependency("numba")
119 @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)
120 def ewm(
121 values: np.ndarray,
122 begin: np.ndarray,
123 end: np.ndarray,
124 minimum_periods: int,
125 ) -> np.ndarray:
126 result = np.empty(len(values))
127 alpha = 1.0 / (1.0 + com)
128 old_wt_factor = 1.0 - alpha
129 new_wt = 1.0 if adjust else alpha
131 for i in numba.prange(len(begin)):
132 start = begin[i]
133 stop = end[i]
134 window = values[start:stop]
135 sub_result = np.empty(len(window))
137 weighted = window[0]
138 nobs = int(not np.isnan(weighted))
139 sub_result[0] = weighted if nobs >= minimum_periods else np.nan
140 old_wt = 1.0
142 for j in range(1, len(window)):
143 cur = window[j]
144 is_observation = not np.isnan(cur)
145 nobs += is_observation
146 if not np.isnan(weighted):
147 if is_observation or not ignore_na:
148 if normalize:
149 # note that len(deltas) = len(vals) - 1 and deltas[i]
150 # is to be used in conjunction with vals[i+1]
151 old_wt *= old_wt_factor ** deltas[start + j - 1]
152 if not adjust and com == 1:
153 # update in case of irregular-interval time series
154 new_wt = 1.0 - old_wt
155 else:
156 weighted = old_wt_factor * weighted
157 if is_observation:
158 if normalize:
159 # avoid numerical errors on constant series
160 if weighted != cur:
161 weighted = old_wt * weighted + new_wt * cur
162 if normalize:
163 weighted = weighted / (old_wt + new_wt)
164 if adjust:
165 old_wt += new_wt
166 else:
167 old_wt = 1.0
168 else:
169 weighted += cur
170 elif is_observation:
171 weighted = cur
173 sub_result[j] = weighted if nobs >= minimum_periods else np.nan
175 result[start:stop] = sub_result
177 return result
179 return ewm
182@functools.cache
183def generate_numba_table_func(
184 func: Callable[..., np.ndarray],
185 nopython: bool,
186 nogil: bool,
187 parallel: bool,
188):
189 """
190 Generate a numba jitted function to apply window calculations table-wise.
192 Func will be passed an M window size x N number of columns array, and
193 must return a 1 x N number of columns array.
195 1. jit the user's function
196 2. Return a rolling apply function with the jitted function inline
198 Parameters
199 ----------
200 func : function
201 function to be applied to each window and will be JITed
202 nopython : bool
203 nopython to be passed into numba.jit
204 nogil : bool
205 nogil to be passed into numba.jit
206 parallel : bool
207 parallel to be passed into numba.jit
209 Returns
210 -------
211 Numba function
212 """
213 numba_func = jit_user_function(func)
214 if TYPE_CHECKING:
215 import numba
216 else:
217 numba = import_optional_dependency("numba")
219 @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)
220 def roll_table(
221 values: np.ndarray,
222 begin: np.ndarray,
223 end: np.ndarray,
224 minimum_periods: int,
225 *args: Any,
226 ):
227 result = np.empty((len(begin), values.shape[1]))
228 min_periods_mask = np.empty(result.shape)
229 for i in numba.prange(len(result)):
230 start = begin[i]
231 stop = end[i]
232 window = values[start:stop]
233 count_nan = np.sum(np.isnan(window), axis=0)
234 nan_mask = len(window) - count_nan >= minimum_periods
235 if nan_mask.any():
236 result[i, :] = numba_func(window, *args)
237 min_periods_mask[i, :] = nan_mask
238 result = np.where(min_periods_mask, result, np.nan)
239 return result
241 return roll_table
244# This function will no longer be needed once numba supports
245# axis for all np.nan* agg functions
246# https://github.com/numba/numba/issues/1269
247@functools.cache
248def generate_manual_numpy_nan_agg_with_axis(nan_func):
249 if TYPE_CHECKING:
250 import numba
251 else:
252 numba = import_optional_dependency("numba")
254 @numba.jit(nopython=True, nogil=True, parallel=True)
255 def nan_agg_with_axis(table):
256 result = np.empty(table.shape[1])
257 for i in numba.prange(table.shape[1]):
258 partition = table[:, i]
259 result[i] = nan_func(partition)
260 return result
262 return nan_agg_with_axis
265@functools.cache
266def generate_numba_ewm_table_func(
267 nopython: bool,
268 nogil: bool,
269 parallel: bool,
270 com: float,
271 adjust: bool,
272 ignore_na: bool,
273 deltas: tuple,
274 normalize: bool,
275):
276 """
277 Generate a numba jitted ewm mean or sum function applied table wise specified
278 by values from engine_kwargs.
280 Parameters
281 ----------
282 nopython : bool
283 nopython to be passed into numba.jit
284 nogil : bool
285 nogil to be passed into numba.jit
286 parallel : bool
287 parallel to be passed into numba.jit
288 com : float
289 adjust : bool
290 ignore_na : bool
291 deltas : tuple
292 normalize: bool
294 Returns
295 -------
296 Numba function
297 """
298 if TYPE_CHECKING:
299 import numba
300 else:
301 numba = import_optional_dependency("numba")
303 @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)
304 def ewm_table(
305 values: np.ndarray,
306 begin: np.ndarray,
307 end: np.ndarray,
308 minimum_periods: int,
309 ) -> np.ndarray:
310 alpha = 1.0 / (1.0 + com)
311 old_wt_factor = 1.0 - alpha
312 new_wt = 1.0 if adjust else alpha
313 old_wt = np.ones(values.shape[1])
315 result = np.empty(values.shape)
316 weighted = values[0].copy()
317 nobs = (~np.isnan(weighted)).astype(np.int64)
318 result[0] = np.where(nobs >= minimum_periods, weighted, np.nan)
319 for i in range(1, len(values)):
320 cur = values[i]
321 is_observations = ~np.isnan(cur)
322 nobs += is_observations.astype(np.int64)
323 for j in numba.prange(len(cur)):
324 if not np.isnan(weighted[j]):
325 if is_observations[j] or not ignore_na:
326 if normalize:
327 # note that len(deltas) = len(vals) - 1 and deltas[i]
328 # is to be used in conjunction with vals[i+1]
329 old_wt[j] *= old_wt_factor ** deltas[i - 1]
330 if not adjust and com == 1:
331 # update in case of irregular-interval time series
332 new_wt = 1.0 - old_wt[j]
333 else:
334 weighted[j] = old_wt_factor * weighted[j]
335 if is_observations[j]:
336 if normalize:
337 # avoid numerical errors on constant series
338 if weighted[j] != cur[j]:
339 weighted[j] = (
340 old_wt[j] * weighted[j] + new_wt * cur[j]
341 )
342 if normalize:
343 weighted[j] = weighted[j] / (old_wt[j] + new_wt)
344 if adjust:
345 old_wt[j] += new_wt
346 else:
347 old_wt[j] = 1.0
348 else:
349 weighted[j] += cur[j]
350 elif is_observations[j]:
351 weighted[j] = cur[j]
353 result[i] = np.where(nobs >= minimum_periods, weighted, np.nan)
355 return result
357 return ewm_table