Suggestions for Jlog3 and Jlog6

[simeon-ned]

@brentyi Hello!

First of all, thank you for your awesome library! I often use jaxlie in my work and research whenever I'm dealing with rotations and translations in JAX.

Right now, I'm working on a project that needs to compute derivatives of log functions for differences between SO3/SE3 elements, involving Jlog3/Jlog6 functions. For example, Jlog3 is defined as:

$$\text{Jlog}_3(R) = \frac{\partial \log_3(R)}{\partial R}$$

And similarly, $\text{Jlog}_6(T)$ for SE3.

These Jacobians admit closed-form solutions, making them particularly efficient and useful. This blog post and its references provide some useful context. These functions are also implemented in libraries like Pinocchio.

I think it would be great to have these as methods in the SO3/SE3 classes or as utility functions in jaxlie. If you agree, I'd love to help by creating a PR and doing the implementation myself, if you have time to review it.

Thanks again for your work :)

[brentyi]

Hello @simeon-ned! Thanks for the nice words, I'm glad the library's working for you!

This sounds useful and I'm happy to review, especially if you have time for tests. I appreciate you offering to help!! For Jacobians specifically, benchmarking both compilation time and runtime against autodiff would also be useful as justification. Often I'm surprised by how comparable runtimes are, but manual Jacobians can accelerate compilation significantly.

If it's possible to also support SO(2) and SE(2) that would also be great for consistency. The 2D support does make the "Jlog3" terminology slightly confusing though, since both SE(2) and SO(3) have 3D log spaces.

[simeon-ned]

@brentyi Sure, I will add tests and consider supporting SO(2) and SE(2). Regarding naming, if we include these functions as methods or properties in dedicated classes, we can keep the names concise, such as jlog, as the class names will clearly indicate the groups they belong to.

[brentyi]

Sounds great, thanks! If you don't have time for SO(2) and SE(2) we could also just implement those using jax.{jacfwd/jacrev}.

[simeon-ned]

Sure, I will create draft of this PR by the end of this week or the beginning of next week.

[simeon-ned]

@brentyi Hello!

I've created a draft PR that includes all jlog functions implemented, along with a draft test suite measuring compilation time, first call time, and average time per call. Generally, performance is comparable between analytical and automatic differentiation approaches, with analytical Jacobians showing faster compilation and first call execution times. The most significant improvement is for SE3, where we observe a 3-4 times speed increase for both first call time and average time per call.

The test suite currently needs to be rewritten for pytest, but I'm encountering some batching issues. Additionally, I'm uncertain if storing Jacobians within the group class is the optimal approach. For future development, we could consider implementing all analytical Jacobians as detailed in the micro Lie theory paper, for instance equations 143-147 for the SO(3) group. This might involve creating a dedicated object for managing these Jacobians.

I'd greatly appreciate your review of the PR draft when you have the opportunity. Once we align on the approach, I'll proceed with rewriting the tests using pytest. Please let me know your thoughts on how we should move forward.

Thanks!