如何加速Pandas系列的“衰减”ffill函数？

Question

我想加速我的“Decage”forward fill函数的以下实现，它用最后一个非零值乘以衰减因子alpha ** (t-T)来填充零值，其中0<alpha<1和(t-T)是距离最后一个非零值的距离：

def decay_series(s):
    decay_fac = 0.9
    for i in range(1, len(s)):
        if abs(s.iloc[i]) < 1e-6:
            s.iloc[i] = s.iloc[i - 1] * decay_fac
    return s
s = pd.Series([0,0, 1, 2, 0,0,1,0,0,1])
s
Out[24]: 
0    0
1    0
2    1
3    2
4    0
5    0
6    1
7    0
8    0
9    1
dtype: int64
decay_series(s)
Out[25]: 
0    0.00
1    0.00
2    1.00
3    2.00
4    1.80
5    1.62
6    1.00
7    0.90
8    0.81
9    1.00
dtype: float64

但是，由于使用了纯python for循环，这太慢了。是否有任何方法可以加快这一速度，例如，通过一些巧妙地应用pandas的本机组件？不幸的是，fillna或replace方法似乎不支持要应用的自定义用户方法

Answer 1

使用mask和广播

(s.mask(s.eq(0)).ffill() * decay_fac ** s.groupby(s.ne(0).cumsum()).cumcount()).fillna(0)

---

0    0.00
1    0.00
2    1.00
3    2.00
4    1.80
5    1.62
6    1.00
7    0.90
8    0.81
9    1.00
dtype: float64

timings

9.62毫秒与10000行的1.12秒相比

%timeit (s.mask(s.eq(0)).ffill() * 0.9 ** s.groupby(s.ne(0).cumsum()).cumcount()).fillna(0)
9.62 ms ± 206 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit decay_series(s)
1.12 s ± 161 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)