{ "cells": [ { "cell_type": "markdown", "id": "491df26c-a2fa-478e-abb8-a8687aa582f0", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ "(example-tn-training-circuits)=\n", "\n", "# Tensor Network Training of Quantum Circuits\n", "\n", "Here we'll run through constructing a tensor network of an ansatz quantum circuit, then training certain 'parametrizable' tensors representing quantum gates in that tensor network to replicate the behaviour of a target unitary." ] }, { "cell_type": "code", "execution_count": 1, "id": "987866b0-f875-44ee-a85e-c0ef5d08829a", "metadata": {}, "outputs": [], "source": [ "%config InlineBackend.figure_formats = ['svg']\n", "import quimb as qu\n", "import quimb.tensor as qtn" ] }, { "cell_type": "markdown", "id": "5c07157a-a166-4995-8e71-f08216ad1f42", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ "## The Ansatz Circuit\n", "\n", "First we set up the ansatz circuit and extract the tensor network. Key here is that when we supply `parametrize=True` to the `'U3'` gate call, it injects a {class}`~quimb.tensor.tensor_core.PTensor` into the network, which lazily represents its data array with a function and set of parameters. Later, when the optimizer\n", "sees this it then knows to optimize the parameters rather than the array itself." ] }, { "cell_type": "code", "execution_count": null, "id": "d67ed341-9bcc-4962-9102-5650170af3e7", "metadata": {}, "outputs": [], "source": [ "def single_qubit_layer(circ, gate_round=None):\n", " \"\"\"Apply a parametrizable layer of single qubit ``U3`` gates.\n", " \"\"\"\n", " for i in range(circ.N):\n", " # initialize with random parameters\n", " params = qu.randn(3, dist='uniform')\n", " circ.apply_gate(\n", " 'U3', *params, i,\n", " gate_round=gate_round, parametrize=True)\n", "\n", "def two_qubit_layer(circ, gate2='CZ', reverse=False, gate_round=None):\n", " \"\"\"Apply a layer of constant entangling gates.\n", " \"\"\"\n", " regs = range(0, circ.N - 1)\n", " if reverse:\n", " regs = reversed(regs)\n", "\n", " for i in regs:\n", " circ.apply_gate(\n", " gate2, i, i + 1, gate_round=gate_round)\n", "\n", "def ansatz_circuit(n, depth, gate2='CZ', **kwargs):\n", " \"\"\"Construct a circuit of single qubit and entangling layers.\n", " \"\"\"\n", " circ = qtn.Circuit(n, **kwargs)\n", "\n", " for r in range(depth):\n", " # single qubit gate layer\n", " single_qubit_layer(circ, gate_round=r)\n", "\n", " # alternate between forward and backward CZ layers\n", " two_qubit_layer(\n", " circ, gate2=gate2, gate_round=r, reverse=r % 2 == 0)\n", "\n", " # add a final single qubit layer\n", " single_qubit_layer(circ, gate_round=r + 1)\n", "\n", " return circ" ] }, { "cell_type": "markdown", "id": "a7d48eb2-b6a7-4d4c-bf7a-0e795a7c7da0", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ "The form of the `'U3'` gate (which generalizes all possible single qubit gates) can be seen here - {func}`~quimb.gen.operators.U_gate`. Now we are ready to instantiate a circuit:" ] }, { "cell_type": "code", "execution_count": 3, "id": "e63d50a3-e3a2-4607-beed-72aa3fb5b3e3", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "n = 6\n", "depth = 9\n", "gate2 = 'CZ'\n", "\n", "circ = ansatz_circuit(n, depth, gate2=gate2)\n", "circ" ] }, { "cell_type": "markdown", "id": "f64b30a4-7dec-44ad-b227-83a2d7722e32", "metadata": {}, "source": [ "We can extract just the unitary part of the circuit as a tensor network like so:" ] }, { "cell_type": "code", "execution_count": 4, "id": "528068ef-d197-4d08-9c4e-2b089d155bef", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/media/johnnie/Storage2TB/Sync/dev/python/quimb/quimb/tensor/circuit.py:1194: FutureWarning: In future the tensor network returned by ``circ.uni`` will not be transposed as it is currently, to match the expectation from ``U = circ.uni.to_dense()`` behaving like ``U @ psi``. You can retain this behaviour with ``circ.get_uni(transposed=True)``.\n", " warnings.warn(\n" ] } ], "source": [ "V = circ.uni" ] }, { "cell_type": "markdown", "id": "263eebea-6830-4d48-b948-d9078899b866", "metadata": {}, "source": [ "You can see it already has various ``tags`` simultaneously identifying different structures:" ] }, { "cell_type": "code", "execution_count": 5, "id": "f97e9e16-b2fb-4bcc-b373-a5dbafa8b9e0", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# types of gate\n", "V.draw(color=['U3', gate2], show_inds=True)" ] }, { "cell_type": "code", "execution_count": 6, "id": "48d405e7-943e-4ca8-a7b4-5957bca93cd4", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# layers of gates\n", "V.draw(color=[f'ROUND_{i}' for i in range(depth + 1)], show_inds=True)" ] }, { "cell_type": "code", "execution_count": 7, "id": "8a3613be-ec12-4394-8b74-968cbf6baed0", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# what register each tensor is 'above'\n", "V.draw(color=[f'I{i}' for i in range(n)], show_inds=True)" ] }, { "cell_type": "code", "execution_count": 8, "id": "12ef219a-ff39-4d42-b539-1d9d9788271c", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# a unique tag for per gate applied\n", "V.draw(color=[f'GATE_{i}' for i in range(circ.num_gates)], legend=False)" ] }, { "cell_type": "markdown", "id": "aa688638-9d49-4e8e-9c11-ede8b2353132", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ "## The Target Unitary\n", "\n", "Next we need a target unitary to try and digitially replicate. Here we'll take an Ising Hamiltonian and a short time evolution. Once we have the dense (matrix) form of the target unitary \\$U\\$ we need to convert it to a tensor which we can put in a tensor network:" ] }, { "cell_type": "code", "execution_count": 9, "id": "6ba5256b-decf-477e-b4ee-8629850e1559", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# the hamiltonian\n", "H = qu.ham_ising(n, jz=1.0, bx=0.7, cyclic=False)\n", "\n", "# the propagator for the hamiltonian\n", "t = 2\n", "U_dense = qu.expm(-1j * t * H)\n", "\n", "# 'tensorized' version of the unitary propagator\n", "U = qtn.Tensor(\n", " data=U_dense.reshape([2] * (2 * n)),\n", " inds=[f'k{i}' for i in range(n)] + [f'b{i}' for i in range(n)],\n", " tags={'U_TARGET'}\n", ")\n", "U.draw(color=['U3', gate2, 'U_TARGET'])" ] }, { "cell_type": "markdown", "id": "97390c6e-e2da-4982-a92c-28419ef8496b", "metadata": {}, "source": [ "The core object describing how similar two unitaries are is: $\\mathrm{Tr}(V^{\\dagger}U)$, which we can naturally visualize at a tensor network:" ] }, { "cell_type": "code", "execution_count": 10, "id": "035cfb1a-841b-4201-ac1e-15eed4c46a0e", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "(V.H & U).draw(color=['U3', gate2, 'U_TARGET'])" ] }, { "cell_type": "markdown", "id": "53beb5e8-4c05-4dea-962b-57d1e67fa390", "metadata": {}, "source": [ "For our loss function we'll normalize this and negate it (since the optimizer *minimizes*)." ] }, { "cell_type": "code", "execution_count": 11, "id": "561c6a7a-ea07-4805-b328-c9be202ebefb", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0.9881326237593226" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "def loss(V, U):\n", " return 1 - abs((V.H & U).contract(all, optimize='auto-hq')) / 2**n\n", "\n", "# check our current unitary 'infidelity':\n", "loss(V, U)" ] }, { "cell_type": "markdown", "id": "90912e67-4118-4d30-893e-2e079271437a", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ "So as expected currently the two unitaries are not similar *at all*.\n", "\n", "## The Tensor Network Optimization\n", "\n", "Now we are ready to construct the {class}`~quimb.tensor.optimize.TNOptimizer` object, with options detailed below:" ] }, { "cell_type": "code", "execution_count": 12, "id": "7e247b1f-fc1d-425b-855a-25b67e4c72f7", "metadata": {}, "outputs": [], "source": [ "# use the autograd/jax based optimizer\n", "\n", "tnopt = qtn.TNOptimizer(\n", " V, # the tensor network we want to optimize\n", " loss, # the function we want to minimize\n", " loss_constants={'U': U}, # supply U to the loss function as a constant TN\n", " tags=['U3'], # only optimize U3 tensors\n", " autodiff_backend='jax', # use 'autograd' for non-compiled optimization\n", " optimizer='L-BFGS-B', # the optimization algorithm\n", ")" ] }, { "cell_type": "markdown", "id": "fc6f8a27-9527-4d21-b48b-23eb1f3b06f6", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ ":::{note}\n", "If `tags` is not specified the default is to optimize **all** tensors.\n", "In that case, instead of specifying tensor tags to opt-in you can use\n", "`constant_tags` to opt-out tensors tags you *don't* want to optimize,\n", "which may be more convenient.\n", ":::\n", "\n", "We could call `optimize` for pure gradient based optimization, but since\n", "unitary circuits can be tricky we'll use `optimize_basinhopping` which\n", "combines gradient descent with 'hopping' to escape local minima:" ] }, { "cell_type": "code", "execution_count": 13, "id": "d6cec3e9-ec6f-49bc-a453-49a996a72363", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "+0.005301594734 [best: +0.005261898041] : 45%|████▌ | 2258/5000 [00:40<00:49, 55.26it/s]\n" ] } ], "source": [ "# allow 10 hops with 500 steps in each 'basin'\n", "V_opt = tnopt.optimize_basinhopping(n=500, nhop=10)" ] }, { "cell_type": "markdown", "id": "b7e8efc4-1d84-413b-ac05-727eefa9cdf3", "metadata": {}, "source": [ "The optimized tensor network still contains ``PTensor`` instances but now with optimized parameters.\n", "For example, here's the tensor of the ``U3`` gate acting on qubit-2 in round-4:" ] }, { "cell_type": "code", "execution_count": 14, "id": "d617745c-1a11-4c8e-83da-9c1da8142c65", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "PTensor(shape=(2, 2), inds=('_8d3684AAABr', '_8d3684AAABh'), tags=oset(['GATE_46', 'ROUND_4', 'U3', 'I2']))" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "V_opt['U3', 'I2', 'ROUND_4']" ] }, { "cell_type": "markdown", "id": "bf18482c-2df8-4c23-8e2e-5634b65339a3", "metadata": {}, "source": [ "We can see the parameters have been updated by the training:" ] }, { "cell_type": "code", "execution_count": 15, "id": "c4f1311f-f04f-49c1-b8ee-8f391be18509", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([0.25927184, 0.10598236, 0.42593174])" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# the initial values\n", "V['U3', 'ROUND_4', 'I2'].params" ] }, { "cell_type": "code", "execution_count": 16, "id": "51bd0057-80b5-45f8-963b-c7a703027ff0", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([ 0.52709323, 0.6036498 , -0.23340225], dtype=float32)" ] }, "execution_count": 16, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# the optimized values\n", "V_opt['U3', 'ROUND_4', 'I2'].params" ] }, { "cell_type": "markdown", "id": "37264774-cbd9-4aaa-8806-216af54ff80a", "metadata": {}, "source": [ "We can see what gate these parameters would generate:" ] }, { "cell_type": "code", "execution_count": 17, "id": "d0010749-3256-443b-8c76-2049a0c6ed85", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[[ 0.965472-0.j -0.253443+0.060252j]\n", " [ 0.214467+0.147877j 0.90005 +0.349352j]]" ] }, "execution_count": 17, "metadata": {}, "output_type": "execute_result" } ], "source": [ "qu.U_gate(*V_opt['U3', 'ROUND_4', 'I2'].params)" ] }, { "cell_type": "markdown", "id": "f8c59dd1-fe7b-46d2-9b2e-1a368bb6a118", "metadata": {}, "source": [ "A final sanity check we can perform is to try evolving a random state with the target unitary and trained circuit and check the fidelity between the resulting states.\n", "\n", "First we turn the tensor network version of $V$ into a dense matrix:" ] }, { "cell_type": "code", "execution_count": 18, "id": "3f71f260-1aae-4843-929d-967c398cb050", "metadata": {}, "outputs": [], "source": [ "V_opt_dense = V_opt.to_dense([f'k{i}' for i in range(n)], [f'b{i}' for i in range(n)])" ] }, { "cell_type": "markdown", "id": "586f5f7c-55f7-47f7-901e-f7ed84776c7e", "metadata": {}, "source": [ "Next we create a random initial state, and evolve it with the" ] }, { "cell_type": "code", "execution_count": 19, "id": "3803d6a7-4783-4d29-8ec3-7fd4e040a3ff", "metadata": {}, "outputs": [], "source": [ "psi0 = qu.rand_ket(2**n)\n", "\n", "# this is the exact state we want\n", "psif_exact = U_dense @ psi0\n", "\n", "# this is the state our circuit will produce if fed `psi0`\n", "psif_apprx = V_opt_dense @ psi0" ] }, { "cell_type": "markdown", "id": "0aef4aac-4f30-45a1-b644-8f6ce9e56d4d", "metadata": {}, "source": [ "The (in)fidelity should broadly match our training loss:" ] }, { "cell_type": "code", "execution_count": 20, "id": "14729a0a-dab8-4a11-9d7c-3719d5a68aa7", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "'Fidelity: 99.57 %'" ] }, "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ "f\"Fidelity: {100 * qu.fidelity(psif_apprx, psif_exact):.2f} %\"" ] }, { "cell_type": "markdown", "id": "0e996238-ead7-485a-98ff-cc0077ecee17", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ "## Extracting the New Circuit\n", "\n", "We can extract the trained circuit parameters by updating the original\n", "{class}`~quimb.tensor.circuit.Circuit` object from the trained TN:" ] }, { "cell_type": "code", "execution_count": 21, "id": "5c136708-4214-40ea-ab2f-033c9d121531", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ,\n", " ]" ] }, "execution_count": 21, "metadata": {}, "output_type": "execute_result" } ], "source": [ "circ.update_params_from(V_opt)\n", "\n", "# the basic gate specification\n", "circ.gates" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "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" }, "vscode": { "interpreter": { "hash": "6132c5c0a7d26b7c311caf7f55df83b87474b489906668e67e2d71a3b39ab16a" } } }, "nbformat": 4, "nbformat_minor": 4 }