Create new pandas dataframe columns from function that returns tuple

Question

I have a function that returns tuple:

def pwrs(x):
   return x*x, x*x*x, x*x*x*x

I would like to apply this function to a single column dataframe named data:

+-------+
|    x  |
+-------+
|    1  |
|    2  |
|    3  |
|    4  |
+-------+

and get back a new dataframe with new columns based on the function return tuples:

+---+------+------+------+
| x |  x^2 |  x^3 |  x^4 |
+---+------+------+------+
| 1 |    1 |    1 |    1 |
| 2 |    4 |    8 |   16 |
| 3 |    9 |   27 |   81 |
| 4 |   16 |   64 |  256 |
+---+------+------+------+

Got as far as iterating through the rows and applying the function:

for _, row in data.iterrows():
    print(pwrs(row['x']))

Not sure how to progress from here....

np.vander might be of use

ALollz
– ALollz

2019-11-15 21:29:51 +00:00
Commented Nov 15, 2019 at 21:29 — ALollz
– ALollz, Commented Nov 15, 2019 at 21:29

piRSquared · Accepted Answer · 2019-11-15 21:26:39Z

4

Generalized

c = np.arange(1, 5)
pd.DataFrame(df.to_numpy() ** c, df.index, c).add_prefix('x^')

   x^1  x^2  x^3  x^4
0    1    1    1    1
1    2    4    8   16
2    3    9   27   81
3    4   16   64  256

answered Nov 15, 2019 at 21:26

piRSquared

296k68 gold badges509 silver badges654 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

BENY · Accepted Answer · 2019-11-15 21:24:43Z

3

IIUC

pd.DataFrame(df.x.values[:,None]**np.array([1,2,3,4]))
Out[290]: 
   0   1   2    3
0  1   1   1    1
1  2   4   8   16
2  3   9  27   81
3  4  16  64  256

answered Nov 15, 2019 at 21:24

BENY

324k22 gold badges176 silver badges250 bronze badges

Comments

sacuL · Accepted Answer · 2019-11-15 21:25:38Z

2

If you want to use your function, you can simply use the dataframe column as the argument:

df['x^2'], df['x^3'], df['x^4'] = pwrs(df['x'])


   x  x^2  x^3  x^4
0  1    1    1    1
1  2    4    8   16
2  3    9   27   81
3  4   16   64  256

answered Nov 15, 2019 at 21:25

sacuL

51.6k9 gold badges88 silver badges115 bronze badges

Comments

Dani Mesejo · Accepted Answer · 2019-11-15 21:28:29Z

1

You could do:

df[['x^2', 'x^3', 'x^4']] = df.x.apply(lambda x: pd.Series(pwrs(x)))
print(df)

Output

   x  x^2  x^3  x^4
0  1    1    1    1
1  2    4    8   16
2  3    9   27   81
3  4   16   64  256

Notice that this works for any function that returns a tuple, not just this mathematical operations.

answered Nov 15, 2019 at 21:28

Dani Mesejo

62.2k6 gold badges56 silver badges86 bronze badges

1 Comment

Leon Adams Over a year ago

This approach avoided a tuple unpacking error I was encountering in the other suggestion

Collectives™ on Stack Overflow

Create new pandas dataframe columns from function that returns tuple

4 Answers 4

Generalized

Comments

Comments

Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

Generalized

Comments

Comments

Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related