Как применить функцию, которая возвращает фрейм данных к групповым объектам

Question

Функция, которую я хочу применить к каждому объекту groupby, является функцией ATR , она принимает три ndarrays и возвращает объект dataframe. И я также хочу поместить результат в новый столбец исходного кадра данных. Вот код, который я написал:

# the original dataframe:
 col1    col2    col3    group
4762.1  4762.9  4762.9     1
4762.3  4762.1  4762.9     2
4762.3  4762.9  4762.1     1
4762.7  4762.5  4762.5     2
4762.5  4762.5  4762.5     1
  ...     ...     ...     ...

def get_atr(df):
    return ta.ATR(df['col1'].values, df['col2'].values, df['col3'].values)

df['atr'] = df.groupby('group').apply(get_atr)

Ожидаемый результат должен быть

 col1    col2    col3    group       atr
4762.1  4762.9  4762.9     1    (the result)
4762.3  4762.1  4762.9     2    (the result)
4762.3  4762.9  4762.1     1    (the result)
4762.7  4762.5  4762.5     2    (the result)
4762.5  4762.5  4762.5     1    (the result)
  ...     ...     ...

Но я получаю все значения NaN

Answer 1

Вы можете попробовать создать новый столбец в функции:

def get_atr(x):
    x['atr'] = ta.ATR(x['col1'].values, x['col2'].values, x['col3'].values)
    return x

df = df.groupby('group').apply(get_atr)