How to forward fill row values with function in pa

2019-08-21 02:18发布

问题:

I have the following MultiIndex dataframe:

                   Close  ATR  condition
Date       Symbol                                                    
1990-01-01 A          24    1       True
           B          72    1      False
           C          40    3      False
           D          21    5       True
1990-01-02 A          65    4       True
           B          19    2       True
           C          43    3       True
           D          72    1      False
1990-01-03 A          92    5      False
           B          32    3       True
           C          52    2      False
           D          33    1      False

I perform the following calculation on this dataframe:

data.loc[data.index.levels[0][0], 'Shares'] = 0
data.loc[data.index.levels[0][0], 'Closed_P/L'] = 0

data = data.reset_index()

Equity = 10000

def calcs(x):

    global Equity

    # Skip first date
    if x.index[0]==0: return x 
    # calculate Shares where condition is True
    x.loc[x['condition'] == True, 'Shares'] = np.floor((Equity * 0.02 / x['ATR']).astype(float))
    # other calulations
    x['Closed_P/L'] = x['Shares'] * x['Close']
    Equity += x['Closed_P/L'].sum()

    return x

data = data.groupby('Date').apply(calcs)

data['Equity'] = data.groupby('Date')['Closed_P/L'].transform('sum')
data['Equity'] = data.groupby('Symbol')['Equity'].cumsum() + Equity

data = data.set_index(['Date','Symbol'])

The output is:

                   Close  ATR  condition  Shares  Closed_P/L   Equity
Date       Symbol                                                    
1990-01-01 A          24  1.2       True     0.0         0.0  10000.0
           B          72  1.4      False     0.0         0.0  10000.0
           C          40    3      False     0.0         0.0  10000.0
           D          21    5       True     0.0         0.0  10000.0
1990-01-02 A          65    4       True    50.0      3250.0  17988.0
           B          19    2       True   100.0      1900.0  17988.0
           C          43    3       True    66.0      2838.0  17988.0
           D          72    1      False     NaN         NaN  17988.0
1990-01-03 A          92    5      False     NaN         NaN  21796.0
           B          32    3       True   119.0      3808.0  21796.0
           C          52    2      False     NaN         NaN  21796.0
           D          33    1      False     NaN         NaN  21796.0

I want to forward fill Shares values - grouped by Symbol - in case condition evaluates to False (except for first date). So the Shares value on 1990-01-02 for D should be 0 (because on 1990-01-01 the Shares value for D was 0 and the condition on 1990-01-02 is False). Also values for Shares on 1990-01-03 for A, C and D should be 50, 66 and 0 respectively based on the logic described above. How can I do that?