Pandas : Create multiple new variables by applying function to dataframe

August 13, 2022

Consider the following simple function:

def Powers(x):
    return [x, x**2, x**3, x**4, x**5]

and input dataframe:

df = pd.DataFrame({ 'x':(1, 2, 3, 4, 5) })

I would like to generate new variables: ['Exp_1', 'Exp_2', 'Exp_3', 'Exp_4', 'Exp_5']

When I apply the function to the dataframe as follows:

df[['Exp_1', 'Exp_2', 'Exp_3', 'Exp_4', 'Exp_5']] = df.apply(lambda x: Powers(x.x), axis=1)

I get:

In other words, the values are transposed. That is, the 5th exponent of 1 is 1 not 5 and the 1st exponent of 5 is 5 and not 1.

I have tried axis=0, in the call above and this does not work either. I also know I have a problem because if the input dataframe is of a different length I get errors.

How do I fix this?

>Solution :

You can return Series in Powers function

def Powers(x):
    return pd.Series([x, x**2, x**3, x**4, x**5])

df[['Exp_1', 'Exp_2', 'Exp_3', 'Exp_4', 'Exp_5']] = df.apply(lambda x: Powers(x.x), axis=1)

print(df)

   x  Exp_1  Exp_2  Exp_3  Exp_4  Exp_5
0  1      1      1      1      1      1
1  2      2      4      8     16     32
2  3      3      9     27     81    243
3  4      4     16     64    256   1024
4  5      5     25    125    625   3125

Or use result_type in DataFrame.apply

def Powers(x):
    return [x, x**2, x**3, x**4, x**5]

df[['Exp_1', 'Exp_2', 'Exp_3', 'Exp_4', 'Exp_5']] = df.apply(lambda x: Powers(x.x), axis=1, result_type='expand')
# or
df[['Exp_1', 'Exp_2', 'Exp_3', 'Exp_4', 'Exp_5']] = df.apply(lambda x: Powers(x.x), axis=1).tolist()