{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallelizing sampling using MPI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To take advantage of modern high performance computing facilities such as clusters with hundreds of CPUs we recommend to use ``MPI`` instead of ``multiprocessing``.\n",
    "\n",
    "To do this we will use the ``ChainManager`` included in ``zeus``.\n",
    "\n",
    "In order to run this example, copy and paste the following script into a file called 'test_mpi.py' and run the following command in the terminal:\n",
    "\n",
    "```\n",
    "mpiexec -n 8 python3 test_mpi.py\n",
    "```\n",
    "\n",
    "This will spawn 8 ``MPI`` processes and divide them into 2 independent chains of 10 walkers each. Unfortunately ``MPI`` is not compatible with ``Jupyter`` notebooks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save this as 'test_mpi.py'\n",
    "\n",
    "```python\n",
    "\n",
    "import numpy as np\n",
    "import zeus\n",
    "from zeus import ChainManager\n",
    "\n",
    "ndim = 5\n",
    "nwalkers = 2 * ndim\n",
    "nsteps = 100\n",
    "nchains = 2\n",
    "\n",
    "def log_prob(x):\n",
    "    return -0.5 * np.sum(x**2.0)\n",
    "\n",
    "start = np.random.randn(nwalkers, ndim)\n",
    "\n",
    "\n",
    "with ChainManager(nchains) as cm:\n",
    "    rank = cm.get_rank\n",
    "\n",
    "    sampler = zeus.EnsembleSampler(nwalkers, ndim, log_prob, pool=cm.get_pool)\n",
    "    sampler.run_mcmc(start, nsteps)\n",
    "    chain = sampler.get_chain(flat=True, discard=0.5)\n",
    "        \n",
    "    np.save('chain_'+str(rank)+'.npy', chain)\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.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "288px"
   },
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}