Python pandas: 如何向量化这个函数

Question

我有两个 DataFrame df 和 evol 如下（为示例进行了简化）：

In[6]: df
Out[6]:
   data  year_final  year_init
0    12        2023       2012
1    34        2034       2015
2     9        2019       2013
...

In[7]: evol
Out[7]: 
      evolution
year           
2000   1.474946
2001   1.473874
2002   1.079157
...
2037   1.463840
2038   1.980807
2039   1.726468

我想以向量化方式进行以下操作（当我有 Gb 数据时，当前的 for 循环实现太长）：

for index, row in df.iterrows():
    for year in range(row['year_init'], row['year_final']):
        factor = evol.at[year, 'evolution']
        df.at[index, 'data'] += df.at[index, 'data'] * factor

复杂性是因为每一行的年份范围都不相同... 在上面的示例中，输出将是：

        data  year_final  year_init
0     163673        2023       2012
1  594596046        2034       2015
2       1277        2019       2013

（用于测试目的的完整 evol 数据框:)

      evolution
year           
2000   1.474946
2001   1.473874
2002   1.079157
2003   1.876762
2004   1.541348
2005   1.581923
2006   1.869508
2007   1.289033
2008   1.924791
2009   1.527834
2010   1.762448
2011   1.554491
2012   1.927348
2013   1.058588
2014   1.729124
2015   1.025824
2016   1.117728
2017   1.261009
2018   1.705705
2019   1.178354
2020   1.158688
2021   1.904780
2022   1.332230
2023   1.807508
2024   1.779713
2025   1.558423
2026   1.234135
2027   1.574954
2028   1.170016
2029   1.767164
2030   1.995633
2031   1.222417
2032   1.165851
2033   1.136498
2034   1.745103
2035   1.018893
2036   1.813705
2037   1.463840
2038   1.980807
2039   1.726468

Answer 1

一种仅使用 pandas 的矢量化方法是在两个帧和子集之间进行笛卡尔连接。开始时会像：

df['dummy'] = 1
evol['dummy'] = 1
combined = df.merge(evol, on='dummy')
# filter date ranges, multiply etc

这可能会比您正在做的更快，但内存效率低下并且可能会破坏您的真实数据。

如果你能接受 numba 依赖，像这样的事情应该会非常快——本质上是你现在正在做的事情的编译版本。在 cython 中也可能有类似的东西。请注意，这要求 evol 数据框按年份排序和连续，可以通过修改放宽。

import numba

@numba.njit
def f(data, year_final, year_init, evol_year, evol_factor):
    data = data.copy()
    for i in range(len(data)):
        year_pos = np.searchsorted(evol_year, year_init[i])
        n_years = year_final[i] - year_init[i]
        for offset in range(n_years):
            data[i] += data[i] * evol_factor[year_pos + offset]            
    return data

f(df['data'].values, df['year_final'].values, df['year_init'].values, evol.index.values, evol['evolution'].values)
Out[24]: array([   163673, 594596044,      1277], dtype=int64)

编辑：测试数据的一些时间

In [25]: %timeit f(df['data'].values, df['year_final'].values, df['year_init'].values, evol.index.values, evol['evolution'].values)
15.6 µs ± 338 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)


In [26]: %%time
    ...: for index, row in df.iterrows():
    ...:     for year in range(row['year_init'], row['year_final']):
    ...:         factor = evol.at[year, 'evolution']
    ...:         df.at[index, 'data'] += df.at[index, 'data'] * factor
Wall time: 3 ms

Python pandas: 如何向量化这个函数

Python pandas: how to vectorize this function

python

numpy

vectorization

pandas

numba