Random sampling using other distributions on SO(3) (ex. Von Mises Fisher)

[femtomc]

Hi -- first off, this package is incredible. Thank you.

Second, have you considered adding further random sampling functionality (for other distributions on spaces of rotations). In the title, I referred to the VMF distribution.

If you think this functionality should live elsewhere, totally fine! If you're curious and open to it, I may setup a PR and move some of the functionality from a package I'm working on to here.

[brentyi]

Thanks for the kind words!!

I did think about including distributions like VMF and Bingham at one point and would be open to it, the main reason this never happened is that tangent-space Gaussians have been lower-effort and sufficient for my own use cases.

Two things I'm wondering:

• I noticed that TensorFlow Probability has a JAX substrate with a high-level VMF implementation; does it still make sense to include one in jaxlie?
• If it does: did you have a particular API in mind? jaxlie.SO3.sample_vmf()? jaxlie.distributions.sample_vmf()?

[femtomc]

@brentyi Re (1) -- that's a fair point.

Honestly, I don't know how large a dependency TFP is. I'd really prefer to not have to pull in all of TensorFlow for my own package. This is something I'll have to investigate.

For (2) -- the main functionality that I'm interested in is sampling and scoring. I would be wrapping distributions for use in a probabilistic programming context, so evaluating logpdf is also valuable to me.

[femtomc]

(Also, re -- TFP: I don't fully understand the "substrate" layer. My stuff is all in pure JAX anyways, so I've been reaching for distributions that are also in pure JAX -- I didn't really investigate integration with TFP, although the documentation seems to indicate a high quality for their distributions. If I could somehow figure out how to just depend on distributions + the JAX substrate, it would basically solve my motivating problem)

[femtomc]

Just to be fully transparent -- I'm interested in functionality like this: https://github.com/probcomp/GenDirectionalStats.jl

But in JAX. If these distributions seem interesting and in scope, I can likely contribute a PR with my colleague.

[brentyi]

On tensorflow_probability: I don't know much about it either, but the required dependencies seem light. I also just tried and was able to install/use it without installing TensorFlow. (turns out they also emphasize this in an example)

If a contribution to jaxlie still makes sense despite that and the you have time for a PR with a sensible API + types + tests, I'm definitely happy to review / merge!

[femtomc]

Cool! -- I'll play a bit with tfp, and see if it still makes sense. Thanks for the discussion nonetheless!