我需要帮助加速这个循环,我不知道该怎么做

import numpy as np
import pandas as pd
import timeit

n = 1000
df = pd.DataFrame({0:np.random.rand(n),1:np.random.rand(n)})

def loop():
    result = pd.DataFrame(index=df.index,columns=['result'])
    for i in df.index:
        last_index_to_consider = df.index.values[::-1][i]
        tdf = df.loc[:last_index_to_consider] - df.shift(-i).loc[:last_index_to_consider]
        tdf = tdf.apply(lambda x: x**2)
        tsumdf = tdf.sum(axis=1)
        result.loc[i,'result'] = tsumdf.mean()
    return result

print(timeit.timeit(loop, number=10))

是否可以调整for循环以使其更快,或者是否有使用numba的选项,或者我可以继续使用多个线程来加速这个循环?
有什么比直接评估代码更明智的方法来获得更高的性能呢?

最佳答案

每次迭代都要进行大量计算保持这种方式,我们可以利用底层阵列数据和np.einsum一起,因为squared-sum-reductions可以带来加速。这是一个遵循这些原则的实现-

def array_einsum_loop(df):
    a = df.values
    l = len(a)
    out = np.empty(l)
    for i in range(l):
        d = a[:l-i] - a[i:]
        out[i] = np.einsum('ij,ij->',d,d)
    df_out = pd.DataFrame({'result':out/np.arange(l,0,-1)})
    return df_out

运行时测试-
In [153]: n = 1000
     ...: df = pd.DataFrame({0:np.random.rand(n),1:np.random.rand(n)})

In [154]: %timeit loop(df)
1 loop, best of 3: 1.43 s per loop

In [155]: %timeit array_einsum_loop(df)
100 loops, best of 3: 5.61 ms per loop

In [156]: 1430/5.61
Out[156]: 254.9019607843137

在不破坏任何环路或气缸组的情况下加速也不错!

09-10 11:10
查看更多