{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ROEehen5W770"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sy63eV30ehE3"
   },
   "source": [
    "# Exercise 1: BBP Phase Transition in the Cards Game"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TyaygwNpenat"
   },
   "source": [
    "1. Write a function to generate the *cards game dataset*. This function should generate a $N$-dimensional vector $s$ with random components $s_i \\in \\{\\pm1\\}$ and return it together with the matrix $Y = \\frac{1}{\\sqrt{N}}ss^\\top + W$, where W is a symmetric noise matrix with i.i.d. entries sampled from a Gaussian distribution $N(0,\\Delta)$. Be careful when generating $W$!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lIYdn1woOS1n"
   },
   "outputs": [],
   "source": [
    "# Generate the cards game dataset\n",
    "\n",
    "def card_game_dataset(N, Delta):\n",
    "\n",
    "    # Your code here\n",
    "\n",
    "    return s, Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9TmqfdWtfK97"
   },
   "source": [
    "2. For $\\Delta \\in \\{0.001,\\, 0.01,\\, 0.1,\\, 1,\\,10\\}$, diagonalize the matrix $Y$ and plot a histogram with the spectrum. What do you observe? *Hint*: Use numpy.linalg.eigh to solve the diagonalization problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "-2n3ZOstfKMt",
    "outputId": "edeb4a00-ff47-4033-9dc6-dc14ff84fdec"
   },
   "outputs": [],
   "source": [
    "# Plot the spectrum for N = 1000, delta = {.001, .01, .1, 10, 100}\n",
    "\n",
    "N = 1000\n",
    "\n",
    "Deltas = np.logspace(-3, 1, 5)\n",
    "\n",
    "for Delta in Deltas:\n",
    "\n",
    "    # your code here\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1qL2HNCjfeA-"
   },
   "source": [
    "3. We are going to infer the cards' values $s_i$ using the spectral estimator $\\hat{s}_i = \\text{sign}(e_i^{\\max})$, with $e_i^{\\max}$ the $i$-th component of eigenvector corresponding to the largest eigenvalue. Compute the alignment $$Q = {\\rm abs}\\left(\\frac{1}{N} \\sum_{i=1}^N s_i \\hat{s}_i \\right).$$\n",
    "Notice the absolute value in order to take into account the symmetry $s \\to -s$ in the problem. Compute $Q$ for different values of $\\Delta$ between 0 and 3. Average over 5 realizations for each considered value of $\\Delta$, and plot $Q$ vs $\\Delta$.  You should get a hint that there is a phase transition at a certain value of $\\Delta$ (which one?). How is the spectrum of Y related to the recovery of the two groups?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Deltas = np.linspace(0, 3, 20)\n",
    "\n",
    "Q = []\n",
    "exps = 5\n",
    "\n",
    "for Delta in Deltas:\n",
    "    #YOUR CODE HERE\n",
    "\n",
    "    Q.append(q.mean())\n",
    "\n",
    "plt.plot(Deltas, Q, 'o')\n",
    "plt.xlabel('$\\Delta$')\n",
    "plt.ylabel('Q')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise 2: Cards Game with Monte Carlo\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Recovering $s^*$\n",
    "\n",
    "The objective is to reconstruct the initial signal $ s^*$ via Monte Carlo Sampling.\n",
    "\n",
    "We consider a slightly different setting, where our ground truth signal is again $s^* =(s^*_1,\\dots,s^*_N)$, $s^*_i=\\pm 1$ with equal probability.\n",
    "Then we build the matrix $Y_{ij}=\\frac{s_i^*s_j^*}{\\sqrt{N}}+\\sqrt{\\Delta}W_{ij}$, for $i\\neq j$,  and  $Y_{ii}=0\\;\\, \\forall\\; i\\in[N]$, where $W_{ij}\\sim N(0,1)$ come from a gaussian distribution and are all independent over $i<j$. Moreover $W_{ij}=W_{ji}$ so that $Y$ is symmetric.\n",
    "$\\{Y_{ij}\\}_{i\\in [N],j\\in[N]}$ are our data.\n",
    "\n",
    "Note that the difference is in the $Y_{ii}$ set to zero!\n",
    "\n",
    "You saw in the lectures that the estimator $\\hat s(Y)$ that minimizes the expected (over the posterior) square loss\n",
    "\n",
    "$$\n",
    "\\text{MSE}=\\mathbb E_{s^*|Y}\\left[\\frac{1}{N}\\sum_i \\left(\\hat s(Y)_i-s^*_i\\right)^2\\right]\n",
    "$$\n",
    "\n",
    "is the Bayes estimator (or minimum mean square error estimator (mmse)) given by the mean of the posterior\n",
    "\n",
    "$$\n",
    "\\hat s_{mmse}(Y)_i=\\mathbb E_{s|Y} \\left[ s_i\\right].\n",
    "$$\n",
    "\n",
    "### Connection to systems in statistical physics\n",
    " Now consider a system of $N$ binary spins $s_1,\\dots,s_N$, whose joint probability distribution is given by the posterior\n",
    " $$\n",
    " P(s|Y)=\\frac{1}{Z(Y)} e^{-H(s;Y)}, \\qquad H(s;Y)=-\\frac{1}{\\Delta\\sqrt{N}}\\sum_{i<j}Y_{ij}s_is_j\n",
    " $$\n",
    " Then the Bayes estimator is nothing but the magnetizations of the spins $s_i$.\n",
    " Denoting the average with respect to $P(s|Y)$ with $\\langle\\cdot\\rangle$ we   can write\n",
    " $$\n",
    "m_i=\\langle s_i \\rangle=\\hat s_{mmse}(Y)_i.\n",
    " $$\n",
    "\n",
    "\n",
    "For the single configuration $s$ we define its overlap with the ground truth $s^*$ as $\\text{ov}(s,s^*)=\\frac{1}{N}\\sum_i s_is^*_i$\n",
    "\n",
    " In order to compute the magnetizations we implement a Markov Chain Monte Carlo algorithm that samples from the posterior.\n",
    "\n",
    " First we define the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generating the planted signal (or ground truth)\n",
    "N = 128  # Number of spins\n",
    "Delta = 0.5  # Noise variance\n",
    "\n",
    "np.random.seed(1234)\n",
    "\n",
    "s_star, Y = card_game_dataset(N, Delta)\n",
    "\n",
    "s_star = s_star.squeeze() # to remove all dimensions of size 1 from the array\n",
    "\n",
    "for i in range(N): # Set to zero the diagonal elements\n",
    "    Y[i, i] = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MCMC Implementation\n",
    "\n",
    "We're given $Y$ and we want to sample from the posterior $P(s|Y)$.\n",
    "MCMC methods are based on designing a Markov chain whose stationary distribution is the distribution we want to sample from.\n",
    "A Markov chain is a stochastic process consisting in a sequence of random variables $\\{s^{(t)}\\}_{t\\in\\mathbb N}$.\n",
    "We will initialize the Markov chain on a random state and then update its state $s^{(t)}$.\n",
    "There are many different Markov chains that have the same stationary probability. \n",
    "Here we will first implement the Metropolis rule for updating the state $s^{(t)}$.\n",
    "\n",
    "### Metropolis\n",
    "\n",
    "Let $s^{(t)}$ be the current state fo the Markov chain.\n",
    "We draw an index $i\\sim Unif([N])$ uniformly at random and propose to flip the spin $s_i$.\n",
    "Let $s'=(s^{(t)}_1,\\dots,s^{(t)}_{i-1},-s^{(t)}_i,s^{(t)}_{i+1},\\dots,s^{(t)}_N)$ be the new proposed state.\n",
    "We then compute the acceptance probability $P_{acc}(s^{(t)},s')=\\min\\{1, P(s'|Y)/P(s^{(t)}|Y)\\}$.\n",
    "We accept the proposal ($s^{(t+1)}=s'$) with probability $P_{acc}(s^{(t)},s')$. In practice, we draw a random variable $u$ uniformly distributed in $[0,1]$ and if $u<P_{acc}(s^{(t)},s')$, we set $s^{(t+1)}=s'$, otherwise we set $s^{(t+1)}=s^{(t)}$.\n",
    "\n",
    "\n",
    "### Special case of Boltzmann measures\n",
    "When we want to sample from a probability of the form $P(s) \\propto e^{-H(s)}$, the acceptance probability can be written as $P_{acc}(s^{(t)},s')=\\min\\{1, \\exp(H(s^{(t)})-H(s'))\\}$.\n",
    "In the particular case of the planted spin glass we have \n",
    "$$\n",
    "    H(s')-H(s^{(t)})=\\frac{2 s_i^{(t)}}{\\Delta\\sqrt{N}}\\sum_k Y_{ik} s^{(t)}_k \n",
    "$$\n",
    "Notice that the Hamiltonian is invariant under the transformation $s\\to-s$, so when running the MCMC we might end up correlating either with $s$ or with $-s$. As a consequence $m$ might either be positive or negative.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def energy_difference(s, i, Y, Delta):\n",
    "    '''\n",
    "    Arguments:\n",
    "    (s) current state, (i) flipped index, (Y) measurement matrix, (Delta) noise variance.\n",
    "    Returns: \n",
    "    H(s')-H(s^t), where s' is obtained by flipping spin i in s^t.\n",
    "    '''\n",
    "    return # YOUR code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uniformly random initial configuration to start the MCMC\n",
    "\n",
    "np.random.seed(0)\n",
    "s0 = 2*np.random.binomial(n=1, p=0.5, size=[N])-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here are some quantities we want to keep track of during the dynamics\n",
    "\n",
    "overlaps = [] # Keeping track of the overlap through time\n",
    "energies = [-np.einsum('ij,i,j->', Y, s0, s0)/(2*np.sqrt(N)*Delta)] # Keep track of all the states visited by the Markov chain\n",
    "state_history = [np.copy(s0)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 10000\n",
    "s = np.copy(s0) # Current state\n",
    "np.random.seed(10)\n",
    "for t in range(T):\n",
    "\n",
    "    # Returns a uniformly distributed integer between 0 and N-1, extremes included\n",
    "    i = np.random.randint(low=0, high=N)\n",
    "    delta_H = energy_difference(s, i, Y, Delta)\n",
    "    # YOUR code here (append the energy values in the array energies)\n",
    "    prob_accept = ...\n",
    "\n",
    "    overlaps.append(np.mean(s*s_star))\n",
    "    state_history.append(np.copy(s))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once your Monte Carlo finally runs, make some plots by running the cells you see below (up to the section \"Computing averages with MCMC\").\n",
    "\n",
    "You will plot overlaps, energies and configurations as a function of time. Are you seeing what you would expect?\n",
    "\n",
    "If you have time you can change the value of $\\Delta$ (for example to something larger than 1) and look for the differences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the overlap against the time\n",
    "\n",
    "plt.figure(figsize=(12, 9))\n",
    "plt.plot(overlaps)\n",
    "plt.xlabel(\"t\")\n",
    "plt.ylabel(\"overlap\")\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the energy against the time\n",
    "plt.figure(figsize=(12, 9))\n",
    "plt.plot(energies)\n",
    "plt.xlabel(\"t\")\n",
    "plt.ylabel(\"energy (H)\")\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing averages from the MCMC\n",
    "\n",
    "Suppose we want to compute the average over the posterior of an observable $A(s)$\n",
    "To do this we use the fact that (under the ergodicity assumption) the time averages over the Markov chain converge to the ensemble averages. In formulas we have\n",
    "$$\n",
    "\\lim_{T\\to\\infty} \\frac{1}{T}\\sum_{t=0}^T A(s^{(t)})=\\langle A \\rangle=\\sum_{s} P(s|Y) A(s)\n",
    "$$\n",
    "\n",
    "In practice things are more involved.\n",
    "\n",
    "As you saw, the Monte Carlo Markov chain needs some time to thermalize, before it starts sampling from the equilibrium measure. Configurations encountered in this initial *burn in* phase are not representative of the posterior, thus they should be discarded.\n",
    "Let $T_{therm}$ be the time needed by the Markov chain to reach equilibrium, and $T$ the maximum time of the simulation.\n",
    "\n",
    "Then we approximate\n",
    "$$\n",
    "\\langle A \\rangle \\approx \\frac{1}{T-T_{therm}}\\sum_{t=T_{therm}}^T A(s^{(t)})\n",
    "$$\n",
    "\n",
    "### Obtaining the MMSE estimator\n",
    "Remember that $\\hat s_{mmse}(Y)_i=m_i=\\langle s_i\\rangle=\\mathbb E_{s|Y} \\left[ s_i\\right]$ is the minimum mean square error estimator. We can approximate the average by the procedure explained above.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_history = np.stack(state_history) # to stack the sequence of states together into a single NumPy array along a new axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the mmse estimator from state_history using the procedure explained above.\n",
    "# What do you think is a reasonable value for T_therm?\n",
    "\n",
    "T_therm = 3000\n",
    "# YOUR code here\n",
    "mmse_estimator = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mmse_estimator[:20] # to check that the result makes sense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the mean squared error between the estimated signal and the true signal, considering both possible signs of the true signal, and keep the smaller one\n",
    "\n",
    "# YOUR code here\n",
    "mse = ...\n",
    "print(\"mean square error =\", mse)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimating $T_{therm}$\n",
    " How do we know when our Markov chain has thermalized? In general there is no good answer to this question, however several heuristics can be used.\n",
    " \n",
    "\n",
    "*   Look at when some observables (e.g. the energy) are stationary, i.e. their value stops drifting and fluctuates around a constant level.\n",
    "*   Initialize two independent Markov chains, one random $s^{(0)}$ independent of $s^*$, and the other at $s^*$. Run the two Markov chains in parallel: $T_{therm}$ is the time at which the observables from the two chains have roughly the same value. This approach has the drawback of requiring the knowledge of $s^*$, which we normally don't have. However in some cases we know the structure of the ordered state (e.g. Ising, where $s_i=1 \\;\\forall i$), so the method can be applied.\n",
    "\n",
    "Let's try this second approach: we run a Monte Carlo Markov chain initialized from the planted configuration $s^*$ and compare the energy and overlap.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First we save the overlaps and energies obtained with the random initialization\n",
    "\n",
    "overlaps_rand_s0 = overlaps.copy()\n",
    "energies_rand_s0 = energies.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Re-execute the MCMC code, this time initializing s0=s_star. Remember to re-initialize all the variables!\n",
    "\n",
    "s0 = np.copy(s_star)\n",
    "\n",
    "overlaps = [] # Keeping track of the overlap through time\n",
    "energies = [-np.einsum('ij,i,j->', Y, s0, s0)/(2*np.sqrt(N)*Delta)]\n",
    "state_history = [np.copy(s0)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 10000\n",
    "s = np.copy(s0)  # Current state\n",
    "\n",
    "for t in range(T):\n",
    "\n",
    "    # Returns a uniformly distributed integer between 0 and N-1, extremes included\n",
    "    i = np.random.randint(low=0, high=N)\n",
    "    delta_H = energy_difference(s, i, Y, Delta)\n",
    "    \n",
    "    # USE YOUR PREVIOUS CODE HERE\n",
    "    prob_accept = ...\n",
    "    overlaps.append(np.mean(s*s_star))\n",
    "    state_history.append(np.copy(s))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot here the overlaps from the Markov chain with random initialization and the Markov with planted initialization, respectively\n",
    "\n",
    "plt.figure(figsize=(12, 9))\n",
    "plt.plot(overlaps_rand_s0, label='random initialization')\n",
    "plt.plot(overlaps, label='planted initialization')\n",
    "plt.grid()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('overlap')\n",
    "plt.legend()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}