Sampling a 1D gaussian mixture

[jecampagne]

Hi,
Here is an exercise using Numpyro JaxNs wrapper that I would like to validate. I have followed
http://num.pyro.ai/en/stable/examples/gaussian_shells.html?highlight=Nested

class MixtureModel(dist.Distribution):
    support = dist.constraints.real_vector

    def __init__(self, locs, scales, weights):
        self.locs, self.scales, self.weights= locs, scales, weights
        super().__init__(batch_shape=locs.shape[:-1], event_shape=locs.shape[-1:])

        self.loc = locs.T
        self.scale = scales.T
        self.weights = weights.T
        norm = jnp.sum(self.weights)
        self.weights = self.weights/norm

        self.num_distr = len(locs)

        
    def sample(self, key, sample_shape=()):
        return jnp.zeros(
            sample_shape + self.shape()
        )  # a dummy sample to initialize the samplers

    def log_prob(self, x):
        log_probs = jax.scipy.stats.norm.logpdf(x,loc=self.loc, scale=self.scale)
        return jax.scipy.special.logsumexp(jnp.log(self.weights) + log_probs, axis=0)[0]

def model():
    x=numpyro.sample("x", dist.Uniform(-3.0, 3.0))
    numpyro.sample("val", MixtureModel(jnp.array([0,1.5]),
                                       jnp.array([0.5,0.1]),
                                       jnp.array([8,2])), obs=x)

ns = NestedSampler(model)
ns.run(random.PRNGKey(2))

samples = ns.get_samples(random.PRNGKey(3), num_samples=100_000)

Well, first the distribution of 'x' looks like I was expecting (although I do not anderstand why it is the 'x' variable and not the 'val' wihich is the output, but I guess it more related to Numpyro wrapping)

But the summary is

--------
# likelihood evals: 10059
# samples: 8354
# likelihood evals / sample: 1.2
--------
logZ=-1.851 +- 0.027
ESS=3209
and then crash for computing percentile

Now I know the Integral value it is 0.63167, and so I do not anderstand the logZ value, also for 100_000 samples wanted I only get 8300 or so.

Do you see any reason that the sampler is not so good here (nb. HMC(NUTS) is doing well) and why the logZ seems out also?

Thanks

[jecampagne]

Add: I have increased significantly the num_live_points from 1_000 (default) to 10_000 and the sampling looks better

--------
# likelihood evals: 2908628
# samples: 82529
# likelihood evals / sample: 35.2
--------

logZ=-1.7929 +- 0.0083
ESS=31622

But,

1. why I need to increase so much the number of live points as the rule of thumb is
   num_live_points=50*(# posterior modes)*(D+1) which here gives 50 x 2 x (1+1) = 200.
2. why LogZ is not Log(value of Integral ) ?
   (nb. there are still summary crashes)

[Joshuaalbert]

Hi @jecampagne let me look into this.

• logZ is the log of the Bayesian evidence, or marginal likelihood. Jaxns is also supposed to be arbitrarily accurate and asymptote to an unbiased estimate. If you're not getting the right number then I'd like to put this into a regression test and solve it. Try again with 0.0.6 and/or 0.0.5 and let me know.

pip install --ignore-installed jaxns==0.0.6
or
pip install --ignore-installed jaxns==0.0.5
Do you get the same results?

• Basically, I've noticed a performance hit to jaxns recently, in my own experiments, which I introduced in the last commit.
• I am in the process of revamping Jaxns to be much more powerful and efficient with likelihood samples, but it will break the API and require numpyro to upgrade their wrappers.

[jecampagne]

Hello @Joshuaalbert
Thanks for interest, currently I use

Name: jaxns
Version: 0.0.7
Summary: Nested Sampling in JAX
Home-page: https://github.com/joshuaalbert/jaxns
Author: Joshua G. Albert
Author-email: [email protected]
License: UNKNOWN

Do you mean that I should downgrade JaxNS?

[Joshuaalbert]

Yes, for these tests let's downgrade to 0.0.6 until 1.0 is released.

[jecampagne]

Dear @Joshuaalbert I cannot downgrade to jaxns 0.0.6 as when I import jaxns I get this error message

ImportError: cannot import name 'squared_norm' from 'jaxns.gaussian_process.utils' (...python3.8/site-packages/jaxns/gaussian_process/utils.py)

[Joshuaalbert]

@jecampagne can you pull the latest and try again. The API has changed slightly, see the tutorial here: https://github.com/Joshuaalbert/jaxns/blob/master/examples/simple/mvn_data_mvn_prior.ipynb

It's currently in release candidate for 1.0. You'll be happy to know that one of the features that I'm releasing is a dynamic arbitrary precision inference (the longer you wait the more precise the evidence and posterior evidence gets). It will also allow refinement of past runs so that you can improve iteratively until you're satisfied. In release candidate I'm not 100% ready, but here's a sneak peak at tutorial https://github.com/Joshuaalbert/jaxns/blob/master/examples/physics/Jones_scalar_modelling.ipynb