Retain in Pandas, faster way than for loop or apply

Question

I have a question like this one (retain function in python). When dealing with time series data, very often the algorithm needs to dynamically referencing the last computed record.

For example, I have some stock trading records, and would like to calculate the 'average cost' for the shares I'm holding. The only solution I can think of is iterate over the target DataFrame. It doesn't feels like using the strength of the Pandas DataFrame.

Fake Data:

import numpy as np
import pandas as pd

aapl = pd.read_csv('https://raw.githubusercontent.com/ktc312/pandas-questions/master/AAPL_exmaple.csv', parse_dates=['Date'])
print aapl

        Date  Quantity     Price
0 2017-01-10      1000  117.2249
1 2017-02-10      -500  130.5928
2 2017-03-10      1500  137.5316
3 2017-04-10     -2000  141.5150
4 2017-05-10       500  151.4884
5 2017-06-09       500  147.8657
6 2017-07-10       500  143.9750
7 2017-08-10     -1000  154.7636
8 2017-09-11      -200  160.9215
9 2017-10-10      1000  155.3416

Some needed variable:

aapl['Total Shares'] = aapl['Quantity'].cumsum()
aapl['Cost'] = aapl['Quantity']*aapl['Price']
print apple

        Date  Quantity     Price  Total Shares       Cost
0 2017-01-10      1000  117.2249          1000  117224.90
1 2017-02-10      -500  130.5928           500  -65296.40
2 2017-03-10      1500  137.5316          2000  206297.40
3 2017-04-10     -2000  141.5150             0 -283030.00
4 2017-05-10       500  151.4884           500   75744.20
5 2017-06-09       500  147.8657          1000   73932.85
6 2017-07-10       500  143.9750          1500   71987.50
7 2017-08-10     -1000  154.7636           500 -154763.60
8 2017-09-11      -200  160.9215           300  -32184.30
9 2017-10-10      1000  155.3416          1300  155341.60

loop through data to get the average cost:

def get_ave_cost(df):
    for index, row in df.iterrows():
        if index == 0:
            df.loc[index,'Ave Cost'] = row['Price']
        elif row['Total Shares'] == 0:
            df.loc[index,'Ave Cost'] = 0.0
        else:
            if row['Quantity'] > 0:
                df.loc[index,'Ave Cost'] = \
                    ((df.loc[index - 1,'Ave Cost'] * \
                      df.loc[index - 1,'Total Shares']) + \
                      row['Cost'])/row['Total Shares']
            else:
                df.loc[index,'Ave Cost'] =  df.loc[index - 1,'Ave Cost']
    return df

get_ave_cost(stock_trading_records_df)

Wanted results:

        Date  Quantity     Price  Total Shares       Cost    Ave Cost
0 2017-01-10      1000  117.2249          1000  117224.90  117.224900
1 2017-02-10      -500  130.5928           500  -65296.40  117.224900
2 2017-03-10      1500  137.5316          2000  206297.40  132.454925
3 2017-04-10     -2000  141.5150             0 -283030.00    0.000000
4 2017-05-10       500  151.4884           500   75744.20  151.488400
5 2017-06-09       500  147.8657          1000   73932.85  149.677050
6 2017-07-10       500  143.9750          1500   71987.50  147.776367
7 2017-08-10     -1000  154.7636           500 -154763.60  147.776367
8 2017-09-11      -200  160.9215           300  -32184.30  147.776367
9 2017-10-10      1000  155.3416          1300  155341.60  153.595777

[See Notebook]

Is there any other way to do this more efficient or simpler?

Thank you!

do you mean second else in get_ave_cost? but this function (for loop) returns the correct answers, just doesn't seems an ideal solution. — KT Chang, Dec 01 '17 at 23:56
For future questions, please provide an easily copy-pasteable code example that includes your mock data. — tobsecret, Mar 07 '18 at 21:36
@tobsecret it's just a fake data, but if you are interested, it's available here https://github.com/ktc312/pandas-questions/blob/master/AAPL_exmaple.csv — KT Chang, Mar 08 '18 at 00:01

Retain in Pandas, faster way than for loop or apply

0 Answers0