python - 如何使用 Pandas MultiIndex DataFrame 中的先前值进行计算？

Question

我有以下多索引数据框。

                         Close     ATR     
Date          Symbol     
1990-01-01    A          24        2       
1990-01-01    B          72        7      
1990-01-01    C          40        3.4 

1990-01-02    A          21        1.5     
1990-01-02    B          65        6        
1990-01-02    C          45        4.2   

1990-01-03    A          19        2.5    
1990-01-03    B          70        6.3       
1990-01-03    C          51        5 

我想计算三列:

• 股数 = 前一天的股本 * 0.02/ATR，向下舍入到整数

• 利润 = 股票 * 平仓

• 净值 = 前一天的净值 + 每个代码的利润总和

权益的初始值为 10,000。

预期输出是:

                         Close     ATR     Shares     Profit     Equity
Date          Symbol     
1990-01-01    A          24        2       0          0          10000
1990-01-01    B          72        7       0          0          10000
1990-01-01    C          40        3.4     0          0          10000

1990-01-02    A          21        1.5     133        2793       17053
1990-01-02    B          65        6       33         2145       17053
1990-01-02    C          45        4.2     47         2115       17053

1990-01-03    A          19        2.5     136        2584       26885
1990-01-03    B          70        6.3     54         3780       26885
1990-01-03    C          51        5       68         3468       26885

我想我需要将一个for循环或一个函数应用于每一行。有了这些我有两个问题。一是我不确定如何在 MultiIndex 数据帧的情况下为此逻辑创建一个 for 循环。第二个是我的数据帧非常大(大约 1000 万行)，所以我不确定 for 循环 是否是一个好主意。但是我怎样才能创建这些列呢？

Answer 1

这个解决方案肯定可以清理，但会产生您想要的输出。我已将您的初始条件包含在示例数据框的构造中:

import pandas as pd
import numpy as np

df = pd.DataFrame({'Date': ['1990-01-01','1990-01-01','1990-01-01','1990-01-02','1990-01-02','1990-01-02','1990-01-03','1990-01-03','1990-01-03'],
    'Symbol': ['A','B','C','A','B','C','A','B','C'],
    'Close': [24, 72, 40, 21, 65, 45, 19, 70, 51],
    'ATR': [2, 7, 3.4, 1.5, 6, 4.2, 2.5, 6.3, 5],
    'Shares': [0, 0, 0, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
    'Profit': [0, 0, 0, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]})

给予:

         Date Symbol  Close  ATR  Shares  Profit
0  1990-01-01      A     24  2.0     0.0     0.0
1  1990-01-01      B     72  7.0     0.0     0.0
2  1990-01-01      C     40  3.4     0.0     0.0
3  1990-01-02      A     21  1.5     NaN     NaN
4  1990-01-02      B     65  6.0     NaN     NaN
5  1990-01-02      C     45  4.2     NaN     NaN
6  1990-01-03      A     19  2.5     NaN     NaN
7  1990-01-03      B     70  6.3     NaN     NaN
8  1990-01-03      C     51  5.0     NaN     NaN

然后将 groupby() 与 apply() 结合使用，并在全局范围内跟踪您的Equity。我花了一秒钟才意识到这个问题的本质要求您分别对两个单独的列进行分组(Symbol 和 Date):

start = 10000
Equity = 10000

def calcs(x):

    global Equity

    if x.index[0]==0: return x #Skip first group

    x['Shares'] = np.floor(Equity*0.02/x['ATR'])
    x['Profit'] = x['Shares']*x['Close']
    Equity += x['Profit'].sum()

    return x

df = df.groupby('Date').apply(calcs)
df['Equity'] = df.groupby('Date')['Profit'].transform('sum')
df['Equity'] = df.groupby('Symbol')['Equity'].cumsum()+start

这会产生:

         Date Symbol  Close  ATR  Shares  Profit   Equity
0  1990-01-01      A     24  2.0     0.0     0.0  10000.0
1  1990-01-01      B     72  7.0     0.0     0.0  10000.0
2  1990-01-01      C     40  3.4     0.0     0.0  10000.0
3  1990-01-02      A     21  1.5   133.0  2793.0  17053.0
4  1990-01-02      B     65  6.0    33.0  2145.0  17053.0
5  1990-01-02      C     45  4.2    47.0  2115.0  17053.0
6  1990-01-03      A     19  2.5   136.0  2584.0  26885.0
7  1990-01-03      B     70  6.3    54.0  3780.0  26885.0
8  1990-01-03      C     51  5.0    68.0  3468.0  26885.0