将公式应用于熊猫 Dataframe

Question

I have a very simple question. I have a dataframe like this

In [19]: df = DataFrame(randn(10,2),columns=list('A'))

In [20]: df

Out[20]: 
          A  
0  0.958465  
1 -0.769077  
2  0.598059  
3  0.290926 
4 -0.248910 
5 -1.352096 
6  0.009125
7 -0.993082
8 -0.593704
9  0.523332

I would like to create a new column B with the following information:

          A              B
0  0.958465  
1 -0.769077  A1*A1+2*A0*A2
2  0.598059  A2*A2+2*A1*A3
3  0.290926  A3*A3+2*A2*A4
4 -0.248910  A4*A4+2*A3*A5
5 -1.352096  ...
6  0.009125  ...
7 -0.993082  ...
8 -0.593704  ...
9  0.523332  ...

It is a sort of convolution or autocorrelation but using everytime a different window. How can I define such a formula in Pandas?

Second question: how can I make variable the number of points involved in the formula (in the example I am just using the previous and the next point to make the calculation, but how can I pass a variable to say to pandas the number of points I want to use for the calculation)?

Answer 1

You can make a function like this to allow a variable number of lags.

def func(s, lags=1):
    return sum(s.shift(lag) * s.shift(-lag) for lag in range(lags+1))

df = pd.DataFrame({"A": [0.958465, -0.769077, 0.598059, 0.290926, -0.248910, -1.352096, 0.009125, 0.993082, -0.593704, 0.523332]})
df["B"] = func(df["A"], 1) # takes 1 point on either side
df["C"] = func(df["A"], 2) # takes 2 points on either side

Answer 2

df['B'] df['A']**2 + 2 * df['A'].shift() * df['A'].shift(-1)
df
          A         B
0  0.958465       NaN
1 -0.769077  1.737917
2  0.598059 -0.089814
3  0.290926 -0.213088
4 -0.248910 -0.724764
5 -1.352096  1.823621
6  0.009125  2.685568
7 -0.993082  0.975377
8 -0.593704 -0.686939
9  0.523332       NaN

Answer 3

I think, it can be done without apply too, in case:

import pandas as pd
import numpy as np

df = pd.DataFrame(np.random.randn(10, 1), columns=['A'])

df['B'] = df['A'] ** + 2 * df['A'].shift(periods=1) * df['A'].shift(periods=-1)

print(df)

output:

          A         B
0  0.383006       NaN
1 -1.240859 -0.469964
2 -0.796920 -0.393244
3  0.499011  0.358630
4 -1.807221 -0.701899
5 -0.430667  0.296360
6 -0.884149  0.475958
7 -1.413762  3.492830
8 -1.976511  5.939588
9 -1.075428       NaN