Skip to content

Commit 54dda90

Browse files
authored
PERF: release gil for ewma_time (#37389)
1 parent e38e987 commit 54dda90

File tree

1 file changed

+29
-20
lines changed

1 file changed

+29
-20
lines changed

pandas/_libs/window/aggregations.pyx

+29-20
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
11
# cython: boundscheck=False, wraparound=False, cdivision=True
22

33
import cython
4-
from cython import Py_ssize_t
54

65
from libcpp.deque cimport deque
76

87
import numpy as np
98

109
cimport numpy as cnp
11-
from numpy cimport float32_t, float64_t, int64_t, ndarray, uint8_t
10+
from numpy cimport float32_t, float64_t, int64_t, ndarray
1211

1312
cnp.import_array()
1413

@@ -1398,7 +1397,7 @@ def roll_weighted_var(float64_t[:] values, float64_t[:] weights,
13981397
# ----------------------------------------------------------------------
13991398
# Exponentially weighted moving average
14001399

1401-
def ewma_time(ndarray[float64_t] vals, int minp, ndarray[int64_t] times,
1400+
def ewma_time(const float64_t[:] vals, int minp, ndarray[int64_t] times,
14021401
int64_t halflife):
14031402
"""
14041403
Compute exponentially-weighted moving average using halflife and time
@@ -1416,30 +1415,40 @@ def ewma_time(ndarray[float64_t] vals, int minp, ndarray[int64_t] times,
14161415
ndarray
14171416
"""
14181417
cdef:
1419-
Py_ssize_t i, num_not_nan = 0, N = len(vals)
1418+
Py_ssize_t i, j, num_not_nan = 0, N = len(vals)
14201419
bint is_not_nan
1421-
float64_t last_result
1422-
ndarray[uint8_t] mask = np.zeros(N, dtype=np.uint8)
1423-
ndarray[float64_t] weights, observations, output = np.empty(N, dtype=np.float64)
1420+
float64_t last_result, weights_dot, weights_sum, weight, halflife_float
1421+
float64_t[:] times_float
1422+
float64_t[:] observations = np.zeros(N, dtype=float)
1423+
float64_t[:] times_masked = np.zeros(N, dtype=float)
1424+
ndarray[float64_t] output = np.empty(N, dtype=float)
14241425

14251426
if N == 0:
14261427
return output
14271428

1429+
halflife_float = <float64_t>halflife
1430+
times_float = times.astype(float)
14281431
last_result = vals[0]
14291432

1430-
for i in range(N):
1431-
is_not_nan = vals[i] == vals[i]
1432-
num_not_nan += is_not_nan
1433-
if is_not_nan:
1434-
mask[i] = 1
1435-
weights = 0.5 ** ((times[i] - times[mask.view(np.bool_)]) / halflife)
1436-
observations = vals[mask.view(np.bool_)]
1437-
last_result = np.sum(weights * observations) / np.sum(weights)
1438-
1439-
if num_not_nan >= minp:
1440-
output[i] = last_result
1441-
else:
1442-
output[i] = NaN
1433+
with nogil:
1434+
for i in range(N):
1435+
is_not_nan = vals[i] == vals[i]
1436+
num_not_nan += is_not_nan
1437+
if is_not_nan:
1438+
times_masked[num_not_nan-1] = times_float[i]
1439+
observations[num_not_nan-1] = vals[i]
1440+
1441+
weights_sum = 0
1442+
weights_dot = 0
1443+
for j in range(num_not_nan):
1444+
weight = 0.5 ** (
1445+
(times_float[i] - times_masked[j]) / halflife_float)
1446+
weights_sum += weight
1447+
weights_dot += weight * observations[j]
1448+
1449+
last_result = weights_dot / weights_sum
1450+
1451+
output[i] = last_result if num_not_nan >= minp else NaN
14431452

14441453
return output
14451454

0 commit comments

Comments
 (0)