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

Question

I have the following MultiIndex dataframe.

                         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 

I want to calculate three columns:

• Shares = previous day's Equity * 0.02 / ATR , rounded down to whole number

• Profit = Shares * Close

• Equity = previous day's Equity + sum of Profit for each Symbol

Equity has an initial value of 10,000.

The expected output is:

                         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

I suppose I need a for loop or a function to be applied to each row. With these I have two issues. One is that I'm not sure how I can create a for loop for this logic in case of a MultiIndex dataframe. The second is that my dataframe is pretty large (something like 10 million rows) so I'm not sure if a for loop would be a good idea. But then how can I create these columns?

Answer 1

This solution can surely be cleaned up, but will produce your desired output. I've included your initial conditions in the construction of your sample dataframe:

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]})

Gives:

         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

Then use groupby() with apply() and track your Equity globally. Took me a second to realize that the nature of this problem requires you to group on two separate columns individually ( Symbol and 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

This yields:

         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

Answer 2

can you try using shift and groupby? Once you have the value of the previous line, all columns operations are straight forward.

table2['previous'] = table2['close'].groupby('symbol').shift(1)

table2

date    symbol      close   atr     previous

1990-01-01  A   24  2   NaN
            B   72  7   NaN
            C   40  3.4     NaN
1990-01-02  A   21  1.5     24
            B   65  6   72
            C   45  4.2     40
1990-01-03  A   19  2.5     21
            B   70  6.3     65
            C   51  5   45