Runs quartz sample analysis with refl1d+dream

gistfile1.txt

#!/usr/bin/env python
"""
DREAM-MCMC analysis of a quartz sample, using refl1d + bumps.

This is a refl1d/bumps re-implementation of the analysis in
`jax_quartz_example.ipynb` (which used refnx + pymc/NUTS). The model
(a quartz slab on Si, measured from air) and the data are the same;
only the modelling package (refl1d) and the MCMC sampler (bumps' DREAM)
differ.

Run directly to reproduce the whole analysis (DE fit, then DREAM,
then plots)::

    python refl1d_quartz_example.py

This file also defines a module level ``problem``, so it can alternatively
be driven from the bumps command line, e.g.::

    bumps refl1d_quartz_example.py --fit=dream --burn=1000 --samples=20000 --store=T1
"""

import time

import numpy as np
import matplotlib.pyplot as plt
import arviz as az

from refl1d.names import SLD, QProbe, Experiment, FitProblem, Parameter

from bumps.fitters import fit
from bumps.dream.views import plot_traces, plot_all
from bumps.dream.stats import var_stats, format_vars

DATA_FILE = "../refnx/reflect/tests/.Quartz_data.txt"

# ---------------------------------------------------------------------------
# Model construction
#
# refl1d builds stacks substrate-first, fronting-medium-last, i.e. the
# reverse order of refnx's `air | quartz(...) | si(...)`. A layer's
# `interface` describes the roughness of the interface between that layer
# and the *next* one in the list (towards the fronting medium), which is
# the same physical convention as refnx's `rough`.
# ---------------------------------------------------------------------------
air = SLD(name="air", rho=0)
quartz = SLD(name="quartz", rho=5)
si = SLD(name="Si", rho=2.07)

quartz.rho.range(0, 5.0)
# si.rho is left fixed at 2.07, matching the notebook (si.real.setp(vary=True, ...)
# is commented out there).

sample = si(0, 5.0) | quartz(1500, 5.0) | air

sample["quartz"].thickness.range(1400.0, 1500.0)
sample["quartz"].interface.range(2.0, 20.0)  # roughness between quartz and air
sample["Si"].interface.range(2.0, 20.0)  # roughness between Si and quartz

bkg = Parameter(1e-7, name="bkg")
bkg.range(1e-20, 1)
scale = Parameter(1.0, name="scale")
scale.range(0.9, 1.5)

# ---------------------------------------------------------------------------
# Data / probe
#
# The data file stores Q, R, dR and the 1-sigma Q resolution. refl1d's
# QProbe expects dQ as a 1-sigma value directly (unlike refnx, which wants
# the resolution as a FWHM), so no 2.3548 conversion is needed here.
# ---------------------------------------------------------------------------
data = np.loadtxt(DATA_FILE, delimiter=",")
data = data[:, 1:]
Q, R, dR, dQ = data.T

probe = QProbe(Q, dQ / 2.3548, R=R, dR=dR, intensity=scale, background=bkg)

experiment = Experiment(probe=probe, sample=sample)
problem = FitProblem(experiment)


def main():
    print(problem.summarize())

    # 1) Differential evolution to find a good starting point, same role as
    #    `fitter.fit("differential_evolution")` in the refnx notebook.
    # Note: verbose=True is avoided here because this bumps version's
    # show_table() unconditionally calls `results.state.draw()`, which only
    # exists for the DREAM fitter's state (DE's state is a plain history dict).
    de_result = fit(problem, method="de", steps=1000, pop=10, verbose=False)
    print(f"DE chisq = {problem.chisq():.4g}")

    # 2) MCMC sampling with bumps' DREAM sampler, starting near the DE
    #    optimum (fit() leaves `problem` set to the DE best-fit point).
    dream_start = time.perf_counter()
    dream_result = fit(
        problem,
        method="dream",
        samples=20000,
        burn=0,
        thin=1,
        verbose=True,
    )
    dream_elapsed = time.perf_counter() - dream_start
    print(f"DREAM sampling took {dream_elapsed:.1f} s")
    state = dream_result.state

    # Parameter summary (mean/median/68%/95% credible intervals), the DREAM
    # equivalent of `az.summary(idata, var_names=["p"], filter_vars="like")`.
    draw = state.draw()
    print(format_vars(var_stats(draw)))

    # Effective sample size per parameter, computed with arviz on the
    # per-chain traces (generation, chain, var) kept by DREAM's population.
    _, chains = state.traces(portion=state.portion)
    ess = az.ess(chains, chain_axis=1, draw_axis=0)
    for label, n_eff in zip(state.labels, ess):
        print(f"ess[{label}] = {n_eff:.0f}")

    # Trace plots, equivalent to `az.plot_trace(idata, ...)`.
    # plt.figure()
    # plot_traces(state)
    #
    # # Corner-style histogram/correlation grid, equivalent to `objective.corner()`.
    # plt.figure()
    # plot_all(state)
    #
    # # Push the posterior mean back into the model and plot data vs. fit,
    # # equivalent to `process_trace(objective, idata)` + `objective.plot()`.
    # problem.setp(draw.points.mean(axis=0))
    # plt.figure()
    # experiment.plot()

    print(problem.summarize())
    # plt.show()


if __name__ == "__main__":
    main()