emcee icon indicating copy to clipboard operation
emcee copied to clipboard

Sampling quaternion variables constrained to a manifold

Open flothesof opened this issue 2 years ago • 3 comments

Hi,

I have a question that may or may not be outside the scope of emcee. I am trying to replicate work found here using emcee: http://asa.scitation.org/doi/10.1121/1.5017840 This work tries to infer elastic constants as well as a 3D angle of rotation using Hamiltonian Monte Carlo. To parametrize this rotation unambiguously, they avoid Euler angles and choose unit quaternions defined by four scalar values (x, y, z, w) which completely describe the rotation that is sought to best match experimental data. I tried to define this problem using emcee but then realized that I am missing something important: the ability to constrain values given by the proposal distribution. Concretely, if the proposal distribution generates 4 new values (x, y, z, w), they have to satisfy x^2 + y^2 + z^2 + w^2 = 1 to lie on the required manifold for rotations. Browsing the documentation, I was wondering if it is possible to use the Move interface to achieve something like this? Or do you see an alternative package I should move to?

Thank you for any pointers. Regards, Florian

flothesof avatar Sep 29 '21 12:09 flothesof

You can reparameterize to sample in an unconstrained 4-D space, e.g.:

theta_1 ~ N(0, 1), theta_2 ~ N(0, 1), theta_3 ~ N(0, 1), theta_4 ~ N(0, 1)

then normalize these to get your quaternions before evaluating your model:

x = theta_1 / sqrt(theta_1^2 + theta_2^2 + theta_3^2 + theta_4^2), ...

You can check that these are uniformly distributed on your target manifold.

In emcee, I would write this as:

def logprob(theta):
    # We're going to sample in the unconstrained space, then normalize those
    # parameters 
    norm = np.sum(theta ** 2)
    x, y, z, w = theta / np.sqrt(norm)

    # Then, we need to put a prior on that space so that we can sample. As long
    # as it's spherical, anything would be fine; let's use a standard normal 
    logprior = -0.5 * norm

    # Use x, y, z, and w, which will now be correctly normalized)...
    loglike = 0.0  # Some function fo x, y, z, and w

    return logprior + loglike

dfm avatar Sep 29 '21 13:09 dfm

Hi Dan,

thank you very much for your quick answer. I’ve tried out the solution you suggest. Here’s what I found for the first part: sampling unconstrained standard normals and then norming them do indeed uniformly distribute when projected to 4D axis-angle representation: image I then tried to setup an inverse problem using your idea:

from scipy.stats import norm
import emcee
import numpy as np

def make_unit_quats(N):
    rvs = [norm() for _ in range(4)]
    values = np.stack([rv.rvs(N) for rv in rvs]).T
    unit_quaternions = values / np.sqrt(np.sum(values ** 2, axis=1)[:, None])
    return unit_quaternions

q0 = np.array([0.9884, 1., -0.1512, 0.4001])
q0 = q0 / np.linalg.norm(q0)
print(q0)

pos = make_unit_quats(20)
nwalkers, ndim = pos.shape

def logprob(theta):
    norm = np.sum(theta ** 2)
    x, y, z, w = theta / np.sqrt(norm)
    q = np.array([x, y, z, w])
    logprior = -0.5 * norm
    sigma = 1
    loglike = -1/sigma**2 * np.sum((q - q0) ** 2 / sigma**2)
    return logprior + loglike

sampler = emcee.EnsembleSampler(nwalkers, ndim, logprob)
sampler.run_mcmc(pos, 5000, progress=True);

With sigma = 1, I get: image

However, if I change sigma to 0.1 (to get "sharper" distributions), I get: image

I would have expected to get a ball-like distribution around the right target value like with sigma = 1. Is there an explanation for why this is not the case?

Thank you for your help!

Florian

flothesof avatar Sep 30 '21 12:09 flothesof

After second thoughts, I’m thinking my loglikelihood is problematic it seems to skew the distributions in a very unspherical way.

I guess the procedure you suggested does indeed do what it should, I’m just not using it right in this simplified problem :)

flothesof avatar Sep 30 '21 12:09 flothesof