我需要在日期时间索引的数据帧上计算从 h1 小时到 h2 小时的每日回报。 我可以使用 asfreq 和 resample 调用来做到这一点，但速度相当慢。例如，假设 h1=0 且 h2=3：
import numpy as np
import pandas as pd
df_index = pd.date_range("1/1/2000", periods=2 * 365 * 24, freq="H")
df_data = np.random.randn(2 * 365 * 24, 250)
df = pd.DataFrame(df_data, index=df_index)
df.asfreq("3H").pct_change().resample("D", closed="right", label="right", offset="3H").last()
我尝试使用 numba 并不成功，因为我观察到时间较慢。
我尝试了两个版本，一个处理完整的 df，一列接一列，参数 idx、col，数据是 numpy 数组，h1 和 h2 是 int64 标量。
我做错了什么？
@numba.njit(error_model="numpy") # expected to be multiples faster
def pct_change_h1_h2_njit(idx, data, h1, h2):
out = np.full_like(data, np.nan)
zeroes = np.full_like(out[0], 0)
denominator = out[0]
for i, v in enumerate(data):
if idx[i] == h1:
denominator = v
elif idx[i] == h2:
if np.isnan(denominator).all():
out[i] = zeroes
else:
out[i] = np.divide(v, denominator) - 1
return out
@numba.njit
def pct_change_col_h1_h2_njit(idx, col, h1=0, h2=3):
out = np.full_like(col, np.nan)
exceptions = np.array([np.nan, 0, np.inf, -np.inf])
denominator = np.nan
for i, v in enumerate(col):
if idx[i] == h1:
denominator = v
elif idx[i] == h2:
if denominator in exceptions:
out[i] = 0
else:
out[i] = v / denominator - 1
return out
在三个等效方法上运行 timeit 显示 njit 版本慢两倍以上。
%timeit for x in range(1000):
df.asfreq("3H").pct_change().resample(
"D", closed="right", label="right", offset="3H"
).last()
# 25 s ± 5.72 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit for x in range(1000):
pd.DataFrame(
pct_change_h1_h2_njit(df.index.hour.values, df.values, 0, 3),
index=df.index, columns=df.columns
).dropna(how="all")
# 59.9 s ± 180 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit
for x in range(1000):
pd.concat(
[
pd.Series(
pct_change_col_h1_h2_njit(
idx=np.array(df.index.hour, dtype=np.int64),
col=df[c].values,
h1=0,
h2=3,
)
)
for c in df.columns
],
axis=1,
).set_axis(df.columns, axis=1).set_axis(df.index, axis=0).dropna(how="all")
# 86.3 s ± 180 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Numba 代码相当“受内存限制”。事实上，最糟糕的是：它甚至是页面错误绑定。这当然就是为什么它在我的机器上更快而在你的机器上更慢（我的机器被设计为实现高吞吐量）：如果 RAM 吞吐量足够高，它会更高效。由于内存墙，如今 RAM 吞吐量非常宝贵（而且这种情况不会很快变得更好）。 因此，这里的关键是
减少内存压力。 np.full_like(data, np.nan)可以替换为
np.empty_like(data)outpct_change_h1_h2_njit初始化的。但在实践中，Pandas 数据框通常每列包含一组多个 Numpy 数组，因此它们更以列为主。这意味着这个基准肯定不能代表您将得到的实际时间。 此外，在 Numba 中，循环通常更好，因为它们
避免创建昂贵的临时数组（这里不多，因为它们可以存储在 CPU L1 缓存中）。更具体地说，v / denominator[i]可以用循环代替。这在我的机器上稍微快一些。
请注意，
pct_change_h1_h2_njit仅占整个计算的 60%。
dropna的开销。 或者，您可以通过预分配输出数组、将其传递到参数并将其写入函数来回收输出数组。这是关键的优化。与之前的版本相比，它在我的机器上的功能速度提高了 2.4 倍，并且比初始版本快了近 3 倍
。这是代码：
@numba.njit(error_model="numpy") # expected to be multiples faster
def pct_change_h1_h2_njit(idx, data, h1, h2, out):
zeroes = np.full_like(out[0], 0)
denominator = out[0]
for i, v in enumerate(data):
if idx[i] == h1:
denominator = v
out[i].fill(np.nan)
elif idx[i] == h2:
if np.isnan(denominator).all():
out[i].fill(0)
else:
for j in range(out[i].shape[0]):
out[i,j] = v[j] / denominator[j] - 1
else:
out[i].fill(np.nan)
return out # Not really needed
我知道创建新数组更方便，但这在性能方面是一种不好的做法。此时，dropna
bug
：根据定义，
np.nan == np.nan始终是False，因此即使np.nan in exceptions
位于
True中，
np.Nan也永远不会是
exceptions。您必须使用
np.isnan。