Bayesian analysis of QENS data#
This notebook assumes that you have already completed the previous notebook.
We will here replicate some of the analysis and show how to get more information about the correlations between the parameters using Bayesian analysis. We have remove most of the explanations from the previous notebook.
import os
import numpy as np
import scipp as sc
import easydynamics as edyn
# Quiz questions for this notebook
from jupyterquiz import display_quiz
from quizlib import bayesian_qens as quiz
# Make the plots interactive
%matplotlib widget
Load and prepare the data#
def prepare_data(experiment):
"""Give bins with too small variance a variance of 1, so that their weight in the fit stays finite."""
data = experiment.data
indices = data.variances <= 0.5
data.variances[indices] = 1.0
experiment.data = data
filename = "../4-reduction/energy_QENS_elastic.h5"
filename2 = "../4-reduction/energy_QENS_sample.h5"
# If you have not completed the reduction exercise, uncomment the following lines to fetch the pre-prepared data
# import utils
# filename = utils.fetch_data('4-reduction/energy_QENS_elastic.h5'
# filename2 = utils.fetch_data('4-reduction/energy_QENS_sample.h5')
elastic_experiment = edyn.Experiment(display_name="Elastic")
elastic_experiment.load_hdf5(filename=filename)
prepare_data(elastic_experiment)
qe_experiment = edyn.Experiment(display_name="QuasiElastic")
qe_experiment.load_hdf5(filename=filename2)
prepare_data(qe_experiment)
Step 1: Determine the resolution from the elastic sample#
resolution_sample_model = edyn.SampleModel(y_unit="counts")
res_gauss = edyn.Gaussian(width=0.002, area=1, name="Res. Gauss", y_unit="counts")
resolution_sample_model.append_component(res_gauss)
background_model_res = edyn.BackgroundModel(y_unit="counts")
poly_res = edyn.Polynomial(coefficients=[0.0], name="Background", y_unit="counts")
poly_res.coefficients[0].min = 0.0
poly_res.coefficients[0].fixed = True
background_model_res.append_component(poly_res)
instrument_model_res = edyn.InstrumentModel(
energy_offset=1e-3,
background_model=background_model_res,
)
elastic_analysis = edyn.Analysis(
display_name="Elastic / Resolution",
experiment=elastic_experiment,
sample_model=resolution_sample_model,
instrument_model=instrument_model_res,
)
elastic_analysis.fit(fit_method="independent")
elastic_analysis.plot_data_and_model()
Now it’s time to investigate the parameters with Bayesian analysis.
Bounds are the prior#
In DREAM, each parameter’s min and max define a uniform prior, so every free parameter must have finite bounds before sampling. Most parameters start with at least one infinite bound, so bayesian.sample() would refuse to run.
bayesian.suggest_bounds() proposes bounds from the fitted values and uncertainties. It is advisory: it changes nothing until you call .apply(), and it only ever fills in an infinite bound, so physical limits you have already set (an area that cannot go below zero, say) are left alone.
We will at first work with just a single Q index. We select an arbitrary one:
analysis_Q4 = elastic_analysis.analysis_list[4]
suggestions = analysis_Q4.bayesian.suggest_bounds()
print(suggestions)
BoundsSuggestions
parameter current suggested
----------------------------------------------------------------------
Res. Gauss area (0, inf) (0, 8.074)
Res. Gauss width (1e-10, inf) (1e-10, 0.002085)
energy_offset (-inf, inf) (0.0008946, 0.002423)
The defaults are deliberately generous: 10 standard deviations plus 20% of the value. Because the bounds are a uniform prior, being too narrow is the dangerous mistake: it truncates the posterior and makes the uncertainty look smaller than it is. The 20% term is there for parameters whose fitted uncertainty comes back as zero. All three settings (n_sigma, relative_pad, absolute_floor) can be adjusted, and you can always set min and max by hand.
It is worth reading the table before applying it. A suggestion many orders of magnitude larger than the parameter itself is a useful warning sign: it means the fit returned a huge uncertainty, which usually happens because two parameters are degenerate, i.e. the data determines only some combination of them, so one can grow while the other shrinks with no effect on the fit. That is a problem to fix in the model, not with the sampler.
Let us apply the bounds:
changed = suggestions.apply()
print(f"Applied bounds to: {[parameter.name for parameter in changed]}")
Applied bounds to: ['Res. Gauss area', 'Res. Gauss width', 'energy_offset']
display_quiz(quiz.q1)
Sample the posterior#
bayesian.sample() runs the chains. The three numbers that matter are:
samples: how many draws to collect in total. More is better, at linear cost.burn: generations discarded at the start, while the chains are still travelling towards the bulk of the posterior.thin: keep only every n-th generation, which reduces the correlation between neighbouring draws.
Sampling never moves your parameters: their values are restored afterwards, so the model is left exactly as the fit left it.
results = analysis_Q4.bayesian.sample(samples=8000, burn=300, thin=2)
print(
f"Collected {results.draws.shape[0]} draws for {results.draws.shape[1]} parameters."
)
Collected 3750 draws for 3 parameters.
Did the chains converge?#
Always look at the traces before trusting the numbers. A converged chain looks like a “hairy caterpillar”: noisy, but flat and stationary. A visible drift or slow wander means the chain has not settled and needs a longer burn-in or more samples.
fig = analysis_Q4.bayesian.plot_trace()
display_quiz(quiz.q2)
Summarize the posterior#
bayesian.summary() reports the median and the 68% credible interval of each parameter, under the parameter’s own name and unit. The interval is asymmetric in general, which is precisely the information a single symmetric error bar throws away.
analysis_Q4.bayesian.summary()
PosteriorSummary
parameter unit median - + current
-----------------------------------------------------------------------------------
Res. Gauss width meV 0.0015239 7.681e-06 8.648e-06 0.0015249
Res. Gauss area meV*counts 5.4226 0.05634 0.05601 5.4255
energy_offset meV 0.0016589 1.504e-05 1.423e-05 0.0016589
display_quiz(quiz.q3)
Correlations between parameters#
The corner plot is the part least available from a least-squares fit. The diagonal shows each parameter’s own distribution; each off-diagonal panel shows a pair. A round blob means the two are independent, while a tilted, narrow ridge means they are correlated and the data constrains only a combination of them.
display_quiz(quiz.q4)
fig = analysis_Q4.bayesian.plot_corner()
Finally, we can show the data and the model with a 68% confidence interval for the model:
fig = analysis_Q4.bayesian.plot_posterior_predictive(n_draws=100)
display_quiz(quiz.q5)
We can also show the correlations as a matrix, where each entry is color-coded to indicate how correlated the parameters are:
fig = analysis_Q4.bayesian.plot_correlations()
Optional: do the analysis for all Q values#
We can also do this analysis for all Q values. Note that it will take 8 times as long, since we have 8 Q values to analyse. Feel free to skip this step if you find that it takes too long.
suggestions = elastic_analysis.bayesian.suggest_bounds()
changed = suggestions.apply()
results = elastic_analysis.bayesian.sample(samples=8000, burn=300, thin=2)
elastic_analysis.bayesian.summary()
elastic_analysis.bayesian.plot_corner()
elastic_analysis.bayesian.plot_posterior_predictive(n_draws=100)
Normalise to the resolution area#
We will again normalise the quasi-elastic data to the resolution area, so we divide it out, just as one would normalise to a vanadium measurement in a real experiment. sc.values(...) drops the uncertainties of the fitted areas, so that they act as a fixed normalisation.
⚠️ Run the next cell only once. Running it a second time would divide the data again.
norm = elastic_analysis.parameters_to_dataset()["Res. Gauss area"].copy()
norm.unit = "dimensionless" # to keep the units of the data consistent
qe_experiment.data = qe_experiment.data / sc.values(norm)
elastic_experiment.data = elastic_experiment.data / sc.values(norm)
Now we set up the model for the quasielastic sample again.
sample_model_qe = edyn.SampleModel(y_unit="counts")
delta_function = edyn.DeltaFunction(name="DeltaFunction", area=0.15, y_unit="counts")
lorentzian = edyn.Lorentzian(name="Lorentzian", area=1.0, width=0.015, y_unit="counts")
sample_model_qe.append_component(delta_function)
sample_model_qe.append_component(lorentzian)
background_model_qe = edyn.BackgroundModel(y_unit="counts")
poly_qe = edyn.Polynomial(
coefficients=[0.0], name="Background", y_unit="counts", suppress_warnings=True
)
poly_qe.coefficients[0].min = 0.0
poly_qe.coefficients[0].fixed = True
background_model_qe.append_component(poly_qe)
instrument_model_qe = edyn.InstrumentModel(
background_model=background_model_qe,
resolution_model=elastic_analysis.sample_model,
energy_offset=1e-3,
)
qe_analysis = edyn.Analysis(
display_name="Quasi-elastic per-Q",
experiment=qe_experiment,
sample_model=sample_model_qe,
instrument_model=instrument_model_qe,
)
Before fitting, it is a good idea to check the start guesses by plotting the data together with the model. If the model does not at all resemble the data, then go back up and change the default values for the SampleModel components and try again. This process may take a few steps to get right.
qe_analysis.plot_data_and_model()
Adjust the start guesses until they look reasonable, then fit every Q independently and plot the result like we did above.
qe_analysis.fit(fit_method="independent")
qe_analysis.plot_data_and_model()
First, inspect the fit. Does it look good at all Q? If not, go back and adjust the start guesses and try again. If yes, please continue.
display_quiz(quiz.q6)
Exercise
Repeat the Bayesian analysis for Q number 1. Start with selecting the right analysis object: qe_analysis_Q1=qe_analysis.analysis_list[1]
Solution:
Optional: do this for all Q#
Again, if you have the time, you can also do this analysis for all Q by outcommenting the code below.
qe_suggestions = qe_analysis.bayesian.suggest_bounds()
changed = qe_suggestions.apply()
results = qe_analysis.bayesian.sample(samples=8000, burn=300, thin=2)
qe_analysis.bayesian.summary()
qe_analysis.bayesian.plot_corner()
Exercise
Create a new SampleModel and Analysis and recreate the analysis we just did for the quasielastic sample, but this time remove all constraints for the coefficients of the BackgroundModel. How do the parameter correlations change?
Solution
Copy/paste the code from above and outcomment or delete the lines that fix and set bounds for the coefficients.
display_quiz(quiz.q7)
Step 3: Fit a jump-diffusion model to the widths#
We will now fit the jump-diffusion model to the widths and carry out Bayesian analysis. We first fix any missing variances (there shouldn’t be any, but better safe than sorry.)
def fix_missing_variances(analysis):
"""Fix missing variances on fitted parameters in an analysis."""
for param in analysis.get_all_parameters():
if param.variance == 0.0:
param.variance = 0.00001 * abs(param.value)
fix_missing_variances(qe_analysis)
jump_diffusion_model = edyn.JumpTranslationalDiffusion(
name="Jump Translational Diffusion",
diffusion_coefficient=4.6e-10, # m^2/s
relaxation_time=22.0, # ps
scale=0.5,
y_unit="counts",
)
binding = edyn.FitBinding(
model=jump_diffusion_model,
targets={
"area": "Lorentzian area",
"width": "Lorentzian width",
},
)
parameter_analysis = edyn.ParameterAnalysis(
parameters=qe_analysis,
bindings=binding,
)
We fit and plot the results
parameter_analysis.fit()
parameter_analysis.plot(names=["Lorentzian width"], xmin=0, xmax=2.1, vmin=0, vmax=0.02)
parameter_analysis.plot(names=["Lorentzian area"], xmin=0, xmax=2.1, vmin=0, vmax=1.2)
We can read off the fitted jump-diffusion parameters, with uncertainties.
parameter_analysis.get_all_parameters()
[<Parameter 'diffusion_coefficient': 4.288e-10 ± 3.148e-11 m^2/s, bounds=[0.0:inf]>,
<Parameter 'scale': 0.8618 ± 0.0132 meV*counts, bounds=[0.0:inf]>,
<Parameter 'relaxation_time': 29.5564 ± 1.4116 ps, bounds=[0.0:inf]>]
Now it’s time for the Bayesian analysis. First a question
display_quiz(quiz.q8)
Exercise
Carry out the Bayesian analysis on the ParameterAnalysis.
Solution:
Conclusions#
This is the end of this analysis. We have shown how to use Bayesian analysis to investigate and understand your fit, the parameters and their correlations. We can use it to justify some of the choices made in the previous notebook. For example, you were told to not fit the background, because the parameter is essentially zero. Now we can see what happens when you do fit it: it correlates with the width of the Lorentzian and makes it even hard to determine correctly. We can also see that the diffusion constant and the relaxation rate are strongly correlated. This could help us provide more realistic bounds on the fitted results.
EasyDynamics is work in progress, and we would very much like to hear your feedback - positive and negative. If you would like to fit your own data with EasyDynamics then please reach out. We would be delighted to help you get started.
Please write to henrik.jacobsen@ess.eu.
Optional: do Bayesian analysis on the global fit#
The global fit obviously has a lot more parameters than the individual fits. The Bayesian analysis will therefore be much more time consuming. Feel free to run the analysis lines below during a break or similar, but be warned - it might take a couple of hours. We are looking into ways to speed this up.
delta_function_diff = edyn.DeltaFunction(name='DeltaFunction', area=0.15, y_unit='counts')
diffusion_model = edyn.JumpTranslationalDiffusion(
name='Jump Translational Diffusion',
diffusion_coefficient=2.5e-10,
relaxation_time=18.0, # ps
scale=1.0,
y_unit='counts',
)
sample_model_diff = edyn.SampleModel(
components=delta_function_diff,
diffusion_models=diffusion_model,
y_unit='counts',
)
poly_diff = edyn.Polynomial(coefficients=[0.0], name='Background', y_unit='counts')
poly_diff.coefficients[0].min = 0.0
poly_diff.coefficients[0].fixed = True
background_model_diff = edyn.BackgroundModel(y_unit='counts')
background_model_diff.append_component(poly_diff)
instrument_model_diff = edyn.InstrumentModel(
background_model=background_model_diff,
resolution_model=elastic_analysis.sample_model,
energy_offset=1e-3,
)
diffusion_analysis = edyn.Analysis(
display_name='Jump Diffusion Full Analysis',
experiment=qe_experiment,
sample_model=sample_model_diff,
instrument_model=instrument_model_diff,
)
diffusion_analysis.fit(fit_method='simultaneous')
diffusion_analysis.plot_data_and_model(plot_residuals=True, autoscale=False)
diff_suggestions = diffusion_analysis.bayesian.suggest_bounds()
print(diff_suggestions)
diff_changed = diff_suggestions.apply()
print(f'Applied bounds to: {[parameter.name for parameter in diff_changed]}')
diff_results = diffusion_analysis.bayesian.sample(samples=1000, burn=200, thin=2,fit_method='simultaneous')
print(f'Collected {diff_results.draws.shape[0]} draws for {diff_results.draws.shape[1]} parameters.')
fig=diffusion_analysis.bayesian.plot_trace()
print(diffusion_analysis.bayesian.summary())
fig=diffusion_analysis.bayesian.plot_corner()