Using different MC packages for Bayesian sampling

refnx can work with a variety of MC packages for inference. This notebook will demonstrate the use of various packages:

An excellent reference to see how to use a wide range of packages for statistical inference is https://mattpitkin.github.io/samplers-demo/pages/samplers-samplers-everywhere.

We’re also going to use the arviz package to calculate effective sample size.

import os

# Set these FIRST before any other imports
os.environ["JAX_ENABLE_X64"] = "1"

import warnings
from importlib import resources
import numpy as np
import matplotlib.pyplot as plt
import scipy
import multiprocessing as mp

import pymc as pm
import dynesty
import arviz as az

import refnx
from refnx.dataset import ReflectDataset, Data1D
from refnx.analysis import (
    Transform,
    CurveFitter,
    Objective,
    Model,
    Parameter,
    pymc_model,
    process_chain,
)
from refnx.reflect import SLD, Slab, ReflectModel

It’s important to note down the versions of the software that you’re using, in order for the analysis to be reproducible.

print(
    f"refnx: {refnx.version.version}\n"
    f"scipy: {scipy.version.version}\n"
    f"numpy: {np.version.version}"
)

refnx: 0.1.67.dev0+git20260730.676d0ea
scipy: 1.18.0
numpy: 2.4.6

The dataset we’re going to use as an example is distributed with every install. The following cell determines its location.

pth = resources.files(refnx.analysis)
DATASET_NAME = "c_PLP0011859_q.txt"
file_path = pth / f"tests/{DATASET_NAME}"

data = ReflectDataset(file_path)

The Structure

si = SLD(2.07, name="Si")
sio2 = SLD(3.47, name="SiO2")
film = SLD(2.0, name="film")
d2o = SLD(6.36, name="d2o")

# first number is thickness, second number is roughness
# a native oxide layer
sio2_layer = sio2(30, 3)

# the film of interest
film_layer = film(250, 3)

# layer for the solvent
d2o_layer = d2o(0, 3)

sio2_layer.thick.setp(bounds=(15, 50), vary=True)
sio2_layer.rough.setp(bounds=(1, 15), vary=True)

film_layer.thick.setp(bounds=(200, 300), vary=True)
film_layer.sld.real.setp(bounds=(0.1, 3), vary=True)
film_layer.rough.setp(bounds=(1, 15), vary=True)

d2o_layer.rough.setp(vary=True, bounds=(1, 15))

structure = si | sio2_layer | film_layer | d2o_layer

print(sio2_layer.parameters)

________________________________________________________________________________
Parameters:     'SiO2'     
<Parameter:'SiO2 - thick' , value=30          , bounds=[15.0, 50.0]>
________________________________________________________________________________
Parameters:     'SiO2'     
<Parameter: 'SiO2 - sld'  , value=3.47  (fixed) , bounds=[-inf, inf]>
<Parameter: 'SiO2 - isld' , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter:'SiO2 - rough' , value=3          , bounds=[1.0, 15.0]>
<Parameter:'SiO2 - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>

ReflectModel, Objective, Curvefitter

model = ReflectModel(structure, bkg=3e-6, dq=5.0)
model.scale.setp(bounds=(0.6, 1.2), vary=True)
model.bkg.setp(bounds=(1e-9, 9e-6), vary=True)

# setting model.threads=1 is useful for reducing oversubscription during sampling
model.threads = 1

objective = Objective(model, data, transform=Transform("logY"))

fitter = CurveFitter(objective)
fitter.fit("differential_evolution", target="nlpost");

-564.650962906306: : 45it [00:03, 13.58it/s]  

objective.plot()
plt.xlabel("Q")
plt.ylabel("logR")
plt.legend();
_images/3aff1645c5c5fdf0be2897e1db7e7df4bc968a3002434150d7d8a9f799b690dd.jpg

emcee

Now lets do a MCMC sampling of the curvefitting system.

fitter.sample(6000, pool=-1);

100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6000/6000 [02:17<00:00, 43.64it/s]

print(f"{fitter.sampler.get_autocorr_time() = }")

# calculate statistics on the chain using arviz
idata = az.from_emcee(fitter.sampler, var_names=None)

# Calculate the effective sample size
az.ess(idata)

fitter.sampler.get_autocorr_time() = array([76.17004152, 78.20967154, 78.61493472, 80.36848009, 78.32501251,
       75.96081422, 74.97741882, 77.75627167])

<xarray.DataTree 'posterior'>
Group: /posterior
    Dimensions:  ()
    Data variables:
        var_0    float64 8B 1.367e+04
        var_1    float64 8B 1.34e+04
        var_2    float64 8B 1.162e+04
        var_3    float64 8B 1.326e+04
        var_4    float64 8B 1.182e+04
        var_5    float64 8B 1.392e+04
        var_6    float64 8B 1.41e+04
        var_7    float64 8B 1.332e+04

Out of a total of 6000 * 200 = 1 200 000 samples, there is an effective sample size of ~ 13200. This is due to autocorrelation within the chain. Consequently the effective sample rate is 13200 / 136 = 97 samples / second.

NOTE At the start of the emcee sampling the walkers have to find their equilibrium positions. This period is called “burn-in” and is normally discarded. The plots below indicate that no burn-in is required. If you need to do any burning in, use the process_chain function (this will be demonstrated a few cells below).

az.plot_trace(idata);
_images/6305e22b87aca9ab97ca866cab6706909e4ecbd6b949e882d97f8a7d20465a67.jpg

In the final output of the sampling each varying parameter is given a set of statistics. Parameter.value is the median of the chain samples. Parameter.stderr is half the [15, 85] percentile, representing a standard deviation.

# thin the chain by the autocorrelation time
process_chain(objective, fitter.chain, nthin=80);

print(objective)

________________________________________________________________________________
Objective - 12901699264
Dataset = c_PLP0011859_q
datapoints = 408
chi2 = 919.58994469468
Weighted = True
Transform = Transform('logY')
________________________________________________________________________________
Parameters:       ''       
________________________________________________________________________________
Parameters: 'instrument parameters'
<Parameter:    'scale'    , value=0.879373 +/- 0.00302, bounds=[0.6, 1.2]>
<Parameter:     'bkg'     , value=4.59384e-07 +/- 2.26e-08, bounds=[1e-09, 9e-06]>
<Parameter:'dq - resolution', value=5  (fixed) , bounds=[-inf, inf]>
<Parameter:  'q_offset'   , value=0  (fixed) , bounds=[-inf, inf]>
________________________________________________________________________________
Parameters: 'Structure - ' 
________________________________________________________________________________
Parameters:      'Si'      
<Parameter: 'Si - thick'  , value=0  (fixed) , bounds=[-inf, inf]>
________________________________________________________________________________
Parameters:      'Si'      
<Parameter:  'Si - sld'   , value=2.07  (fixed) , bounds=[-inf, inf]>
<Parameter:  'Si - isld'  , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter: 'Si - rough'  , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter:'Si - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>
________________________________________________________________________________
Parameters:     'SiO2'     
<Parameter:'SiO2 - thick' , value=38.6452 +/- 0.367, bounds=[15.0, 50.0]>
________________________________________________________________________________
Parameters:     'SiO2'     
<Parameter: 'SiO2 - sld'  , value=3.47  (fixed) , bounds=[-inf, inf]>
<Parameter: 'SiO2 - isld' , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter:'SiO2 - rough' , value=5.83482 +/- 0.296, bounds=[1.0, 15.0]>
<Parameter:'SiO2 - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>
________________________________________________________________________________
Parameters:     'film'     
<Parameter:'film - thick' , value=259.037 +/- 0.247, bounds=[200.0, 300.0]>
________________________________________________________________________________
Parameters:     'film'     
<Parameter: 'film - sld'  , value=2.40178 +/- 0.0126, bounds=[0.1, 3.0]>
<Parameter: 'film - isld' , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter:'film - rough' , value=8.83255 +/- 0.356, bounds=[1.0, 15.0]>
<Parameter:'film - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>
________________________________________________________________________________
Parameters:      'd2o'     
<Parameter: 'd2o - thick' , value=0  (fixed) , bounds=[-inf, inf]>
________________________________________________________________________________
Parameters:      'd2o'     
<Parameter:  'd2o - sld'  , value=6.36  (fixed) , bounds=[-inf, inf]>
<Parameter: 'd2o - isld'  , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter: 'd2o - rough' , value=3.78923 +/- 0.113, bounds=[1.0, 15.0]>
<Parameter:'d2o - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>

A corner plot shows the covariance between parameters. You need to install the matplotlib and corner packages to create these graphs.

objective.corner();
_images/c400ebffc668ad2254b3c82f827518a7bf6423322d3fbceba58993b086dec21b.jpg

Once we’ve done the sampling we can look at the variation in the model at describing the data. In this example there isn’t much spread.

objective.plot(samples=True);
_images/43d975dd93748de3708c063d9af7226a1103fa2cfd8a3b23e8fafec220264583.jpg

In a similar manner we can look at the spread in SLD profiles consistent with the data.

structure.plot(samples=True)
plt.ylim(2.2, 6);
_images/8ca85152696597b22c2824be1d340eded5151221e623a26367d95b4691b126c9.jpg

Sampling with pymc

pymc is also an excellent Bayesian package. refnx has some features built in to work with pymc models.

from refnx.reflect.extra import to_pymc_model, process_trace

refnx can now use the NUTS sampler, which is high performance.

_model = to_pymc_model(objective, _dist="potential")
with _model:
    idata = pm.sample(mp_ctx=mp.get_context("forkserver"))

NUTS[nutpie]: [p0, p1, p2, p3, p4, p5, p6, p7]
/Users/andrew/.conda/envs/dev3/lib/python3.14/site-packages/pytensor/link/numba/dispatch/basic.py:234: UserWarning: Numba will use object mode to run _LogLikeValueGradOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details.
  warnings.warn(











az.plot_dist(idata, var_names="p", filter_vars="like");








_images/30b0ee0b7401a6b956ce4b00a9059e3200afdd39318dd9c8cecaba8120ff50a4.jpg









az.plot_autocorr(idata, max_lag=1000, var_names="p", filter_vars="like");








_images/edd5a82e2c918e8b6bb7fd5c8974831b25cd312e3c014b46abb1d7badcd5316d.jpg





The autocorrelation time is 0, so we have a total of 4000 independent samples (1000 on each of the 4 chains).
The time taken on my machine was ~45 sec, which includes tuning time. This is an effective sampling rate of ~90 samples/sec.


Note how all the parameters are labelled p0, p1, ..., pn. Each of those parameters correspond to a Parameter in Objective.varying_parameters(). Let’s do some processing to update the objective with the sampling results.






# bring the pymc traces back into the Objective.
process_trace(objective, idata);













print(objective)









________________________________________________________________________________
Objective - 12901699264
Dataset = c_PLP0011859_q
datapoints = 408
chi2 = 940.0816760603088
Weighted = True
Transform = Transform('logY')
________________________________________________________________________________
Parameters:       ''       
________________________________________________________________________________
Parameters: 'instrument parameters'
<Parameter:    'scale'    , value=0.875408 +/- 0.00301, bounds=[0.6, 1.2]>
<Parameter:     'bkg'     , value=4.07962e-07 +/- 1.89e-08, bounds=[1e-09, 9e-06]>
<Parameter:'dq - resolution', value=5  (fixed) , bounds=[-inf, inf]>
<Parameter:  'q_offset'   , value=0  (fixed) , bounds=[-inf, inf]>
________________________________________________________________________________
Parameters: 'Structure - ' 
________________________________________________________________________________
Parameters:      'Si'      
<Parameter: 'Si - thick'  , value=0  (fixed) , bounds=[-inf, inf]>
________________________________________________________________________________
Parameters:      'Si'      
<Parameter:  'Si - sld'   , value=2.07  (fixed) , bounds=[-inf, inf]>
<Parameter:  'Si - isld'  , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter: 'Si - rough'  , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter:'Si - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>
________________________________________________________________________________
Parameters:     'SiO2'     
<Parameter:'SiO2 - thick' , value=39.0467 +/- 0.346, bounds=[15.0, 50.0]>
________________________________________________________________________________
Parameters:     'SiO2'     
<Parameter: 'SiO2 - sld'  , value=3.47  (fixed) , bounds=[-inf, inf]>
<Parameter: 'SiO2 - isld' , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter:'SiO2 - rough' , value=5.63212 +/- 0.291, bounds=[1.0, 15.0]>
<Parameter:'SiO2 - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>
________________________________________________________________________________
Parameters:     'film'     
<Parameter:'film - thick' , value=258.876 +/- 0.234, bounds=[200.0, 300.0]>
________________________________________________________________________________
Parameters:     'film'     
<Parameter: 'film - sld'  , value=2.41802 +/- 0.0122, bounds=[0.1, 3.0]>
<Parameter: 'film - isld' , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter:'film - rough' , value=8.36809 +/- 0.347, bounds=[1.0, 15.0]>
<Parameter:'film - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>
________________________________________________________________________________
Parameters:      'd2o'     
<Parameter: 'd2o - thick' , value=0  (fixed) , bounds=[-inf, inf]>
________________________________________________________________________________
Parameters:      'd2o'     
<Parameter:  'd2o - sld'  , value=6.36  (fixed) , bounds=[-inf, inf]>
<Parameter: 'd2o - isld'  , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter: 'd2o - rough' , value=3.72803 +/- 0.111, bounds=[1.0, 15.0]>
<Parameter:'d2o - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>









The parameter values from pymc and emcee are effectively the same.






Sampling with dynesty






import dynesty









You can find out how to estimate posteriors with dynesty using this page, https://dynesty.readthedocs.io/en/stable/dynamic.html. It’s best to use the dynamic nested sampler, then you need to reweight the samples with the weights. Dynesty provides a utility function for that.
You can also use dynesty to perform model comparison by looking at the evidence term.






nested_sampler = dynesty.DynamicNestedSampler(
    objective.logl,
    objective.prior_transform,
    ndim=len(objective.varying_parameters()),
)
nested_sampler.run_nested()
# process the samples
chain = nested_sampler.results.samples_equal()

# another way of processing the samples (reweighting is needed)
logZdynesty = nested_sampler.results.logz[-1]  # value of logZ
weights = np.exp(nested_sampler.results.logwt - logZdynesty)
chain = dynesty.utils.resample_equal(nested_sampler.results.samples, weights)

print(chain.shape)









27096it [01:19, 339.22it/s, batch: 5 | bound: 14 | nc: 1 | ncall: 118088 | eff(%): 22.812 | loglstar: 563.368 < 570.812 < 569.117 | logz: 536.702 +/-  0.188 | stop:  0.861]              





(27096, 8)









The size of the chain resulting from the samples_equal method is not equal to the number of effective samples. One can estimate the effective number of posterior samples resulting from a run using the following:






def ess(weights):
    """
    Estimate the effective sample size from the weights.

    Args:
        weights (array_like): an array of weights values for each nested sample

    Returns:
        int: the effective sample size
    """

    N = len(weights)
    w = weights / weights.sum()
    ess = N / (1.0 + ((N * w - 1) ** 2).sum() / N)

    return int(ess)


print("effective number of samples: ", ess(np.exp(weights)))









effective number of samples:  27095









The effective sampling rate for dynesty seems to be ~26048/61 ~ 427 samples/sec.


Let’s process the chain to put the statistics into the objective. Here we’ll use the process_chain utility function that’s designed for use with emcee chains. This function assumes that the chain has shape (nsteps, nwalkers, nvars). The chain from dynesty has shape (nsamples, nvars), so we can fake the dynesty chain into looking like an emcee chain by putting an extra axis in.






process_chain(objective, chain[:, None, :]);













print(objective)









________________________________________________________________________________
Objective - 12901699264
Dataset = c_PLP0011859_q
datapoints = 408
chi2 = 919.5885800417523
Weighted = True
Transform = Transform('logY')
________________________________________________________________________________
Parameters:       ''       
________________________________________________________________________________
Parameters: 'instrument parameters'
<Parameter:    'scale'    , value=0.879373 +/- 0.00302, bounds=[0.6, 1.2]>
<Parameter:     'bkg'     , value=4.5947e-07 +/- 2.22e-08, bounds=[1e-09, 9e-06]>
<Parameter:'dq - resolution', value=5  (fixed) , bounds=[-inf, inf]>
<Parameter:  'q_offset'   , value=0  (fixed) , bounds=[-inf, inf]>
________________________________________________________________________________
Parameters: 'Structure - ' 
________________________________________________________________________________
Parameters:      'Si'      
<Parameter: 'Si - thick'  , value=0  (fixed) , bounds=[-inf, inf]>
________________________________________________________________________________
Parameters:      'Si'      
<Parameter:  'Si - sld'   , value=2.07  (fixed) , bounds=[-inf, inf]>
<Parameter:  'Si - isld'  , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter: 'Si - rough'  , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter:'Si - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>
________________________________________________________________________________
Parameters:     'SiO2'     
<Parameter:'SiO2 - thick' , value=38.6418 +/- 0.37 , bounds=[15.0, 50.0]>
________________________________________________________________________________
Parameters:     'SiO2'     
<Parameter: 'SiO2 - sld'  , value=3.47  (fixed) , bounds=[-inf, inf]>
<Parameter: 'SiO2 - isld' , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter:'SiO2 - rough' , value=5.83723 +/- 0.301, bounds=[1.0, 15.0]>
<Parameter:'SiO2 - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>
________________________________________________________________________________
Parameters:     'film'     
<Parameter:'film - thick' , value=259.039 +/- 0.246, bounds=[200.0, 300.0]>
________________________________________________________________________________
Parameters:     'film'     
<Parameter: 'film - sld'  , value=2.40194 +/- 0.0125, bounds=[0.1, 3.0]>
<Parameter: 'film - isld' , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter:'film - rough' , value=8.83133 +/- 0.357, bounds=[1.0, 15.0]>
<Parameter:'film - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>
________________________________________________________________________________
Parameters:      'd2o'     
<Parameter: 'd2o - thick' , value=0  (fixed) , bounds=[-inf, inf]>
________________________________________________________________________________
Parameters:      'd2o'     
<Parameter:  'd2o - sld'  , value=6.36  (fixed) , bounds=[-inf, inf]>
<Parameter: 'd2o - isld'  , value=0  (fixed) , bounds=[-inf, inf]>
<Parameter: 'd2o - rough' , value=3.78977 +/- 0.115, bounds=[1.0, 15.0]>
<Parameter:'d2o - volfrac solvent', value=0  (fixed) , bounds=[0.0, 1.0]>













Conclusions


Hopefully you’ve found it useful to see how the three sampling packages can be used to obtain posterior distributions for the parameter set. All have high performance.