{ "cells": [ { "cell_type": "markdown", "id": "4065cad3-0d0c-41ff-97e4-120e2af31d54", "metadata": {}, "source": [ "(ex-bayesian-optimize-qaoa)=\n", "\n", "Bayesian Optimizing QAOA Circuit Energy\n", "=======================================\n", "\n", "In this quantum circuit example, we'll optimize a 54-qubit, $p=4$ (depth 12), QAOA\n", "circuit on a random 3-regular graph, with respect to its energy - a sum of local\n", "expectations." ] }, { "cell_type": "code", "execution_count": 1, "id": "12bc290c-58e4-409f-acee-244a65c2083e", "metadata": {}, "outputs": [], "source": [ "%config InlineBackend.figure_formats = ['svg']\n", "\n", "import quimb as qu\n", "import quimb.tensor as qtn" ] }, { "cell_type": "markdown", "id": "230aeec0-9661-40c5-a130-9cb1c137d44e", "metadata": {}, "source": [ "First we instantiate a high quality contraction path finder from ``cotengra``,\n", "because each energy term will be a unique contraction, we'll make a\n", "'reusable' optimizer that can be used on multiple different contractions." ] }, { "cell_type": "code", "execution_count": 2, "id": "3984aaa2-6efd-409c-9144-9303064151a3", "metadata": {}, "outputs": [], "source": [ "import cotengra as ctg\n", "\n", "opt = ctg.ReusableHyperOptimizer(\n", " methods=['greedy'],\n", " reconf_opts={},\n", " max_repeats=32,\n", " max_time=\"rate:1e6\",\n", " parallel=True,\n", " # use the following for persistently cached paths\n", " # directory=True,\n", ")" ] }, { "cell_type": "markdown", "id": "92f37d3f-dbf4-4912-9507-7961f2d5d8e4", "metadata": {}, "source": [ "Setting Up the Circuit\n", "----------------------\n", "\n", "Then we generate a random regular graph of conditions to satisfy. Optimizing\n", "the antiferromagnetic coupling on this graph is equivalent to trying to solve\n", "the MAX-CUT problem." ] }, { "cell_type": "code", "execution_count": 3, "id": "a9babce7-9e46-4fa1-9768-8ffee656b59c", "metadata": {}, "outputs": [], "source": [ "import networkx as nx\n", "\n", "reg = 3\n", "n = 54\n", "seed = 666\n", "G = nx.random_regular_graph(reg, n, seed=seed)\n", "\n", "terms = {(i, j): 1 for i, j in G.edges}" ] }, { "cell_type": "markdown", "id": "b3ca6b6c-2ae4-4c6b-8291-1186c6c65454", "metadata": {}, "source": [ "`quimb` has a built-in QAOA circuit ansatz, which takes the dict of couplings\n", "to weights, as well as the $\\beta$ and $\\gamma$ parameters describing gate\n", "rotations:" ] }, { "cell_type": "code", "execution_count": 4, "id": "ffaae3af-cb4a-4797-b709-ab38e16490cd", "metadata": {}, "outputs": [], "source": [ "p = 4\n", "gammas = qu.randn(p)\n", "betas = qu.randn(p)\n", "circ_ex = qtn.circ_qaoa(terms, p, gammas, betas)" ] }, { "cell_type": "markdown", "id": "20b674f9-7dc6-4354-ae44-8d698a482f02", "metadata": {}, "source": [ "The overall circuit this generates is very complex:" ] }, { "cell_type": "code", "execution_count": 5, "id": "1a206d7b-0f19-4bb0-b701-42920203e914", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "circ_ex.psi.draw(color=['PSI0', 'H', 'RZZ', 'RX'])" ] }, { "cell_type": "markdown", "id": "247a3a3f-d6e0-4c3b-9877-56d48ff3e66d", "metadata": {}, "source": [ "But because of the unitary structure of quantum circuits, local quantities can usually\n", "be significantly simplified (automatically by ``quimb``). Here, e.g., is the simplfied\n", "tensor network describing the reduced density matrix of qubit 0 only:" ] }, { "cell_type": "code", "execution_count": 6, "id": "b2c5307e-061e-4aae-8d4d-00f40c631955", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "circ_ex.get_rdm_lightcone_simplified([0]).draw(color=['PSI0', 'H', 'RZZ', 'RX'], highlight_inds=['k0', 'b0'])" ] }, { "cell_type": "markdown", "id": "b8a47c63-1e31-4cfd-a553-047b0c538140", "metadata": {}, "source": [ "Rehearsing the Computation\n", "--------------------------" ] }, { "cell_type": "markdown", "id": "624d6ce1-8fdc-4865-ae30-a6d63cd038f3", "metadata": {}, "source": [ "Before we actually compute the QAOA energy, its usually worth 'rehearsing' it -\n", "finding the contraction widths and costs of each energy term to check they are\n", "not too big. The contraction paths found (which can take some time), will be\n", "cached by the `ReusableHyperOptimizer` for the actual computation later." ] }, { "cell_type": "code", "execution_count": 7, "id": "96d70a8a-a529-441b-ba04-ac245124f372", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 81/81 [05:55<00:00, 4.39s/it]\n" ] } ], "source": [ "import tqdm\n", "\n", "ZZ = qu.pauli('Z') & qu.pauli('Z')\n", "\n", "local_exp_rehs = [\n", " circ_ex.local_expectation_rehearse(weight * ZZ, edge, optimize=opt)\n", " for edge, weight in tqdm.tqdm(list(terms.items()))\n", "]" ] }, { "cell_type": "markdown", "id": "de50ce34-ef4f-4a84-a2e5-b857166609c9", "metadata": {}, "source": [ "If we plot these we can see that the size of each and the number of ops is\n", "very moderate:" ] }, { "cell_type": "code", "execution_count": 8, "id": "e071f92a-5d94-42ca-8b37-1180c9413a7f", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "import matplotlib.pyplot as plt\n", "\n", "with plt.style.context(qu.NEUTRAL_STYLE):\n", " fig, ax1 = plt.subplots()\n", " ax1.plot([rehs['W'] for rehs in local_exp_rehs], color='green')\n", " ax1.set_ylabel('contraction width, $W$, [log2]', color='green')\n", " ax1.tick_params(axis='y', labelcolor='green')\n", "\n", " ax2 = ax1.twinx()\n", " ax2.plot([rehs['C'] for rehs in local_exp_rehs], color='orange')\n", " ax2.set_ylabel('contraction cost, $C$, [log10]', color='orange')\n", " ax2.tick_params(axis='y', labelcolor='orange')" ] }, { "cell_type": "markdown", "id": "963c701d-7cef-4367-8f4a-796e50d9a6a6", "metadata": {}, "source": [ "In fact, the actual contractions will probably not be the most time consuming\n", "part of the computation." ] }, { "cell_type": "markdown", "id": "3bb0bcf5-4cbb-4357-bc2a-252c3443721d", "metadata": {}, "source": [ "Bayesian Optimizing the Energy\n", "------------------------------\n", "\n", "Since we only have $2 p=8$ variational parameters we can easily apply something\n", "like Bayesian optimization to minimize the energy.\n", "\n", "We'll use ``scikit-optimize`` here, but many gradient free optimizers could be\n", "used instead. Most simply require casting the problem function as taking a\n", "single vector of parameters - ``x``:" ] }, { "cell_type": "code", "execution_count": 9, "id": "551c7066-d616-4fc7-94c9-53b7a9b2e77d", "metadata": {}, "outputs": [], "source": [ "def energy(x):\n", " p = len(x) // 2\n", " gammas = x[:p]\n", " betas = x[p:]\n", " circ = qtn.circ_qaoa(terms, p, gammas, betas)\n", "\n", " ZZ = qu.pauli('Z') & qu.pauli('Z')\n", "\n", " ens = [\n", " circ.local_expectation(weight * ZZ, edge, optimize=opt, backend=\"jax\")\n", " for edge, weight in terms.items()\n", " ]\n", "\n", " return sum(ens).real" ] }, { "cell_type": "markdown", "id": "cd785d08-20d1-4cbe-8ac7-6a7d2c190af3", "metadata": {}, "source": [ "Now we can setup some bounds for our parameters and the optimizer object itself:" ] }, { "cell_type": "code", "execution_count": 10, "id": "5992b99e-8294-48fc-a7cd-19d9c05070a9", "metadata": {}, "outputs": [], "source": [ "from skopt import Optimizer\n", "from skopt.plots import plot_convergence, plot_objective" ] }, { "cell_type": "code", "execution_count": 11, "id": "14db3ece-3793-4096-8598-06e4188b969a", "metadata": {}, "outputs": [], "source": [ "eps = 1e-6\n", "bounds = (\n", " [(0.0 + eps, qu.pi / 2 - eps)] * p +\n", " [(-qu.pi / 4 + eps, qu.pi / 4 - eps)] * p\n", ")\n", "\n", "bopt = Optimizer(bounds)" ] }, { "cell_type": "markdown", "id": "7733ab69-e6a8-4071-aced-0b26231a1568", "metadata": {}, "source": [ "For the actual minimization, we ask the optimizer to suggest a vector of\n", "parameters to sample, then simply report the corresponding energy:" ] }, { "cell_type": "code", "execution_count": 12, "id": "563598a3-4423-472a-a872-37cb154723c7", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [1:06:39<00:00, 40.00s/it]\n" ] } ], "source": [ "for i in tqdm.trange(100):\n", " x = bopt.ask()\n", " res = bopt.tell(x, energy(x))" ] }, { "cell_type": "code", "execution_count": 13, "id": "5c25b9f1-3141-464c-990a-79ad2b2952f0", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "with plt.style.context(qu.NEUTRAL_STYLE):\n", " plot_convergence(res);" ] }, { "cell_type": "markdown", "id": "8d3807cf-18e8-4946-931c-5446924ee4b4", "metadata": {}, "source": [ "One of the advantages of using Bayesian optimization and ``scikit-optimize`` is\n", "that we get an estimate of the actual energy landscape we can visualize:" ] }, { "cell_type": "code", "execution_count": 14, "id": "85e2c9c6-1c9e-4ad7-86c0-1318aaf48ef2", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "