build_your_own_gibbs_sampler.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7c808935-d50b-4dcc-b2db-45cd9464a946",
   "metadata": {},
   "source": [
    "# Component Separation Using Bayesian Methods\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f54ba897-65a6-4f19-83d5-76ecb6117456",
   "metadata": {},
   "source": [
    "## Chapter 1 - Simple Test Case\n",
    "Before we move on to any real data, let us take a step back and look at a couple ways to solve for parameters from a distribution. \n",
    "\n",
    "First let us draw 1000 samples from a Gaussian distribution and look at a histogram of the result. \n",
    "\n",
    "As a reminder, a Gaussian distribution for parameter $x$ with mean $\\mu$ and standard deviation $\\sigma$ is given by\n",
    "\n",
    "\\begin{equation}\n",
    "p(x) = \\frac{1}{\\sqrt{2 \\pi \\sigma^2}} e^{-\\frac{(x-\\mu)^2}{2\\sigma^2}}.\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "57ce5b05-e867-4262-979c-1192e82d7b9c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYMAAAELCAYAAAA7h+qnAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAAbEklEQVR4nO3de7xVZb3v8c83EAXvAl4SFc1e5DlWZhxN3Zml52h5wfKS5i4p9+Z0emma7leBdra29SidjTs9lewoEyoTyEsKmhfYoqWJgrLTRM9WQ4MQliioXEV/548x8KwW6zLnZI75rLHG9/16rdeac8yx5vgyXM7fep5njOdRRGBmZtX2ntQBzMwsPRcDMzNzMTAzMxcDMzPDxcDMzID+qQM0asiQITF8+PDUMawHzz77LAAjRoxInKQ+zm191fz581+JiKEdt5e2GAwfPpx58+aljmE9OOqoowCYM2dO0hz1cm7rqyS92Nl2dxOZmVl5WwZWDt/+9rdTR2iIc1vVqKx3II8cOTLcTWRmVh9J8yNiZMft7iayQi1YsIAFCxakjlE357aqccvAClXWAU3ntr7KLQMzM+uSi4GZmbkYmJmZi4GZmeH7DKxgV155ZbJjDx97Z037LRp//GbbUubeEmXNbem5GFihDj/88NQRGuLcVjUuBlaohx9+GOjdH1KdtSDWLV4IwDbDDnh3W2ctiN6mDOfbeicXAyvUxRdfDJTvuveVD04BYPcvjE+cpD5lPd+WngeQzczMxcDMzFwMzMwMFwMzM8MDyFawa665JnWEhuxy9JjUERpS1vNt6bkYWKEOOuig1BEaMmC3/VJHaEhZz7el524iK9SsWbOYNWtW6hh1W7toAWsXLUgdo25lPd+WnlsGVqgrrrgCgGOOOSZxkvqsengqAAOHH5Q2SJ3Ker4tPbcMzMzMxcDMzFwMzMwMFwMzM8MDyFawH/3oR6kjNGTwseemjtCQsp5vS8/FwAo1YsSI1BEastXgYakjNKSs59vSczeRFWrGjBnMmDEjdYy6rXluLmuem5s6Rt3Ker4tPbcMrFBXX301ACeeeGLiJPV5/dHbABi0/6GJk9SnrOfb0nPLwMzMXAzMzMzFwMzMcDEwMzM8gGwF+/nPf546QkOGnHBRwz87fOydNe23aPzxDR+jK2U935aei4EVaq+99kodoSH9dxiaOkJDynq+Lb2WdxNJ+oakP0p6StJNkraRtK+kuZKekzRN0oBW57JiTJs2jWnTpqWOUbfVCx9k9cIHU8eoW1nPt6XX0mIgaU/g68DIiDgQ6AecAXwX+F5E7A+8BpzTylxWnIkTJzJx4sTUMer2xhN38cYTd6WOUbeynm9LL8UAcn9goKT+wCBgKfAp4Ob89SnAyQlymZlVVkuLQUQsASYAL5EVgVXAfGBlRGzMd1sM7NnZz0saI2mepHltbW2tiGxmVgmt7ibaGRgF7Au8F9gWOK7Wn4+ISRExMiJGDh1azgE+M7PeqNVXEx0D/Cki2gAk3QocAewkqX/eOhgGLGlxLrMe1XrJqFkZtboYvAR8TNIgYC1wNDAPuB84FZgKnA3c3uJcVpCbb7655516oaEnj0sdoSFlPd+WXkuLQUTMlXQz8DiwEXgCmATcCUyVdEW+7fpW5rLiDBkyJHWEhvQbtGPqCA0p6/m29Fp+01lEXApc2mHzC8Ahrc5ixZs8eTIAo0ePTpqjXm8+OQuA7T54TOIk9Snr+bb0PDeRFWry5MnvfkCVyZtPznq3IJRJWc+3pediYGZmLgZmZuZiYGZmuBiYmRmewtoKdtddzZ/srRU3f+162mWFH6MIRZxvqwYXAyvUoEGDUkdoyHu22iZ1hIaU9Xxbei4GVqjrrrsOgK997WuJk9Tnjcez1sf2Bzd/NbJN6mnh1LoqWlnPt6XnMQMr1PTp05k+fXrqGHVb/cxvWf3Mb1PHqFtZz7el52JgZmYuBmZm5mJgZma4GJiZGb6ayAo2Z86c1BEasvsXxqeO0JCynm9Lzy0DMzNzMbBiTZgwgQkTJqSOUbdVc29l1dxbU8eoW1nPt6XnYmCFmjlzJjNnzkwdo25rn3+Utc8/mjpG3cp6vi09FwMzM/MAsvUOrZh8zsy65paBmZm5ZWDFGjhwYOoIDVH/rVNHaEhZz7el52JghfrNb36TOkJDdjv9O6kjNKSs59vSczeRmZm5GFixLr/8ci6//PLUMeq28qGbWPnQTalj1K2s59vSczGwQs2ePZvZs2enjlG3dS/+O+te/PfUMepW1vNt6XnMwKwPeeSFFUDPl+rWunKaVYdbBmZm5mJgZmY1FgNJX5I0uIvXdpH0pebGsr5i8ODBDB7c6a9Or9Zv4A70G7hD6hh1K2tuS6/WMYMbgMOAFZ28tm/++s+aFcr6jltuuSV1hIYM/ezFqSM0pKy5Lb1au4nUzWuDgdebkMXMzBLpsmUgaRQwqt2m/ymprcNu2wAfBx4rIJv1AePGjQPgqquuSpykPq89MBmAnT8xOmmOepU1t6XXXTfRrsAH2z1/H7B7h302APcCVzQ5l/URv//971NHaMj6Jc+kjtCQsua29LosBhHxY+DHAJLuB/5HRPg3zcysD6ppADkiPll0EDMzS6fmO5AlvRc4ARhGNlbQXkTEt2p8n52AnwAHAgF8BXgWmAYMBxYBp0fEa7VmMzOzLVNTMZD0WeAmoB+wnGysoL0AaioGwLXA3RFxqqQBwCDgYmB2RIyXNBYYW8f7WS82bNiw1BEa0n/7IakjNKSsuS29WlsGV5INFI+OiFcbPZikHYEjgdEAEbEB2JBfuXRUvtsUYA4uBn3CL37xi9QRGjLkxH9IHaEhZc1t6dVaDPYCztuSQpDbF2gDbpD0YWA+cD6wW0Qszfd5Gditsx+WNAYYA7D33ntvYRRrBa9tbFYOtd509jAwognH6w8cDEyMiI8Aq8m6hN4VEUHW7bSZiJgUESMjYuTQoUObEMeK9uqsSbw6a1LqGHVzbquaWlsGFwI3SnoTuA9Y2XGHiFhTw/ssBhZHxNz8+c1kxWCZpD0iYqmkPcjGJawP2LD8hdQRGuLcVjW1FoM/5N9voIu/2skGl7sVES9L+rOkERHxLHA08HT+dTYwPv9+e425zMysCWotBl+h6yJQr/PIWhkDgBeAL5N1V02XdA7wInB6k45lZmY1qPWms8nNOmBELABGdvLS0c06hpmZ1cfLXlqhttplz9QRGuLcVjW13nTWRg/dRBGxa1MSWZ8y+LjzUkdoiHNb1dTaMvghmxeDncm6dnYAftrMUGZm1lq1jhlc1tl2SQKmA281MZP1ISvu/j5Qvr9YnduqZovGDCIiJP2E7JLTK5sTyfqSt15dkjpCQ5zbqqbWO5C7sx8woAnvY2ZmidQ6gPy1TjYPAA4AzgJ+1cxQZmbWWrV2E/2gk23ryaaXuA74TtMSmZlZy9U6gNyM7iSroAG77pc6QkOc26rGN51ZoXY5ZkzqCA1xbquamv/il7SfpImSnpS0JP9+nST/KWJmVnK1DiB/FLgfWAfMBJaRLUBzCnCWpE9GxOOFpbTSemXGBKB8K3A5t1VNrd1EE4AngE+3X7dA0iDgrvz1TzU/npXdxjdeSR2hIc5tVVNrN9EhwP/uuIBN/nwCcGizg5mZWevU2jJYCwzu4rVdyLqPrEK8tnFr+Xxb0WptGdwJjJf0N+035s+vAmY0O5iZmbVOPWsg3w48IGk52RrFu+ZfvwcuKiaeld3We34gdYSGOLdVTa03na0A/kbSccB/AfYAlgJzI+LeAvNZye38idGpIzTEua1q6rrpLCLuBu4uKIuZmSXS5ZiBpD0k3SLp2G72OTbfx6ucWafabruSttvKN7u5c1vVdDeA/A9k01N31w10L7AvHjOwLry99nXeXvt66hh1c26rmu6KwQnAv0ZEl2sf56/9CBjV7GBmZtY63RWDfYCna3iPhcDwpqQxM7MkuisGa8kWu+/Jdvm+ZmZWUt1dTfQ4cBLZDWfdGZXva7aZbfb5cOoIDXFuq5ruisF1wDRJD0fElM52kPQl4MvA54sIZ+W30xFnpo7QEOe2qumyGETELZKuBW6QdC7Z/QUvAQHsDRwLjAS+FxG3tSKsmZkVo9ubziLiIklzgAvILjXdOn9pPfAQMCoiZhYZ0Mpt2fRLAdjt9HItk+3cVjU93oEcETOAGZL68/9nLl0RERsLTWZ9QmxcnzpCQ5zbqqbm6SjyD/9lBWYxM7NEal4D2czM+i4XAzMzq2/WUrN6DXzfIakjNMS5rWpcDKxQOx76udQRGuLcVjXuJjIzs9paBpK2As4HPgcMA7bpuE9EeE0D28zLvxwLwO5fGJ84SX2c26qm1m6i7wH/HZgJ3A9s2JKDSuoHzAOWRMQJkvYFppLdxzAf+GJEbNExzMysdrUWg9OAsRFxdZOOez7Z1NebZkX9Ltm0FlMl/StwDjCxSccyM7Me1DpmIOAPzTigpGHA8cBP8ucCPgXcnO8yBTi5GccyM7Pa1FoMfgw0azrEa4BvAu/kzwcDK9tNb7EY2LOzH5Q0RtI8SfPa2tqaFMfMzGrtJloGnCXpfuA+YGWH1yMieuzWkXQCsDwi5ks6qo6cmw4yCZgEMHLkyC6X47TeY9sPfDx1hIY4t1VNrcXgmvz73sAnOnk9qK2P/wjgJEmfIbsiaQfgWmAnSf3z1sEwYEmNuayX2/7g41NHaIhzW9XU1E0UEe/p4atfje8zLiKGRcRw4Azg3yLiLLIrlE7NdzsbuL2Bf4v1Qu+8tY533lqXOkbdnNuqprfcdPYt4EJJz5GNIVyfOI81yfJfXcbyX12WOkbdnNuqpstuIkn/CXg+Itbnj7sVEU/Xc+CImAPMyR+/AHhSFTOzRLobM3gK+BjwaP64qwFb5a/V1FVkZma9T3fF4JPA0+0em5lZH9VlMYiIBzp7bGZmfU9DU1hL2gl4H/DniFje1ESW1PCxdzb1/bb74DFNfb9W6eu5a/3vvGi8L1Wtim6LgaQzyKaG2Aq4NSJulPSPwDhgQL7Pr4EvRcTqYqNaGfX1D9Xepqy5Lb0uLy2V9PfAL4F9gR2BGyR9D/gGcDHZ/EJjgaOBS4qPamX09ppVvL1mVeoYdXNuq5ruWgbnAddExIUAkv6WbBK58yPiB/k+d0vaCHyVrECY/ZW2X18FlG9+fee2qunuprP3ATPaPb+d7DLS+R32mwfs0+RcZmbWQt0Vg4FA+3GANfn39R3220A2pmBmZiXV03QUnd1o5tlCzcz6mJ4uLb0nHxNob3aHbQ1dnmpmZr1Hdx/k32lZCuuztv/IZ1JHaIhzW9V0dweyi4FtsW0PODJ1hIY4t1VNb5nC2vqoja+3sfH18i1R6txWNS4GVqhXZl7NKzOvTh2jbs5tVeNiYGZmLgZmZuZiYGZmuBiYmRm+YcwKtsMhn00doSHObVXjYmCFGrT/oakjNMS5rWrcTWSFemvFYt5asTh1jLo5t1WNi4EVasU9P2DFPT/oecdexrmtalwMzMzMxcDMzFwMzMwMFwMzM8OXllrBdjz8jNQRGuLcVjUuBlaogcMPSh2hIc5tVeNuIivUhmUvsGHZC6lj1M25rWpcDKxQr86exKuzJ6WOUTfntqpxMTAzMxcDMzNzMTAzM3w1kZl1Y/jYO2vab9H44wtOYkVraTGQtBfwM2A3IIBJEXGtpF2AacBwYBFwekS81spsVoydjjw7dYSGOLdVTatbBhuBiyLicUnbA/Ml3QeMBmZHxHhJY4GxwLdanM0KsM2wA1JHaIhzW9W0dMwgIpZGxOP54zeAhcCewChgSr7bFODkVuay4qxbvJB1ixemjlE357aqSTaALGk48BFgLrBbRCzNX3qZrBups58ZI2mepHltbW2tCWpbZOWDU1j54JSed+xlnNuqJkkxkLQdcAtwQUS83v61iAiy8YTNRMSkiBgZESOHDh3agqRmZtXQ8quJJG1FVghujIhb883LJO0REUsl7QEsb3Wuvq7Wq0LMiuSrk3qvlrYMJAm4HlgYEf/S7qU7gE2XQZwN3N7KXGZmVdfqlsERwBeBJyUtyLddDIwHpks6B3gROL3FuczMKq2lxSAifgeoi5ePbmUWa41djh6TOkJDnNuqxncgW6EG7LZf6ggNcW6rGhcDK9TaRQuA8i264tz18QUK5ediYIVa9fBUoHwfqs5tVeNZS83MzMXAzMxcDMzMDBcDMzPDA8hWsMHHnps6QkOc26rGxaCX6iuX6m01eFjqCA1xbqsadxNZodY8N5c1z81NHaNuzm1V45aBFer1R28DYND+hyZOUh/ntqpxy8DMzFwMzMzMxcDMzHAxMDMzPIBsBRtywkWpIzTEua1qXAysUP13GJo6QkOc26rG3URWqNULH2T1wgdTx6ibc1vVuGVghXrjibsA2PaAIxMnqY9zW9W4ZWBmZi4GZmbmYmBmZrgYmJkZHkC2gg09eVzqCA1x7vKodbr3ReOPLzhJubkYWKH6DdoxdYSGOLdVjbuJrFBvPjmLN5+clTpG3ZzbqsbFwApV1g8n57aqcTEwMzOPGTSDB7DMmquvrAFeJm4ZmJmZWwat5L92zKy3cjGwQu162mWpIzTEua1qKlkM3MffOu/ZapvUERri3FY1HjOwQr3x+J288Xj5usec26qmki0Da53Vz/wWgO0PLlcry7n7nt7eI5A6n1sGZmbWe4qBpOMkPSvpOUljU+cxM6uSXtFNJKkf8EPgvwKLgcck3RERT6fM5UtBzaqn2f/fl+VClN7SMjgEeC4iXoiIDcBUYFTiTGZmlaGISJ0BSacCx0XE3+XPvwgcGhHndthvDDAmfzoCeHYLDjsEeGULfj6VMuYuY2Zw7lYqY2YoZ+59ImJox429opuoVhExCZjUjPeSNC8iRjbjvVqpjLnLmBmcu5XKmBnKm7szvaWbaAmwV7vnw/JtZmbWAr2lGDwGvF/SvpIGAGcAdyTOZGZWGb2imygiNko6F7gH6Af8NCL+WPBhm9LdlEAZc5cxMzh3K5UxM5Q392Z6xQCymZml1Vu6iczMLCEXAzMzq14xkDRC0oJ2X69LuiB1rp5I+oakP0p6StJNkkoxV7Gk8/PMf+zN51nSTyUtl/RUu227SLpP0n/k33dOmbGjLjKflp/rdyT1ykseu8j9z5KekfQHSbdJ2ilhxE51kfvyPPMCSfdKem/KjFuicsUgIp6NiIMi4iDgo8Aa4La0qbonaU/g68DIiDiQbJD9jLSpeibpQODvye4w/zBwgqT906bq0mTguA7bxgKzI+L9wOz8eW8ymc0zPwV8Dniw5WlqN5nNc98HHBgRHwL+LzCu1aFqMJnNc/9zRHwo/zyZCfxjq0M1S+WKQQdHA89HxIupg9SgPzBQUn9gEPCXxHlqcQAwNyLWRMRG4AGyD6peJyIeBF7tsHkUMCV/PAU4uZWZetJZ5ohYGBFbcmd+4brIfW/+OwLwCNm9Rr1KF7lfb/d0W6C0V+RUvRicAdyUOkRPImIJMAF4CVgKrIqIe9OmqslTwMclDZY0CPgMf31zYW+3W0QszR+/DOyWMkyFfAX4TeoQtZL0vyT9GTgLtwzKJ7+57STgV6mz9CTvqx4F7Au8F9hW0t+mTdWziFgIfBe4F7gbWAC8nTJToyK7Bru0f/WVhaRLgI3Ajamz1CoiLomIvcgyn9vT/r1VZYsB8Gng8YhYljpIDY4B/hQRbRHxFnArcHjiTDWJiOsj4qMRcSTwGll/cFksk7QHQP59eeI8fZqk0cAJwFlRzhugbgROSR2iUVUuBmdSgi6i3EvAxyQNkiSysY6FiTPVRNKu+fe9ycYLfpk2UV3uAM7OH58N3J4wS58m6Tjgm8BJEbEmdZ5aSXp/u6ejgGdSZdlSlbwDWdK2ZB+w+0XEqtR5aiHpO8DnyZrQTwB/FxHr06bqmaTfAoOBt4ALI2J24kidknQTcBTZlMTLgEuBXwPTgb2BF4HTI6LjIHMyXWR+Ffg+MBRYCSyIiGMTRexUF7nHAVsDK/LdHomIryYJ2IUucn+GbDr9d8h+R76aj/GVTiWLgZmZ/bUqdxOZmVnOxcDMzFwMzMzMxcDMzHAxMDMzXAysD5H0J0nR6snwJA3Pj3tCu22LJE1o9/z0/KaqZh53nqTJzXxPqy4XA+sTJB0GDM+fnpkwyiafBf5Pu+enA6PTRDHrmYuB9RVnAquBufSCYhART0TES6lzmNXKxcBKT1I/sr+87wB+Chwg6cPtXh+dd+McLGmOpDX5YiQHS9pW0g2SVkl6QdKZHd57jqSbJY3Ju37WSrozX2Oiu0zvdhPlXTmnAJ/Ic4Skyzru10ne7dptO1DSQ5LWSVoo6aQujvtxSQ/k/8YVkn4safvaz6ZVlYuB9QWfJJteeipwM9nUF521DqaQzUd1CqB83+vJ1oY4laxV8TNJHefSPww4D7gQOAf4ENlUFbW6HLifbBqRw/Kvn9T6w5IGAvcA2wFfAK4AriGbJqP9fkcAs8im2z4VuIBsuoQb6shqFdU/dQCzJjiTbB6euyNig6R7gTMkjesw++WEiJgCkE/4dycwJyIuybc9SvYheiIwsd3P7QoctqnbR9KLwO8kHRcRd/cULiKel/Qq8J6IeKSBf9+X8wyHRsTiPMMi4Hcd9hsPPBwRn9+0QdISYLakAyPiKcy64JaBlVq+LsXngNsiYkO+eSqwD9lf4O21nyTvufz7v23akE9a2AZ07AJ6vH3/f0Q8RDad9SFb/A+ozSHA/E2FoEMGAPLFgw4Dpkvqv+mLrGC8RbbEq1mXXAys7D4N7ATcJWmnfCH1OcB6Nu8qWtnu8YZOtm3avk2HbZ2tY7Ac2KPesA3avZsMm+xMtjb2dWQf/pu+1gNbUa4V5iwBdxNZ2W36wO9sxbrTJF3QhGPs2sW2pZ1sr9c6YECHbTt3eP4y8IEecq0kW4ntMuCuTvYtw5rZlpCLgZVWvi7FiWSDwpM6vPwR4F+ATzXhUAdL2rvdmMERZB/Ej9bxHp21OAAWAwd02PbfOjx/DDhL0rB2YwabMgAQEaslPQKMiIh/qiOXGeBiYOU2ChgEXBsRc9u/IOkh4BKylsODW3icNuBOSZeSfaB/l2wcocfB43aeAUZJOpmsAPwlIv4C3AZ8X9LFZB/6pwD/ucPP3gB8O89wGTCQ7AqlVzrs902yweJ3yK6UeoPsiqPjgUsiokxLjlqLeczAyuxM4D86FgKAfK3o6WSDy1tv4XEeBn5Idjnn9cBTwMl1vsd1wL1k90E8BozJt0/K3/fred71ZJeOvitfBvJYspvqppKtsHUR2cpa7ff7HXAk2SpnPwdmkBWIP5OtzGXWJa90ZtYNSXOAVyLi1NRZzIrkloGZmbkYmJmZu4nMzAy3DMzMDBcDMzPDxcDMzHAxMDMzXAzMzAz4f2xFoBXq6hl2AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "from corner import corner\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "ndraws = 1000\n",
    "mu     = 10.\n",
    "sigma  = 1.\n",
    "\n",
    "data = np.zeros(ndraws)\n",
    "for i in range(ndraws):\n",
    "    data[i] = np.random.normal(mu,sigma)\n",
    "\n",
    "plt.hist(data,bins=30)\n",
    "plt.xlabel('Amplitude',size=15)\n",
    "plt.ylabel('Bin Count',size=15)\n",
    "plt.axvline(mu,color='k',linestyle='--')\n",
    "plt.axvline(mu-sigma,color='k',linestyle='--')\n",
    "plt.axvline(mu+sigma,color='k',linestyle='--')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47484222-b704-462f-b925-3ac12ff9bc50",
   "metadata": {},
   "source": [
    "Now let us suppose we only have the `data` vector defined above, and want to find the mean value of that data. How do we go about that?\n",
    "\n",
    "The first step is always to craft a data model. Let us switch to an observational standpoint. We start by assuming that the true value of the data $\\mu'$ is a single value who's obserbed value has been obscured by some underlying Gaussian noise, i.e. $x = \\mu' + n$ where $x$ is the data, and $n$ is some Gaussian noise. We will generalize this to the following equation, which ought to be baked into your memory\n",
    "\n",
    "\\begin{equation}\n",
    "d = s + n\n",
    "\\end{equation}\n",
    "\n",
    "where $d$ is the observed data, $s$ is the true underlying signal, and $n$ is again a noise term. This formalism is the basis of our component separation methods and will be expanded upon further down the road."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd2e86f4-4b97-435d-80ee-d5551dcf8bfc",
   "metadata": {},
   "source": [
    "### Maximum-Likelihood Solution\n",
    "\n",
    "Now that we have a simple data model, how does one go about actually determining the true underlying signal $s$? A good starting point is by finding what is called the maximum likelihood solution, using maximum likelihood estimation (MLE). A cursory Google search for maximum likelihood estimation will provide a mix of math heavy resources (like Wikipedia), and some practical examples (like towardsdatascience). \n",
    "\n",
    "Here we will stick with a Gaussian assumption, i.e. that the noise inherent in the data has a Gaussian form. \n",
    "\n",
    "For any observed point $x_i$, the probability of observing that data point, given the underlying mean ($\\mu$) and standard deviation ($\\sigma$) of the distribution is given by\n",
    "\n",
    "$p(x_i|\\mu,\\,\\sigma) = \\frac{1}{\\sqrt{2 \\pi \\sigma^2}} e^{-\\frac{(x_i-\\mu)^2}{2\\sigma^2}}$,\n",
    "\n",
    "and the probability of observing say 3 data points $x_1$, $x_2$, and $x_3$ is given by \n",
    "\n",
    "$p(x_1,x_2,x_3|\\mu,\\,\\sigma) = \\frac{1}{\\sqrt{2 \\pi \\sigma^2}} e^{-\\frac{(x_1-\\mu)^2}{2\\sigma^2}}\\times\\frac{1}{\\sqrt{2 \\pi \\sigma^2}} e^{-\\frac{(x_2-\\mu)^2}{2\\sigma^2}}\\times\\frac{1}{\\sqrt{2 \\pi \\sigma^2}} e^{-\\frac{(x_3-\\mu)^2}{2\\sigma^2}}$.\n",
    "\n",
    "Let us take all of our observations into account and condense our probability expression (the likelihood) to be\n",
    "\n",
    "\\begin{equation}\n",
    "p(\\vec{x}|\\mu,\\,\\sigma) = {\\displaystyle \\prod_{i=1}^{1000}}\\frac{1}{\\sqrt{2 \\pi \\sigma^2}} e^{-\\frac{(x_i-\\mu)^2}{2\\sigma^2}}.\n",
    "\\end{equation}\n",
    "\n",
    "From our Calculus experience, we may note that a sure fire way to determine the maximum or minimum of a function is by taking its derivative. However we can tell by looking at the above equation that taking the derivative of this thing will be a cumbersome operation.\n",
    "\n",
    "Those versed in MLE will be quick to point out that taking the natural logarithm of the above expression gives us something nice to work with!\n",
    "\n",
    "To be explicit, finding the maximum likelihood solution involves finding $\\mu$ where\n",
    "   \n",
    "\\begin{equation}\n",
    "\\frac{\\partial \\ln (p(\\vec{x}|\\mu,\\,\\sigma))}{\\partial \\mu} = 0.\n",
    "\\end{equation}\n",
    "\n",
    "Before we move on, let us touch on a couple final pieces of information which will be important for the future. First let's consider the equation we are solving here $p(\\vec{x}|\\mu,\\,\\sigma)$. This equation reads \"probability of the data given the parameters $\\mu$ and $\\sigma$. Since we already know the data, and don't know the values of $\\mu$ and $\\sigma$ a priori, it appears that this equation itself isn't very useful. However, the <span style=\"background-color:yellow\">likelihood</span> of having $\\mu$ and $\\sigma$ given some data $\\vec{x}$ is defined as\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathcal{L}(\\mu,\\,\\sigma|\\vec{x}) = p(\\vec{x}|\\mu,\\,\\sigma).\n",
    "\\end{equation}\n",
    "\n",
    "Though the equations are equal, they ask different questions. The left-hand side is concerned about the likelihood of the parameters given the data, while the right-hand side is asking about what the probability of observing the data is, given the parameters. Hence the phrasing *Maximum Likelihood Solution*.\n",
    "\n",
    "Finally, as we will see often later on, let us clearly define the log-likelihood as\n",
    "\n",
    "\\begin{equation}\n",
    "\\ln\\mathcal{L}(\\mu,\\,\\sigma|\\vec{x}) = \\ln\\,p(\\vec{x}|\\mu,\\,\\sigma).\n",
    "\\end{equation}\n",
    "\n",
    "Compiled with help from [towardsdatascience.com](https://towardsdatascience.com/probability-concepts-explained-maximum-likelihood-estimation-c7b4342fdbb1)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6aadd0b-c68a-4699-9b93-bb30de30a2cf",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Task 1: \n",
    "\n",
    "- Derive an expression $\\ln\\mathcal{L}(\\mu,\\,\\sigma|\\vec{x})$. Hold on to this result as we will come back to it later.\n",
    "\n",
    "- Derive an expression for $\\frac{\\partial \\ln\\mathcal{L}(\\mu,\\,\\sigma|\\vec{x})}{\\partial \\mu} = 0$.\n",
    "\n",
    "- Given the above expression for $\\frac{\\partial \\ln\\mathcal{L}(\\mu,\\,\\sigma|\\vec{x})}{\\partial \\mu}$, find an expression for the mean of the data.\n",
    "\n",
    "- Create a function which takes in the data and a noise ($\\sigma$) estimate and returns the log-likelihood as a function of the mean. Plot the result.\n",
    "\n",
    "- Create a function which solves for the ML solution given the data, and compare your answer to the value of $\\mu$ given in the first cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4414ad2-c941-4a12-840e-cfedc5a4f245",
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_log_likelihood(data,sigma):\n",
    "    samples = np.shape(data)[0]\n",
    "    \n",
    "    # Make a grid of uniform possible mean values\n",
    "    mu_grid = np.linspace(np.min(data),np.max(data),1000)\n",
    "\n",
    "    # Calculate the log-likelihood given the data for each mu in mu_grid\n",
    "    \n",
    "    return mu_grid,lnL\n",
    "    \n",
    "# Get the values\n",
    "x, lnL = return_log_likelihood(data,sigma)\n",
    "# Plot\n",
    "plt.plot(x,lnL)\n",
    "plt.show()\n",
    "\n",
    "def solve_for_ML(data,sigma):\n",
    "    \n",
    "    return amplitude\n",
    "\n",
    "print(solve_for_ML(data,sigma))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4be45303-01f3-49c9-b34a-5d96a280c9b1",
   "metadata": {},
   "source": [
    "### Multi-frequency example\n",
    "\n",
    "As we will see further on, determining the mean of an observable can be aided by (or in practice requires) the addition of multiple data points. As we move more towards a realistic example, we will say that we have observations at multiple frequencies. For a simple example, let's consider some observable which scales linearly with frequency. \n",
    "\n",
    "Let's create a new data set, where our observable has a mean value $a_{obs} = \\mu$, and each observation frequency has its own noise characterisitic $n_{\\nu}$. So instead of multiple observations of a single point $x_i$, we have one data point at each frequency $\\nu$, giving us a vector of data $d_{\\nu}$.\n",
    "\n",
    "Explicitly, we can now write our likelihood expression as\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathcal{L}(\\mu,\\,\\sigma|\\vec{d_{\\nu}}) = p(\\vec{d_{\\nu}}|\\mu,\\,\\sigma) = {\\displaystyle \\prod_{\\nu}}\\frac{1}{\\sqrt{2 \\pi \\sigma_{\\nu}^2}} e^{-\\frac{(d_{\\nu}-\\mu)^2}{2\\sigma_{\\nu}^2}}.\n",
    "\\end{equation}\n",
    "\n",
    "When carrying out this multi-frequency analysis, we need to choose a reference frequency $\\nu_{\\rm ref}$, at which we evaluate the recovered amplitude $a_{obs}$.\n",
    "\n",
    "Therefore, our new data model will be written as\n",
    "\n",
    "$d_{\\nu} = s_{\\nu} + n_{\\nu} = a_{obs}(\\nu/\\nu_{\\rm ref}) + n_{\\nu}$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9ebaa8c-9a85-47fe-a2cb-a64002eb059f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# New variables\n",
    "nfreq  = 5\n",
    "nus    = np.asarray([10.,20., 30.,40.,50.])\n",
    "sigmas = np.asarray([1., 1.5, 1.25, 2., 0.75])\n",
    "\n",
    "# Remake our data array\n",
    "data   = np.empty((ndraws,nfreq))\n",
    "\n",
    "# Reference frequency\n",
    "nu_ref = 20.\n",
    "\n",
    "# Define a frequency scaling function for the data\n",
    "def gen_spectrum(nu,nu_ref):\n",
    "    return (nu/nu_ref)\n",
    "\n",
    "# Simulate some data\n",
    "for i in range(ndraws):\n",
    "        for j in range(nfreq):\n",
    "            data[i][j] = mu*gen_spectrum(nus[j],nu_ref)+np.random.normal(0.0,sigmas[j])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba4a45e9-3032-4be1-abb4-1ec2e5089554",
   "metadata": {},
   "source": [
    "What does this data look like for us?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "080867e6-b095-48e2-a563-2a5a38e28b48",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_mean  = np.empty(nfreq)\n",
    "data_error = np.empty(nfreq)\n",
    "for j in range(nfreq):\n",
    "    data_mean[j]  = np.mean(data.T[j])\n",
    "    data_error[j] = np.std(data.T[j])\n",
    "    \n",
    "plt.errorbar(nus,data_mean,yerr=data_error,fmt='.',color='k')\n",
    "plt.xlabel('\"Frequency\"',size=15)\n",
    "plt.ylabel('Amplitude',size=15)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e5c3344-6e79-410b-bfd8-f7e35a91ac99",
   "metadata": {},
   "source": [
    "We can see that it would be simple to draw a straight line through the points, though many straight lines may fit through all of the error bars. It is vital that we make an estimate of the *best* fit to the data, not just *a* fit."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db54d87c-a1da-46be-8f7b-6048ddac90fb",
   "metadata": {},
   "source": [
    "### Task 2:\n",
    "\n",
    "- Generalize the `return_log_likelihood` function to work for the multi-frequency data and plot $\\mathcal{L}(a_{obs}|\\boldsymbol d)$ as a function of $a_{obs}$.\n",
    "\n",
    "- Generalize the `solve_for_ML` function to work for the multi-frequency data and find the ML value of $a_{obs}$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8706a1de-4ea5-47f8-a244-918d83a7bdb0",
   "metadata": {},
   "source": [
    "# Interjection - Some Important Information\n",
    "\n",
    "## Generalization\n",
    "\n",
    "Up until now we have taken a very simple example, where we have a data vector containing samples, and we find the MLS. Before we move to drawing samples from noisy data, we will define some more general nomenclature. Our original data model looked something like\n",
    "\n",
    "\\begin{equation}\n",
    "d = s + n.\n",
    "\\end{equation}\n",
    "\n",
    "We then expanded the actual signal into an amplitude and a frequency scaling\n",
    "\n",
    "\\begin{equation}\n",
    "s = a_{obs}(\\nu/\\nu_{\\rm ref}).\n",
    "\\end{equation}\n",
    "\n",
    "This expansion of the signal term can be generalized in the following way\n",
    "\n",
    "\\begin{equation}\n",
    "s = \\mathsf{T} a,\n",
    "\\end{equation}\n",
    "\n",
    "where $\\mathsf{T}$ is a matrix containing the scaling relation information (in this simple case $\\mathsf{T}=(\\nu/\\nu_{\\rm ref})$), and $a$ is a vector corresponding to the amplitudes we are looking for $a=a_{ obs}$.\n",
    "\n",
    "In essence, by finding the ML solution, we are solving the equation\n",
    "\n",
    "\\begin{equation}\n",
    "(\\mathsf{T}^t \\mathsf{N}^{-1} \\mathsf{T})a = \\mathsf{T}^t \\mathsf{N}^{-1} d,\n",
    "\\end{equation}\n",
    "\n",
    "where $\\mathsf{N}^{-1}$ is the noise covariance matrix.\n",
    "\n",
    "To be explicit, through the remainder of this notebook we will be looking and multi-frequency observations (vital in Cosmological Component Separation) and aim to solve for the true signal sky signal $s_{\\nu}$ with the data model\n",
    "\n",
    "\\begin{equation}\n",
    "d_{\\nu} = s_{\\nu} + n_{\\nu},\n",
    "\\end{equation}\n",
    "\n",
    "by solving the linear equation  \n",
    "\n",
    "\\begin{equation}\n",
    "{\\displaystyle \\sum_{\\nu}(\\mathsf{T}_{\\nu}^t \\mathsf{N}_{\\nu}^{-1} \\mathsf{T}_{\\nu})a = \\sum_{\\nu}\\mathsf{T}_{\\nu}^t \\mathsf{N}_{\\nu}^{-1} d_{\\nu}}.\n",
    "\\end{equation}\n",
    "\n",
    "This equation will be expanded upon to include sampling and prior terms throughout this notebook. \n",
    "\n",
    "\n",
    "Note that when we move to more complex systems (many data sets with many pixels), filling and multiplying these matricies can be very expensive both computationally and for memory usage. A student interested in should read \"An Introduction to the Conjugate Gradient Method withouth the Agonizing Pain\" by Jonathan Shewchuk. It is excellent. Additionally, Numerical Recipes is always your friend when trying to solve complex linear systems computationally.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f605f15-9fda-4ba9-b03c-05995ba42be1",
   "metadata": {},
   "source": [
    "## Bayes Theorem\n",
    "\n",
    "Before getting too far into the weeds here, we will introduce Bayes' Theorem as well as a practical definition for our purposes. Given a set of parameters $\\vec{\\omega} = \\{\\omega_1,...,\\omega_k\\}$, we will define the posterior distribution of $\\boldsymbol \\omega$ with respect to the data as \n",
    "\n",
    "\\begin{equation}\n",
    "P(\\vec{\\omega}\\mid\\boldsymbol{d}) = \\frac{P(\\boldsymbol{d} \\mid \\vec{\\omega}) P(\\vec{\\omega})}{P(\\boldsymbol{d})} \\propto \\mathcal{L}(\\vec{\\omega}) P(\\vec{\\omega}),\n",
    "\\end{equation}\n",
    "where $ P(\\boldsymbol{d}\\mid\\vec{\\omega}) \\equiv \\mathcal{L}(\\vec{\\omega})$ is the likelihood function (as we have talked about above), $P(\\vec{\\omega})$ is the prior, and $P(\\boldsymbol{d})$ is a normalization factor which we disregard here as it is independent of the parameters $\\vec{\\omega}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54cf9798-416c-4487-8868-c2e7263745b5",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Chapter 2 - Sampling Methods\n",
    "\n",
    "In Chapter 1 we have simply drawn Gaussian samples of the data using `numpy`, having already known the mean and standard deviation of our data. In map-space analysis, we work with sky maps that are characterized by a total sky signal amplitude $s_{\\nu}$, with an estimate of the noise $n_{\\nu}$, giving us our data\n",
    "\n",
    "\\begin{equation}\n",
    "d_{\\nu} = s_{\\nu} + n_{\\nu}.\n",
    "\\end{equation}\n",
    "\n",
    "This has been stated before and will be stated again in hopes of making you see the above equation when you close your eyes. We now wish to sample for $s_{\\nu}$, the *true* sky signal, seeing as we know the data $d_{\\nu}$ and have an estimate of the noise $n_{\\nu}$.\n",
    "\n",
    "Here we introduce three methods which we can use for sampling parameters. The Metropolis and Metropolis-Hastings methods are broadly used in MCMC methods and will be introduced here. These methods are usable in the general sense. Inversion Sampling is also introduced, which is used in the Commander/\\texttt{Cosmoglobe} framework for sampling non-linear parameters, i.e., spectral indices. Finally from a Gaussian distribution is introduced. This is a core method to learn for sampling linear parameters and will be expanded upon in the next Chapter to the multi-variate case. In the Cosmoglobe/Commander framework this method is used to sample for Galactic foreground amplitudes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3af1bf66-b13d-42cd-b4fc-c2bf92d49b8a",
   "metadata": {},
   "source": [
    "### Metropolis and Metropolis-Hastings Methods\n",
    "\n",
    "The Metropolis method is by far the most well-known of all of the MCMC sampling methods due to it's ease of implementation, tunability, and effectiveness. The Metropolis and Metropolis-Hastings algorithms are effective as all samples with a higher likelihood are accepted, while some samples with lower likelihoods will be accepted from time to time (this can be helpful to avoid becoming stuck around a local likelihood peak). \n",
    "\n",
    "Let us say we wish to sample for some parameter $\\omega$. Let $\\omega_j$ represent te $j$th sample of $\\omega$ in a Markov chain. We also need to define a stochastic proposal probablility distribution, $T(\\omega_{j+1}\\mid \\omega_{j})$, with which we tune in order to help with parameter convergence. In the Metropolis method, the proposal distribution is symmetric, i.e., $T(\\omega_{j+1}\\mid \\omega_{j})=T(\\omega_{j}\\mid \\omega_{j+1})$. In the Metropolis-Hastings method, this proposal distribution need not be symmetric.\n",
    "\n",
    "The next thing we need to define is the acceptance criterion, which tells us which samples we will accept and which ones we will not. This acceptance criterion is defined by the ratio of the densities, which is given by the ratio of the posterior probabilities as defined by Bayes' theorem ($\\frac{P(\\boldsymbol d \\mid \\omega_{j+1})}{P(\\boldsymbol d \\mid \\omega_{j})}$). If a prior is used to inform our sampling, this ratio is weighted by the ratio of the prior probabilities. Therefore our total definition of the acceptance probability $q$ is defined as\n",
    "\n",
    "\\begin{equation}\n",
    "q = \\frac{P(\\boldsymbol d \\mid \\omega_{j+1})}{P(\\boldsymbol d \\mid \\omega_{j})}\\frac{P(\\omega_{j+1})}{P(\\omega_j)}.\n",
    "\\end{equation}\n",
    "\n",
    "With these defined, we can now outline the algorithm itself. The algorthim is given by:\n",
    "\n",
    "1. Initialie chain at some parameter value $\\omega_0$.\n",
    "2. Draw a random sample from the proposal probability distribution, i.e., $\\omega_{j+1} \\leftarrow T(\\omega_{j+1}\\mid\\omega_j)$.\n",
    "3. Compute the acceptance probability $q$.\n",
    "4. Draw a random number $\\eta$ from a uniform distribution $U[0,1]$.\n",
    "5. Accept the proposed $\\omega_{j+1}$ if $\\eta < q$. Otherwise, set $\\omega_{j+1} = \\omega_{j}$.\n",
    "6. Repeate steps 2-5 until convergence.\n",
    "\n",
    "Aside:\n",
    "Tuning of the Metropolis and Metropolis-Hastings step size is vital to properly exploring the distribution within a realistic number of samples. Having a step size which is too small can make exploring the full distribution too expensive as many steps are accepted without the sampled values moving much. When the step size is too large, very few samples are accepted, wasting a large amount of time drawing samples, again not exploring the distribution well.\n",
    "\n",
    "Tuning the step size though testing to accept between 20 and 80 % of the samples will provide you a sampling procedure that allows for good exploration of the underlying distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b5b8281-ffae-401e-9399-a2de145ef28f",
   "metadata": {},
   "source": [
    "### Inversion sampling\n",
    "Now that we have sampled the two amplitude parameters, we want to sample spectral parameters, which do not depend linearly on the data $d$, and are not Gaussian distributions, hence demand a different approach. \n",
    "To sample these two conditional distributions, we employ the \"inversion sampler\", which is completely general and works for all univariate distributions.\n",
    "In our case, we have the distribution\n",
    "\\begin{align}\n",
    "  P(\\beta\\mid d, a) &\\propto P(d\\mid a, \\beta) P(\\beta)\\\\\n",
    "  &\\propto \\left[\\prod_{\\nu}\n",
    "    \\mathrm \\exp\\left(-\\frac{1}{2}\\left(d_{\\nu}-\\mathsf{T}(\\beta)a\\right)^t \\mathsf{N}_{\\nu}^{-1}\\left(d_{\\nu}-\\mathsf{T}(\\beta)a\\right)\\right)\\right]P(\\beta).\n",
    "\\end{align}\n",
    "\n",
    "Using this for $P(\\beta)$ we can sample using the inversion sampler by following the procedure:\n",
    "1. Compute $P(\\beta)$ over a grid in $\\beta$, making sure to probe the tails to sufficient accuracy.\n",
    "2. Compute the cumulative probability distribution, ${F(\\beta) = \\int_{-\\infty}^{\\beta} P(\\beta')\\,\\mathrm d\\beta'}$.\n",
    "3. Draw a random uniform variate, $\\eta \\sim U[0,1]$.\n",
    "4. Solve the nonlinear equation $\\eta = F(\\beta)$ for $\\beta$.\n",
    "\n",
    "Note that we can extend this sampling scheme for any parameter, amplitude or spectral parameter. To do so, we simply grid out the distrubution $P(x)$ for parameter $x$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a76edaf-61db-4af0-a555-63f62693ded6",
   "metadata": {},
   "source": [
    "### Sampling a Gaussian\n",
    "We may write $d = \\mathsf{T}a+n = f(\\omega)\\cdot A + n$ where $\\omega$ is the set of spectral indices (such as $\\beta_{\\rm s}$ or $T_{\\rm d}$ for dust).\n",
    "If we rewrite this and assume that the noise is Gaussian $n = d - \\mathsf{T}a$ we can express the distribution as\n",
    "\\begin{align}\n",
    "P(a\\mid d, \\omega\\setminus a) &\\propto P(d\\mid\\omega)P(a)\\\\\n",
    "  &\\propto P(d\\mid a)P(a)\\\\\n",
    "  &\\propto \\left[\\prod_{\\nu} \\exp\\left(-\\frac{1}{2}(d_{\\nu}-\\mathsf{T}_{\\nu}a)^t \\mathsf{N}_{\\nu}^{-1}(d_{\\nu}-\\mathsf{T}_{\\nu}a)\\right)\\right]\\cdot P(a)\n",
    "\\end{align}\n",
    "where $P(a)$ is a prior on the amplitude. This equation which can be sampled by solving for $a$ in \n",
    "$$\\biggl(S^{-1} + \\sum_{\\nu}\\mathsf{T}^t_{\\nu}\\mathsf{N}_{\\nu}^{-1}\\mathsf{T}_{\\nu}\\biggr)\\,a = \\sum_\\nu \\mathsf{T}_{\\nu}^t \\mathsf{N}_\\nu^{-1}d_{\\nu}+\\sum_\\nu \\mathsf{T}_{\\nu}^t \\mathsf{N}_\\nu^{-1/2}m + \\sum_{\\nu}\\mathsf{T}_{\\nu}^t\\mathsf{N}_{\\nu}^{-1/2}\\eta + \\mathsf{S}^{-1/2}\\eta_{0},$$\n",
    "where $\\mathsf{S}$ is a prior standarad deviation, $\\mathsf{N}$ is the noise, $m$ is the prior mean, and $\\eta$ and $\\eta_0$ are random normal variates (i.e. $\\eta \\leftarrow \\mathcal{N}(0,1)$). For a full explanation on this, look at [BeyondPlanck 1](https://arxiv.org/abs/2011.05609), Appendix 2. For a proper breakdown of what is going on here, check out [this blog post](https://dncnwtts.github.io/blg/docs/journal/sampling-gaussians.html) by Duncan Watts.\n",
    "\n",
    "Note here that the simplified case (no prior) \n",
    "\n",
    "\\begin{equation}\n",
    "\\biggl(\\sum_{\\nu}\\mathsf{T}^t_{\\nu}\\mathsf{N}_{\\nu}^{-1}\\mathsf{T}_{\\nu}\\biggr)\\,a = \\sum_\\nu \\mathsf{T}_{\\nu}^t \\mathsf{N}_\\nu^{-1}d_{\\nu} + \\sum_{\\nu}\\mathsf{T}_{\\nu}^t\\mathsf{N}_{\\nu}^{-1/2}\\eta.\n",
    "\\end{equation}\n",
    "\n",
    "This equation is identical to the equation at the end of the previous section with an additional sampling term.\n",
    "\n",
    "We can also think of the above equation as \n",
    "\n",
    "\\begin{equation}\n",
    "a = \\biggl(\\sum_{\\nu}\\mathsf{T}^t_{\\nu}\\mathsf{N}_{\\nu}^{-1}\\mathsf{T}_{\\nu}\\biggr)^{-1}\\sum_\\nu \\mathsf{T}_{\\nu}^t \\mathsf{N}_\\nu^{-1}d_{\\nu} + \\biggl(\\sum_{\\nu}\\mathsf{T}^t_{\\nu}\\mathsf{N}_{\\nu}^{-1}\\mathsf{T}_{\\nu}\\biggr)^{-1}\\sum_{\\nu}\\mathsf{T}_{\\nu}^t\\mathsf{N}_{\\nu}^{-1/2}\\eta_{\\nu},\n",
    "\\end{equation}\n",
    "which is nice for the most simple cases.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a5a5775-bded-4f88-ac4b-02b69f737073",
   "metadata": {},
   "source": [
    "### Task 3: \n",
    "- With the new data below, again make new versions of the `return_log_likelihood` and `solve_for_ML` functions\n",
    "    (Hint: for our new version of `solve_for_ML`, let's solve the equation $\\biggl(\\sum_{\\nu}\\mathsf{T}^t_{\\nu}\\mathsf{N}_{\\nu}^{-1}\\mathsf{T}_{\\nu}\\biggr)\\,a = \\sum_\\nu \\mathsf{T}_{\\nu}^t \\mathsf{N}_\\nu^{-1}d_{\\nu}$).\n",
    "- Create implementations for each of these sampling algorithms to sample for the amplitude of the data set defined in the cell below (no priors).\n",
    "- Plot chains for each of these sampling algorithms (Overplot the input amplitude and the ML solution).\n",
    "- Plot histograms of the samples for each of the algorithms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "306e6110-653e-4877-b76b-73126fe0e0bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "nbands      = 4\n",
    "newamp      = 7.5\n",
    "frequencies = np.asarray([10.,40.,70.,100.])\n",
    "ref_freq    = 40.\n",
    "noise       = np.asarray([0.75, 1.25, 2.0, 0.5])\n",
    "\n",
    "data        = newamp*gen_spectrum(frequencies,ref_freq)+np.random.normal(0.0,noise)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "607cc6ff-e251-4a35-a390-5a490a6bb147",
   "metadata": {},
   "source": [
    "# Disclaimer\n",
    "Below is some skeleton code that may be confusing. The goal below is to define some functions that will be re-usable later in the notebook. We will define the arguments as follows:\n",
    "\n",
    "- data:        This is quite clear, this is the data which we are fitting the model to\n",
    "- noise:       The uncertainty in the data\n",
    "- spectrum:    The rule defined for our spectral behavior. For our first example, this is given by `general_spectrum`, or a power-law\n",
    "- amplitude:   The amplitude applied to our spectrum - necessary to evaluate how good our model fits\n",
    "- frequencies: The frequencies being used in our model fit, to be put into our spectrum\n",
    "\n",
    "This all being said, one can define their likelihood function as they please. For example, one could just evaluate the log likelihood by passing in the data, the noise, and the model with which we are evaluating. There is no single way to do this correctly! Feel free to be creative!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1dfc7c8-1623-45ca-9b70-0955d9b49fe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_log_likelihood(data,noise,spectrum,amplitude,frequencies):\n",
    "    return lnl\n",
    "\n",
    "def metropolis_sample(data,noise,spectrum,amplitude,frequencies,niter):\n",
    "    samples    = np.zeros(niter)\n",
    "    amp_sample = amplitude # Initial value of the Metropolis chain\n",
    "    # Define the step length of the proposal distribution\n",
    "    # step_size = \n",
    "    \n",
    "    # Evaluate the inital likelihood P(data|amp_sample)\n",
    "    # lnL_old = eval_loglikelihood(data,noise,spectrum,amp_sample)\n",
    "    # We will adopt a Gaussian distribution for our proposal distribution\n",
    "    for i in range(niter):\n",
    "        sample = amp_sample + np.random.normal(0.0,step_size)\n",
    "        \n",
    "        # Evaluate the probability of new sample\n",
    "        \n",
    "        # Calculate the ratio of the likelihoods\n",
    "        \n",
    "        # Calculate acceptance probability\n",
    "        \n",
    "        # If q >= U[0,1] accept, else, try again\n",
    "        \n",
    "        # Save sample so we can look at the chain afterwards\n",
    "        samples[i] = amp_sample\n",
    "    \n",
    "    return samples\n",
    "\n",
    "def inversion_sample(data,noise,spectrum,xs,frequencies,niter):\n",
    "    samples = np.zeros(niter)\n",
    "\n",
    "    # Compute P(x) over a grid in x ^^(xs)\n",
    "    \n",
    "    # Compute the cumulative probability distribution\n",
    "    \n",
    "    for i in range(niter):\n",
    "        # draw a random number\n",
    "        eta        = np.random.uniform(0,np.max(Fx))\n",
    "\n",
    "        # Find the value of the parameter which corresponds to the random number (i.e. solve for x where eta = F(x))\n",
    "\n",
    "    return samples\n",
    "\n",
    "def gaussian_sample(data,noise,spectrum,frequencies,niter):\n",
    "    # In this single pixel, single foreground case, the matrix T_nu is given\n",
    "    # by the spectrum at each frequency, as we attempted to described previously \n",
    "        \n",
    "    # Compute the static parts of the equation (no eta)\n",
    "\n",
    "    # For each iteration draw a random normal for each entry of eta\n",
    "    for i in range(niter):\n",
    "        for j in range(nbands):\n",
    "            eta[j] = np.random.normal(0.0,1.0)    \n",
    "        # Compute the sampling portion and save the sample\n",
    "        \n",
    "    return samples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26965237-caab-42c0-ba7e-0fd8c4565b61",
   "metadata": {},
   "source": [
    "# Chapter 3 - A Step Closer to the Microwave Sky - Introducing More Foregrounds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12d0effc-b68d-4619-8747-d440c441fbc0",
   "metadata": {},
   "source": [
    "In this chapter we will simulate a simplified example of the Microwave Sky. We will consider a single-pixel example that includes the CMB, synchrotron, and thermal dust radiation, with frequencies similar to that of the Planck satellite. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20c2c30d-20b7-451d-85fc-785fe373b74d",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Simulate a Single Pixel Data Set"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e15cad8-83ae-45dd-9391-3087dbc91a3b",
   "metadata": {},
   "source": [
    "Let's begin by first importing a couple useful functions and constants. We will also introduce an array of frequencies corresponding to the *Planck* satellite band's central frequencies, as well as their corresponding polarization sensitivities. Will use these for the remainder of the exercises. You are of course free to add your own frequencies and sensitivities if you wish."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a88d032-8b87-42ef-8e3a-afe63dca1c2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constants\n",
    "kb     = 1.38e-23      # Boltzmanns constant\n",
    "Tcmb   = 2.7255        # CMB temperature\n",
    "h      = 6.626e-34     # Plancks constant\n",
    "\n",
    "#Pixel size. Small number => large pixels => little noise pr pixel\n",
    "nside  = 64\n",
    "\n",
    "#Frequency bands for Planck\n",
    "nbands = 9\n",
    "freq   = np.array([30., 44., 70., 100., 143., 217., 353., 545., 857.])\n",
    "\n",
    "#Noise level pr pixel, pr frequency\n",
    "sigma  = np.array([1.2, 2.7, 3.2, 5.6, 2.1, 3.8, 1.3, 2.9, 1.7])*nside/512 \n",
    "\n",
    "# Unit conversion between Rayleigh–Jeans units and cmb units\n",
    "def Kcmb_to_Krj(frequency):\n",
    "    xx = h*frequency*1.0e9/(kb*Tcmb)\n",
    "    return xx**2*np.exp(xx)/((np.exp(xx)-1.)**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "374d12bb-a37d-4a9d-b2dc-a8c6dd62df3b",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Simulated sky data\n",
    "Now we can simulate our data using this sky model for a given set of frequencies. Below we set the model parameters, the set of frequencies, the level of noise for each frequency and which frequency we use as pivot frequency (reference frequency) for dust and syncrotron. \n",
    "\n",
    "We have added all the Planck frequencies for to illuminate the coverage of the experiment. We will begin with an example similar to that of our last. Synchrotron radiation is, with both physical and observation support, assume to have a power-law frequency scaling (in $\\mu \\mathrm{K_{RJ}}$ units, see Synchrotron lecture). This is essentially the same model we just ran with, with just an additional term. Explicitly the synchrotron sky signal at frequency $\\nu$, $s_{\\rm s,\\,\\nu}$ is given by \n",
    "\n",
    "\\begin{equation}\n",
    "s_{\\rm s,\\,\\nu} = A_{\\rm s} \\Big(\\frac{\\nu}{\\nu_{\\rm 0,\\,s}}\\Big)^{\\beta_{\\rm s}},\n",
    "\\end{equation}\n",
    "where $\\beta_{\\rm s}$ has typical values between -3.0 and -3.2. Let us begin by defining a simple class to store our synchrotron information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d397ad4-3706-4f6a-8047-2ee06e7c528c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class synch_comp:\n",
    "    def __init__(self,nu_ref,amplitude,beta):\n",
    "        self.par    = np.empty(2)\n",
    "        self.nu_ref = nu_ref\n",
    "        self.par[0] = amplitude\n",
    "        self.par[1] = beta\n",
    "    def spectrum(self,nu,beta=None):\n",
    "        beta = self.par[1] if beta is None else beta\n",
    "        return (nu/self.nu_ref)**beta\n",
    "    def model(self,nu):\n",
    "#         self.model = self.par[0]*self.spectrum(self.nu_ref)\n",
    "        return self.par[0]*self.spectrum(nu)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e295a7cc-92a2-4b10-9862-11adcb3e85c2",
   "metadata": {},
   "source": [
    "Now we will simulate the synchrotron sky signal across the *Planck* frequency channels, and sample the corresponding parameters (i.e. $A_{\\rm s}$ and $\\beta_{\\rm s}$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49c18042-2287-4f0a-af68-086737760a0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will adopt a reference frequency of 30 GHz as this is where we have the highest synchrotron signal-to-noise\n",
    "# We Will also take a fiducial value of 12 uK_{RJ} for the amplitude\n",
    "# Finally we will also assume that beta_s = -3.1\n",
    "\n",
    "synch = synch_comp(30.0,12.,-3.1)\n",
    "d_synch = np.empty(nbands)\n",
    "s_synch = np.empty(nbands)\n",
    "for j in range(nbands):\n",
    "    s_synch[j] = synch.par[0]*synch.spectrum(freq[j],synch.par[1])\n",
    "    d_synch[j] = synch.par[0]*synch.spectrum(freq[j],synch.par[1]) + np.random.normal(0.0,sigma[j])\n",
    "    \n",
    "plt.scatter(freq,s_synch)\n",
    "plt.scatter(freq,d_synch)\n",
    "plt.yscale('log')\n",
    "plt.xscale('log')\n",
    "\n",
    "print(d_synch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9296e9b5-3382-4e99-988a-0d2b6b07bc6f",
   "metadata": {
    "tags": []
   },
   "source": [
    "If we print out the data we can see that we quickly hit the \"noise floor\", where the magnitude of the synchrotron emission is lower than the uncertainty in our measurement. Gathering high signal-to-noise observations is vital in CMB experiment, specifically in polarization as the CMB signal is very weak. Hopefully we will see some of the difficulties this can cause us through these exercises.\n",
    "\n",
    "### Task 4:\n",
    "- Find the ML solution for $A_{\\rm s}$.\n",
    "- With the simulated data, use the Gaussian sampling method to make an estimate of $A_{\\rm s}$ given the data, and assuming $\\beta_{\\rm s} = -3.1$. Plot the chain of samples. Along with the samples, also plot the $\\chi^2 = \\sum_{\\nu}(data_{\\nu}-model_{\\nu})^2/\\sigma_{\\nu}^2$.\n",
    "- Using some intuition, try removing some frequency bands from the fit and see how it affects the ML solution and the sampling chain. Comment on how/if making data cuts affects your results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d721cbf8-6b32-45e4-8122-062ab698e155",
   "metadata": {},
   "source": [
    "### Task 5\n",
    "- Adjust the likelihood evaluation within the inversion sampler to accept arbitrary data models to fit to the data.\n",
    "- Assume that $A_{\\rm s} = 12.0$. Using the inversion sampler and all of the data, sample for $\\beta_{\\rm s}$.\n",
    "- Again, try making a data cut to see how it affects the fit."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "209b2d28-2fb4-426e-bc27-f07a3bac2f1c",
   "metadata": {},
   "source": [
    "## Gibbs Sampling \n",
    "\n",
    "Now that we have sampled two parameters separately, while holding the other constant, we can introduce Gibbs sampling. Gibbs sampling is the core of the Commander and Cosmoglobe frameworks, and is a useful tool for mapping out multi-dimensional parameter spaces. Consider a data model with some set of parameters $\\vec{\\omega} = \\{\\omega_0,..,\\omega_k\\}$. We utilize Bayes Theorem, as defined earlier, to explore this parameter space. Gibbs sampling functions by sampling each of these parameters $\\omega_{i}$ individually, while holding all of the parameters constant. Therefore, in each Gibbs chain there are $k$ steps, such that we draw samples in the following way\n",
    "\n",
    "\\begin{align}\n",
    "\\omega_{0} &\\leftarrow P(\\omega_{0}\\mid \\vec{d},\\,\\omega_{1},\\,\\omega_{2},\\cdots,\\,\\omega_{k})\\\\\n",
    "\\omega_{1} &\\leftarrow P(\\omega_{1}\\mid \\vec{d},\\,\\omega_{0},\\,\\omega_{2},\\,\\cdots,\\,\\omega_{k})\\\\\n",
    "&\\vdots\\\\\n",
    "\\omega_{k} &\\leftarrow P(\\omega_{k}\\mid \\vec{d},\\,\\omega_{0},\\,\\omega_{1},\\,\\cdots,\\,\\omega_{k-1}).\n",
    "\\end{align}\n",
    "\n",
    "Let's try a simple example here with our two synchrotron parameters. Seeing as we have simulated the data ourselves, we already know what the correct answer should be, but in practice we don't know what the true sky signal looks like *a priori*. Let's say that we think that the realy synchrotron amplitude $A_{\\rm s} = 10.0$, and that the synchrotron spectral index $\\beta_{\\rm s} = -3.0$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49e6c247-1148-407d-9d39-5e244173a39a",
   "metadata": {},
   "source": [
    "### Task 6\n",
    "- Initialize with $A_{\\rm s} = 5.0$, $\\beta_{\\rm s} = -2.5$. Remember to set your inversion sampler grid to a wider range.\n",
    "- Instead of many samples per iteration, draw just a few (3-5) samples of each parameter, then switch to the next. I.e. sample $A_{\\rm s}\\mid\\beta_{\\rm s}$, then sample $\\beta_{\\rm s}\\mid A_{\\rm s}$, and save both values, then repeat. Do this for 250 Gibbs samples. Also plot the $\\chi^2$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2a4f518-ed73-45a1-bc04-4289aff3b8e5",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Chapter 4 - Now with More Components - Expanding the Gibbs sampler\n",
    "\n",
    "### Excellent! Well done!\n",
    "\n",
    "Now we're really getting some practical stuff done! Hopefully you've gained some understanding of the theory behind sampling for parameters, and how it is done in practice as well. Now that we have the base infrastructure built, we want to expand to more and more realistic examples. Let's introduce some more foregrounds which you can create classes for and simulate a more complex data set. Let's add two more components, thermal dust emission, and the CMB. These are two good components to introduce as they are prevalent in total intensity, as well as in polarization. Recall that we are working in Rayleigh-Jeans units, and will introduce the parametric models for these components in such units. And as a consistent reminder, details on microwave sky foregrounds can be found in the online lectures.\n",
    "\n",
    "### CMB\n",
    "For the CMB radiation we have \n",
    "$$s_{\\mathrm{CMB},\\,\\nu} = a_{\\mathrm{CMB}} \\frac{x^2e^{x}}{\\left(e^{x}-1\\right)^2}$$\n",
    "where $a_{\\mathrm{CMB}}$ is the strength (amplitude) of the signal and the rest is a converstion factor between Rayleigh-Jeans units and CMB units where  $x=h\\nu/kT_0$. (For the interested reader, take a look [here](https://dncnwtts.github.io/blg/docs/journal/unit-conversion.html) to see what the heck CMB units are).\n",
    "\n",
    "### Thermal dust\n",
    "Thermal dust dominates the higher frequencies. Dust is modelled with a modified blackbody function, \n",
    "$$s_{\\mathrm{d},\\,\\nu} = a_{\\mathrm{d}}\\,\\left(\\frac{\\nu}{\\nu_{0,\\mathrm{d}}}\\right)^{\\beta_{\\mathrm{d}}+1}\\frac{e^{h\\nu_{\\mathrm{0,\\mathrm{d}}}/kT_{\\mathrm{d}}}-1}{e^{h\\nu /kT_{\\mathrm{d}}}-1}$$\n",
    "where $a_{\\mathrm{d}}$ is the strength of the signal at the reference frequency $\\nu_{0,\\mathrm{d}}$, $\\beta_{\\mathrm{d}}$ is the spectral index and $T_{\\mathrm{d}}$ is the dust temperature. \n",
    "\n",
    "## Full sky model \n",
    "The sky model $s_{\\mathrm{RJ},\\,\\nu}$ that we will use is thus given by this model \n",
    "$$s_{\\mathrm{RJ},\\,\\nu} =s_{\\mathrm{CMB},\\,\\nu} + s_{\\mathrm{s},\\,\\nu}+ s_{\\mathrm{d},\\,\\nu}$$containing a total of seven \n",
    "parameters to be sampled. Three amplitudes ($a_{\\mathrm{CMB}}$, $a_{\\mathrm{s}}$ and $a_{\\mathrm{d}}$), two spectral indexes ($\\beta_{\\mathrm{s}}$ and $\\beta_{\\mathrm{s}}$) and thermal dust ($T_{\\mathrm{d}}$). Given the total degrees of freedom, we need at least 7 frequency bands to constrain this model. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97c3a2ac-7c37-483b-8b2b-f40284e25ebc",
   "metadata": {},
   "source": [
    "### Task 7\n",
    "- Create classes for both thermal dust emission and the CMB component (Hint: The CMB is constant in $\\mathrm{K}_{\\rm CMB}$.\n",
    "- Create a set of simulated *Planck* data, with $A_{\\rm CMB} = 2.00$, and $A_{\\rm d} = 8.0$ @ 353 GHz, with $\\beta_{\\rm d} = 1.6$ and $T_{\\rm d} = 19.6 \\mathrm{K}$, as well as the appropraite noise realizations for each data band.\n",
    "- Plot the SED of the simulated data. At what frequencies do each of the three components dominate? Does the CMB dominate over the sum of foregrounds in any frequency range?\n",
    "- Add noise to the simulations, and add a plot the sum of all components compared to the noiseless SEDs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60b58910-2958-41c4-88e0-31b955a265bf",
   "metadata": {},
   "source": [
    "Now that we have simulated a more complex data model, something that better represents the real microwave sky, we can Gibbs sample the components!\n",
    "\n",
    "### Task 8\n",
    "- Gibbs sample all of the components ($A_{\\rm cmb}$, $A_{\\rm s}$, $A_{\\rm d}$, $\\beta_{\\rm s}$, $\\beta_{\\rm d}$, $T_{\\rm d}$).\n",
    "- Plot the chains for each of the parameters, as well as the residual ($\\sum_{\\nu}(d_{\\nu}-s_{\\nu})$) and the $\\chi^2$.\n",
    "- Use `corner` to make a corner plot of all of the parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "936c32f3-b080-4604-8cab-2839754c6cc6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}