использовать результат функции, примененной к groupby, для расчета на оригинальном df

Question

У меня есть некоторые данные, которые выглядят так, как показано ниже df. Я пытаюсь сначала вычислить средний угол для каждой группы, используя функцию mean_angle. Вычисленный средний угол затем используется для другого расчета для группы с использованием функции fun.

import pandas as pd
import numpy as np

генерация данных выборки

a = np.array([1,2,3,4]).repeat(4)
x1 = 90 + np.random.randint(-15, 15, size=a.size//2 - 2 )
x2 = 270 + np.random.randint(-50, 50, size=a.size//2 + 2 )

b = np.concatenate((x1, x2))
np.random.shuffle(b) 

df = pd.DataFrame({'a':a, 'b':b})

Возвращенный кадр данных напечатан ниже.

    a   b
0   1   295
1   1   78
2   1   280
3   1   94
4   2   308
5   2   227
6   2   96
7   2   299
8   3   248
9   3   288
10  3   81
11  3   78
12  4   103
13  4   265
14  4   309
15  4   229

Мои функции mean_angle и fun

def mean_angle(deg):

    deg = np.deg2rad(deg)

    deg = deg[~np.isnan(deg)]

    S = np.sum(np.sin(deg))
    C = np.sum(np.cos(deg))
    mu = np.arctan2(S,C)
    mu = np.rad2deg(mu)

    if mu <0:
        mu = 360 + mu

    return mu

def fun(x, mu):

    return  np.where(abs(mu - x) < 45, x, np.where(x+180<360, x+180, x-180))

что я пробовал

mu = df.groupby(['a'])['b'].apply(mean_angle)

df2 = df.groupby(['a'])['b'].apply(fun, args = (mu,)) #this function should be element wise

Я знаю, что это совершенно неправильно, но я не мог придумать лучшего способа.

Требуемый результат выглядит примерно так , где mu среднее значение для каждой группы

 a   b   c
0   1   295 np.where(abs(mu - 295) < 45, 295, np.where(295 +180<360, 295 +180, 295 -180))
1   1   78 np.where(abs(mu - 78) < 45, 78, np.where(78 +180<360, 78 +180, 78 -180))
2   1   280 np.where(abs(mu - 280  < 45, 280, np.where(280  +180<360, 280  +180, 280  -180))
3   1   94                ...
4   2   308               ...
5   2   227                .
6   2   96                 .
7   2   299                .
8   3   248                .
9   3   288                .
10  3   81                 .
11  3   78                 .
12  4   103                .
13  4   265                .
14  4   309                .
15  4   229                .

Любая помощь приветствуется

Answer 1

Вам не нужна ваша вторая функция, просто передайте необходимые столбцы в np.where(). Таким образом, создавая ваш фрейм данных таким же образом и не изменяя вашу функцию mean_angle, у нас есть следующий пример фрейма данных:

    a    b
0   1  228
1   1  291
2   1   84
3   1  226
4   2  266
5   2  311
6   2   82
7   2  274
8   3   79
9   3  250
10  3  222
11  3   88
12  4   80
13  4  291
14  4  100
15  4  293

Затем создайте столбец c (содержащий значения mu), используя groupby() и transform(), и, наконец, примените логику np.where():

df['c'] = df.groupby(['a'])['b'].transform(mean_angle)
df['c'] = np.where(abs(df['c'] - df['b']) < 45, df['b'], np.where(df['b']+180<360, df['b']+180, df['b']-180))

Урожайность:

    a    b    c
0   1  228  228
1   1  291  111
2   1   84  264
3   1  226  226
4   2  266  266
5   2  311  311
6   2   82  262
7   2  274  274
8   3   79  259
9   3  250   70
10  3  222   42
11  3   88  268
12  4   80  260
13  4  291  111
14  4  100  280
15  4  293  113