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