{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "06ddf709",
   "metadata": {},
   "source": [
    "## Solving Gaussian Weighted Integrals using GBS\n",
    "\n",
    "In the previous talk, we discussed how Gaussian Boson Sampling (GBS) can be used to approximate the value of Gaussian-weighted integrals. We saw that there exists a large class of problems for which it outperforms plain Monte Carlo (MC) integration. Remember that MC is the best classical baseline that does not directly suffer from the curse of dimensionality. In this notebook, we will look at a specific example of an integral where GBS outperforms MC. Specifically, you will do the following:\n",
    "\n",
    "1. Compute the integral exactly.\n",
    "2. Approximate the integral using the MC estimator.\n",
    "3. Approximate the integral using the GBS-P estimator.\n",
    "4. Compare the convergence of the two approximation methods."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a38a040",
   "metadata": {},
   "source": [
    "**The Integral:** We must specify the dimension of the integral $N$, the covariance matrix $B$, and the function $f$ to integrate over. In Sec. 6 of the paper [[*1*](#references:)], the main theorems are used to find specific problems where GBS outperforms MC. It is crucial to emphasize that this is not a forced example; the space where GBS outperforms MC is generally large [[*2*](#references:)]! Let $N = 3$ and choose the covariance matrix to be:\n",
    "\n",
    "$$ B = \\begin{bmatrix} 0.3421 & 0.3364 & 0.3244 \\\\ 0.3364 & 0.3392 & 0.3225 \\\\ 0.3244 & 0.3225 & 0.3520 \\end{bmatrix}. $$\n",
    "\n",
    "We choose to integrate over a function $f$ with a maximal degree of $2 \\cdot 15 = 30$:\n",
    "\n",
    "$$f(x) = \\sum_{k=0}^{15} \\sum_{|I|=2k} a_I x^{I}.$$\n",
    "\n",
    "Notice how we only care about even-degree terms, since all odd Gaussian moments vanish. The coefficients $a_I$ are defined as in Sec. 6.3 of [[*1*](#references:)], but for the sake of simplicity, we will omit the mathematical details here. Instead, the logic for the coefficients $a_I$ has been allocated to the following code cell. When you run it, it will find all indices with $a_I \\neq 0$ and pre-compute $a_I$, $\\text{Haf}(B_I)$, and $P(I)$. This will allow us to scale the simulations to larger sample sizes than what would otherwise be possible during a 30-minute hands-on exercise session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1117d180",
   "metadata": {},
   "outputs": [],
   "source": [
    "# All libraries used throughout notebook. You might need to 'pip install thewalrus'.\n",
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from thewalrus import hafnian_repeated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5db1224a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pre-defined constants of the problem; the covariance matrix B, the degree K, and the number d\n",
    "# appearing in the GBS distribution, which is determined by the eigenvalues of B.\n",
    "B = np.array([[0.3421, 0.3364, 0.3244], \n",
    "     [0.3364, 0.3392, 0.3225], \n",
    "     [0.3244, 0.3225, 0.3520]])\n",
    "d = 1/223.7037\n",
    "K = 15\n",
    "\n",
    "# Find all patterns I with a_I != 0. For each of them, compute a_I, Haf(B_I), and P(I).\n",
    "patterns = [[0, 0, 0]]\n",
    "coeffs = [1] # Definition of a_0\n",
    "hafnians = [1]\n",
    "probabilities = [d]\n",
    "for k in range(1, K+1):\n",
    "    # Generate all multi-indices I = [i, j, l] such that i + j + l = 2k\n",
    "    for i in range(2*k + 1):\n",
    "        for j in range(2*k - i + 1):\n",
    "            l = 2*k - i - j\n",
    "            I = [i, j, l]\n",
    "            \n",
    "            # Compute m_k.\n",
    "            N = 3\n",
    "            s_k = (2 * k) // N\n",
    "            r_k = (2 * k) % N\n",
    "            m_k = (math.factorial(s_k)**N) * ((s_k + 1)**r_k)\n",
    "\n",
    "            # Compute factorial of I.\n",
    "            I_factorial = math.factorial(I[0]) * math.factorial(I[1]) * math.factorial(I[2])\n",
    "            \n",
    "            # Check a_I != 0.\n",
    "            if I_factorial != m_k:\n",
    "                continue\n",
    "            patterns.append(I)\n",
    "            \n",
    "            # Compute a_I.\n",
    "            coeffs.append((np.sqrt(k) * (1.4368**k)) / (math.comb(N, r_k) * math.factorial(k)))\n",
    "            \n",
    "            # Compute Haf(B_I).\n",
    "            hafnian = hafnian_repeated(B, I)\n",
    "            hafnians.append(hafnian)\n",
    "             \n",
    "            # Compute P(I).property\n",
    "            probabilities.append((d/I_factorial)*abs(hafnian)**2)\n",
    "             \n",
    "# Note: All patterns with |I| > 2K will have a_I = 0. We therefore collect\n",
    "# the probability of all such outcomes into one [2K+1, 0 , 0], since it does not\n",
    "# matter for the computation which one we drew. This simplifies computational complexity.\n",
    "patterns.append([2*K+1,0,0])\n",
    "probabilities.append(1-sum(probabilities))\n",
    "coeffs.append(0) # Index has a_I = 0\n",
    "hafnians.append(0) # Contribution vanishes anyways"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "870b93ac",
   "metadata": {},
   "source": [
    "**1. Exact Computation:**\n",
    "\n",
    "We will use Wick's theorem to express the integral as a sum of Hafnians, each of which we compute classically, yielding the exact value of the integral. Notice that this approach is not viable for large problems due to the #P-hard complexity of computing Hafnians. By using Wick's theorem and the linearity of the integral, we get:\n",
    "\n",
    "$$\\mu_{\\text{Haf}} = \\int_{\\mathbb{R}^N} \\left(\\sum_{k=0}^K \\sum_{|I|=2k} a_Ix^I\\right) g(x) dx = \\sum_{k=0}^K \\sum_{|I|=2k} a_I \\int_{\\mathbb{R}^N} x^I g(x) dx = \\sum_{k=0}^K \\sum_{|I|=2k} a_I \\text{Haf}(B_I).$$\n",
    "\n",
    "This is the expression that will be used to compute the integral exactly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e343d3d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Sum over contributions for I in the precomputed 'patterns' with non-zero a_Is.\n",
    "# Remember: The hafnians Haf(B_I), the coefficients a_I, and the patterns with non-zero\n",
    "# a_I are precomputed and stored in 'hafnians', 'coeffs', and 'patterns' respectively.\n",
    "exact_value = 0\n",
    "\n",
    "# Uncomment the code below to print the exact value.\n",
    "# print(f\"Exact value of integral: {exact_value:.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e7e8f80",
   "metadata": {},
   "source": [
    "**2. The MC Estimator:**\n",
    "\n",
    "Here we implement our classical baseline, the plain Monte Carlo estimator. The covariance matrix $B$ defines our multivariate normal distribution centered at zero. Your task is to draw i.i.d. samples $X_1, X_2, \\dots, X_n$ from this distribution, evaluate the provided polynomial function $f(X_i)$ for each sample, and calculate the sample mean:\n",
    "\n",
    "$$\\varepsilon_n^{\\text{MC}} = \\frac{1}{n} \\sum_{i = 1}^n f(X_i).$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ad7bce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(xs):\n",
    "    \"\"\"\n",
    "    Input: list of points x = [x_1, x_2, ...] in R^3. Output is a list of evaluations [f(x_1), f(x_2), ...].\n",
    "    \"\"\"\n",
    "    # TODO: Evaluate f(x) at each x in 'xs'. For fixed x, iterate over patterns I with a_I != 0\n",
    "    # and add the contribution a_I * x_{I[0]}^{I[0]} * x_{I[1]}^{I[1]} * x_{I[2]}^{I[2]} to the total sum.\n",
    "    # Hint: Vectorize computation over 'xs' using numpy for efficiency.\n",
    "    evaluations = np.zeros(xs.shape[0])\n",
    " \n",
    "    return evaluations\n",
    "\n",
    "def mc_estimator(num_samples):\n",
    "    \"\"\"\n",
    "    Approximates the integral using standard Monte Carlo sampling.\n",
    "    \"\"\"\n",
    "    # TODO: Draw `num_samples` from a multivariate normal distribution \n",
    "    # with mean 0 and covariance matrix B.\n",
    "    # Hint: Use np.random.multivariate_normal\n",
    "    # samples = ... \n",
    "    \n",
    "    # TODO: Evaluate the function `f(x)` for each sample and return the mean.\n",
    "    # evaluations = ...\n",
    "    \n",
    "    # return np.mean(evaluations)\n",
    "    return 0.0 # Placeholder\n",
    "\n",
    "# Test your estimator. How does it compare to the exact value?\n",
    "# print(f\"MC Estimate (100.000 samples): {mc_estimator(100000):.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b50c2e77",
   "metadata": {},
   "source": [
    "**3. The GBS-P Estimator:**\n",
    "\n",
    "Here we implement the GBS-P estimator as discussed in the talk. A function `get_gbs_samples` is provided to simulate the output of a GBS device for our specific problem. Your task is to draw i.i.d samples $I_1, \\dots, I_n$ from this GBS simulator and compute the GBS-P estimator using those samples:\n",
    "\n",
    "$$ \\varepsilon_n^{\\text{GBS-P}} = \\sum_J a_J \\sqrt{\\frac{J!}{d}} \\sqrt{\\frac{S_n^{(J)}}{n}}, \\quad S_n^{(𝐽)} = \\text{\\# occurences of } J $$\n",
    "\n",
    "where $d$, $a_J$, and $J!$ has been pre-computed for all patterns $J$ with $a_J \\neq 0$. Notice that you only need to sum over unique patterns you have observed, since any pattern you have not observed will have $S_n^{(J)} = 0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a13a5f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_gbs_samples(num_samples):\n",
    "    \"\"\"\n",
    "    Generates simulated GBS samples using The Walrus.\n",
    "    Uses the input covariance matrix to program the GBS \n",
    "    device as discussed during the talk.\n",
    "    \"\"\"\n",
    "    \n",
    "    sampled_indices = np.random.choice(len(patterns), size=num_samples, p=probabilities)\n",
    "    samples = [patterns[idx] for idx in sampled_indices]\n",
    "\n",
    "    return samples\n",
    "\n",
    "def gbs_p_estimator(num_samples):\n",
    "    \"\"\"\n",
    "    Approximates the integral using the GBS-P estimator.\n",
    "    \"\"\"\n",
    "    \n",
    "    # 1. Fetch the samples from the (simulated) GBS device\n",
    "    samples = get_gbs_samples(num_samples)\n",
    "    \n",
    "    # TODO: Evaluate the GBS-P estimator derived in the talk.\n",
    "    # (1) Find number of occurences S_n^{(J)} of each unique sample J.\n",
    "    # (2) Compute the factorial J!.\n",
    "    # (3) Retrieve the coefficient a_J.\n",
    "    # (4) Sum over contributions from unique patterns J. \n",
    "    # Hint: Use np.unique() to find unique samples and their counts.\n",
    "    estimate = 0\n",
    "        \n",
    "    return estimate\n",
    "\n",
    "# Test your estimator. How does it compare to the exact value?\n",
    "# print(f\"GBS-P Estimate (100.000 samples): {gbs_p_estimator(100000):.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82451e3f",
   "metadata": {},
   "source": [
    "**4. Convergence of Methods:**\n",
    "\n",
    "Let's compare the sample efficiency of both estimators. We will compute the estimates for increasing sample sizes and plot the convergence. Notice how the scaling behavior differs between the two methods. For comparisons with larger sample sizes, we refer to Fig. 3 in the paper [[*1*](#references:)]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fa902a3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nplt.figure(figsize=(6, 5))\\n\\nplt.axhline(6.8134, color=\\'grey\\', linestyle=\\'--\\', linewidth=2, label=\\'Ground Truth\\', zorder=1)\\nplt.plot(sample_sizes, mc_results, label=\\'MC\\', color=\"skyblue\", marker=\\'o\\', linewidth=2, zorder=2)\\nplt.plot(sample_sizes, gbs_results, label=\\'GBS-P\\', color=\"tomato\", marker=\\'s\\', linewidth=2, zorder=3)\\n\\nplt.xlabel(\\'Number of Samples\\', fontsize=12)\\nplt.xticks(np.arange(0, 400001, 100000, dtype=int))\\nplt.ticklabel_format(style=\\'scientific\\', axis=\\'x\\', scilimits=(0, 0))\\nplt.ylabel(\\'Estimated Value of Integral\\', fontsize=12)\\nplt.title(\\'Convergence Comparison: MC vs GBS-P\\', fontsize=14)\\nplt.legend(fontsize=11, loc=\\'best\\')\\nplt.grid(True, which=\"both\", ls=\"--\", alpha=0.5, zorder=0)\\n\\nplt.tight_layout()\\nplt.show()\\n'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_sizes = np.arange(1000, 401001, 100000, dtype=int)\n",
    "mc_results = []\n",
    "gbs_results = []\n",
    "\n",
    "# TODO: Loop over `sample_sizes`, compute the estimate for both methods, \n",
    "# and append the results to the lists above.\n",
    "\n",
    "# TODO: Uncomment the code below to plot the results once the lists are populated.\n",
    "\"\"\"\n",
    "plt.figure(figsize=(6, 5))\n",
    "\n",
    "plt.axhline(exact_value, color='grey', linestyle='--', linewidth=2, label='Ground Truth', zorder=1)\n",
    "plt.plot(sample_sizes, mc_results, label='MC', color=\"skyblue\", marker='o', linewidth=2, zorder=2)\n",
    "plt.plot(sample_sizes, gbs_results, label='GBS-P', color=\"tomato\", marker='s', linewidth=2, zorder=3)\n",
    "\n",
    "plt.xlabel('Number of Samples', fontsize=12)\n",
    "plt.xticks(np.arange(0, 400001, 100000, dtype=int))\n",
    "plt.ticklabel_format(style='scientific', axis='x', scilimits=(0, 0))\n",
    "plt.ylabel('Estimated Value of Integral', fontsize=12)\n",
    "plt.title('Convergence Comparison: MC vs GBS-P', fontsize=14)\n",
    "plt.legend(fontsize=11, loc='best')\n",
    "plt.grid(True, which=\"both\", ls=\"--\", alpha=0.5, zorder=0)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e958462",
   "metadata": {},
   "source": [
    "**Extra Exercise:** If you managed to vectorize the computation of $f(x)$, we can scale this up to larger sample sizes to make the convergence more apparent. In the code below, you should see the GBS-P estimator converge smoothly toward the exact value, while the MC estimator still jumps around."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae36bee6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_sizes = np.arange(1000, 1201001, 100000, dtype=int)\n",
    "mc_results = []\n",
    "gbs_results = []\n",
    "\n",
    "# TODO: Loop over `sample_sizes`, compute the estimate for both methods, \n",
    "# and append the results to the lists above.\n",
    "\n",
    "# TODO: Uncomment the code below to plot the results once the lists are populated.\n",
    "'''\n",
    "plt.figure(figsize=(6, 5))\n",
    "\n",
    "plt.axhline(exact_value, color='grey', linestyle='--', linewidth=2, label='Ground Truth', zorder=1)\n",
    "plt.plot(sample_sizes, mc_results, label='MC', color=\"skyblue\", marker='o', linewidth=2, zorder=2)\n",
    "plt.plot(sample_sizes, gbs_results, label='GBS-P', color=\"tomato\", marker='s', linewidth=2, zorder=3)\n",
    "\n",
    "plt.xlabel('Number of Samples', fontsize=12)\n",
    "plt.xticks(np.arange(0, 1200001, 300000, dtype=int))\n",
    "plt.ticklabel_format(style='scientific', axis='x', scilimits=(0, 0))\n",
    "plt.ylabel('Estimated Value of Integral', fontsize=12)\n",
    "plt.title('Convergence Comparison: MC vs GBS-P', fontsize=14)\n",
    "plt.legend(fontsize=11, loc='best')\n",
    "plt.grid(True, which=\"both\", ls=\"--\", alpha=0.5, zorder=0)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1526988b",
   "metadata": {},
   "source": [
    "### References:\n",
    "\n",
    "[1]: Jørgen Ellegaard Andersen, & Shan Shan. \"Using Gaussian Boson Samplers to Approximate Gaussian Expectation Problems.\" arXiv preprint arXiv:2502.19336 (2025).\n",
    "\n",
    "[2]: Jørgen Ellegaard Andersen, & Shan Shan. \"Estimating the Percentage of GBS Advantage in Gaussian Expectation Problems.\" arXiv preprint arXiv:2502.19362 (2025)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.13.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}