改变熊猫多头/空头头寸的更有效方法

Question

有没有比我在下面提出的更有效的方法来改变 pandas 的多头/空头头寸？

逻辑要点：

• 在出现买入 (+1) 或卖出 (-1) 信号之前，仓位为 np.nan。
• 第一次出现买入/卖出 (+/-1) 信号时，将仓位设置为该数字。
• 设置第一个位置后，继续前进。仅在相同符号的情况下添加到该位置。否则，更改符号并设置符号乘以 1（即，如果当前仓位为 3，然后您收到卖出 (-1) 信号，则仓位从 3 变为 -1。反之亦然，如果仓位为 -3 并且有买入信号，仓位从 -3 变为 1)。

我的代码

import numpy as np
import pandas as pd

df = pd.DataFrame({
    'buy_sell': [np.nan, 1, np.nan, 1, np.nan, np.nan, 1, -1, np.nan, -1, np.nan, 1, np.nan, -1],
    'position': np.nan
})

for i, r, in df.iterrows():
    buy_sell = r['buy_sell']

    # Check if first index
    if i != 0:
        last_position = df.loc[i-1, 'position']

        if np.isnan(buy_sell):
            df.loc[i, 'position'] = last_position
        else:
            if np.isnan(last_position) or last_position * buy_sell <= 0:
                df.loc[i, 'position'] = buy_sell
            else:
                df.loc[i, 'position'] = last_position + buy_sell
    else:
        df.loc[i, 'position'] = buy_sell

预期的解决方案

df_expected = pd.DataFrame({
    'buy_sell': [np.nan, 1, np.nan, 1, np.nan, np.nan, 1, -1, np.nan, -1, np.nan, 1, np.nan, -1],
    'position': [np.nan, 1, 1, 2, 2, 2, 3, -1, -1, -2, -2, 1, 1, -1],
})

    buy_sell  position
0        NaN       NaN
1        1.0       1.0
2        NaN       1.0
3        1.0       2.0
4        NaN       2.0
5        NaN       2.0
6        1.0       3.0
7       -1.0      -1.0
8        NaN      -1.0
9       -1.0      -2.0
10       NaN      -2.0
11       1.0       1.0
12       NaN       1.0
13      -1.0      -1.0

注意：我上面的代码会产生预期的数据帧。我在问是否有更有效/更好的方式来做我上面所做的事情。

Answer 1

你可以这样做：

df["position"] = df.buy_sell.groupby(\
                      df.buy_sell.bfill()\
                      .diff()\
                      .abs()\
                      .cumsum()\
                 ).cumsum().ffill()
                 

输出

df
     buy_sell   position
0         NaN        NaN
1      1.0000     1.0000
2         NaN     1.0000
3      1.0000     2.0000
4         NaN     2.0000
5         NaN     2.0000
6      1.0000     3.0000
7     -1.0000    -1.0000
8         NaN    -1.0000
9     -1.0000    -2.0000
10        NaN    -2.0000
11     1.0000     1.0000
12        NaN     1.0000
13    -1.0000    -1.0000

---

说明

基本上，您想要某种cumsum 由列中的符号更改创建的组，那么，首先我们需要找到这些更改：

>>> df.buy_sell.bfill().diff().abs()
0           NaN
1        0.0000
2        0.0000
3        0.0000
4        0.0000
5        0.0000
6        0.0000
7        2.0000
8        0.0000
9        0.0000
10       2.0000
11       0.0000
12       2.0000
13       0.0000

如您所见，每一个2 都是一个符号变化。现在，使用pandas.Series.cumsum 我们可以为每一行分配一个组

>>> df.buy_sell.bfill().diff().abs().cumsum()
0           NaN
1        0.0000
2        0.0000
3        0.0000
4        0.0000
5        0.0000
6        0.0000
7        2.0000
8        2.0000
9        2.0000
10       4.0000
11       4.0000
12       6.0000
13       6.0000

因此，0 的每一行都是一个组，2 的每一行都是一个组，依此类推。有了这个，我们可以直接使用pandas.Series.groupby和pandas.core.groupby.GroupBy.cumsum：

>>> df.buy_sell.groupby(df.buy_sell.bfill().diff().abs().cumsum()).cumsum()
0           NaN
1        1.0000
2           NaN
3        2.0000
4           NaN
5           NaN
6        3.0000
7       -1.0000
8           NaN
9       -2.0000
10          NaN
11       1.0000
12          NaN
13      -1.0000

我们现在只需要使用pandas.Series.ffill 将最后一个有效观察值向前传播到下一个有效值：

>>> df.buy_sell.groupby(df.buy_sell.bfill().diff().abs().cumsum()).cumsum().ffill()
0           NaN
1        1.0000
2        1.0000
3        2.0000
4        2.0000
5        2.0000
6        3.0000
7       -1.0000
8       -1.0000
9       -2.0000
10      -2.0000
11       1.0000
12       1.0000
13      -1.0000