Being careful is great, but I think these `assert` statements could be omitted. If `np.arctan2` ever returns something outside the range [-pi, pi] something is very wrong.