{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88e5466b",
   "metadata": {},
   "source": [
    "# Boson Sampling as a hardware accelerator for Monte Carlo integration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9716945",
   "metadata": {},
   "source": [
    "This notebook demonstrates how a boson sampler can be used to accelerate Monte Carlo integration <cite data-footcite=\"correa2025bosonsampling\"></cite>. For an introduction to boson sampling, please refer to the boson sampling tutorial.\n",
    "\n",
    "## Background\n",
    "\n",
    "#### Algorithm\n",
    "\n",
    "We would like to compute a quantity via Monte Carlo integration of a function $f$ over some multidimensional domain $D$:\n",
    "\n",
    "$$\n",
    "F = \\int_{{D}} \\, d \\textbf{X} \\, f (\\textbf{X}).\n",
    "$$\n",
    "\n",
    "The integral can be estimated by drawing a set of samples from a probability distribution $p(\\textbf{X})$ and rewriting the integral as an estimator:\n",
    "$$\n",
    "F = \\int_{{D}} \\, d \\textbf{X} \\, \\frac{f (\\textbf{X})}{p(\\textbf{X})}p(\\textbf{X}) \\approx \\frac{1}{N} \\sum_{i=1}^{N} \\frac{f (\\textbf{X}^{(i)})}{p(\\textbf{X}^{(i)})}\n",
    "$$\n",
    "with $\\textbf{X}^{(i)} \\sim p(\\textbf{X})$.\n",
    "\n",
    "Then, we consider integrands that may be factorized as $f (\\textbf{X}) = g (\\textbf{X}) h (\\textbf{X}) $ such that $g (\\textbf{X})$ is a probability distribution and $h (\\textbf{X})$ encodes a quantity that we would like to be computed. In our case, $g (\\textbf{X})$ gives the sampling distribution produced by a boson sampler.\n",
    "\n",
    "#### Specific application: perturbation theory\n",
    "\n",
    "Computing special types of integrals as described above can be usefully applied in perturbation theory to study complicated quantum systems. In time-independent perturbation theory, writing the trapped-boson Hamiltonian as $H = H_0 + V$, where $H_0$ is a sum of single-particle harmonic-trap Hamiltonians whose eigenfunctions define the occupied orbitals, and $V(X)$ is a small three-body perturbing potential evaluated in position space.\n",
    "\n",
    "The first-order energy correction is then the expectation value\n",
    "$$\n",
    "E^{(1)}=\\langle \\Psi_0|V|\\Psi_0\\rangle=\\int_D |\\Psi_0(\\textbf{X})|^2\\,V(\\textbf{X})\\,d\\textbf{X},\\\n",
    "$$\n",
    "which can be estimated by sampling configurations $\\textbf{X}$ from $g(\\textbf{X})=|\\Psi_0(\\textbf{X})|^2$ using a boson sampler and then classically averaging $h(\\textbf{X})=V(\\textbf{X})$ over those samples (importance sampling).\n",
    "\n",
    "For $V(\\textbf{X})$ the three-body potential of Efimov physics is considered:\n",
    "$$\n",
    "V_{Ef}(\\textbf{X}) = - \\frac{C + 1/4}{R^2},\n",
    "$$\n",
    "where $C$ is a constant and $R^2 = \\frac{2}{3} ( r^2_{12} + r^2_{13} + r^2_{23})$, with $r_{ij} = |x_i − x_j|$ denoting pairwise distances.\n",
    "\n",
    "For this demonstration, we work with the wavefunctions of Fock state. The normalized wave function of Fock states is:\n",
    "$$\n",
    "\\psi_n(x) = \\frac{1}{\\sqrt{\\sqrt{\\pi} 2^{n} n!}} H_n(x) e^{-\\frac{x^2}{2}},\n",
    "$$\n",
    "where $n$ is the number of particles in the Fock state, $x$ is the position of the state, and $H_n$ are the physicist's Hermite polynomials.\n",
    "\n",
    "#### Encoding the orbitals\n",
    "\n",
    "In practice, to accelerate the estimation of the integral discussed before, we would like to define a boson sampling experiment. To do that, we consider a photonic quantum computer running a circuit where the input states to the circuit are Fock states, an interferometer is applied, and then photon number resolving detection is performed. To perform the desired sampling, we have to encode the orbitals in our boson sampling experiment. For this, we construct the unitary matrix of the interferometer by evaluating each eigenfunction at the chosen positions. Hence, three vectors are obtained (one for each of the three eigenfunctions) which are then orthogonalized using a singular value decomposition technique (Löwdin ortogonalization) and then further orthonormal vectors are gathered to make up a $12 × 12$ unitary matrix that can be fit on a 12-mode photonic quantum chip."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b402dde",
   "metadata": {},
   "source": [
    "## Demonstration\n",
    "\n",
    "#### Outline\n",
    "\n",
    "1. Import libraries\n",
    "2. Define single-particle wavefunctions and build orbital matrix\n",
    "3. Orthogonalize rows (Löwdin) and extend to a full unitary\n",
    "4. Build a Piquasso program\n",
    "5. Define the classical observable (Efimov-inspired 3-body potential)\n",
    "6. Run the computation and compare to analytical value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b39e8456-3ba4-49d2-9564-72ab77303837",
   "metadata": {},
   "source": [
    "#### 1. Imports\n",
    "As described, we first import all relevant libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0bc7ab2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import piquasso as pq\n",
    "from scipy.special import factorial\n",
    "from scipy.constants import pi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04f15c71-f5ef-42d7-8e43-3d4d81417db9",
   "metadata": {},
   "source": [
    "#### 2. Wavefunction and building the orbital matrix\n",
    "\n",
    "We then define the wavefunction of Fock states and the procedure to evaluate the single-particle eigenfunctions. The values of the eigenfunctions at specific positions will be later encoded in the unitary of the interferometer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1ac7d6f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def wavefunction(x, n):\n",
    "    in_sqrt = np.sqrt(pi) * 2**n * factorial(n)\n",
    "    a = 1 / np.sqrt(in_sqrt)\n",
    "    H1 = np.polynomial.hermite.hermval(x, [0] * n + [1])\n",
    "    exp = np.exp(-(x**2) / 2)\n",
    "    return a * H1 * exp\n",
    "\n",
    "\n",
    "def build_single_particle_orbitals(bin_positions, quantum_numbers=None):\n",
    "    if quantum_numbers is None:\n",
    "        quantum_numbers = list(range(3))\n",
    "    n_orbitals = len(quantum_numbers)\n",
    "    n_modes = bin_positions.shape[0]\n",
    "    orbital_matrix = np.zeros((n_orbitals, n_modes), dtype=complex)\n",
    "    for i, n_qn in enumerate(quantum_numbers):\n",
    "        for j, x in enumerate(bin_positions):\n",
    "            val = wavefunction(x, n_qn)\n",
    "            orbital_matrix[i, j] = val[0] if np.ndim(val) > 0 else val\n",
    "    return orbital_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18e11093-fa2b-444c-978d-989d2b4f69bb",
   "metadata": {},
   "source": [
    "#### 3. Creating the unitary\n",
    "\n",
    "Since we are encoding only 3 eigenfunctions, but need a $12 × 12$ unitary matrix, we perform an SVD-like method (that involves Löwdin ortogonalization) to create a unitary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a39e761e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lowdin(P):\n",
    "    U, _, Vh = np.linalg.svd(P, full_matrices=False)\n",
    "    return U @ Vh\n",
    "\n",
    "\n",
    "def extend_to_orthonormal_rows(mat_ortogonal_rows, seed=None, max_tries=10000):\n",
    "    tol = 1e-12\n",
    "    rng = np.random.default_rng(seed)\n",
    "    basis = [mat_ortogonal_rows[i].copy() for i in range(len(mat_ortogonal_rows))]\n",
    "\n",
    "    def hermitian_proj_coeff(b, w):\n",
    "        bb = np.vdot(b, b)\n",
    "        if np.abs(bb) <= tol:\n",
    "            return 0.0j\n",
    "        return np.vdot(b, w) / bb\n",
    "\n",
    "    def orthogonalize(v, basis_list):\n",
    "        w = v.copy()\n",
    "        for b in basis_list:\n",
    "            w = w - hermitian_proj_coeff(b, w) * b\n",
    "        return w\n",
    "\n",
    "    n_modes = mat_ortogonal_rows.shape[1]\n",
    "    tries = 0\n",
    "    while len(basis) < n_modes:\n",
    "        if tries >= max_tries:\n",
    "            msg = (\n",
    "                \"Could not generate enough independent vectors; \"\n",
    "                \"increase max_tries or check input.\"\n",
    "            )\n",
    "            raise RuntimeError(msg)\n",
    "        tries += 1\n",
    "        v = rng.standard_normal(n_modes) + 1j * rng.standard_normal(n_modes)\n",
    "        w = orthogonalize(v, basis)\n",
    "        if np.linalg.norm(w) > tol:\n",
    "            basis.append(w)\n",
    "    Q = np.stack(basis, axis=0)\n",
    "    norms = np.sqrt(np.real(np.einsum(\"ij,ij->i\", Q.conj(), Q)))\n",
    "    Q = Q / norms[:, None]\n",
    "    return Q"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad7000b8-a0ca-498f-849f-93c4a5dfdb46",
   "metadata": {},
   "source": [
    "#### 4. Creating the Piquasso program\n",
    "\n",
    "Once we have are unitary matrix, we can use Piquasso to create a boson sampling program for our experiment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cf1bac7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_piquasso_program(unitary, input_occupation):\n",
    "    m = unitary.shape[0]\n",
    "    instructions = []\n",
    "    instructions.append(pq.Vacuum().on_modes(*range(m)))\n",
    "    for mode, photon_count in enumerate(input_occupation):\n",
    "        if photon_count == 1:\n",
    "            instructions.append(pq.Create().on_modes(mode))\n",
    "    instructions.append(pq.Interferometer(unitary))\n",
    "    instructions.append(pq.ParticleNumberMeasurement())\n",
    "    return pq.Program(instructions=instructions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cdc31fd-1880-403a-a565-b2a5cc380e09",
   "metadata": {},
   "source": [
    "#### 5. Efimov potential and randomization\n",
    "\n",
    "To compute the integral that's being considered, we define the Efimov potential. Integrating the function of the Efimov potential leads to divergence in certain points (three particles cannot spatially overlap completely in position space), furthermore discretizing the domain of the function introduces a bias. To mitigate the impact of this discretization bias on the final estimate of the integral we also perform a randomization procedure<cite data-footcite=\"correa2025bosonsampling\"></cite> for the exact positions considered when evaluating the Efimov potential function. Note that the lower and upper limits used for the discretization of the domain were obtained by evaluating $f (\\textbf{X})$ over points in the domain and considering such contributions to the expectation value that we're estimating (energy correction)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "976267d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def efimov_potential_3body(positions_xyz, C=1.0):\n",
    "    x = positions_xyz\n",
    "    eps = 1e-10\n",
    "    r12 = abs(x[0] - x[1]) + eps\n",
    "    r13 = abs(x[0] - x[2]) + eps\n",
    "    r23 = abs(x[1] - x[2]) + eps\n",
    "    R2 = (2.0 / 3.0) * (r12**2 + r13**2 + r23**2)\n",
    "    return -(C + 1.0 / 4.0) / R2\n",
    "\n",
    "\n",
    "def efimov_normalized(bin_positions, C=1.0):\n",
    "    lower, upper = -4, 4\n",
    "    num_pos = len(bin_positions)\n",
    "    bin_borders = np.linspace(lower, upper, num_pos + 1)\n",
    "    N = 1000\n",
    "    s = 0.0\n",
    "    for _ in range(N):\n",
    "        x_bin, y_bin, z_bin = bin_positions\n",
    "        x = y = z = 0.0\n",
    "        c = 2 / 3\n",
    "        while abs(x - y) < c or abs(x - z) < c or abs(y - z) < c:\n",
    "            x = np.random.uniform(bin_borders[x_bin], bin_borders[x_bin + 1])\n",
    "            y = np.random.uniform(bin_borders[y_bin], bin_borders[y_bin + 1])\n",
    "            z = np.random.uniform(bin_borders[z_bin], bin_borders[z_bin + 1])\n",
    "        s += efimov_potential_3body([x, y, z], C)\n",
    "    return s / N"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d2db8e7-7313-4808-990f-42bd69a49ab1",
   "metadata": {},
   "source": [
    "#### 6. Running the simulation\n",
    "\n",
    "We combine the Piquasso program execution and Efimov potential computation to build our boson sampling procedure. We leverage Piquasso's boson sampling backend (``SamplingSimulator``) for efficiency. Only those samples are kept in which at most a single photon is detected on each mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ef397897-de9d-42a6-b49a-b80984e737a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def importance_sampling_with_bs(positions, input_occupation=None, n_samples=1000):\n",
    "    m = positions.shape[0]\n",
    "    input_occupation = [1, 1, 1] + [0] * (m - 3)\n",
    "    orbital_matrix = build_single_particle_orbitals(positions)\n",
    "    orthogonal_rows_mat = lowdin(orbital_matrix)\n",
    "    unitary = extend_to_orthonormal_rows(orthogonal_rows_mat, seed=0)\n",
    "\n",
    "    values = []\n",
    "    sim = pq.SamplingSimulator(d=m)\n",
    "    prog = build_piquasso_program(unitary, input_occupation)\n",
    "    result = sim.execute(prog, shots=n_samples)\n",
    "    for pattern in result.samples:\n",
    "        idx = [i for i, n in enumerate(pattern) if n == 1]\n",
    "        X = positions[idx, :]\n",
    "        if len(X) != 3:\n",
    "            continue\n",
    "        values.append(efimov_potential_3body(X))\n",
    "    if len(values) == 0:\n",
    "        raise RuntimeError(\"No valid 3-photon samples collected.\")\n",
    "    return np.mean(values)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ac814d-1eec-459d-b856-7fd8869364e8",
   "metadata": {},
   "source": [
    "Now we have every building block to perform the importance sampling procedure that includes boson sampling to estimate the integral. The analytic value for the energy correction has been pre-computed via numerical integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "38cf3bb3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Piquasso estimated E^(1):       -0.2677\n",
      "Analytic value for E^(1):       -0.2600\n"
     ]
    }
   ],
   "source": [
    "n_modes = 12\n",
    "xs = np.arange(n_modes)\n",
    "positions = np.stack([xs], axis=1)\n",
    "n_samples = 10000  # reduce for quick runs\n",
    "analytical_expval = -0.259961\n",
    "\n",
    "boson_sampling_estimate = importance_sampling_with_bs(positions, n_samples=n_samples)\n",
    "print(f\"Piquasso estimated E^(1):       {boson_sampling_estimate:.4f}\")\n",
    "print(f\"Analytic value for E^(1):       {analytical_expval:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "820caa58-fc93-4f4e-a8eb-0e656f0c4724",
   "metadata": {},
   "source": [
    "What we observe is that the quantity we get with importance sampling is close enough to the analytic value, so we have successfully estimated the integral (i.e., the first order energy correction). 🎉"
   ]
  }
 ],
 "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}