{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b10488aa-dfe0-4594-939c-14a215a3c793",
   "metadata": {},
   "source": [
    "<h1>Introduction to quantum science and technology - QUANT 400</h1>\n",
    "\n",
    "<p> <b>Lecturer:</b> <i>Prof. N. Macris</i> <p>\n",
    "    \n",
    "<p> <b>Assistant: </b> <i>gian.gentinetta@epfl.ch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af426b4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade pip\n",
    "%pip install qiskit pylatexenc networkx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46b0c9b1-3617-43e1-bfe5-2a76237204f9",
   "metadata": {},
   "source": [
    "## Exercise 3 - Variational quantum algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cfbb3da",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from qiskit import QuantumCircuit\n",
    "from qiskit.primitives import StatevectorEstimator, StatevectorSampler\n",
    "from qiskit.quantum_info import Statevector, SparsePauliOp, Operator\n",
    "from qiskit.circuit import ParameterVector\n",
    "import networkx as nx\n",
    "from qiskit.visualization import plot_histogram"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caee92ab-871c-4b84-964f-4d1c2e75a463",
   "metadata": {},
   "source": [
    "### Parameterized quantum circuits in Qiskit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8b36d4d",
   "metadata": {},
   "source": [
    "We have already encountered examples of parameterized gates in Qiskit like the `rx` or `rzz` gates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89d775c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "circ = QuantumCircuit(2)\n",
    "circ.rx(0.12, 0)\n",
    "circ.rzz(0.34, 0, 1)\n",
    "circ.rz(0.56, 1)\n",
    "circ.draw(\"mpl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6f64505",
   "metadata": {},
   "source": [
    "In the example above we specified the value of the parameter (the angle) directly when constructing the circuit. However, there might be cases where we want to keep the parameters as free variables. For example, imagine you want to run the same circuit many times with different parameters. Then, it makes sense to define the circuit only once with parameters unspecified and only assign the actual values to the parameters when executing the circuit.\n",
    "\n",
    "We can define parameterized circuits in Qiskit using a parameter placeholder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c2cf2a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "dummy_params = ParameterVector(\"theta\", 3)  # 3 parameters\n",
    "dummy_params_iter = iter(dummy_params)\n",
    "circ = QuantumCircuit(2)\n",
    "circ.rx(next(dummy_params_iter), 0)\n",
    "circ.rzz(next(dummy_params_iter), 0, 1)\n",
    "circ.rz(next(dummy_params_iter), 1)\n",
    "circ.draw(\"mpl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f559584c",
   "metadata": {},
   "source": [
    "When we run the circuit to, for example, compute an expecation value, we need to pass the actual parameter values to the `Estimator().run` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56db0d79",
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = StatevectorEstimator()\n",
    "\n",
    "params = [0.12, 0.34, 0.56]\n",
    "\n",
    "observable = SparsePauliOp.from_list([(\"ZZ\", 1.0)])\n",
    "\n",
    "pub = (circ, [observable], [params])\n",
    "result = estimator.run(pubs=[pub]).result()\n",
    "result[0].data.evs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9b57e48",
   "metadata": {},
   "source": [
    "We can also conveniently execute the circuit in parallel for different parameter settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d5b27da",
   "metadata": {},
   "outputs": [],
   "source": [
    "params2 = [1.2, 3.4, 5.6]\n",
    "\n",
    "pub = (circ, [observable], [params, params2])\n",
    "result = estimator.run(pubs=[pub]).result()\n",
    "result[0].data.evs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5484c81",
   "metadata": {},
   "source": [
    "If you want to assign parameters to a predefined circuit yourself, you can use the `assign_parameters` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f8fe81",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Statevector(circ) # this does not work because the circuit is abstractely defined without parameters\n",
    "\n",
    "# first assign the parameters to the circuit\n",
    "circ2 = circ.assign_parameters(params2)\n",
    "Statevector(circ2).data.reshape(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c4b2e08",
   "metadata": {},
   "outputs": [],
   "source": [
    "circ2.draw(\"mpl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f4cf267",
   "metadata": {},
   "source": [
    "### Problem: Finding ground states with the Variational Quantum Eigensolver (VQE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef1dc72d",
   "metadata": {},
   "source": [
    "In this exercise our goal is to implement the VQE algorithm which allows us to prepare the ground state of a given Hamiltonian on a quantum device. So far, we have always considered the Ising model as an example. To mix things up this time, we will study another paradigmatic spin Hamiltonian: the spin-1/2 Heisenberg chain\n",
    "\n",
    "$$ H = J \\sum_{i=1}^{N-1} \\vec{S}_i \\cdot \\vec{S}_{i+1} =  \\frac{J}{4} \\sum_{i=1}^{N-1} X_i X_{i+1} + Y_i Y_{i+1} + Z_i Z_{i+1} .$$\n",
    "\n",
    "The Hamiltonian above is actually the isotropic Heisenberg model, since the coupling constant is the same in all three spatial directions. Therefore, the model above is also called the XXX Heisenberg chain.\n",
    "\n",
    "We will study the XXX model on an open chain with positive exchange interaction $J$ for which the system's ground state is antiferromagnetic."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a61dac6",
   "metadata": {},
   "source": [
    "a) Define a function that returns the Heisenberg Hamiltonian as a `SparsePauliOp`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e93ca346",
   "metadata": {},
   "outputs": [],
   "source": [
    "def heisenberg(N, J):\n",
    "    h_terms = [(\"XX\", [i, i + 1], J / 4) for i in range(N - 1)]\n",
    "    h_terms += [(\"YY\", [i, i + 1], J / 4) for i in range(N - 1)]\n",
    "    h_terms += [(\"ZZ\", [i, i + 1], J / 4) for i in range(N - 1)]\n",
    "    return SparsePauliOp.from_sparse_list(h_terms, num_qubits=N)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a344739",
   "metadata": {},
   "source": [
    "VQE is based on a parameterized quantum circuit $U(\\theta)$, whose parameters are optimized such that the energy is minimized, i.e.,\n",
    "\n",
    "$$\\min_\\theta E(\\theta)= \\min_\\theta \\langle 0| {U}^{\\dagger}(\\theta) H {U}(\\theta)|0\\rangle .$$\n",
    "\n",
    "If we optimize all the way to the exact ground state energy then the quantum state at the end of the circuit $|\\psi\\rangle = {U}(\\theta_{opt})|0\\rangle$ is the exact ground state of $H$.\n",
    "\n",
    "The parameterized quantum circuit is also often referred to as an ansatz (for the ground state in this case). Coming up with a good parameterized quantum circuit ansatz is a difficult task and often done heuristically. A common choice is the hardware-efficient ansatz which is composed of a layer of parameterized single-qubit rotations and a layer of fixed, two-qubit entangling gates."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d10a78a",
   "metadata": {},
   "source": [
    "b) Define a function that implements the following circuit. `reps` refers to the number of repetitions of the subcircuit.\n",
    "\n",
    "<div>\n",
    "<img src=\"ansatz.png\" width=\"500\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6adfb080",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hardware_efficient_ansatz(N, reps=1):\n",
    "    circ = QuantumCircuit(N)\n",
    "    dummy_params = iter(ParameterVector(\"theta\", 3 * N * (reps + 1)))\n",
    "    for _ in range(reps):\n",
    "        for i in range(N):\n",
    "            circ.rz(next(dummy_params), i)\n",
    "            circ.ry(next(dummy_params), i)\n",
    "            circ.rz(next(dummy_params), i)\n",
    "        for i in range(N - 1):\n",
    "            circ.cx(i, i + 1)\n",
    "    for i in range(N):\n",
    "        circ.rz(next(dummy_params), i)\n",
    "        circ.ry(next(dummy_params), i)\n",
    "        circ.rz(next(dummy_params), i)\n",
    "    return circ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "863087f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ansatz = hardware_efficient_ansatz(4, reps=3)\n",
    "ansatz.draw(\"mpl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d2fd234",
   "metadata": {},
   "source": [
    "c) Next define the loss function we want to minimize, i.e., \n",
    "\n",
    "$$ \\mathcal{L}(\\theta)=E(\\theta)= \\langle 0| {U}^{\\dagger}(\\theta) H {U}(\\theta)|0\\rangle .$$\n",
    "\n",
    "Given a vector of parameters $\\theta$, the circuit ansatz $U(\\theta)$, the Hamiltonian $H$, and an instance of the `Estimator`, this function should return the value of the loss function (the energy)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85922cd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function(params, ansatz, hamiltonian, estimator):\n",
    "    pub = (ansatz, [hamiltonian], [params])\n",
    "    result = estimator.run(pubs=[pub]).result()\n",
    "    energy = result[0].data.evs[0]\n",
    "    return energy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93611a1c",
   "metadata": {},
   "source": [
    "We want to optimize the circuit parameters $\\theta$ via gradient descent. For that we need to compute the gradient of our loss function w.r.t. the parameters. In the lecture you learned that you can compute gradients via the parameter shift rule\n",
    "\n",
    "$$\\partial_{\\theta_k} \\mathcal{L}\\left(\\theta_1, \\ldots, \\theta_k, \\ldots \\theta_l\\right)=\\frac{\\mathcal{L}\\left(\\theta_1, \\ldots, \\theta_k+\\frac{\\pi}{2}, \\ldots \\theta_l\\right)-\\mathcal{L}\\left(\\theta_1, \\ldots, \\theta_k-\\frac{\\pi}{2}, \\ldots \\theta_l\\right)}{2}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4c628c1",
   "metadata": {},
   "source": [
    "d) Compute the gradient of the loss function using the parameter shift rule for each parameter. The function below should return the gradient vector `grad` having the same shape as the parameter vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "738922bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gradient(params, ansatz, hamiltonian, estimator):\n",
    "    grad = np.zeros_like(params)\n",
    "    for j in range(len(params)):\n",
    "        params_shifted = params.copy()\n",
    "        params_shifted[j] += np.pi / 2.0\n",
    "        l_p = loss_function(params_shifted, ansatz, hamiltonian, estimator)\n",
    "\n",
    "        params_shifted = params.copy()\n",
    "        params_shifted[j] -= np.pi / 2.0\n",
    "        l_m = loss_function(params_shifted, ansatz, hamiltonian, estimator)\n",
    "\n",
    "        grad[j] = (l_p - l_m) / 2.0\n",
    "    return grad"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a5cadf7",
   "metadata": {},
   "source": [
    "e) Finally write a function which performs the actual optimization via gradient descent with learning rate `eta`. The number of optimization/iteration steps is specified by `n_iter`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0cebebd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimize(params, ansatz, hamiltonian, estimator, eta=0.2, n_iter=100):\n",
    "    for i in range(n_iter):\n",
    "        grad = gradient(params, ansatz, hamiltonian, estimator)\n",
    "        params = params - eta * grad\n",
    "\n",
    "        energy = loss_function(params, ansatz, hamiltonian, estimator)\n",
    "        opt_data_dict[\"iters\"] += 1\n",
    "        opt_data_dict[\"params\"] = params\n",
    "        opt_data_dict[\"energies\"].append(energy)\n",
    "        print(f\"Iteration: {opt_data_dict['iters']} [Current energy: {energy}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6b44498",
   "metadata": {},
   "source": [
    "Let's see all of this in action.\n",
    "\n",
    "f) Run the optimization for the parameters specified below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c1d6e7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 4  # number of spins/qubits\n",
    "J = 1.0  # coupling strength\n",
    "reps = 2  # number of circuit repetitions\n",
    "\n",
    "n_params = 3 * (reps + 1) * N  # number of parameters in the ansatz\n",
    "# technically the initial parameters should be random, but for reproducibility we fix them\n",
    "initial_params = np.ones(n_params)  # initial parameters\n",
    "# initial_params = 2 * np.pi * np.random.rand(n_params) # initial parameters\n",
    "\n",
    "# leave this unchanged\n",
    "opt_data_dict = {\n",
    "    \"params\": None,\n",
    "    \"iters\": 0,\n",
    "    \"energies\": [],\n",
    "}\n",
    "\n",
    "hamiltonian = heisenberg(N, J)\n",
    "ansatz = hardware_efficient_ansatz(N, reps=reps)\n",
    "estimator = StatevectorEstimator()\n",
    "\n",
    "optimize(initial_params, ansatz, hamiltonian, estimator, eta=0.2, n_iter=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2ca56cf",
   "metadata": {},
   "source": [
    "g) Compute the exact ground state energy for comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55526e24",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.sparse.linalg import eigsh\n",
    "\n",
    "H = hamiltonian.to_matrix(sparse=True)\n",
    "E_ed, v = eigsh(H, k=4, which=\"SA\", return_eigenvectors=True)\n",
    "psi_exact = v[:, 0]\n",
    "E_exact = E_ed[0]\n",
    "E_exact"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23b55e6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(range(opt_data_dict[\"iters\"]), opt_data_dict[\"energies\"])\n",
    "ax.axhline(E_exact, color=\"red\", linestyle=\"--\")\n",
    "ax.set_xlabel(\"Iterations\")\n",
    "ax.set_ylabel(\"Energy\")\n",
    "plt.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b883887c",
   "metadata": {},
   "source": [
    "Is the energy minimized and converges to the exact energy? Do the optimized quantum state and the exact ground state coincide?\n",
    "\n",
    "h) Compute the fidelity between the quantum state at the end of the optimized circuit w.r.t. the exact state.\n",
    "\n",
    "$$ F = |\\langle \\psi_{exact} | U(\\theta) |0 \\rangle|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73756a56",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_ansatz = ansatz.assign_parameters(opt_data_dict[\"params\"])\n",
    "psi_vqe = Statevector(optimized_ansatz).data.reshape(-1)\n",
    "np.abs(psi_vqe.conj().T @ psi_exact)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e65a1f9",
   "metadata": {},
   "source": [
    "i) In case the energy/state did not converge to the ground state, we need to do some hyperparameter tuning. What hyperparameters are there? Which hyperparameter(s) do you think we could change to improve convergence to the ground state? Try it out!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bdab9eb",
   "metadata": {},
   "source": [
    "Hyperparameters are the learning rate `eta`, the numer of iterations `n_iter`, the number of circuit repetitions `reps`.\n",
    "\n",
    "One reason we might not be able to learn the ground state to a good accurcay is that the ansatz is not expressive enough, i.e., there is no parameter setting for which the quantum circuit ansatz state corresponds to the true ground state. The expressivity is controlled by the number of parameters which in turn is controlled by the number of circuit repetitions `reps`. Increasing `reps` to 6 already gives a very good approximation to the ground state."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "593a404b",
   "metadata": {},
   "source": [
    "Below I optimize the circuit for different number of circuit repetitions and plot the final energy as a function of `reps`. You clearly see that the accuracy improves as we add more parameters/gates to the circuit.\n",
    "\n",
    "(Note that the code below takes a minute to run.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5baa5b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "optimized_energies = []\n",
    "\n",
    "for reps in range(1, 6):\n",
    "    n_params = 3 * (reps + 1) * N  # number of parameters in the ansatz\n",
    "    initial_params = np.ones(n_params)  # initial parameters\n",
    "    # initial_params = 2 * np.pi * np.random.rand(n_params) # initial parameters\n",
    "\n",
    "    opt_data_dict = {\n",
    "        \"params\": None,\n",
    "        \"iters\": 0,\n",
    "        \"energies\": [],\n",
    "    }\n",
    "\n",
    "    ansatz = hardware_efficient_ansatz(N, reps=reps)\n",
    "\n",
    "    optimize(initial_params, ansatz, hamiltonian, estimator, eta=0.2, n_iter=100)\n",
    "\n",
    "    ax.plot(\n",
    "        range(opt_data_dict[\"iters\"]), opt_data_dict[\"energies\"], label=f\"Reps: {reps}\"\n",
    "    )\n",
    "\n",
    "    optimized_energies.append(opt_data_dict[\"energies\"][-1])\n",
    "\n",
    "ax.axhline(E_exact, color=\"red\", linestyle=\"--\")\n",
    "ax.set_xlabel(\"Iterations\")\n",
    "ax.set_ylabel(\"Energy\")\n",
    "ax.legend()\n",
    "plt.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c059e76",
   "metadata": {},
   "source": [
    "You see that the purple line, which corresponds to using 5 circuit repetitions, converges very close to the exact ground state energy. Using fewer circuit repetitions leads to a larger deviation to the ground state energy.\n",
    "\n",
    "Let's plot the final energy at the end of each optimization against the number of circuit repetitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab6f68da",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(range(1, 6), optimized_energies, marker=\"o\")\n",
    "ax.axhline(E_exact, color=\"red\", linestyle=\"--\")\n",
    "ax.set_xlabel(\"repetitions\")\n",
    "ax.set_ylabel(\"Final energy\")\n",
    "plt.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b9a9ed1",
   "metadata": {},
   "source": [
    "Indeed, with 5 circuit repetitions we get a very good approximation to the ground state energy at the end of the optimization. Hence, the 5-reps circuit is more expressive than a circuit with less repetitions/parameters.\n",
    "\n",
    "Finally, let's look at the fidelity of the optimized state (of the reps=5 circuit) to the exact ground state. As expected, the state overlap is large and we indeed prepared a state that is very close to the exact ground state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f93bc28",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_ansatz = ansatz.assign_parameters(opt_data_dict[\"params\"])\n",
    "psi_vqe = Statevector(optimized_ansatz).data.reshape(-1)\n",
    "np.abs(psi_vqe.conj().T @ psi_exact)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem: Approximating Max-Cut with the Quantum Approximate Optimization Algorithm (QAOA)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this exercise we implement the QAOA algorithm for the Max-Cut application (see [Wikipedia](https://en.wikipedia.org/wiki/Quantum_optimization_algorithms#QAOA_for_finding_the_minimum_vertex_cover_of_a_graph)).\n",
    "\n",
    "The classical Max-Cut problem asks us to divide the vertices of a graph into two sets so that the number of edges crossing the cut is maximized.\n",
    "We can promote this objective to a quantum cost Hamiltonian where each computational basis state encodes a partition, and lower energies correspond to larger cuts.\n",
    "For an edge $(i, j)$ the projector onto vertices being separated is $\\tfrac{1}{2}(\\mathbb{1} - Z_i Z_j)$; summing these projectors over all edges yields the cost operator.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The figure below shows an example of a graph for which we want to solve the Max-Cut problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = nx.Graph()\n",
    "graph.add_edges_from([(0, 1), (1, 2), (2, 0), (2, 3)])\n",
    "pos = nx.spring_layout(graph)\n",
    "nx.draw(graph, pos, with_labels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "a) Implement a function `maxcut_hamiltonian(graph)` that returns the cost operator as a `SparsePauliOp` for a given graph.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def maxcut_hamiltonian(graph: nx.Graph) -> SparsePauliOp:\n",
    "    \"\"\"Return the Max-Cut cost Hamiltonian for a ring graph.\"\"\"\n",
    "    edges = graph.edges()\n",
    "    num_qubits = len(graph.nodes())\n",
    "    terms = [(\"I\", (0,), len(edges) / 2)]\n",
    "    for edge in edges:\n",
    "        i, j = edge\n",
    "        terms.append((\"ZZ\", (i, j), -0.5))\n",
    "\n",
    "    return SparsePauliOp.from_sparse_list(terms, num_qubits=num_qubits)\n",
    "\n",
    "\n",
    "two_qubit_h = maxcut_hamiltonian(nx.Graph([(0, 1)]))\n",
    "assert np.allclose(two_qubit_h.to_matrix(), np.diag([0, 1, 1, 0]))\n",
    "\n",
    "hamiltonian = maxcut_hamiltonian(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab21b504",
   "metadata": {},
   "source": [
    "b) Show that, up to a global phase, the two-qubit gate $\\exp\\left(-i\\beta/2(1-Z_iZ_j) \\right)$ is equivalent to a $ZZ$ rotation. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "c) Build a function `qaoa_ansatz(graph, p=1)` that prepares the usual QAOA circuit. Use a `ParameterVector` of length `2 * p`. Initialize all qubits in $|+\\rangle$, apply the cost and mixer layers in order, and return the parameterized circuit.\n",
    "```math\n",
    "\n",
    "|\\psi(\\gamma, \\beta)\\rangle = \\left(\\prod_{i=k}^p \\exp(-i\\gamma_kH_x)\\exp(-i\\beta_k H_Z)\\right)|+\\rangle^{\\otimes n},\n",
    "```\n",
    "\n",
    "where $H_x = \\sum_i^N \\hat X_i$ and $H_z = \\sum_{(i,j)\\in G} \\frac{1}{2}(1 - \\hat Z_i \\hat Z_j)$\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def qaoa_ansatz(graph: nx.Graph, p: int = 1) -> QuantumCircuit:\n",
    "    \"\"\"Return a depth-p QAOA ansatz for the ring Max-Cut problem.\"\"\"\n",
    "    edges = graph.edges()\n",
    "    num_qubits = len(graph.nodes())\n",
    "    circuit = QuantumCircuit(num_qubits)\n",
    "    dummy_params = iter(ParameterVector(\"theta\", 2 * p))\n",
    "\n",
    "    # Initialize all qubits in |+>\n",
    "    for qubit in range(num_qubits):\n",
    "        circuit.h(qubit)\n",
    "    for layer in range(p):\n",
    "        # Cost layer\n",
    "        gamma = next(dummy_params)\n",
    "        beta = next(dummy_params)\n",
    "        for edge in edges:\n",
    "            i, j = edge\n",
    "            circuit.rzz(beta, i, j)\n",
    "        for qubit in range(num_qubits):\n",
    "            circuit.rx(gamma, qubit)\n",
    "    return circuit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "d) For p=1, plot a 2d grid of the qaoa loss (you can use the VQE loss function defined earlier) as a function of $\\gamma$ and $\\beta$. Find the optimal values of $\\gamma$ and $\\beta$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = 1\n",
    "n_params = 2 * p  # number of parameters in the ansatz\n",
    "# technically the initial parameters should be random, but for reproducibility we fix them\n",
    "initial_params = np.ones(n_params)  # initial parameters\n",
    "# initial_params = 2 * np.pi * np.random.rand(n_params) # initial parameters\n",
    "\n",
    "gammas = np.linspace(0, 2 * np.pi, 50)\n",
    "betas = np.linspace(0, 2 * np.pi, 50)\n",
    "loss_grid = np.zeros((len(gammas), len(betas)))\n",
    "for i, gamma in enumerate(gammas):\n",
    "    for j, beta in enumerate(betas):\n",
    "        params = [gamma, beta]\n",
    "        ansatz = qaoa_ansatz(graph, p=p)\n",
    "        loss = loss_function(params, ansatz, hamiltonian, estimator)\n",
    "        loss_grid[i, j] = loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b52959a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Maximal loss:\", np.max(loss_grid))\n",
    "indices = np.unravel_index(np.argmax(loss_grid, axis=None), loss_grid.shape)\n",
    "gamma_opt = gammas[indices[0]]\n",
    "beta_opt = betas[indices[1]]\n",
    "print(f\"Optimal gamma: {gamma_opt}, Optimal beta: {beta_opt}\")\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "c = ax.pcolormesh(betas, gammas, loss_grid, shading=\"auto\")\n",
    "ax.set_xlabel(r\"$\\beta$\")\n",
    "ax.set_ylabel(r\"$\\gamma$\")\n",
    "fig.colorbar(c, ax=ax, label=\"QAOA Loss\")\n",
    "plt.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39ff4df0",
   "metadata": {},
   "source": [
    "e) Sample from the optimal circuit to find the optimal cut."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7f98096",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimal_circ = qaoa_ansatz(graph, p=1)\n",
    "optimal_circ.measure_all()  # we first need to specify which qubits are measured\n",
    "sampler = StatevectorSampler()\n",
    "result = sampler.run([(optimal_circ, [gamma_opt, beta_opt])], shots=10000).result()[0]\n",
    "counts = result.data.meas.get_counts()\n",
    "\n",
    "plot_histogram(counts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac3473c8",
   "metadata": {},
   "source": [
    "f) Plot the graph with edges coloured according to the most sampled bitstring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57af73fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_counts = 0\n",
    "best_bitstring = None\n",
    "for bitstring, count in counts.items():\n",
    "    if count > max_counts:\n",
    "        max_counts = count\n",
    "        best_bitstring = bitstring\n",
    "\n",
    "bits = [int(b) for b in best_bitstring[::-1]]  # reverse the bitstring\n",
    "nx.draw(graph, pos, with_labels=True, node_color=bits)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}