Calibration of a line#

Simple demo demonstrating the workflow of rxmc.

from collections import OrderedDict

import corner
import numpy as np
from matplotlib import pyplot as plt
from scipy import stats

import rxmc
Using database version X4-2024-12-31 located in: /home/kyle/db/exfor/unpack_exfor-2024/X4-2024-12-31
rng = np.random.default_rng(49)

define parameters#

help(rxmc.params.Parameter)
Help on class Parameter in module rxmc.params:

class Parameter(builtins.object)
 |  Parameter(name, dtype=<class 'float'>, unit='', latex_name=None, bounds=(-inf, inf))
 |
 |  A single scalar model parameter.
 |
 |  Parameters
 |  ----------
 |  name : str
 |      Human-readable name of the parameter.
 |  dtype : type, optional
 |      Data type of the parameter value.  Defaults to ``float``.
 |  unit : str, optional
 |      Physical unit string (e.g. ``"MeV"``).  Defaults to ``""``.
 |  latex_name : str, optional
 |      LaTeX representation used in plots and documentation.  Defaults to
 |      ``name`` when not supplied.
 |  bounds : tuple of float, optional
 |      ``(lower, upper)`` bounds for the parameter.  Defaults to
 |      ``(-np.inf, np.inf)``.
 |
 |  Methods defined here:
 |
 |  __eq__(self, other)
 |      Return self==value.
 |
 |  __init__(self, name, dtype=<class 'float'>, unit='', latex_name=None, bounds=(-inf, inf))
 |      Initialize self.  See help(type(self)) for accurate signature.
 |
 |  ----------------------------------------------------------------------
 |  Data descriptors defined here:
 |
 |  __dict__
 |      dictionary for instance variables
 |
 |  __weakref__
 |      list of weak references to the object
 |
 |  ----------------------------------------------------------------------
 |  Data and other attributes defined here:
 |
 |  __hash__ = None

make the model#

help(rxmc.physical_model.PhysicalModel)
Help on class PhysicalModel in module rxmc.physical_model:

class PhysicalModel(builtins.object)
 |  PhysicalModel(params: list[rxmc.params.Parameter], transform=None)
 |
 |  Abstract base class for parametric physical models.
 |
 |  Represents an arbitrary parametric model
 |  $y_{\mathrm{model}}(x;\,\alpha)$ for comparison to an experimental
 |  measurement $\{x_i,\, y(x_i)\}$ encapsulated in an
 |  :class:`~rxmc.observation.Observation`.
 |
 |  Subclasses implement :meth:`evaluate` in physical space.  An optional
 |  *parametric* ``transform`` (see :mod:`rxmc.transforms`) is applied on top by
 |  :meth:`__call__`; its parameters are appended to :attr:`params` so they flow
 |  through the ordinary model-parameter machinery (priors, ``split_parameters``).
 |  Typical uses are a latent normalisation :func:`rxmc.transforms.scale` or one
 |  per dataset via :func:`rxmc.transforms.per_observation_scaling`.  Comparison-
 |  space transforms (e.g. comparing in log space) are *not* the model's
 |  business: declare them on the :class:`~rxmc.observation.Observation`.
 |
 |  Parameters
 |  ----------
 |  params : list of Parameter
 |      Physical parameters of the model.  Each entry should carry a name
 |      and a data type.
 |  transform : Transform or callable, optional
 |      Model-side transform ``y -> transform(y, *values)`` applied after
 |      :meth:`evaluate`.  Its parameters (if any) are appended to ``params``.
 |
 |  Methods defined here:
 |
 |  __call__(self, observation: rxmc.observation.Observation, *params) -> numpy.ndarray
 |      Physical-space :meth:`evaluate` followed by the model transform.
 |
 |  __init__(self, params: list[rxmc.params.Parameter], transform=None)
 |      Initialize self.  See help(type(self)) for accurate signature.
 |
 |  apply_transform(self, observation, y, transform_values=())
 |      Apply the model-side transform to a physical-space prediction.
 |
 |  evaluate(self, observation: rxmc.observation.Observation, *params) -> numpy.ndarray
 |      Evaluate the model at the given parameter values.
 |
 |      Must be overridden by subclasses.
 |
 |      Parameters
 |      ----------
 |      observation : Observation
 |          Observation containing the independent-variable grid.
 |      *params : float
 |          Model parameter values.
 |
 |      Returns
 |      -------
 |      np.ndarray
 |          Predicted observable values on the observation grid.
 |
 |      Raises
 |      ------
 |      NotImplementedError
 |          Always — subclasses must implement this method.
 |
 |  split_params(self, params)
 |      Split a full parameter tuple into ``(base_params, transform_values)``.
 |
 |  ----------------------------------------------------------------------
 |  Data descriptors defined here:
 |
 |  __dict__
 |      dictionary for instance variables
 |
 |  __weakref__
 |      list of weak references to the object

Clearly to make a model, we need to understand these things called Observations.#

This is an important detail, the whole point of rxmc is to compare predictions of a PhysicalModel to data contained in an Observation, to calibrate the parameters of the PhysicalModel.

help(rxmc.observation.Observation)
Help on class Observation in module rxmc.observation:

class Observation(builtins.object)
 |  Observation(x: numpy.ndarray, y: numpy.ndarray, y_stat_err=None, y_sys_err_normalization=None, y_sys_err_offset=None, label=None, transform=None, mask=None)
 |
 |  Experimental data: ``x``, ``y``, and the statistical error on ``y``.
 |
 |  Parameters
 |  ----------
 |  x : np.ndarray
 |      Independent-variable data.
 |  y : np.ndarray
 |      Dependent-variable data, same shape as ``x``.
 |  y_stat_err : np.ndarray, optional
 |      Statistical (uncorrelated) error on ``y``.  Defaults to zeros.
 |  y_sys_err_normalization : float or np.ndarray, optional
 |      Reported *fractional* (dimensionless) normalisation uncertainty —
 |      inert metadata; see :meth:`systematic_terms`.
 |  y_sys_err_offset : float or np.ndarray, optional
 |      Reported *absolute* offset uncertainty, in the same units as ``y`` —
 |      inert metadata; see :meth:`systematic_terms`.
 |  label : str, optional
 |      Human-readable dataset identifier used in error messages.
 |  transform : Transform or callable, optional
 |      Parameter-free *comparison-space* transform (see :mod:`rxmc.transforms`).
 |      Pass **raw** ``y``: the observation stores ``y = transform(y_raw)`` and
 |      propagates ``y_stat_err`` by the delta method, and the
 |      :class:`~rxmc.constraint.Constraint` applies the same transform to the
 |      model prediction — so ``transform=rxmc.transforms.log`` compares in log
 |      space with the model written once, in physical space.
 |  mask : array_like of bool, optional
 |      Which points are *active* in a likelihood (default all).  Inactive
 |      points stay in the block (supports/terms are authored over all points)
 |      but are excluded from the residual; use :meth:`masked` /
 |      :meth:`masked_where` to derive fit/held-out views.
 |
 |  Attributes
 |  ----------
 |  x, y : np.ndarray
 |      The data, ``y`` in comparison space.
 |  y_raw, y_stat_err_raw : np.ndarray
 |      ``y`` and its statistical error as given (physical space).
 |  y_stat_err : np.ndarray
 |      Statistical error on ``y`` in comparison space (raw, not squared).
 |  transform : Transform
 |      The comparison-space transform (identity by default).
 |  mask : np.ndarray of bool
 |      Active points.
 |  identity : Observation
 |      The root observation this one is a view of.  Views made by
 |      :meth:`masked` share it, so anything routing by observation (e.g.
 |      :func:`rxmc.transforms.per_observation_scaling`) treats a masked view
 |      and its root as the same dataset.
 |  y_sys_err_normalization : float or np.ndarray or None
 |      Fractional normalisation uncertainty (dimensionless).
 |  y_sys_err_offset : float or np.ndarray or None
 |      Absolute offset uncertainty (units of ``y``).
 |  label : str or None
 |      Human-readable dataset identifier.
 |  n_data_pts : int
 |      Number of data points.
 |
 |  Methods defined here:
 |
 |  __init__(self, x: numpy.ndarray, y: numpy.ndarray, y_stat_err=None, y_sys_err_normalization=None, y_sys_err_offset=None, label=None, transform=None, mask=None)
 |      Initialize self.  See help(type(self)) for accurate signature.
 |
 |  masked(self, mask, label=None)
 |      A shallow copy of this observation with a new point mask.
 |
 |      No data or pre-computed workspaces are rebuilt: the copy shares them and
 |      only changes which points enter a likelihood.
 |
 |  masked_where(self, predicate, label=None)
 |      :meth:`masked` with ``mask = predicate(x)`` (points where it is True).
 |
 |  num_pts_within_interval(self, ylow: numpy.ndarray, yhigh: numpy.ndarray, xlim=None)
 |      Number of active points of ``y`` that fall within ``[ylow, yhigh)``.
 |
 |      Useful for empirical-coverage diagnostics.  ``ylow``/``yhigh`` are in
 |      comparison space and indexed over *all* points of the block.
 |
 |      Parameters
 |      ----------
 |      ylow, yhigh : np.ndarray
 |          Interval bounds, same shape as ``y``.
 |      xlim : tuple, optional
 |          ``(x_min, x_max)`` range to restrict the count.
 |
 |  statistical_term(self, support=None)
 |      The always-on, genuinely uncorrelated statistical diagonal.
 |
 |      Parameters
 |      ----------
 |      support : np.ndarray, optional
 |          Indices of this observation's block in the stacked vector
 |          (``None`` for a single-observation constraint).
 |
 |      Returns
 |      -------
 |      Term
 |          ``diag(y_stat_err**2)`` on ``support`` (comparison space).
 |
 |  systematic_terms(self, support=None) -> list
 |      This dataset's reported correlated systematics as fixed rank-one terms.
 |
 |      Opt-in — **not** added to any covariance automatically.  Pass the result
 |      via ``Constraint(extra_terms=[*obs.systematic_terms(), ...])``.
 |      Zero magnitudes are skipped, so an observation without reported
 |      systematics yields an empty list.  Magnitudes are reported in physical
 |      space and propagated to the comparison space by the delta method
 |      (``|t'| * omega`` for an offset, ``|t'(ym_raw)| * eta * ym_raw`` for a
 |      normalisation).
 |
 |      Parameters
 |      ----------
 |      support : np.ndarray, optional
 |          Indices of this observation's block in the stacked vector
 |          (``None`` for a single-observation constraint).
 |
 |      Returns
 |      -------
 |      list of Term
 |          The absolute offset mode (``outer(omega, omega)``) first, then the
 |          fractional, prediction-scaled normalisation mode
 |          (``eta**2 * outer(ym, ym)``).
 |
 |  ----------------------------------------------------------------------
 |  Readonly properties defined here:
 |
 |  log_jacobian
 |      ``sum(log |t'(y_raw)|)`` over the active points.
 |
 |      The log-Jacobian of the comparison-space transform: a constant in the
 |      parameters, needed only to compare marginal likelihoods (log Z) across
 |      different comparison spaces (``log Z_raw = log Z_transformed +
 |      log_jacobian``).  Zero for the identity.
 |
 |  n_active
 |      Number of active (unmasked) points.
 |
 |  ----------------------------------------------------------------------
 |  Data descriptors defined here:
 |
 |  __dict__
 |      dictionary for instance variables
 |
 |  __weakref__
 |      list of weak references to the object

Ok, so basically it’s just some x and y, ans some information about the errors of y. We can think of Observations like experimental measurements of an observable, and we want to build a model to make uncertainty-quantified predictions of that observable.

class LinearModel(rxmc.physical_model.PhysicalModel):
    def __init__(self):
        params = [
            rxmc.params.Parameter("m", float, "no-units"),
            rxmc.params.Parameter("b", float, "y-units"),
        ]
        super().__init__(params)

    def evaluate(self, observation, m, b):
        return self.y(observation.x, m, b)

    def y(self, x, m, b):
        # useful to have a function hat takes in an array-like x
        # rather than an Observation, e.g. for plotting
        return m * x + b

Well that’s not too complicated.

my_model = LinearModel()

Let’s test it out:

# true params are obviously m = 1, b = 0
observation = rxmc.observation.Observation(x=np.array([1, 2, 3]), y=np.array([1, 2, 3]))
my_model(observation, 1, 0)
array([1., 2., 3.])
my_model.y(observation.x, 2, 1)
array([3, 5, 7])
# now with parameters that are obviously wrong
my_model(observation, 2, 1)
array([3., 5., 7.])
# just to show some ways that may be convenient to pass around params
# note they must be in the same order as in my_model.evaluate,
# which should be in the same order as my_model.params

p = [2, 1]
print(my_model(observation, *p))
p = np.array([2, 1])
print(my_model(observation, *p))
[3. 5. 7.]
[3. 5. 7.]

define a prior#

Let’s imagine this line corresponds to some physics, and we have some backround knowledge to inform us what we expect $m$ and $b$ to be. We can encode that into a prior distribution:

prior_mean = OrderedDict(
    [
        ("m", 1),
        ("b", 1),
    ]
)
prior_std_dev = OrderedDict(
    [
        ("m", 1),
        ("b", 1),
    ]
)
covariance = np.diag(list(prior_std_dev.values())) ** 2
mean = np.array(list(prior_mean.values()))
n_prior_samples = 1000
prior_distribution = stats.multivariate_normal(mean, covariance)
prior_samples = prior_distribution.rvs(size=n_prior_samples, random_state=rng)

Let’s plot some samples of lines from this prior

x = np.linspace(0, 1, 10)

# array to hold the lines
y = np.zeros((n_prior_samples, len(x)))

# propagate the prior through to the observation
for i in range(n_prior_samples):
    sample = prior_samples[i, :]
    y[i, :] = my_model.y(x, *sample)

# grab confidence intervals for plotting
upper, lower = np.percentile(y, [5, 95], axis=0)
for i in np.random.choice(n_prior_samples, 1000):
    plt.plot(x, y[i, :], zorder=1, alpha=0.2)
    # pass

plt.fill_between(
    x, lower, upper, alpha=0.5, zorder=2, label=r"inner 90$^\text{th}$ pctl"
)
plt.xlabel("x")
plt.ylabel("y")
plt.legend()
plt.title("prior")
Text(0.5, 1.0, 'prior')
../_images/4913e47ca66b632836065f9056042c06d63faf51dd2dbf047cec6f8c08d3cc2a.png
fig = corner.corner(prior_samples, labels=[p.name for p in my_model.params])
fig.suptitle("prior")
Text(0.5, 0.98, 'prior')
../_images/77034736fd9d05b1a2df160173542fee64cbd2b40774260724aee2cf1e1fa226.png

Now let’s update our prior by comparing to some data#

This will require learning how rxmc encodes our assumptions about the error on an experimental Observation: as a covariance built from explicit terms, evaluated under a likelihood functional. These are a necessary ingredient for comparing to the predictions of a PhysicalModel.

In our case we will mock experimental data by synthetically generate some data with noise about a “true” $m$ and $b$. Our calibration posterior should converge to be centered about this true point.

Let us assume that the experimentalists made a perfect estimate of the experimental noise in their setup. That is, the error bars they report will correspond exactly to the true distribution from which we sample.

This noise will correspond to statistical noise. Later on we will look at systematic experimental error.

true_params = OrderedDict(
    [
        ("m", 0.6),
        ("b", 2),
    ]
)

x = np.linspace(0, 1, 10)
noise = 0.1
y_exp = my_model.y(x, *list(true_params.values())) + rng.normal(
    scale=noise, size=len(x)
)
y_stat_err = noise * np.ones_like(y_exp)  # noise is just a constant fraction of y
obs1 = rxmc.observation.Observation(x=x, y=y_exp, y_stat_err=y_stat_err)
plt.errorbar(
    x,
    y_exp,
    y_stat_err,
    color="k",
    marker=".",
    linestyle="none",
    label="obs1",
)
plt.plot(x, my_model.y(x, *list(true_params.values())), "k--", label="truth")
plt.xlabel("x")
plt.ylabel("y")

plt.fill_between(
    x,
    lower,
    upper,
    alpha=0.5,
    zorder=2,
    label=r"prior inner 90$^\text{th}$ pctl",
)

plt.legend()
plt.title("experimental constraint")
Text(0.5, 1.0, 'experimental constraint')
../_images/95e3cae7f3a3004eab51d7f6808984d70a56cbec2898b07b27c5c3fcdd12dd26.png

Clearly, our prior is at odds with our observation. We will now determine a posterior distribution of $m$ and $b$ that takes obs1 into account. To do this, we will need to think about a likelihood for obs1 — in rxmc that means a covariance model (here, just the reported statistical errors) and a likelihood functional (the default GaussianLikelihood).

set up the likelihood and Constraint#

We will use the simplest assumption about the error on y: that the experimentalists exactly reported the statistical error, and there is no systematic error at all. This implies that each data point in obs1.y, say obs1.y[i] can be modeled as being an random variate, each sampled independently from normal distributions with mean obs1.y[i] and with standard deviation obs1.y_stat_err[i].

This is exactly the default behavior: a Constraint automatically includes each observation’s statistical diagonal in its covariance, and evaluates it under the default GaussianLikelihood:

help(rxmc.likelihood_model.GaussianLikelihood)
Help on class GaussianLikelihood in module rxmc.likelihood_model:

class GaussianLikelihood(Likelihood)
 |  Multivariate-normal likelihood over the stacked residual.
 |
 |  Parameter-free — all uncertainty lives on the covariance terms.
 |
 |  Method resolution order:
 |      GaussianLikelihood
 |      Likelihood
 |      builtins.object
 |
 |  Methods defined here:
 |
 |  log_likelihood(self, d2, logdet, n, *like_params)
 |
 |  ----------------------------------------------------------------------
 |  Data and other attributes defined here:
 |
 |  __annotations__ = {}
 |
 |  ----------------------------------------------------------------------
 |  Methods inherited from Likelihood:
 |
 |  chi2(self, d2, logdet, n, *like_params)
 |
 |  ----------------------------------------------------------------------
 |  Data descriptors inherited from Likelihood:
 |
 |  __dict__
 |      dictionary for instance variables
 |
 |  __weakref__
 |      list of weak references to the object
 |
 |  ----------------------------------------------------------------------
 |  Data and other attributes inherited from Likelihood:
 |
 |  n_params = 0
 |
 |  params = ()

With a purely statistical (diagonal) covariance, the likelihood is proportional to the familiar form:

\begin{equation} \mathcal{L}(\alpha|y) \propto e^{ - \chi^2(\alpha,y) } \end{equation}

where the Chi-squared is simply

\begin{equation} \chi^2(\alpha,y) = \sum_i \frac{(y(x_i) - y_m(x_i;\alpha) )^2}{\sigma^2_i} \end{equation}

In a later tutorial we will use more complicated Observations that include systematic error, and look at other LikelihoodModels.

likelihood_model = rxmc.likelihood_model.GaussianLikelihood()
constraint = rxmc.constraint.Constraint(
    [obs1],
    my_model,
    likelihood_model,
)

Let’s test this constraint thing out. What is the reduced $\chi^2$ for the prior mean?

constraint.chi2(prior_distribution.mean) / constraint.n_data_pts
64.18631209285807

Here is proof that, in this case, we reduce to the form described above:

y = my_model(obs1, *prior_distribution.mean)
np.sum((y - obs1.y) ** 2 / y_stat_err**2) / constraint.n_data_pts
np.float64(64.18631209285807)

running the calibration#

help(rxmc.walker.Walker)
Help on class Walker in module rxmc.walker:

class Walker(builtins.object)
 |  Walker(model_sampler: rxmc.param_sampling.Sampler, evidence: rxmc.evidence.Evidence, likelihood_samplers: list[rxmc.param_sampling.Sampler] | None = None, rng: numpy.random._generator.Generator | None = None)
 |
 |  Gibbs-style MCMC coordinator for a Bayesian calibration problem.
 |
 |  Manages one sampler for the physical-model parameters and, optionally,
 |  per-constraint samplers for parametric likelihood parameters.  The samplers
 |  alternate in a Gibbs framework: model parameters are updated with the
 |  likelihood parameters held fixed, then each set of likelihood parameters is
 |  updated with the model parameters held fixed.
 |
 |  Parameters
 |  ----------
 |  model_sampler : Sampler
 |      Sampler for the physical-model parameters.
 |  evidence : Evidence
 |      Evidence object containing the observations and likelihood models.
 |  likelihood_samplers : list of Sampler, optional
 |      One sampler per entry in ``evidence.parametric_constraints``.
 |  rng : np.random.Generator, optional
 |      Random number generator.  Defaults to ``default_rng(42)``.
 |
 |  Raises
 |  ------
 |  ValueError
 |      If the physical-model parameters in *evidence* and *model_sampler* do
 |      not match.
 |  ValueError
 |      If the number of *likelihood_samplers* does not equal the number of
 |      parametric constraints in *evidence*.
 |  ValueError
 |      If any likelihood sampler's parameters do not match those of the
 |      corresponding parametric constraint.
 |
 |  Methods defined here:
 |
 |  __init__(self, model_sampler: rxmc.param_sampling.Sampler, evidence: rxmc.evidence.Evidence, likelihood_samplers: list[rxmc.param_sampling.Sampler] | None = None, rng: numpy.random._generator.Generator | None = None)
 |      Initialize self.  See help(type(self)) for accurate signature.
 |
 |  log_likelihood(self, model_params, likelihood_params)
 |
 |  log_posterior(self, model_params, likelihood_params)
 |
 |  log_prior(self, model_params, likelihood_params)
 |      Log prior probability of model and likelihood parameters.
 |
 |      Parameters
 |      ----------
 |      model_params : tuple
 |          Physical-model parameter values.
 |      likelihood_params : list of tuple
 |          One tuple of likelihood parameter values per parametric constraint.
 |
 |      Returns
 |      -------
 |      float
 |          Sum of log prior densities for model and likelihood parameters.
 |
 |  run_likelihood_batches(self, n_steps, starting_locations, model_params, burn=False)
 |      Sample each set of likelihood parameters for fixed model parameters.
 |
 |      Parameters
 |      ----------
 |      n_steps : int
 |          Number of MCMC steps per likelihood sampler.
 |      starting_locations : list of np.ndarray
 |          Starting locations for each likelihood sampler.
 |      model_params : tuple
 |          Fixed physical-model parameter values.
 |      burn : bool, optional
 |          If ``True``, treat as burn-in (samples are not recorded).
 |
 |  run_model_batch(self, n_steps, x0, likelihood_params=None, burn=False)
 |      Sample model parameters for fixed likelihood parameters.
 |
 |      Parameters
 |      ----------
 |      n_steps : int
 |          Number of MCMC steps.
 |      x0 : np.ndarray
 |          Starting location for the model parameters.
 |      likelihood_params : list of tuple, optional
 |          Fixed values of the likelihood parameters for each parametric
 |          constraint.  Defaults to ``[]``.
 |      burn : bool, optional
 |          If ``True``, treat as burn-in (samples are not recorded).
 |
 |  walk(self, n_steps: int, burnin: int = 0, batch_size: int = None, verbose: bool = True)
 |      Run the full MCMC chain.
 |
 |      Updates the internal state of ``model_sampler`` and each entry of
 |      ``likelihood_samplers`` with the accumulated chain, log posteriors,
 |      and acceptance statistics.
 |
 |      Parameters
 |      ----------
 |      n_steps : int
 |          Total number of active (post-burn-in) steps.
 |      burnin : int, optional
 |          Number of burn-in steps discarded before recording.
 |          Defaults to ``0``.
 |      batch_size : int, optional
 |          Steps per batch.  If ``None`` the entire chain is one batch.
 |      verbose : bool, optional
 |          Print batch completion messages.  Defaults to ``True``.
 |
 |  ----------------------------------------------------------------------
 |  Data descriptors defined here:
 |
 |  __dict__
 |      dictionary for instance variables
 |
 |  __weakref__
 |      list of weak references to the object

First we need to put together our Evidence. With one constraint this seems trivial, but this will be useful down the road when we may want to combine multiple constraints together.

evidence = rxmc.evidence.Evidence([constraint])

Another chore we have to do is configure how we will sample our parameter space.

help(rxmc.param_sampling.MetropolisHastingsSampler)
Help on class MetropolisHastingsSampler in module rxmc.param_sampling:

class MetropolisHastingsSampler(Sampler)
 |  MetropolisHastingsSampler(params: list[rxmc.params.Parameter], prior, starting_location: numpy.ndarray, proposal: rxmc.proposal.ProposalDistribution)
 |
 |  Metropolis-Hastings sampler with a fixed proposal distribution.
 |
 |  Parameters
 |  ----------
 |  params : list of Parameter
 |      Parameters to sample.
 |  prior : object
 |      Prior distribution with a callable ``logpdf(x)`` method.
 |  starting_location : np.ndarray, shape (ndim,)
 |      Initial parameter vector.
 |  proposal : ProposalDistribution
 |      Callable proposal distribution.  Must accept ``(x, rng)`` and return
 |      a proposed parameter vector.
 |
 |  Method resolution order:
 |      MetropolisHastingsSampler
 |      Sampler
 |      builtins.object
 |
 |  Methods defined here:
 |
 |  __init__(self, params: list[rxmc.params.Parameter], prior, starting_location: numpy.ndarray, proposal: rxmc.proposal.ProposalDistribution)
 |      Initialize self.  See help(type(self)) for accurate signature.
 |
 |  ----------------------------------------------------------------------
 |  Methods inherited from Sampler:
 |
 |  batch_acceptance_fractions(self) -> numpy.ndarray
 |      Acceptance fraction for each completed batch.
 |
 |      Returns
 |      -------
 |      np.ndarray
 |          Per-batch acceptance fractions, or ``[0.0]`` if no batches run.
 |
 |  most_recent_batch_acceptance_fraction(self) -> float
 |      Acceptance fraction of the most recent batch.
 |
 |      Returns
 |      -------
 |      float
 |          Fraction of proposals accepted in the last batch, or ``0.0`` if
 |          no batches have been run.
 |
 |  overall_acceptance_fraction(self) -> float
 |      Overall acceptance fraction across all completed batches.
 |
 |      Returns
 |      -------
 |      float
 |          Total accepted / total proposed, or ``0.0`` if no batches run.
 |
 |  record_batch(self, n_steps: int, n_accepted: int, chain: numpy.ndarray, logp_chain: numpy.ndarray)
 |      Append a completed batch to the running chain.
 |
 |      Parameters
 |      ----------
 |      n_steps : int
 |          Number of steps in the batch.
 |      n_accepted : int
 |          Number of accepted proposals in the batch.
 |      chain : np.ndarray, shape (n_steps, ndim)
 |          Sampled parameter vectors.
 |      logp_chain : np.ndarray, shape (n_steps,)
 |          Log posterior values for the batch.
 |
 |  sample(self, n_steps: int, starting_location: numpy.ndarray, rng: numpy.random._generator.Generator, log_posterior: Callable[[numpy.ndarray], float], burn: bool = False)
 |      Run the sampling algorithm for one batch.
 |
 |      Updates ``self.state`` to the last sample; records the batch unless
 |      *burn* is ``True``.
 |
 |      Parameters
 |      ----------
 |      n_steps : int
 |          Number of steps to run.
 |      starting_location : np.ndarray, shape (ndim,)
 |          Starting parameter vector for this batch.
 |      rng : np.random.Generator
 |          Random number generator.
 |      log_posterior : callable
 |          Function ``f(x) -> float`` returning the log posterior at ``x``.
 |      burn : bool, optional
 |          If ``True``, discard samples (burn-in); only ``self.state`` is
 |          updated.  Defaults to ``False``.
 |
 |  ----------------------------------------------------------------------
 |  Data descriptors inherited from Sampler:
 |
 |  __dict__
 |      dictionary for instance variables
 |
 |  __weakref__
 |      list of weak references to the object

This means we have to decide on a proposal distribution. We will use a simple form: just sampling from a scaled down version of the prior about the previos point:

def proposal_distribution(x, rng):
    return stats.multivariate_normal.rvs(
        mean=x, cov=prior_distribution.cov / 100, random_state=rng
    )
sampling_config = rxmc.param_sampling.MetropolisHastingsSampler(
    params=my_model.params,
    starting_location=prior_distribution.mean,
    proposal=proposal_distribution,
    prior=prior_distribution,
)
walker = rxmc.walker.Walker(
    sampling_config,
    evidence,
)
%%time
walker.walk(
    n_steps=20000,
    burnin=1000,
)
Burn-in batch 1/1 completed, 1000 steps.
Batch: 1/1 completed, 20000 steps. 
  Model parameter acceptance fraction: 0.273
CPU times: user 5.89 s, sys: 40.5 ms, total: 5.93 s
Wall time: 5.88 s
chain = walker.model_sampler.chain
logp = walker.model_sampler.logp_chain
chain.shape
(20000, 2)
fig, axes = plt.subplots(chain.shape[1] + 1, 1, figsize=(8, 8), sharex=True)
for i in range(chain.shape[1]):
    axes[i].plot(chain[:, i])
    axes[i].set_ylabel(f"${my_model.params[i].latex_name}$ [{my_model.params[i].unit}]")
    true_value = true_params[my_model.params[i].name]
    axes[i].hlines(true_value, 0, len(chain), "r", linestyle="--")


axes[-1].plot(logp)
axes[-1].set_ylabel(r"$\log{\mathcal{L}(\alpha_i | \mathcal{O})}$")

axes[-1].set_xlabel(r"$i$")
# plt.legend(title="chains", ncol=3,)
Text(0.5, 0, '$i$')
../_images/1b60187568abc13bf277e25f371af35ea746b42959c93ffc1a3c346cfb9b107f.png
posterior_range = np.vstack([np.min(chain, axis=0), np.max(chain, axis=0)]).T
fig = corner.corner(
    chain,
    labels=[p.name for p in my_model.params],
    label="posterior",
    truths=[true_params["m"], true_params["b"]],
)
fig.suptitle("posterior")
Text(0.5, 0.98, 'posterior')
../_images/17f4f5604a626d3f42fd11880c30a07f86610a3276f55119ffa935b195be0205.png
x_full = np.linspace(-1, 2, 10)
n_posterior_samples = chain.shape[0]
y = np.zeros((n_posterior_samples, len(x_full)))
for i in range(n_posterior_samples):
    sample = chain[i, :]
    y[i, :] = my_model.y(x_full, *sample)

upper, median, lower = np.percentile(y, [5, 50, 95], axis=0)
plt.plot(x_full, my_model.y(x_full, *list(true_params.values())), "k--", label="truth")
plt.xlabel("x")
plt.ylabel("y")
plt.errorbar(
    x,
    obs1.y,
    y_stat_err,
    color="k",
    marker="o",
    linestyle="none",
    label="experiment",
)
plt.fill_between(
    x_full,
    lower,
    upper,
    alpha=0.5,
    zorder=2,
    label=r"posterior inner 90$^\text{th}$ pctl",
)
plt.plot(x_full, median, "m:", label="posterior median")

plt.legend()
# plt.title("predictive posterior")
<matplotlib.legend.Legend at 0x7b23082120c0>
../_images/2b34d52e1f2db34470a08a540d462b09b2765594e02d6bccb812454dc574f7f8.png

Nice!#

Hopefully this simple example served to illustrate the basic function of the working pieces of rxmc. The true power is the ability to compose different Constraints, and easily manage and test different model forms for both the PhysicalModel and the covariance terms describing the uncertainty.

Check out the other demos to see how rxmc helps us handle more realistic problems.