{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# MCMC sampling\n\nIn this example, we demonstrate how to perform uncertainty quantification (UQ) using\nparallel tempered MCMC (PTMCMC). We use a Stillinger-Weber (SW) potential for silicon\nthat is archived in OpenKIM_.\n\nFor simplicity, we only set the energy-scaling parameters, i.e., ``A`` and ``lambda`` as\nthe tunable parameters. Furthermore, these parameters are physically constrained to be\npositive, thus we will work in log parameterization, i.e. ``log(A)`` and ``log(lambda)``.\nThese parameters will be calibrated to energies and forces of a small dataset,\nconsisting of 4 compressed and stretched configurations of diamond silicon structure.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "To start, let's first install the SW model::\n\n   $ kim-api-collections-management install user SW_StillingerWeber_1985_Si__MO_405512056662_006\n\n.. seealso::\n   This installs the model and its driver into the ``User Collection``. See\n   `install_model` for more information about installing KIM models.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from multiprocessing import Pool\n\nimport numpy as np\nfrom corner import corner\n\nfrom kliff.calculators import Calculator\nfrom kliff.dataset import Dataset\nfrom kliff.dataset.weight import MagnitudeInverseWeight\nfrom kliff.loss import Loss\nfrom kliff.models import KIMModel\nfrom kliff.models.parameter_transform import LogParameterTransform\nfrom kliff.uq import MCMC, autocorr, mser, rhat\nfrom kliff.utils import download_dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Before running MCMC, we need to define a loss function and train the model. More detail\ninformation about this step can be found in `tut_kim_sw` and\n`tut_params_transform`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Instantiate a transformation class to do the log parameter transform\nparam_names = [\"A\", \"lambda\"]\nparams_transform = LogParameterTransform(param_names)\n\n# Create the model\nmodel = KIMModel(\n    model_name=\"SW_StillingerWeber_1985_Si__MO_405512056662_006\",\n    params_transform=params_transform,\n)\n\n# Set the tunable parameters and the initial guess\nopt_params = {\n    \"A\": [[\"default\", -8.0, 8.0]],\n    \"lambda\": [[\"default\", -8.0, 8.0]],\n}\n\nmodel.set_opt_params(**opt_params)\nmodel.echo_opt_params()\n\n# Get the dataset and set the weights\ndataset_path = download_dataset(dataset_name=\"Si_training_set_4_configs\")\n# Instantiate the weight class\nweight = MagnitudeInverseWeight(\n    weight_params={\n        \"energy_weight_params\": [0.0, 0.1],\n        \"forces_weight_params\": [0.0, 0.1],\n    }\n)\n# Read the dataset and compute the weight\ntset = Dataset(dataset_path, weight=weight)\nconfigs = tset.get_configs()\n\n# Create calculator\ncalc = Calculator(model)\nca = calc.create(configs)\n\n# Instantiate the loss function\nresidual_data = {\"normalize_by_natoms\": False}\nloss = Loss(calc, residual_data=residual_data)\n\n# Train the model\nloss.minimize(method=\"L-BFGS-B\", options={\"disp\": True})\nmodel.echo_opt_params()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "To perform MCMC simulation, we use :class:`~kliff.uq.MCMC`.This class interfaces with\nptemcee_ Python package to run PTMCMC, which utilizes the affine invariance property\nof MCMC sampling. We simulate MCMC sampling at several different temperatures to\nexplore the effect of the scale of bias and overall error bars.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Define some variables that correspond to the dimensionality of the problem\nntemps = 4  # Number of temperatures to simulate\nndim = calc.get_num_opt_params()  # Number of parameters\nnwalkers = 2 * ndim  # Number of parallel walkers to simulate"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We start by instantiating :class:`~kliff.uq.MCMC`. This requires :class:`~kliff.loss.Loss`\ninstance to construct the likelihood function. Additionally, we can specify the prior\n(or log-prior to be more precise) via the ``logprior_fn`` argument, with the default\noption be a uniform prior that is bounded over a finite range that we specify via the\n``logprior_args`` argument.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>When user uses the default uniform prior but doesn't specify the bounds, then the\n   sampler will retrieve the bounds from the model\n   (see :meth:`~kliff.models.KIMModel.set_opt_params`). Note that an error will be\n   raised when the uniform prior extends to infinity in any parameter direction.</p></div>\n\nTo specify the sampling temperatures to use, we can use the arguments ``ntemps`` and\n``Tmax_ratio`` to set how many temperatures to simulate and the ratio of the highest\ntemperature to the natural temperature $T_0$, respectively. The default values of\n``ntemps`` and ``Tmax_ratio`` are 10 and 1.0, respectively. Then, an internal function\nwill create a list of logarithmically spaced points from $T = 1.0$ to\n$T = T_{\\text{max\\_ratio}} \\times T_0$. Alternatively, we can also give a list of\nthe temperatures via ``Tladder`` argument, which will overwrites ``ntemps`` and\n``Tmax_ratio``.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>It has been shown that including temperatures higher than $T_0$ helps the\n   convergence of walkers sampled at $T_0$.</p></div>\n\nThe sampling processes can be parallelized by specifying the pool. Note that the pool\nneeds to be declared after instantiating :class:`~kliff.uq.MCMC`, since the posterior\nfunction is defined during this process.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Set the boundaries of the uniform prior\nbounds = np.tile([-8.0, 8.0], (ndim, 1))\n\n# It is a good practice to specify the random seed to use in the calculation to generate\n# a reproducible simulation.\nseed = 1717\nnp.random.seed(seed)\n\n# Create a sampler\nsampler = MCMC(\n    loss,\n    ntemps=ntemps,\n    logprior_args=(bounds,),\n    random=np.random.RandomState(seed),\n)\n# Declare a pool to use parallelization\nsampler.pool = Pool(nwalkers)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>As a default, the algorithm will set the number of walkers for each sampling\n   temperature to be twice the number of parameters, but we can also specify it via\n   the ``nwalkers`` argument.</p></div>\n\nTo run the MCMC sampling, we use :meth:`~kliff.uq.MCMC.run_mcmc`. This function requires\nus to provide initial states $p_0$ for each temperature and walker. We also need\nto specify the number of steps or iterations to take.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>The initial states $p_0$ need to be an array with shape ``(K, L, N,)``, where\n   ``K``, ``L``, and ``N`` are the number of temperatures, walkers, and parameters,\n   respectively.</p></div>\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Initial starting point. This should be provided by the user.\np0 = np.empty((ntemps, nwalkers, ndim))\nfor ii, bound in enumerate(bounds):\n    p0[:, :, ii] = np.random.uniform(*bound, (4, 4))\n\n# Run MCMC\nsampler.run_mcmc(p0, 5000)\nsampler.pool.close()\n\n# Retrieve the chain\nchain = sampler.chain"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The resulting chains still need to be processed. First, we need to discard the first few\niterations in the beginning of each chain as a burn-in time. This is similar to the\nequilibration time in a molecular dynamic simulation before we can start the\nmeasurement. KLIFF provides a function to estimate the burn-in time, based on the\nMarginal Standard Error Rule (MSER). This can be accessed via\n:func:`~kliff.uq.mcmc_utils.mser`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Estimate equilibration time using MSER for each temperature, walker, and dimension.\nmser_array = np.empty((ntemps, nwalkers, ndim))\nfor tidx in range(ntemps):\n    for widx in range(nwalkers):\n        for pidx in range(ndim):\n            mser_array[tidx, widx, pidx] = mser(\n                chain[tidx, widx, :, pidx], dmin=0, dstep=10, dmax=-1\n            )\n\nburnin = int(np.max(mser_array))\nprint(f\"Estimated burn-in time: {burnin}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>:func:`~kliff.uq.mcmc_utils.mser` only compute the estimation of the burn-in time for\n   one single temperature, walker, and parameter. Thus, we need to calculate the burn-in\n   time for each temperature, walker, and parameter separately.</p></div>\n\nAfter discarding the first few iterations as the burn-in time, we only want to keep\nevery $\\tau$-th iteration from the remaining chain, where $\\tau$ is the\nautocorrelation length, to ensure uncorrelated samples.\nThis calculation can be done using :func:`~kliff.uq.mcmc_utils.autocorr`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Estimate the autocorrelation length for each temperature\nchain_no_burnin = chain[:, :, burnin:]\n\nacorr_array = np.empty((ntemps, nwalkers, ndim))\nfor tidx in range(ntemps):\n    acorr_array[tidx] = autocorr(chain_no_burnin[tidx], c=1, quiet=True)\n\nthin = int(np.ceil(np.max(acorr_array)))\nprint(f\"Estimated autocorrelation length: {thin}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>:func:`~kliff.uq.mcmc_utils.acorr` is a wrapper for emcee.autocorr.integrated_time_,\n   As such, the shape of the input array for this function needs to be ``(L, M, N,)``,\n   where ``L``, ``M``, and ``N`` are the number of walkers, steps, and parameters,\n   respectively. This also implies that we need to perform the calculation for each\n   temperature separately.</p></div>\n\nFinally, after obtaining the independent samples, we need to assess whether the\nresulting samples have converged to a stationary distribution, and thus a good\nrepresentation of the actual posterior. This is done by computing the potential scale\nreduction factor (PSRF), denoted by $\\hat{R}^p$. The value of $\\hat{R}^p$\ndeclines to 1 as the number of iterations goes to infinity. A common threshold is about\n1.1, but higher threshold has also been used.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Assess the convergence for each temperature\nsamples = chain_no_burnin[:, :, ::thin]\n\nthreshold = 1.1  # Threshold for rhat\nrhat_array = np.empty(ntemps)\nfor tidx in range(ntemps):\n    rhat_array[tidx] = rhat(samples[tidx])\n\nprint(f\"$\\hat{{r}}^p$ values: {rhat_array}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>:func:`~kliff.uq.mcmc_utils.rhat` only computes the PSRF for one temperature, so that\n   the calculation needs to be carried on for each temperature separately.</p></div>\n\nNotice that in this case, $\\hat{R}^p < 1.1$ for all temperatures. When this\ncriteria is not satisfied, then the sampling process should be continued. Note that\nsome sampling temperatures might converge at slower rates compared to the others.\n\nAfter obtaining the independent samples from the MCMC sampling, the uncertainty of the\nparameters can be obtained by observing the distribution of the samples. As an example,\nwe will use corner_ Python package to present the MCMC result at sampling\ntemperature 1.0 as a corner plot.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Plot samples at T=1.0\ncorner(samples[0].reshape((-1, ndim)), labels=[r\"$\\log(A)$\", r\"$\\log(\\lambda)$\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>As an alternative, KLIFF also provides a wrapper to emcee_. This can be accessed by\n   setting ``sampler=\"emcee\"`` when instantiating :class:`~kliff.uq.MCMC`. For further\n   documentation, see :class:`~kliff.uq.EmceeSampler`.</p></div>\n\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}