NUTS/BarkerMH not sampling posterior properly

numpyro
1

I’m new at numpyro, and I’m trying to sample the Fourier series parameters that fit my data. The likelihood function uses the misfit between a Fourier series model (curve) and the data at each iteration to discard model that have a misfit higher than 0.5.
This is how I parameterised the model: (Just for the record, this parametrization works perfectly on pymc)

def fold_bayesian_model(w, y=None):

    c0 = numpyro.sample('c0', dist.Normal(0, 1), 
                                                           )
    c1 = numpyro.sample('c1', dist.Normal(0, 1), 
                                                           )
    c2 = numpyro.sample('c2', dist.Normal(0, 1), 
                                                           )
    wl = numpyro.sample('wl', dist.Normal(w, w/4))
    
# fourier parameters
    theta = [c0, c1, c2, wl]

# function that calculate the misfit
    misfitx = misfit(theta, fld, flr)
   
    numpyro.sample('misfit', dist.Normal(0, 0.5), obs=misfitx)

        
kernel = BarkerMH(fold_bayesian_model)

mcmc = MCMC(kernel, num_warmup=3000, num_samples=1000)

rng_key = random.PRNGKey(64)

mcmc.run(rng_key, s0_lwmu, extra_fields=('potential_energy',))

Then I get the summary with r_hat and n_eff are weird:
mean std median 5.0% 95.0% n_eff r_hat
c0 1.04 0.00 1.04 1.04 1.04 6.48 1.07
c1 0.15 0.00 0.15 0.15 0.15 4.52 1.32
c2 -0.88 0.00 -0.88 -0.88 -0.88 4.61 1.14
wl -1.89 0.00 -1.89 -1.89 -1.89 267.23 1.00

wl is the wavelength of the Fourier series curve and seems weird, it should be around 14 000 m while I get -1.89.

Do you know why I get this weird results?

Thank you in advance

2
3

Yes I tried using more samples I used over 100 000 samples in both barkermh and nuts and i get the same problem.
I tried as you suggested 64 bit precision and still the same problem.

4

can you post the rest of your model? maybe something is mis-specified

5

you might also try some of the flags in NUTS like dense_mass=True, max_tree_depth=12, find_heuristic_step_size=True, regularize_mass_matrix=False etc

6

This is the whole model:

@jit
def jax_fourier_series(x, c0, c1, c2, w):
“”"

Parameters
----------
x
c0
c1
c2
w

Returns
-------

"""
v = jnp.array(x.astype(float))
# v.fill(c0)
v = c0 + c1 * jnp.cos(2 * np.pi / w * x) + c2 * jnp.sin(2 * np.pi / w * x)
return jnp.rad2deg(jnp.arctan(v))

fld = jnp.asarray(s0.fold.fold_limb_rotation.fold_frame_coordinate)
flr = jnp.asarray(s0.fold.fold_limb_rotation.rotation_angle)
s0_lwmu = 1.48804364e+04
@jit
def misfit(theta, axis=False):

# fold_limb_rotation = FoldRotationAngle(rotation_angle, fold_frame)
def fold_function(x, theta):
    
    return jnp.rad2deg(jnp.arctan(
       
            jax_fourier_series(x, *theta))
        )

# calculate the observed misfit for for the likelihood function
def calculate_misfit():

    return jnp.tan(jnp.deg2rad(rotation_angle)) - jnp.tan(
        jnp.deg2rad(fold_function(fold_frame, theta)))

return calculate_misfit()

def fold_bayesian_model(w, y=None):

c0 = numpyro.sample('c0', dist.Normal(0., 1.), 
                                                       )
c1 = numpyro.sample('c1', dist.Normal(0., 1.), 
                                                       )
c2 = numpyro.sample('c2', dist.Normal(0., 1.), 
                                                       )
wl = numpyro.sample('wl', dist.Normal(w, w/3))


theta = [c0, c1, c2, wl]

misfitx = misfit(theta)

numpyro.sample('misfit', dist.Normal(0, 5), obs=misfitx)

kernel = NUTS(fold_bayesian_model, dense_mass=True,
max_tree_depth=20, find_heuristic_step_size=True,
regularize_mass_matrix=False)

mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000)

rng_key = random.PRNGKey(64)

mcmc.run(rng_key, s0_lwmu, extra_fields=(‘potential_energy’,))

7

I tried the flags but same problem!

8

@kodi i meant that you should try those non-default settings one-by-one: not all at once

are you aware that the second argument to Normal(loc, scale) is a square root variance and not a variance?