What's missing for sampling on product manifolds?

[sethaxen]

It would be great if this package had the flexibility to run HMC on product manifolds. See for example this paper (by @simonbyrne, who I see filed an issue here), which presents an approach for sampling on embedded manifolds using the geodesic flow on the manifold (see Algorithm 1). The paper also notes that the geodesic flow on a product manifold splits into the flows on the submanifolds, so that each step of the integration is done in parallel. This would be very useful for applying HMC to sample geometric quantities. For example, I'm interested in sampling orientations of rigid bodies (SE(3), implemented as S^3 x R^3). With n Euclidean variables and m rigid bodies, the variable space is the product manifold R^n x (S^3 x R^3)^m = R^(n+3m) x (S^3)^m). The Euclidean part can be updated with all Euclidean variables, while the latter must be considered as a separate unit with their own kinetic energy, leapfrog integrator, and integrator step size (equivalent to scale of the spherical metric; alternatively, scaling factor for step size).

The machinery to do this might be too specialized for this package, but I'm wondering whether you think the building blocks here are sufficiently general to enable this kind of manifold sampling. At first glance, it seems that a reimplementation of much of hamiltonian.jl might be all that's necessary.

[tpapp]

I don't know anything about sampling on product manifolds, so can't say much about this.

That said, I would prefer to keep this package focused on a single method, mainly what NUTS evolved to in practice. This package itself is a building block (just for sampling from a log density on R^n), and the way I envision MCMC in Julia is to provide similar, interchangeable building blocks that interface well.

I am happy to consider suggestions about exposing various parts of the API, or even factoring out pieces to packages. The simplectic integrator itself is the most trivial part — the challenge is providing an efficient and consistent tree-building algorithm, adaptation, and a well-designed API.

I am currently redesigning the API (#30), but I am happy to discuss this in the meantime.

[simonbyrne]

The first challenge would be figuring out what the analogy of the NUTS criterion would be on a manifold: I haven't thought about it a lot, it might be possible to use the same criterion directly, i'm not sure.

I suspect it would also require modifications to the AbstractLogDensity interface: e.g. rather than just expecting a vector of inputs, you would need to pass through a type containing the manifold information, e.g.

EuclideanManifold
SphericalManifold
ProductManifold

[sethaxen]

@simonbyrne

The first challenge would be figuring out what the analogy of the NUTS criterion would be on a manifold: I haven't thought about it a lot, it might be possible to use the same criterion directly, i'm not sure.

I had a brief exchange with @betanalpha a ways back, and he pointed me to this paper, which generalizes the NUTS criterion to any Riemannian manifold. Looks like the implementation of the criterion here is from that paper. That being said, I don't know if there's a fundamental difference between executing the criterion on the momentum or on the velocity from your paper.

@tpapp:

That said, I would prefer to keep this package focused on a single method, mainly what NUTS evolved to in practice. This package itself is a building block (just for sampling from a log density on R^n), and the way I envision MCMC in Julia is to provide similar, interchangeable building blocks that interface well.

I agree that it probably doesn't make sense to add the complication of product manifold sampling to this package, but it would be great if the building blocks were sufficiently general that they could be assembled for product manifold sampling. As far as I understand, the NUTS machinery itself is fairly general, and there's no reason it couldn't be applied to product manifolds, except most existing implementations (e.g. Stan's) assume pure Euclidean/Riemannian manifolds.

[tpapp]

@sdaxen: I understand that certain elements can be shared between implementations. But it would also complicate the code in this package for something which, if I understand correctly, is experimental.

As I said, this is beyond my time & knowledge to implement anything like this. If someone wants to implement it, they should start a separate package. If it works out, I am happy to talk about merging, or making it easier to share code. In the meantime, I am closing this.