Add geodesic distance implementation stable around pi

[HampSwe]

Both roma.utils.rotmat_geodesic_distance and roma.utils.rotmat_geodesic_distance_naive are numerically imprecise around $\alpha=\pi$, in the same way as roma.utils.rotmat_geodesic_distance_naive is imprecise around $\alpha=0$. This PR adds another implementation, roma.utils.rotmat_geodesic_distance_pi_stable, which is numerically precise at $\alpha=\pi$ too (i.e. precise for all $\alpha$), at the price of slower runtime performance around $\alpha=\pi$. See the figures below for the numerical improvement.

For performance reasons, roma.utils.rotmat_geodesic_distance_pi_stable first calls roma.utils.rotmat_geodesic_distance. If the resulting $\alpha$ is too close to $\pi$ given a tolerance, it recalculates the value of $\alpha$ using a more precise, but slower, function called roma.utils._rotmat_geodesic_distance_atan2. roma.utils._rotmat_geodesic_distance_atan2 works by the formulas $\alpha=atan2(sin(\alpha), cos(\alpha))$, $sin(\alpha)=\frac{1}{2}|(R_{21}-R_{12}, R_{02}-R_{20}, R_{10}-R_{01})|_2$ and $cos(\alpha)=\frac{1}{2}(Trace(R)-1)$. See https://en.wikipedia.org/wiki/Rotation_matrix#Axis_and_angle for more details on this implementation.

Note that roma.utils._rotmat_geodesic_distance_atan2 is significantly slower than roma.utils._rotmat_geodesic_distance. In this implementation, which has been optimized for CUDA, it is approximately 3x slower on CUDA tensors and 10x slower on CPU tensors. One can change the implementation of it to reverse the results and instead achieve higher runtime performance on the CPU; a comment can be found in the implementation on how to do this. Of course, one could dispatch based on the given device, but for maintainability I have simply opted for the CUDA-friendly implementation. However, for all $\alpha$ that are not close to $\pi$, roma.utils.rotmat_geodesic_distance_pi_stable produces the same result as roma.utils.rotmat_geodesic_distance with only a minimal overhead cost in runtime performance.

This PR also adds a specific test for this new implementation which roma.utils.rotmat_geodesic_distance fails. It also updates the API docs, and makes a small comment under the "Care for numerical precision" section in the docs.

[rbregier]

Hi, thank you for the pull request (and sorry for not commenting earlier).

Numerical precision close to π was indeed not one of my main concerns here. The relative error remains small with rotmat_geodesic_distance and the gradient direction is mathematically unstable in that region anyway. That said, I agree we could include a function such as rotmat_geodesic_distance_atan2 for cases where extra precision is needed.

However, I would prefer not to add the additional rotmat_geodesic_distance_pi_stable function. In my opinion, the extra if/else branching adds complexity while bringing only limited practical benefit: users can simply use rotmat_geodesic_distance_atan2 instead.

Would you mind updating the pull request to remove it?