apply user defined function to pandas dataframe specific columns and add new columns to dataframe

Question

I created a function called polar(X_relative , Y_relative , Z_relative) that takes these 3 arguments as shown and the result are new arguments which are (azimuth_angle , tilt_angle). I want to apply this function to a pandas dataframe where the function arguments are certain columns in the dataframe and would like to add the output arguments of the functions (azimuth_angle , tilt_angle) as new coulmns in the dataframe where the function is calculated for each row.

dataframe columns: X , Y, Z , X_relative , Y_relative , Z_relative
dataframe coulmns I expect after applying the function: X , Y, Z , X_relative , Y_relative , Z_relative , azimuth_angle , tilt_angle

Answer 1

Wrap the return of polar into a series:

return pd.Series([azimuth_angle, tilt_angle]) 

Then use apply:

df[['azimuth_angle', 'tilt_angle']] = df.apply(lambda x: polar(x.X_relative, x.Y_relative, x.Z_relative), axis = 1)