Dataframe快速展开时在列中向后搜索的方法

Question

这是我的价格数据框

   stock1  stock2
0     2.3    10.1
1     1.9    11.2
2     3.5    10.5
3     2.8    10.8
4     3.1    10.3
5     2.7     9.8
6     3.3    10.2

这是我想要获得的价格支持数据框

   stock1  stock2
0     NaN     NaN
1     1.9    10.1
2     1.9    10.1
3     1.9    10.5
4     2.8    10.1
5     1.9     NaN
6     2.7     9.8

让我们关注价格的第一列。想法是通过这种方式计算下峰

    prices1 = prices['stock1']
    mask = (prices1.shift(1) > prices1) & (prices1.shift(-1) > prices1)
    supports1 = prices1.where(mask, NaN)
    supports1.iloc[0] = min(prices1[0],prices1[1])
    supports1 = supports1.shift(1).fillna(method='ffill')

我们获得

   stock1
0  NaN
1  NaN
2  1.9
3  1.9
4  2.8
5  2.8
6  2.7

[另一个规则是，对于每个价格，支撑都必须更低。在第5行中不会发生这种情况，因为2.8> 2.7。为了更正，我们必须在此支持列中向后看以查找低于当前价格的首次出现的价格（如果存在，否则为NaN）。在这种情况下，正确值为1.9

我发现了两种解决问题的方法，但是我需要进行迭代，并且当数据帧增加时，它变得非常慢。我想快10倍，希望100倍。这是我的代码

from pandas import DataFrame
from numpy import NaN
from numpy.random import uniform
from timeit import timeit

##rows = 5000
##cols = 10
##d={}
##for i in range(cols):
##    d['stock_{}'.format(i)] = 100*uniform(0.95,1.05,rows).cumprod()
##prices = DataFrame(d)

prices = DataFrame({'stock1':[2.3, 1.9, 3.5, 2.8, 3.1, 2.7, 3.3],\
                    'stock2':[10.1, 11.2, 10.5, 10.8, 10.3, 9.8, 10.2]})



#----------------------------------------------------------------
def calc_supports1(prices):
    supports = DataFrame().reindex_like(prices)
    for stock in prices:
        prices1 = prices[stock]
        mask = (prices1.shift(1) > prices1) & (prices1.shift(-1) > prices1)
        supports1 = prices1.where(mask, NaN)
        supports1.iloc[0] = min(prices1[0],prices1[1])
        supports1 = supports1.shift(1).fillna(method='ffill')

        sup = supports1.drop_duplicates()
        for i,v in prices1.loc[prices1 < supports1].iteritems():
            mask = (sup.index < i) & (sup < v)
            sup2 = sup.values[mask.values]
            supports1.at[i] = sup2[-1] if len(sup2) > 0 else NaN

        supports[stock] = supports1
    return supports

#----------------------------------------------------------------
def calc_supports2(prices):
    supports = DataFrame().reindex_like(prices)
    for stock in prices:
        prices1 = prices[stock]
        sup = [min(prices1[0],prices1[1])]
        supports1 = [NaN, sup[0]]

        for i in xrange(2,len(prices1)):
            while len(sup) > 0 and prices1[i] < sup[0]:
                sup.pop(0)

            if prices1[i-1]<prices1[i] and prices1[i-1]<prices1[i-2]:
                sup.insert(0, prices1[i-1])

            supports1.append(sup[0] if len(sup) > 0 else NaN)

        supports[stock] = supports1
    return supports
#----------------------------------------------------------------

print 'fun1', timeit('calc_supports1(prices)', \
            setup='from __main__ import calc_supports1, prices',number = 1)
print 'fun2', timeit('calc_supports2(prices)', \
            setup='from __main__ import calc_supports2, prices',number = 1)

我如何加快速度？

Answer 1

此代码中存在多个性能问题。

• Python循环很慢（使用默认解释器CPython）。最好使用Cython或Numba之类的东西进行此类计算。
• 熊猫手动索引很慢。最好在这里使用Python列表，甚至避免在代码的热门部分进行索引。
• 在列表开头插入/删除元素可能要慢得多。 Python文档建议在末尾添加元素以实现堆栈。

这里是正确的代码：

def calc_supports3(prices):
    supports = DataFrame().reindex_like(prices)
    for stock in prices:
        prices1 = list(prices[stock])
        sup = [min(prices1[0],prices1[1])]
        supports1 = [NaN, sup[-1]]

        # Sliding window
        prices1_im2 = NaN
        prices1_im1 = prices1[0]
        prices1_im0 = prices1[1]

        for i in xrange(2,len(prices1)):
            prices1_im2, prices1_im1, prices1_im0 = prices1_im1, prices1_im0, prices1[i]

            while len(sup) > 0 and prices1_im0 < sup[-1]:
                sup.pop()

            if prices1_im1<prices1_im0 and prices1_im1<prices1_im2:
                sup.append(prices1_im1)

            supports1.append(sup[-1] if len(sup) > 0 else NaN)

        supports[stock] = supports1
    return supports

这是您在我的计算机上的小型数据集上的性能结果：

fun1 0.006461 s
fun2 0.000901 s
fun3 0.000648 s (40% faster than fun2)

以下是在具有5万行的随机生成的数据集上的性能结果：

fun1 3.916947 s
fun2 2.064891 s
fun3 0.034465 s (60 times faster than fun2)