I would like to use the following function on a large dataset, 300000 rows with 1500 unique ID. And create two new columns as in below. The function below works when applied to the data frame as a whole (len(df1)). But I want to apply it based on length of each unique ID.
d = {'ID':['a12', 'a12','a12','b33','b33','b33','v55','v55','v55','v55'], 'Exp_A':[2.2,2.2,2.2,3.1,3.1,3.1,1.5,1.5,1.5,1.5],
'Exp_B':[2.4,2.4,2.4,1.2,1.2,1.2,1.5,1.5,1.5,1.5],
'A':[0,0,0,1,0,1,0,1,0,1], 'B':[0,0,1,0,1,0,1,0,1,0]}
df1 = pd.DataFrame(data=d)
def adj_Apois(row):
i = row.name
if row.A == 1:
return poisson.pmf(row.A, row.Exp_A*(i+1)/len(df1)) * row.Exp_A
else:
return poisson.pmf(row.A, row.Exp_A*(i+1)/len(df1)) * row.Exp_A
def adj_Bpois(row):
i = row.name
if row.B == 1:
return poisson.pmf(row.B, row.Exp_B*(i+1)/len(df1)) * row.Exp_B
else:
return poisson.pmf(row.B, row.Exp_B*(i+1)/len(df1)) * row.Exp_B
df1['New1'] = df1.apply(adj_Apois, axis=1)
df1['New2'] = df1.apply(adj_Bpois, axis=1)
Desired out come
replacing df1 above with length of unique ID yield the following
A B Exp_A Exp_B ID new1 new2
0 0 0 2.2 2.4 a12 1.05667 1.07839
1 0 0 2.2 2.4 a12 0.50752 0.48455
2 0 1 2.2 2.4 a12 0.24377 0.86271
3 1 0 3.1 1.2 b33 1.13981 0.80438
4 0 1 3.1 1.2 b33 0.39248 0.43136
5 1 0 3.1 1.2 b33 0.43292 0.36143
6 0 1 1.5 1.5 v55 1.03093 0.3866
7 1 0 1.5 1.5 v55 0.53141 0.70855
8 0 1 1.5 1.5 v55 0.48698 0.54785
9 1 0 1.5 1.5 v55 0.50204 0.33469
I would appreciate any help.
IIUC, just use poisson.pmf directly
from spicy.stats import poisson
df1['New2a'] = poisson.pmf(df1.B, df1.Exp_B*(df1.index+1)/len(df1)) * df1.Exp_B
df1['New1a'] = poisson.pmf(df1.A, df1.Exp_A*(df1.index+1)/len(df1)) * df1.Exp_A
Yields same results
>>> df1['New1a'].eq(df1['New1']).all()
True
To group:
>>> df1.groupby('ID').apply(lambda s: poisson.pmf(s.B, s.Exp_B*(s.index+1)/len(s)) * s.Exp_B).reset_index()