lookback/shift in pandas dataframes calling functions

Question

If I have the following dataframe:

  date       A     B    M     S
 20150101    8     7    7.5   0
 20150101    10    9    9.5   -1
 20150102    9     8    8.5   1
 20150103    11    11   11    0
 20150104    11    10   10.5  0
 20150105    12    10   11    -1
 ...

If I want to create another column 'cost' by the following rules:

if S < 0, cost = (M-B).shift(1)*S
if S > 0, cost = (M-A).shift(1)*S
if S == 0, cost=0

currently, I am using the following function:

def cost(df):
if df[3]<0:
    return np.roll((df[2]-df[1]),1)*df[3]
elif df[3]>0:
    return np.roll((df[2]-df[0]),1)*df[3]
else:
    return 0
df['cost']=df.apply(cost,axis=0)

Is there any other way to do it? can I somehow use pandas shift function in user defined functions? thanks.

score 5 · Accepted Answer · edited May 23 '17 at 11:46

5

It's generally expensive to do it this way, as you're losing the vector speed advantage when you apply a user defined function. Instead, how about using the numpy version of the ternary operator:

import numpy as np

np.where(df[3] < 0,
    np.roll((df[2]-df[1]),1),
    np.where(df[3] > 0,
        np.roll((df[2]-df[0]),1)*df[3] 
        0))

(of course assign it to df['cost']).

edited May 23 '17 at 11:46

Community

1
1

answered May 10 '15 at 14:06

Ami Tavory

74,578
11
141
185

Thank you so much, the method is very enlightening to me. – user6396 May 10 '15 at 14:30
Thanks, however see unutbu's answer below, which is better. – Ami Tavory May 10 '15 at 14:50

score 5 · Answer 2 · answered May 10 '15 at 14:30

np.where(condition, A, B) is the NumPy elementwise equivalent of

A if condition else B

np.select(conditions, choices) is a generalization of np.where which is useful when there are more than two choices.

So, like Ami Tavory's answer except using np.select, you could use

import numpy as np
import pandas as pd
df = pd.read_table('data', sep='\s+')
conditions = [S < 0, S > 0]
M, A, B, S = [df[col] for col in 'MABS']
choices = [(M-B).shift(1)*S, (M-A).shift(1)*S]
df['cost'] = np.select(conditions, choices, default=0)

which yields

       date   A   B     M  S  cost
0  20150101   8   7   7.5  0   0.0
1  20150101  10   9   9.5 -1  -0.5
2  20150102   9   8   8.5  1  -0.5
3  20150103  11  11  11.0  0   0.0
4  20150104  11  10  10.5  0   0.0
5  20150105  12  10  11.0 -1  -0.5

Ooh, very nice, I did not know this. – Ami Tavory May 10 '15 at 14:48 — Ami Tavory, May 10 '15 at 14:48

lookback/shift in pandas dataframes calling functions

2 Answers2