{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b10488aa-dfe0-4594-939c-14a215a3c793",
   "metadata": {},
   "source": [
    "<h1><center>Introduction to quantum science and technology - QUANT 400</center></h1>\n",
    "\n",
    "<p><center> <b>Lecturer:</b> <i>Prof. G. Carleo</i> </center><p>\n",
    "    \n",
    "<p><center> <b>Assistant: </b> <i>friederike.metz@epfl.ch"
   ]
  },
  {
   "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",
    "from qiskit.visualization import plot_histogram\n",
    "import networkx as nx"
   ]
  },
  {
   "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": "97948145",
   "metadata": {},
   "outputs": [],
   "source": [
    "def heisenberg(N, J):\n",
    "    return  # TODO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6752edd",
   "metadata": {},
   "source": [
    "Upon running `print(heisenberg(4, 1.))` you should get this output:\n",
    "\n",
    "```python\n",
    "SparsePauliOp(['IIXX', 'IXXI', 'XXII', 'IIYY', 'IYYI', 'YYII', 'IIZZ', 'IZZI', 'ZZII'],\n",
    "            coeffs=[0.25+0.j, 0.25+0.j, 0.25+0.j, 0.25+0.j, 0.25+0.j, 0.25+0.j, 0.25+0.j,\n",
    "            0.25+0.j, 0.25+0.j])\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. CNOTs are applied to neighboring pairs of qubits.\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",
    "    # TODO\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": "2da1af3c",
   "metadata": {},
   "source": [
    "If you run the cell above your circuit should look like this:\n",
    "<div>\n",
    "<img src=\"ansatz2.png\" width=\"900\"/>\n",
    "</div>"
   ]
  },
  {
   "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",
    "    # TODO\n",
    "    return energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d4fe586",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test your solution\n",
    "hamiltonian = heisenberg(3, 1.0)\n",
    "ansatz = hardware_efficient_ansatz(3, reps=1)\n",
    "estimator = StatevectorEstimator()\n",
    "params = np.ones(18)\n",
    "np.allclose(\n",
    "    0.346156899, loss_function(params, ansatz, hamiltonian, estimator)\n",
    ")  # should return True"
   ]
  },
  {
   "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",
    "    # TODO\n",
    "    return grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df9063ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test your solution\n",
    "np.allclose(\n",
    "    0.27643390, np.linalg.norm(gradient(params, ansatz, hamiltonian, estimator))\n",
    ")  # should return True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11ed76e5",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\n",
    "\n",
    "There are different ways of implementing the gradient function. Here is one idea:\n",
    "\n",
    "```python\n",
    "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.\n",
    "        l_p = loss_function(params_shifted, ansatz, hamiltonian, estimator)\n",
    "\n",
    "        # TODO same for the -pi/2 shift\n",
    "\n",
    "        grad[j] = # TODO\n",
    "    return grad\n",
    "```\n",
    "\n",
    "</details>"
   ]
  },
  {
   "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",
    "        # TODO\n",
    "\n",
    "        # leave the code below unchanged; we use it for tracking the learning progress\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": "code",
   "execution_count": null,
   "id": "6ad47e07",
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_data_dict = {\n",
    "    \"params\": None,\n",
    "    \"iters\": 0,\n",
    "    \"energies\": [],\n",
    "}\n",
    "optimize(params, ansatz, hamiltonian, estimator, n_iter=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fe63ef8",
   "metadata": {},
   "source": [
    "The line above should print (up to machine precision): `Iteration: 1 [Current energy: 0.33047536425822827]`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2669c3e4",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\n",
    "\n",
    "```python\n",
    "def optimize(params, ansatz, hamiltonian, estimator, eta=0.2, n_iter=100):\n",
    "    for i in range(n_iter):\n",
    "        grad = # get gradient\n",
    "        params = # update parameters via gradient descent\n",
    "```\n",
    "\n",
    "</details>"
   ]
  },
  {
   "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": "e56d07dd",
   "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",
    "# TODO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "719b7897",
   "metadata": {},
   "source": [
    "On my machine the last line that was printed above read (up to machine precision): `Iteration: 100 [Current energy: -1.5283147163580864]`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2ca56cf",
   "metadata": {},
   "source": [
    "g) Compute the exact ground state energy for comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "924852d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO\n",
    "E_exact = # exact ground state energy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9927b16c",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\n",
    "\n",
    "You can use `scipy.sparse.linalg.eigsh` to compute the ground state and ground state energy. The ground state energy is: -1.6160254"
   ]
  },
  {
   "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": [
    "# TODO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2756e0f",
   "metadata": {},
   "source": [
    "On my machine the fidelity computes to `0.9744610946455567`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5aea8674",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\n",
    "\n",
    "You first need to assign the optimized parameters to the ansatz. Then you can use `Statevector().data` to obtain the quantum state at the end of the circuit. After reshaping the output, you can compute the fidelity w.r.t. the exact ground state (obtained via scipy) using conventional numpy methods.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "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": "91a02a99",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\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. How can we increase the expressivity of the ansatz state?\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31b0eb63",
   "metadata": {},
   "source": [
    "<details><summary>✨Solution</summary>\n",
    "\n",
    "Hyperparameters are the learning rate `eta`, the numer of iterations `n_iter`, the number of circuit repetitions `reps`. 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": "3191c4a3",
   "metadata": {},
   "source": [
    "### Problem: Approximating Max-Cut with the Quantum Approximate Optimization Algorithm (QAOA)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28be1d19",
   "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",
   "id": "92399397",
   "metadata": {},
   "source": [
    "The figure below shows an example of a graph for which we want to solve the Max-Cut problem.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19eb8fcf",
   "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, seed=7)\n",
    "nx.draw(graph, pos, with_labels=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25d6cff8",
   "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,
   "id": "b9835bc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def maxcut_hamiltonian(graph: nx.Graph) -> SparsePauliOp:\n",
    "    \"\"\"Return the Max-Cut cost Hamiltonian for a given graph.\"\"\"\n",
    "    num_qubits = len(graph.nodes())\n",
    "    edges = list(graph.edges())\n",
    "    terms = [(\"I\" * num_qubits, len(edges) / 2)]\n",
    "    for i, j in edges:\n",
    "        # TODO: place Z operators on qubits i and j\n",
    "        # TODO: append the ZZ term with coefficient -0.5\n",
    "    return SparsePauliOp.from_sparse_list(terms, num_qubits=num_qubits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f200be3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "hamiltonian = maxcut_hamiltonian(graph)\n",
    "assert hamiltonian.num_qubits == len(graph.nodes())\n",
    "print(\"maxcut_hamiltonian test passed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5c68bd4",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\n",
    "Each edge contributes $\\frac{1}{2}(\\mathbb{1} - Z_i Z_j)$. Build the Pauli string for qubits $i$ and $j$ and append it with coefficient $-0.5$. The identity shift is the sum of the $\\frac{1}{2}$ factors over all edges.\n",
    "</details>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e18dda93",
   "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.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1435a8b6",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\n",
    "Factor the exponential into $e^{-i\beta/2} e^{+i\beta/2 Z_i Z_j}$ and note that the first factor is a global phase while the second equals an $RZZ(-\beta)$ gate.\n",
    "</details>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d505eb6",
   "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,
   "id": "728c0ae2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def qaoa_ansatz(graph: nx.Graph, p: int = 1) -> QuantumCircuit:\n",
    "    \"\"\"Return a depth-p QAOA ansatz for the Max-Cut problem.\"\"\"\n",
    "    num_qubits = len(graph.nodes())\n",
    "    circuit = QuantumCircuit(num_qubits)\n",
    "    dummy_params = iter(ParameterVector(\"theta\", 2 * p))\n",
    "\n",
    "    for qubit in range(num_qubits):\n",
    "        # TODO: prepare |+> on this qubit\n",
    "    for layer in range(p):\n",
    "        gamma = next(dummy_params)\n",
    "        beta = next(dummy_params)\n",
    "        for i, j in graph.edges():\n",
    "            # TODO: apply the cost unitary on edge (i, j)\n",
    "        for qubit in range(num_qubits):\n",
    "            # TODO: apply the mixer rotation on this qubit\n",
    "\n",
    "    return circuit\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6986088d",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ansatz = qaoa_ansatz(graph, p=1)\n",
    "ops = test_ansatz.count_ops()\n",
    "assert len(test_ansatz.parameters) == 2\n",
    "assert ops.get(\"h\", 0) == len(graph.nodes())\n",
    "assert ops.get(\"rzz\", 0) == len(graph.edges())\n",
    "assert ops.get(\"rx\", 0) == len(graph.nodes())\n",
    "print(\"qaoa_ansatz test passed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26f2616e",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\n",
    "Create the parameter vector via `params = ParameterVector(\"theta\", 2 * p)` and iterate over it with `next`. The cost layer applies `RZZ` gates on every edge, while the mixer layer applies `RX` gates on every qubit.\n",
    "</details>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ebbab31",
   "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,
   "id": "f84700ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = 1\n",
    "hamiltonian = maxcut_hamiltonian(graph)\n",
    "estimator = StatevectorEstimator()\n",
    "gammas = np.linspace(0.0, 2 * np.pi, 50)\n",
    "betas = np.linspace(0.0, 2 * np.pi, 50)\n",
    "loss_grid = np.zeros((len(gammas), len(betas)))\n",
    "\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",
    "        # TODO: evaluate the loss for these parameters\n",
    "        # TODO: store the result in loss_grid[i, j]\n",
    "\n",
    "max_index = np.unravel_index(np.argmax(loss_grid, axis=None), loss_grid.shape)\n",
    "gamma_opt = gammas[max_index[0]]\n",
    "beta_opt = betas[max_index[1]]\n",
    "\n",
    "\n",
    "# TODO: print maximal loss and optimal parameters\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.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78e5aeb0",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\n",
    "Use your `loss_function`: `loss_function(params, ansatz, hamiltonian, estimator)`. After filling the grid, `np.argmax` helps locate the optimal pair of angles.\n",
    "</details>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcb13dd6",
   "metadata": {},
   "source": [
    "e) Sample from the optimal circuit to identify the optimal cut.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86eb1ee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimal_circ = qaoa_ansatz(graph, p=1)\n",
    "optimal_circ.measure_all()\n",
    "sampler = StatevectorSampler()\n",
    "# TODO: run the sampler on (optimal_circ, [gamma_opt, beta_opt]) with 10000 shots\n",
    "# TODO: extract the measurement counts dictionary\n",
    "plot_histogram(counts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4f5c53f",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\n",
    "`sampler.run([(optimal_circ, [gamma_opt, beta_opt])], shots=10000)` returns a `SamplerResult`; use `.result()[0]` and `.data.meas.get_counts()` to obtain the histogram.\n",
    "</details>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15b41ba5",
   "metadata": {},
   "source": [
    "f) Plot the graph with edges coloured according to the most sampled bitstring.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0db3afa",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_counts = 0\n",
    "best_bitstring = None\n",
    "for bitstring, count in counts.items():\n",
    "    # TODO: keep track of the bitstring with the largest count\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)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83c59b9c",
   "metadata": {},
   "source": [
    "<details><summary>💡Hint</summary>\n",
    "Use a running maximum over the `counts` dictionary to find the most frequent bitstring before converting it to colours.\n",
    "</details>\n"
   ]
  }
 ],
 "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
}
