Pandas Custom Cumulative Calculation Over Group By in DataFrame

user_6483074 0
Viewed 121

I am trying to run a simple calculation over the values of each row from within a group inside of a dataframe, but I'm having trouble with the syntax, I think I'm specifically getting confused in relation to what data object I should return, i.e. dataframe vs series etc.

For context, I have a bunch of stock values for each product I am tracking and I want to estimate the number of sales via a custom function which essentially does the following:

# Because stock can go up and down, I'm looking to record the difference 
# when the stock is less than the previous stock number from the previous row.
# How do I access each row of the dataframe and then return the series I need?

def get_stock_sold(x):
    # Written in pseudo
    stock_sold = previous_stock_no - current_stock_no if current_stock_no < previous_stock_no else 0
    return pd.Series(stock_sold)

I then have the following dataframe:

# 'Order' is a date in the real dataset.

data = { 
    'id'            : ['1', '1', '1', '2', '2', '2'],
    'order'         : [1, 2, 3, 1, 2, 3],
    'current_stock' : [100, 150, 90, 50, 48, 30]
}

df = pd.DataFrame(data)
df = df.sort_values(by=['id', 'order'])
df['previous_stock'] = df.groupby('id')['current_stock'].shift(1)

I'd like to create a new column (stock_sold) and apply the logic from above to each row within the grouped dataframe object:

df['stock_sold'] = df.groupby('id').apply(get_stock_sold)

Desired output would look as follows:

| id | order | current_stock | previous_stock | stock_sold |
|----|-------|---------------|----------------|------------|
| 1  | 1     | 100           | NaN            | 0          |
|    | 2     | 150           | 100.0          | 0          |
|    | 3     | 90            | 150.0          | 60         |
| 2  | 1     | 50            | NaN            | 0          |
|    | 2     | 48            | 50.0           | 2          |
|    | 3     | 30            | 48             | 18         |
custom-function pandas-groupby pandas dataframe python
1 Answers

Try:

df["previous_stock"] = df.groupby("id")["current_stock"].shift()
df["stock_sold"] = np.where(
    df["current_stock"] > df["previous_stock"].fillna(0),
    0,
    df["previous_stock"] - df["current_stock"],
)
print(df)

Prints:

  id  order  current_stock  previous_stock  stock_sold
0  1      1            100             NaN         0.0
1  1      2            150           100.0         0.0
2  1      3             90           150.0        60.0
3  2      1             50             NaN         0.0
4  2      2             48            50.0         2.0
5  2      3             30            48.0        18.0
user_100359850
Related