Using python array's column as boolean to change another column's values

Question

I am farily new to python, so that may be why I haven't been able to search properly in this site for an answer to my problem, in case someone know about it already.

I am reading several data files consisting in six columns: XYZ coordinates and 3-vector components around a sphere. I am using the X and Z coordinates to find the angle location of my vectors, and using the angle results as a new column to my array, as a fourth column. So, I have a 7 column array.

The angles are calculated with np.arctan and they get (-) or (+) sign depending on the quadrant the XZ coordinates are located.

Now, I want to make all the angles around the upper half of the sphere (Z positive) positive and spanning from 0º to 180º. On the other hand, I want my angles around the lower half to be negative, and going from 0º to -180º.

The operations to change sign and apply the 0º-180º range are easy to implement. But, I want to use the third column of my array (Z coordinates) as the boolean to decide how to change the angles (fourth column of the array). I have read quite a lot of info about slicing arrays, or modifying columns/rows based on arbitrary boolean conditions applied to the same columns/rows.

I am trying to avoid combining for-loops and if-statements, and learn how to use the power of the python :)

Thanks in advance!

Answer 1

np.where(cond, if_true, if_false) solves the problem you think you have:

fixed_angle = np.where(z > 0, angle, -angle)

The problem you actually have is that you are not using atan2, or as it's spelt in numpy, np.arctan2