Jasper
Jasper

Reputation: 371

How to create a column in a pandas dataframe based on a previous state?

I have a pandas Dataframe with a timeseries index. One column contains buy signals, another contains sell signals.

buy     0 0 1 0 1 1 0 1 1 1 0 1 0 0 0 0 0 0 0 0 
sell    0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 0 1 1 0

I want to add a state to this simplistic model: have a maximum of 1 stock at any time: Only buy when you currently have nothing, and sell when you have a stock.

buy     0 0 1 0 1 1 0 1 1 1 0 1 0 0 0 0 0 0 0 0 
sell    0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 0 1 1 0
wallet  0 0 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0

How would I create the df['wallet'] column based on df['buy'] and df['sell']?

Upvotes: 2

Views: 324

Answers (2)

elPastor
elPastor

Reputation: 8966

This isn't the prettiest solution, but I believe it achieves what you're looking for.

import pandas as pd

df = pd.DataFrame({ 'buy':  [0,0,0,1,0,1,1,1,1,1,0,0,0,0,0,0], 
                    'sell': [0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1]   })


df['wallet'] = pd.np.where(df['buy'] - df['sell'] < 0, 0, df['buy'].cumsum().count() - df['sell'].cumsum())
df.loc[df['wallet'] > 1, 'wallet'] = 1

Upvotes: 3

joaoavf
joaoavf

Reputation: 1383

import pandas as pd

wallet = 0


def create_wallet(df_line):
    global wallet

    if df_line['buy']:
        wallet = 1

    elif df_line['sell']:
        wallet = 0

    return wallet


df = pd.DataFrame(
{'buy': [0, 0, 0, 1, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0], 'sell': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1]})
df['wallet'] = df.apply(create_wallet, axis=1)

Upvotes: 2

Related Questions