Есть ли способ pandas для вычисления функции между 2 столбцами?

Question

Я ищу более быстрый способ вычисления какой-либо функции по нескольким столбцам.

мой фрейм данных выглядит так:

c = 12*1000
b = int(c/2)
d = int(b/2)

newdf = {'Class': ['c1']*c+['c2']*c+['c3']*c,
        'Section': ['A']*b+['B']*b+['C']*b+['D']*b+['E']*b+['F']*b,
        'Time': [1,2,3,4,5,6]*d+[3,1,3,4,5,7]*d}

test = pd.DataFrame(newdf)
test['f_x'] = test['Time']**2/5
test['f_x_2'] = test['Time']**2/5+test['f_x']
#working with 1 column
test['section_mean'] = test.groupby(['Class','Section'])['f_x'].transform(lambda x: x.mean())
test['two_col_sum'] = test[['Time','f_x']].apply(lambda x: x.Time+x.f_x,axis=1)
cols = ['f_x','f_x_2']

, и я знаю, как вычислить, например, значение для ряда столбцов для групп:

test['section_mean'] = test.groupby(['Class','Section'])['f_x'].transform(lambda x: x.mean())

Или, в конце концов, выполните простые операции между несколькими столбцами:

test['two_col_sum'] = test[['Time','f_x']].apply(lambda x: x.Time+x.f_x,axis=1)

Однако я пытаюсь выполнить какое-то вычисление над полный столбец сгруппированного экземпляра:

%%time
slopes_df = pd.DataFrame()
grouped = test.groupby(['Class','Section'])

for name, group in grouped:
    nd=[]
    for col in cols:
        ntest = group[['Time',col]]
        x = ntest.Time
        y = ntest[col]
        f=np.polyfit(x,y, deg=1).round(2)
        data = [name[0],name[1],col,f[0],f[1]]
        nd.append(data)

    slopes_df=pd.concat([slopes_df,pd.DataFrame(nd)])

slopes_df.columns=['Class','Section','col','slope','intercept']
slopes_df_p = pd.pivot_table(data=slopes_df,index=['Class','Section'], columns=['col'], values=['slope','intercept']).reset_index()
slopes_df_p.columns = pd.Index(e[0] if e[0] in ['Class','Section'] else e[0]+'_'+e[1] for e in slopes_df_p.columns)
fdf = pd.merge(test, slopes_df_p, on=['Class','Section'])

Я попробовал решение, предложенное следующим образом:

%%time
for col in cols:
    df1 = (test.groupby(['Class','Section'])
              .apply(lambda x: np.polyfit(x['Time'],x[col], deg=1).round(2)[0])
              .rename('slope_'+str(col)))
    df2 = (test.groupby(['Class','Section'])
              .apply(lambda x: np.polyfit(x['Time'],x[col], deg=1).round(2)[1])
              .rename('intercept_'+str(col)))
    df1['col']=col
    df2['col']=col

    test = pd.merge(test,df1, on=['Class','Section'])
    test = pd.merge(test,df2, on=['Class','Section'])

, но на моем p c first l * 1022 оно кажется более медленным * занимает 150 мс и второй код 300 мс

Андреа

Answer 1

Ваше решение l oop не работает по данным групп, поэтому я думаю, что вам нужно GroupBy.apply:

def f(x):
    for col in cols:
        x[f'slope_{col}'], x[f'intercept_{col}'] = np.polyfit(x['Time'],x[col], deg=1).round(2)
    return x
df1 = test.groupby(['Class','Section']).apply(f)