Add an if inside apply in pandas-CodePudding

I have the following function:

d["direction"] = d.apply(lambda x: 180   180/(m.pi*(m.atan2(x.water_u,x.water_v))), axis=1, result_type="expand")

The function adds a column named direction based on the calculation of 2 columns water_u and water_v - As you can see if the angle is zero there will be a zero-division and therefore an error, now I can make a function to add an if, such that it is zero the value of the division remains zero but is there a pythonic way to do this inside my lambda function?

Anyway, I would be grateful, if someone knows a smart succinct way to achieve this.

Regards

CodePudding user response：

You don't have to use lambda function.

Compute your angle separately and use np.where to compute the direction or not:

import numpy as np

angle = (np.pi*(np.arctan2(df.water_u, df.water_v)))
df['direction'] = np.where(angle != 0, 180   180 / angle, 0)

Example:

# Before
>>> df
   water_u  water_v
0        0        0
1       10       20

# After
>>> df
   water_u  water_v   direction
0        0        0    0.000000
1       10       20  303.576135