An exercice both native JaxNS and Numpyro wrapper

[jecampagne]

Here is an exercise with a fit over observations. Here is the snippet un JaxNS (0.0.7). At the end I point some questions:

from jaxns.prior_transforms import UniformPrior, PriorChain, NormalPrior

def log_normal(x, mean, scale):
    dx = (x - mean) / scale
    return -0.5 * jnp.log(2. * jnp.pi) - jnp.log(scale)  - 0.5 * dx * dx

def solve(X_obs, Y_obs):
    
    def log_lik(a0,a1,a2,a3,sigma, **kwargs):
        Y = a0 + a1*X_obs + a2*X_obs**2 + a3*X_obs**3
        LogL = log_normal(Y,Y_obs,sigma)
        return jnp.sum(LogL)
    
    prior_chain = PriorChain(NormalPrior('a0',0.,10.),
                             NormalPrior('a1',0.,10.),
                             NormalPrior('a2',0.,10.),
                             NormalPrior('a3',0.,10.),
                             UniformPrior('sigma',0.,10.))
    
    print('num_live_points:',prior_chain.U_ndims*500)
                             
    ns = jaxns.nested_sampling.NestedSampler(log_lik, prior_chain,
                       num_live_points=prior_chain.U_ndims*500)
    
    print('Go...')
                             
    results = ns(jax.random.PRNGKey(32564))
    
    return results

param_true = np.array([1.0, 0.0, 0.2, 0.5, 1.5])
sample_size = 5_000
sigma_e = param_true[4]          # true value of parameter error sigma
random_num_generator = np.random.RandomState(0)
xi = 5*random_num_generator.rand(sample_size)-2.5
e = random_num_generator.normal(0, sigma_e, sample_size)
yi = param_true[0] + param_true[1] * xi + param_true[2] * xi**2 + param_true[3] *xi**3 +  e  
plt.hist2d(xi, yi, bins=50);

results = solve(xi,yi)

num_live_points: 2500
Go...

Q0 : I am running on a CPU machine (ie not a GPU) and it takes 3min30sec to get the results. This is strange because using Numpyro wrapper takes about 30sec to get running the NestedSampler and get 5000 samples.

Then

jaxns.utils.summary(results)

crashes TypeError: percentile requires ndarray or scalar arguments, got <class 'list'> at position 1. I get this error with other JaxNS exercices.

plot_diagnostics(results)

I do not know if its ok or not, but after

samples = jaxns.utils.resample(random.PRNGKey(43083245),results.samples, results.log_p, S=int(results.ESS))

I cannot use the corner plots given by the library, so I have written mine

def plot_corner_diagnostics(samples, param_true):
    mpl.rcParams['font.size'] = 12
    ndim = len(samples.keys())
    # This is the empirical mean of the sample:
    value2 = np.mean(np.array(list(samples.values())),axis=1)
    #True
    value1 = param_true

    # Make the base corner plot
    # 68% et 95% quantiles 1D et levels in 2D
    figure = corner.corner(samples,quantiles=(0.025, 0.158655, 0.841345, 0.975), levels=(0.68,0.95), 
                        show_titles=True, title_kwargs={"fontsize": 12}, 
                        truths=param_true, truth_color='g', color='b'
                        );

    # Extract the axes
    axes = np.array(figure.axes).reshape((ndim, ndim))

    # Loop over the diagonal
    for i in range(ndim):
        ax = axes[i, i]
        ax.axvline(value2[i], color="r")
    
    # Loop over the histograms
    for idy in range(ndim):
        for idx in range(idy):
            ax = axes[idy, idx]
            ax.axvline(value2[idx], color="r")
            ax.axhline(value2[idy], color="r")
            ax.plot(value2[idx], value2[idy], "sr")

    plt.show()

So,

plot_corner_diagnostics(samples,param_true=param_true)

gives

which looks fine, at least it corresponds to the Numpy Wrapper results and also to NUTS sampling results too.

But (see also Q0 above):
Q1) I wander: from where JaxNS tells to give me 100_000 samples? while sometimes it is few hundred only?
Q2) Also why the percentile crashes in the JaxNS disgnotics (print/corner...)
Q3) why in the https://github.com/Joshuaalbert/jaxns/blob/master/examples/jones_scalar_model.py example there is no needs to use jaxns.utils.resample to get the samples for corner plot ???
R4) By the way a get_samples(random key, num_samples) a-la-Numpyro wrapping would be useful.

Thanks

[Joshuaalbert]

Q0: In your code you're not jit compling. Try:

ns = jaxns.nested_sampling.NestedSampler(log_lik, prior_chain,
                       num_live_points=prior_chain.U_ndims*500)
ns = jax.jit(ns)

Q1: Just to be clear are you asking about results.num_samples? This is the total number of samples generated, but they do not have equal sample weight. results.ESS is the effective number of samples, taking into account weight. When setting max_samples in the NestedSampler constructor, you are saying the maximum num of total samples you are willing to collect. This allocates a maximum memory space for those samples.

Q2: Bug in 0.0.7, that I've fixed, and will release in 0.0.8.

Q3: you never need to resample for corner plot. It does it for you. You just feed in cornerplot(results) and it resamples for you.

R4: With nested sampling you can't really control the number of samples you get out, since they are weighted. I guess it would be possible to recalculate ESS after each run and make a stopping criterion that is based on having enough ESS. That's not too hard. I'll make an issue for it.

[jecampagne]

Thanks @Joshuaalbert for your kind answers.
I guess for R4 you know how Numpyro proceeds from the output of JaxNS (here is the link).

[Joshuaalbert]

There is now a stopping criterion for having a certain number of ESS. It works with dynamic sampling.

[Joshuaalbert]

@jecampagne anything unanswered here?