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 nogil: bool,
25 parallel: bool,
26):
27 """
28 Generate a numba jitted apply function specified by values from engine_kwargs.
30 1. jit the user's function
31 2. Return a rolling apply function with the jitted function inline
33 Configurations specified in engine_kwargs apply to both the user's
34 function _AND_ the rolling apply function.
36 Parameters
37 ----------
38 func : function
39 function to be applied to each window and will be JITed
40 nogil : bool
41 nogil to be passed into numba.jit
42 parallel : bool
43 parallel to be passed into numba.jit
45 Returns
46 -------
47 Numba function
48 """
49 numba_func = jit_user_function(func)
50 if TYPE_CHECKING:
51 import numba
52 else:
53 numba = import_optional_dependency("numba")
55 @numba.jit(nogil=nogil, parallel=parallel)
56 def roll_apply(
57 values: np.ndarray,
58 begin: np.ndarray,
59 end: np.ndarray,
60 minimum_periods: int,
61 *args: Any,
62 ) -> np.ndarray:
63 result = np.empty(len(begin))
64 for i in numba.prange(len(result)):
65 start = begin[i]
66 stop = end[i]
67 window = values[start:stop]
68 count_nan = np.sum(np.isnan(window))
69 if len(window) - count_nan >= minimum_periods:
70 result[i] = numba_func(window, *args)
71 else:
72 result[i] = np.nan
73 return result
75 return roll_apply
78@functools.cache
79def generate_numba_ewm_func(
80 nogil: bool,
81 parallel: bool,
82 com: float,
83 adjust: bool,
84 ignore_na: bool,
85 deltas: tuple[float, ...],
86 normalize: bool,
87):
88 """
89 Generate a numba jitted ewm mean or sum function specified by values
90 from engine_kwargs.
92 Parameters
93 ----------
94 nogil : bool
95 nogil to be passed into numba.jit
96 parallel : bool
97 parallel to be passed into numba.jit
98 com : float
99 adjust : bool
100 ignore_na : bool
101 deltas : tuple
102 normalize : bool
104 Returns
105 -------
106 Numba function
107 """
108 if TYPE_CHECKING:
109 import numba
110 else:
111 numba = import_optional_dependency("numba")
113 @numba.jit(nogil=nogil, parallel=parallel)
114 def ewm(
115 values: np.ndarray,
116 begin: np.ndarray,
117 end: np.ndarray,
118 minimum_periods: int,
119 ) -> np.ndarray:
120 result = np.empty(len(values))
121 alpha = 1.0 / (1.0 + com)
122 old_wt_factor = 1.0 - alpha
123 new_wt = 1.0 if adjust else alpha
125 for i in numba.prange(len(begin)):
126 start = begin[i]
127 stop = end[i]
128 window = values[start:stop]
129 sub_result = np.empty(len(window))
131 weighted = window[0]
132 nobs = int(not np.isnan(weighted))
133 sub_result[0] = weighted if nobs >= minimum_periods else np.nan
134 old_wt = 1.0
136 for j in range(1, len(window)):
137 cur = window[j]
138 is_observation = not np.isnan(cur)
139 nobs += is_observation
140 if not np.isnan(weighted):
141 if is_observation or not ignore_na:
142 if normalize:
143 # note that len(deltas) = len(vals) - 1 and deltas[i]
144 # is to be used in conjunction with vals[i+1]
145 old_wt *= old_wt_factor ** deltas[start + j - 1]
146 if not adjust and com == 1:
147 # update in case of irregular-interval time series
148 new_wt = 1.0 - old_wt
149 else:
150 weighted = old_wt_factor * weighted
151 if is_observation:
152 if normalize:
153 # avoid numerical errors on constant series
154 if weighted != cur:
155 weighted = old_wt * weighted + new_wt * cur
156 if normalize:
157 weighted = weighted / (old_wt + new_wt)
158 if adjust:
159 old_wt += new_wt
160 else:
161 old_wt = 1.0
162 else:
163 weighted += cur
164 elif is_observation:
165 weighted = cur
167 sub_result[j] = weighted if nobs >= minimum_periods else np.nan
169 result[start:stop] = sub_result
171 return result
173 return ewm
176@functools.cache
177def generate_numba_table_func(
178 func: Callable[..., np.ndarray],
179 nogil: bool,
180 parallel: bool,
181):
182 """
183 Generate a numba jitted function to apply window calculations table-wise.
185 Func will be passed an M window size x N number of columns array, and
186 must return a 1 x N number of columns array.
188 1. jit the user's function
189 2. Return a rolling apply function with the jitted function inline
191 Parameters
192 ----------
193 func : function
194 function to be applied to each window and will be JITed
195 nogil : bool
196 nogil to be passed into numba.jit
197 parallel : bool
198 parallel to be passed into numba.jit
200 Returns
201 -------
202 Numba function
203 """
204 numba_func = jit_user_function(func)
205 if TYPE_CHECKING:
206 import numba
207 else:
208 numba = import_optional_dependency("numba")
210 @numba.jit(nogil=nogil, parallel=parallel)
211 def roll_table(
212 values: np.ndarray,
213 begin: np.ndarray,
214 end: np.ndarray,
215 minimum_periods: int,
216 *args: Any,
217 ):
218 result = np.empty((len(begin), values.shape[1]))
219 min_periods_mask = np.empty(result.shape)
220 for i in numba.prange(len(result)):
221 start = begin[i]
222 stop = end[i]
223 window = values[start:stop]
224 count_nan = np.sum(np.isnan(window), axis=0)
225 nan_mask = len(window) - count_nan >= minimum_periods
226 if nan_mask.any():
227 result[i, :] = numba_func(window, *args)
228 min_periods_mask[i, :] = nan_mask
229 result = np.where(min_periods_mask, result, np.nan)
230 return result
232 return roll_table
235# This function will no longer be needed once numba supports
236# axis for all np.nan* agg functions
237# https://github.com/numba/numba/issues/1269
238@functools.cache
239def generate_manual_numpy_nan_agg_with_axis(nan_func):
240 if TYPE_CHECKING:
241 import numba
242 else:
243 numba = import_optional_dependency("numba")
245 # Not jitted directly: this is passed through jit_user_function in
246 # generate_numba_table_func, which wraps it with register_jitable so it
247 # inlines into roll_table's outer jit. A standalone @jit(parallel=True)
248 # here nests parallel regions and crashes the workqueue threading layer
249 # on platforms without tbb/openmp (GH#40454).
250 def nan_agg_with_axis(table):
251 result = np.empty(table.shape[1])
252 for i in numba.prange(table.shape[1]):
253 partition = table[:, i]
254 result[i] = nan_func(partition)
255 return result
257 return nan_agg_with_axis
260@functools.cache
261def generate_numba_ewm_table_func(
262 nogil: bool,
263 parallel: bool,
264 com: float,
265 adjust: bool,
266 ignore_na: bool,
267 deltas: tuple[float, ...],
268 normalize: bool,
269):
270 """
271 Generate a numba jitted ewm mean or sum function applied table wise specified
272 by values from engine_kwargs.
274 Parameters
275 ----------
276 nogil : bool
277 nogil to be passed into numba.jit
278 parallel : bool
279 parallel to be passed into numba.jit
280 com : float
281 adjust : bool
282 ignore_na : bool
283 deltas : tuple
284 normalize: bool
286 Returns
287 -------
288 Numba function
289 """
290 if TYPE_CHECKING:
291 import numba
292 else:
293 numba = import_optional_dependency("numba")
295 @numba.jit(nogil=nogil, parallel=parallel)
296 def ewm_table(
297 values: np.ndarray,
298 begin: np.ndarray,
299 end: np.ndarray,
300 minimum_periods: int,
301 ) -> np.ndarray:
302 alpha = 1.0 / (1.0 + com)
303 old_wt_factor = 1.0 - alpha
304 new_wt = 1.0 if adjust else alpha
305 old_wt = np.ones(values.shape[1])
307 result = np.empty(values.shape)
308 weighted = values[0].copy()
309 nobs = (~np.isnan(weighted)).astype(np.int64)
310 result[0] = np.where(nobs >= minimum_periods, weighted, np.nan)
311 for i in range(1, len(values)):
312 cur = values[i]
313 is_observations = ~np.isnan(cur)
314 nobs += is_observations.astype(np.int64)
315 for j in numba.prange(len(cur)):
316 if not np.isnan(weighted[j]):
317 if is_observations[j] or not ignore_na:
318 if normalize:
319 # note that len(deltas) = len(vals) - 1 and deltas[i]
320 # is to be used in conjunction with vals[i+1]
321 old_wt[j] *= old_wt_factor ** deltas[i - 1]
322 if not adjust and com == 1:
323 # update in case of irregular-interval time series
324 new_wt = 1.0 - old_wt[j]
325 else:
326 weighted[j] = old_wt_factor * weighted[j]
327 if is_observations[j]:
328 if normalize:
329 # avoid numerical errors on constant series
330 if weighted[j] != cur[j]:
331 weighted[j] = (
332 old_wt[j] * weighted[j] + new_wt * cur[j]
333 )
334 if normalize:
335 weighted[j] = weighted[j] / (old_wt[j] + new_wt)
336 if adjust:
337 old_wt[j] += new_wt
338 else:
339 old_wt[j] = 1.0
340 else:
341 weighted[j] += cur[j]
342 elif is_observations[j]:
343 weighted[j] = cur[j]
345 result[i] = np.where(nobs >= minimum_periods, weighted, np.nan)
347 return result
349 return ewm_table