{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Variational Inference using Pathfinder\n",
    "\n",
    "Stan supports the Pathfinder algorithm\n",
    "([Zhang, 2022](https://jmlr.org/papers/v23/21-0889.html)).\n",
    "Pathfinder is a variational method for approximately\n",
    "sampling from differentiable log densities.  Starting from a random\n",
    "initialization, Pathfinder locates normal approximations to the target\n",
    "density along a quasi-Newton optimization path, with local covariance\n",
    "estimated using the negative inverse Hessian estimates produced by the\n",
    "LBFGS optimizer.  Pathfinder returns draws from the Gaussian approximation\n",
    "with the lowest estimated Kullback-Leibler (KL) divergence to the true\n",
    "posterior.\n",
    "\n",
    "There are two Stan implementations of the Pathfinder algorithm:\n",
    "single-path Pathfinder and multi-path Pathfinder.\n",
    "Single-path Pathfinder generates a set of approximate draws from one run of the basic Pathfinder algorithm.\n",
    "Multi-path Pathfinder uses importance resampling over the draws from multiple runs of Pathfinder.\n",
    "This better matches non-normal target densities and also mitigates\n",
    "the problem of L-BFGS getting stuck at local optima or in saddle points on plateaus."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example: variational inference with Pathfinder for model ``bernoulli.stan``\n",
    "\n",
    "The [CmdStanModel pathfinder](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanModel.pathfinder ) method\n",
    "wraps the CmdStan [pathfinder ](https://mc-stan.org/docs/cmdstan-guide/pathfinder-config.html) method.\n",
    "\n",
    "By default, CmdStanPy runs multi-path Pathfinder which returns an importance-resampled set of draws over the outputs of 4 independent single-path Pathfinders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from cmdstanpy.model import CmdStanModel, cmdstan_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bernoulli_dir = os.path.join(cmdstan_path(), 'examples', 'bernoulli')\n",
    "stan_file = os.path.join(bernoulli_dir, 'bernoulli.stan')\n",
    "data_file = os.path.join(bernoulli_dir, 'bernoulli.data.json')\n",
    "# instantiate, compile bernoulli model\n",
    "model = CmdStanModel(stan_file=stan_file)\n",
    "# run CmdStan's pathfinder method, returns object `CmdStanPathfinder`\n",
    "pathfinder = model.pathfinder(data=data_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(pathfinder)\n",
    "print(pathfinder.metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `pathfinder` method returns a [CmdStanPathfinder](https://mc-stan.org/cmdstanpy/api.html#cmdstanpathfinder) object,\n",
    "which provides access to the disparate information from the Stan CSV files.\n",
    "\n",
    "\n",
    "- The [stan_variable](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanPathfinder.stan_variable) and\n",
    "[stan_variables](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanPathfinder.stan_variables) methods \n",
    "return a Python [numpy.ndarray](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html#numpy.ndarray)\n",
    "containing all draws from the sample where the structure of each draw corresponds to the structure of the\n",
    "Stan variable.\n",
    "\n",
    "- The [draws](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanPathfinder.draws) method returns the sample as a numpy.ndarray."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pathfinder.stan_variable(\"theta\").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pathfinder.column_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pathfinder.draws().shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pathfinders as initialization for the MCMC sampler\n",
    "\n",
    "The method [create_inits](https://mc-stan.org/cmdstanpy/api.html#cmdstanpy.CmdStanPathfinder.create_inits) returns a Python Dict containing a set of per-chain initializations for the model parameters.  Each set of initializations is a random draw from the Pathfinder sample.   These initializations can be used as the initial parameter values for Stan's NUTS-HMC sampler, which will reduce the number of warmup iterations needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inits = pathfinder.create_inits()\n",
    "print(inits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `create_inits` takes two arguments:\n",
    "\n",
    "* `seed` - used for random selection.\n",
    "* `chains` - the number of draws to return, default is 4.  This should match the number of sampler chains to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inits = pathfinder.create_inits(chains=3)\n",
    "print(inits)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}