docs/source/tutorials/dqas.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d41f435a",
   "metadata": {},
   "source": [
    "# Differentiable Quantum Architecture Search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50b478cb",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This tutorial demonstrates how to utilize the advanced computational features provided by TensorCircuit such as ``jit`` and ``vmap`` to super efficiently simulate the differentiable quantum architecture search (DQAS) algorithm, where an ensemble of quantum circuits with different structures can be compiled to simulate at the same time.\n",
    "\n",
    "[WIP note]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0704d1f4",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b523ddaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorcircuit as tc\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b35c5c59",
   "metadata": {},
   "outputs": [],
   "source": [
    "K = tc.set_backend(\"tensorflow\")\n",
    "ctype, rtype = tc.set_dtype(\"complex128\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff3654c0",
   "metadata": {},
   "source": [
    "## Problem Description\n",
    "\n",
    "The task is to find the state preparation circuit for GHZ state $\\vert \\text{GHZ}_N\\rangle = \\frac{1}{\\sqrt{2}}\\left(\\vert 0^N\\rangle +\\vert 1^N\\rangle \\right)$. We prepare a gate pool with rx0, rx1, ry0, ry1, rz0, rz1, cnot01, cnot10 for the $N=2$ demo. Amongst the eight gates, six are parameterized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9731c922",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rx0(theta):\n",
    "    return K.kron(\n",
    "        K.cos(theta) * K.eye(2) + 1.0j * K.sin(theta) * tc.gates._x_matrix, K.eye(2)\n",
    "    )\n",
    "\n",
    "\n",
    "def rx1(theta):\n",
    "    return K.kron(\n",
    "        K.eye(2), K.cos(theta) * K.eye(2) + 1.0j * K.sin(theta) * tc.gates._x_matrix\n",
    "    )\n",
    "\n",
    "\n",
    "def ry0(theta):\n",
    "    return K.kron(\n",
    "        K.cos(theta) * K.eye(2) + 1.0j * K.sin(theta) * tc.gates._y_matrix, K.eye(2)\n",
    "    )\n",
    "\n",
    "\n",
    "def ry1(theta):\n",
    "    return K.kron(\n",
    "        K.eye(2), K.cos(theta) * K.eye(2) + 1.0j * K.sin(theta) * tc.gates._y_matrix\n",
    "    )\n",
    "\n",
    "\n",
    "def rz0(theta):\n",
    "    return K.kron(\n",
    "        K.cos(theta) * K.eye(2) + 1.0j * K.sin(theta) * tc.gates._z_matrix, K.eye(2)\n",
    "    )\n",
    "\n",
    "\n",
    "def rz1(theta):\n",
    "    return K.kron(\n",
    "        K.eye(2), K.cos(theta) * K.eye(2) + 1.0j * K.sin(theta) * tc.gates._z_matrix\n",
    "    )\n",
    "\n",
    "\n",
    "def cnot01():\n",
    "    return K.cast(K.convert_to_tensor(tc.gates._cnot_matrix), ctype)\n",
    "\n",
    "\n",
    "def cnot10():\n",
    "    return K.cast(\n",
    "        K.convert_to_tensor(\n",
    "            np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])\n",
    "        ),\n",
    "        ctype,\n",
    "    )\n",
    "\n",
    "\n",
    "ops_repr = [\"rx0\", \"rx1\", \"ry0\", \"ry1\", \"rz0\", \"rz1\", \"cnot01\", \"cnot10\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "be9e5e07",
   "metadata": {},
   "outputs": [],
   "source": [
    "n, p, ch = 2, 3, 8\n",
    "# number of qubits, number of layers, size of operation pool\n",
    "\n",
    "target = tc.array_to_tensor(np.array([1, 0, 0, 1.0]) / np.sqrt(2.0))\n",
    "# target wavefunction, we here use GHZ2 state as the objective target\n",
    "\n",
    "\n",
    "def ansatz(params, structures):\n",
    "    c = tc.Circuit(n)\n",
    "    params = K.cast(params, ctype)\n",
    "    structures = K.cast(structures, ctype)\n",
    "    for i in range(p):\n",
    "        c.any(\n",
    "            0,\n",
    "            1,\n",
    "            unitary=structures[i, 0] * rx0(params[i, 0])\n",
    "            + structures[i, 1] * rx1(params[i, 1])\n",
    "            + structures[i, 2] * ry0(params[i, 2])\n",
    "            + structures[i, 3] * ry1(params[i, 3])\n",
    "            + structures[i, 4] * rz0(params[i, 4])\n",
    "            + structures[i, 5] * rz1(params[i, 5])\n",
    "            + structures[i, 6] * cnot01()\n",
    "            + structures[i, 7] * cnot10(),\n",
    "        )\n",
    "    s = c.state()\n",
    "    loss = K.sum(K.abs(target - s))\n",
    "    return loss\n",
    "\n",
    "\n",
    "vag1 = K.jit(K.vvag(ansatz, argnums=0, vectorized_argnums=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "751beb4d",
   "metadata": {},
   "source": [
    "## Probability Ensemble Approach\n",
    "\n",
    "This approach is more practical and experimental relevant and is the same algorithm described in Ref.1, though we here use advanced vmap to accelerate the simulation of circuits with different structures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c71c5552",
   "metadata": {},
   "outputs": [],
   "source": [
    "structures = K.zeros([p, ch], rtype)\n",
    "\n",
    "\n",
    "def sampling_from_structure(structures, batch=1):\n",
    "    prob = K.softmax(K.real(structures), axis=-1)\n",
    "    return np.array([np.random.choice(ch, p=K.numpy(prob[i])) for i in range(p)])\n",
    "\n",
    "\n",
    "@K.jit\n",
    "def best_from_structure(structures):\n",
    "    return K.argmax(structures, axis=-1)\n",
    "\n",
    "\n",
    "@K.jit\n",
    "def nmf_gradient(structures, oh):\n",
    "    \"\"\"\n",
    "    compute the Monte Carlo gradient with respect of naive mean-field probabilistic model\n",
    "    \"\"\"\n",
    "    choice = K.argmax(oh, axis=-1)\n",
    "    prob = K.softmax(K.real(structures), axis=-1)\n",
    "    indices = K.transpose(K.stack([K.cast(tf.range(p), \"int64\"), choice]))\n",
    "    prob = tf.gather_nd(prob, indices)\n",
    "    prob = K.reshape(prob, [-1, 1])\n",
    "    prob = K.tile(prob, [1, ch])\n",
    "\n",
    "    return tf.tensor_scatter_nd_add(\n",
    "        tf.cast(-prob, dtype=ctype),\n",
    "        indices,\n",
    "        tf.ones([p], dtype=ctype),\n",
    "    )\n",
    "\n",
    "\n",
    "nmf_gradient_vmap = K.vmap(nmf_gradient, vectorized_argnums=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e9072ce1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Using a while_loop for converting GatherNd\n",
      "WARNING:tensorflow:Using a while_loop for converting TensorScatterAdd\n",
      "----------epoch 0-----------\n",
      "batched average loss:  1.4592934704536051\n",
      "best candidates so far: ['rz0', 'rx1', 'ry1']\n",
      "corresponding weights for each gate: [-0.057327405380022844, -0.045601297710585284, -0.04636363032931133]\n",
      "----------epoch 40-----------\n",
      "batched average loss:  1.030357139800294\n",
      "best candidates so far: ['rz1', 'rz1', 'rz1']\n",
      "corresponding weights for each gate: [0.010308291000538396, 0.023200255203119288, 0.027883394318059553]\n",
      "----------epoch 80-----------\n",
      "batched average loss:  0.22124358935767952\n",
      "best candidates so far: ['ry0', 'ry0', 'cnot01']\n",
      "corresponding weights for each gate: [-0.14457921547763192, -0.7551699008410323, 0.0]\n",
      "----------epoch 120-----------\n",
      "batched average loss:  0.023774129899832103\n",
      "best candidates so far: ['rz0', 'ry0', 'cnot01']\n",
      "corresponding weights for each gate: [0.02096313495888724, -0.8123872497459989, 0.0]\n",
      "----------epoch 160-----------\n",
      "batched average loss:  0.004546865740426344\n",
      "best candidates so far: ['cnot01', 'ry0', 'cnot01']\n",
      "corresponding weights for each gate: [0.0, -0.7612616254391652, 0.0]\n",
      "----------epoch 200-----------\n",
      "batched average loss:  0.014621489354652689\n",
      "best candidates so far: ['cnot01', 'ry0', 'cnot01']\n",
      "corresponding weights for each gate: [0.0, -0.765156297260239, 0.0]\n",
      "----------epoch 240-----------\n",
      "batched average loss:  0.006180035538359833\n",
      "best candidates so far: ['cnot01', 'ry0', 'cnot01']\n",
      "corresponding weights for each gate: [0.0, -0.8020375864055442, 0.0]\n",
      "----------epoch 280-----------\n",
      "batched average loss:  0.010882250249743137\n",
      "best candidates so far: ['cnot01', 'ry0', 'cnot01']\n",
      "corresponding weights for each gate: [0.0, -0.7988850988688151, 0.0]\n",
      "----------epoch 320-----------\n",
      "batched average loss:  0.001633935902769458\n",
      "best candidates so far: ['cnot01', 'ry0', 'cnot01']\n",
      "corresponding weights for each gate: [0.0, -0.7962192912123128, 0.0]\n",
      "----------epoch 360-----------\n",
      "batched average loss:  0.000863812265081192\n",
      "best candidates so far: ['cnot01', 'ry0', 'cnot01']\n",
      "corresponding weights for each gate: [0.0, -0.794013318167943, 0.0]\n",
      "----------epoch 399-----------\n",
      "batched average loss:  0.014226146799717954\n",
      "best candidates so far: ['cnot01', 'ry0', 'cnot01']\n",
      "corresponding weights for each gate: [0.0, -0.785398050740091, 0.0]\n"
     ]
    }
   ],
   "source": [
    "verbose = False\n",
    "epochs = 400\n",
    "batch = 256\n",
    "lr = tf.keras.optimizers.schedules.ExponentialDecay(0.06, 100, 0.5)\n",
    "structure_opt = tc.backend.optimizer(tf.keras.optimizers.Adam(0.12))\n",
    "network_opt = tc.backend.optimizer(tf.keras.optimizers.Adam(lr))\n",
    "nnp = K.implicit_randn(stddev=0.02, shape=[p, 6], dtype=rtype)\n",
    "stp = K.implicit_randn(stddev=0.02, shape=[p, 8], dtype=rtype)\n",
    "avcost1 = 0\n",
    "for epoch in range(epochs):  # iteration to update strcuture param\n",
    "    avcost2 = avcost1\n",
    "    costl = []\n",
    "    batched_stuctures = K.onehot(\n",
    "        np.stack([sampling_from_structure(stp) for _ in range(batch)]), num=8\n",
    "    )\n",
    "    infd, gnnp = vag1(nnp, batched_stuctures)\n",
    "    gs = nmf_gradient_vmap(structures, batched_stuctures)  # \\nabla lnp\n",
    "    gstp = [K.cast((infd[i] - avcost2), ctype) * gs[i] for i in range(infd.shape[0])]\n",
    "    gstp = K.real(K.sum(gstp, axis=0) / infd.shape[0])\n",
    "    avcost1 = K.sum(infd) / infd.shape[0]\n",
    "    nnp = network_opt.update(gnnp, nnp)\n",
    "    stp = structure_opt.update(gstp, stp)\n",
    "\n",
    "    if epoch % 40 == 0 or epoch == epochs - 1:\n",
    "        print(\"----------epoch %s-----------\" % epoch)\n",
    "        print(\n",
    "            \"batched average loss: \",\n",
    "            np.mean(avcost1),\n",
    "        )\n",
    "\n",
    "        if verbose:\n",
    "            print(\n",
    "                \"strcuture parameter: \\n\",\n",
    "                stp.numpy(),\n",
    "                \"\\n network parameter: \\n\",\n",
    "                nnp.numpy(),\n",
    "            )\n",
    "\n",
    "        cand_preset = best_from_structure(stp)\n",
    "        print(\"best candidates so far:\", [ops_repr[i] for i in cand_preset])\n",
    "        print(\n",
    "            \"corresponding weights for each gate:\",\n",
    "            [K.numpy(nnp[j, i]) if i < 6 else 0.0 for j, i in enumerate(cand_preset)],\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e2824b3",
   "metadata": {},
   "source": [
    "## Directly Optimize the Structure Parameters\n",
    "\n",
    "Since we are using numerical simulation anyway, we can directly optimize the structure parameter and omit whether the super circuit is unitary or not, this approach can be faster and more reliable for some scenarios."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e709fb2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ansatz2(params, structures):\n",
    "    c = tc.Circuit(n)\n",
    "    params = K.cast(params, ctype)\n",
    "    structures = K.softmax(structures, axis=-1)\n",
    "    structures = K.cast(structures, ctype)\n",
    "    for i in range(p):\n",
    "        c.any(\n",
    "            0,\n",
    "            1,\n",
    "            unitary=structures[i, 0] * rx0(params[i, 0])\n",
    "            + structures[i, 1] * rx1(params[i, 1])\n",
    "            + structures[i, 2] * ry0(params[i, 2])\n",
    "            + structures[i, 3] * ry1(params[i, 3])\n",
    "            + structures[i, 4] * rz0(params[i, 4])\n",
    "            + structures[i, 5] * rz1(params[i, 5])\n",
    "            + structures[i, 6] * cnot01()\n",
    "            + structures[i, 7] * cnot10(),\n",
    "        )\n",
    "    s = c.state()\n",
    "    s /= K.norm(s)\n",
    "    loss = K.sum(K.abs(target - s))\n",
    "    return loss\n",
    "\n",
    "\n",
    "vag2 = K.jit(K.value_and_grad(ansatz2, argnums=(0, 1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "69258140",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------epoch 0-----------\n",
      "batched average loss:  1.3046788213442395\n",
      "strcuture parameter: \n",
      " [[ 0.07621179  0.04934165  0.04669995  0.04737751  0.02036102  0.01170415\n",
      "   0.03786593 -0.05644197]\n",
      " [ 0.01168381  0.0561013   0.02979136  0.03134415  0.03763557  0.03739249\n",
      "   0.03408754 -0.05335854]\n",
      " [ 0.03540374  0.03219197  0.01680129  0.02014464  0.06939972  0.02393527\n",
      "   0.04619596 -0.01844729]] \n",
      " network parameter: \n",
      " [[-0.0584098   0.04281717  0.0642035   0.06008445  0.0357175   0.05512457]\n",
      " [-0.07067937  0.04410743  0.03608519  0.03465959  0.02446072  0.06917318]\n",
      " [-0.01337738  0.04776898  0.04278249  0.04917169  0.0495427   0.01059102]]\n",
      "best candidates so far: ['rx0', 'rx1', 'rz0']\n",
      "corresponding weights for each gate: [-0.058409803714939854, 0.04410743113093344, 0.04954270315507654]\n",
      "----------epoch 70-----------\n",
      "batched average loss:  1.0081966098666586\n",
      "strcuture parameter: \n",
      " [[-0.91750096  0.35057522  0.32585577  0.37681816  1.77239369  1.7734987\n",
      "   1.80143958 -0.38591221]\n",
      " [ 0.30087524  0.28764993  0.36971695  0.36078872  1.79887933  1.47542633\n",
      "   1.79490296 -0.38283427]\n",
      " [ 0.29950339  0.32101711 -0.07372448  0.34959339  1.83486426  1.78887106\n",
      "   1.81320642 -0.34792317]] \n",
      " network parameter: \n",
      " [[ 0.01163284 -0.02749067 -0.00602475  0.46422017 -0.03365732 -0.01443091]\n",
      " [-0.00057541 -0.02624807 -0.03408587  0.43879875 -0.04520759 -0.00055711]\n",
      " [ 0.05673025  0.03099979 -0.02736317  0.45331194 -0.02026327 -0.00559595]]\n",
      "best candidates so far: ['cnot01', 'rz0', 'rz0']\n",
      "corresponding weights for each gate: [0.0, -0.045207589104267296, -0.02026326781055693]\n",
      "----------epoch 140-----------\n",
      "batched average loss:  0.8049806880722175\n",
      "strcuture parameter: \n",
      " [[-3.20900567 -2.18126972  1.96173331  0.3704988   0.75310085  2.01979348\n",
      "   2.47701794 -0.37965676]\n",
      " [-0.78487034 -1.05072503  0.83960507  0.35409074  1.49913186  0.4284363\n",
      "   4.58858068 -0.37664102]\n",
      " [ 0.72348068  0.29661214  0.82121041  0.34328667  4.57946006  3.79373413\n",
      "   2.24252671 -0.3416766 ]] \n",
      " network parameter: \n",
      " [[-5.93268249e-04 -4.03543595e-02 -1.13260135e+00  4.62883177e-01\n",
      "  -3.47753230e-02 -1.57096245e-02]\n",
      " [-1.38210543e-03 -5.03624409e-02  1.02006945e+00  4.37465879e-01\n",
      "  -4.64645263e-02 -1.16956649e-03]\n",
      " [-7.80346264e-02  1.90816551e-02  1.09724554e+00  4.51972627e-01\n",
      "  -2.15345680e-02 -6.84665987e-03]]\n",
      "best candidates so far: ['cnot01', 'cnot01', 'rz0']\n",
      "corresponding weights for each gate: [0.0, 0.0, -0.02153456797370576]\n",
      "----------epoch 210-----------\n",
      "batched average loss:  0.041816948869616476\n",
      "strcuture parameter: \n",
      " [[-3.86458991 -2.84058112  2.62327171  0.26992388  0.09167012  1.35827717\n",
      "   1.81549048 -0.56415243]\n",
      " [-1.16314411 -1.7698344   1.49466411 -0.30614419  0.75064439 -0.31409853\n",
      "   5.25000534 -0.65623059]\n",
      " [ 1.4704075   0.89799938  2.01474589  2.50046978  4.8946084   4.44647834\n",
      "   1.49549043 -0.52420796]] \n",
      " network parameter: \n",
      " [[ 0.00716229  0.0950563  -1.62490102  0.60459966 -0.033863   -0.01472524]\n",
      " [ 0.41329341  0.02296645  1.58326833  0.57927215 -0.04604745 -0.05234586]\n",
      " [-0.07409766  0.08796055 -0.2881097  -0.52346262 -0.02053635 -0.00585734]]\n",
      "best candidates so far: ['ry0', 'cnot01', 'rz0']\n",
      "corresponding weights for each gate: [-1.624901021386579, 0.0, -0.020536353581361112]\n",
      "----------epoch 280-----------\n",
      "batched average loss:  0.04771541732805661\n",
      "strcuture parameter: \n",
      " [[-3.86686803 -2.8436989   2.62698311  0.26657346  0.0879849   1.35457084\n",
      "   1.81178177 -0.67868932]\n",
      " [-1.33926046 -1.75120967  1.4909566  -0.18080598  0.75530859 -0.30731494\n",
      "   5.25370236 -0.67797002]\n",
      " [ 1.47961761  0.93984054  2.12875762  2.49693907  4.78860574  4.56031727\n",
      "   1.50105437 -3.90111934]] \n",
      " network parameter: \n",
      " [[ 0.01376387  0.10062571 -1.62306954  0.60290458 -0.0321182  -0.01292187]\n",
      " [ 0.41014329  0.10543278  1.58481867  0.66429872 -0.04488787 -0.0548457 ]\n",
      " [-0.07315826  0.09047629 -0.35272068 -0.52529957 -0.01871782 -0.00404166]]\n",
      "best candidates so far: ['ry0', 'cnot01', 'rz0']\n",
      "corresponding weights for each gate: [-1.6230695444527814, 0.0, -0.018717818357922293]\n",
      "----------epoch 350-----------\n",
      "batched average loss:  0.0484244468649333\n",
      "strcuture parameter: \n",
      " [[-3.86889078 -2.84635759  2.63008504  0.26367828  0.08490232  1.35147264\n",
      "   1.80868181 -0.68180282]\n",
      " [-1.61872015 -1.73606594  1.48784963 -0.16422792  0.75890549 -0.30250011\n",
      "   5.25680704 -0.70078063]\n",
      " [ 1.48953381  0.96377125  2.13183873  2.49380904  4.79193079  4.5635497\n",
      "   1.5054478  -5.6630325 ]] \n",
      " network parameter: \n",
      " [[ 0.02014668  0.10330406 -1.62102814  0.60086296 -0.03016668 -0.01090957]\n",
      " [ 0.40649692  0.12322429  1.58685415  0.67753112 -0.04356077 -0.05752607]\n",
      " [-0.0730366   0.09201651 -0.35067966 -0.52733454 -0.01668937 -0.00201604]]\n",
      "best candidates so far: ['ry0', 'cnot01', 'rz0']\n",
      "corresponding weights for each gate: [-1.6210281393366774, 0.0, -0.01668936553578817]\n",
      "----------epoch 420-----------\n",
      "batched average loss:  0.0490371292665724\n",
      "strcuture parameter: \n",
      " [[-3.8707214  -2.84868677  2.63274431  0.26116667  0.08225763  1.34881617\n",
      "   1.80602401 -0.68447817]\n",
      " [-2.17422677 -1.72638998  1.48517837 -0.14718937  0.76130497 -0.30044208\n",
      "   5.25948165 -0.72889795]\n",
      " [ 1.50031397  0.986817    2.13447811  2.49111713  4.79485557  4.56643982\n",
      "   1.50815833 -6.42399688]] \n",
      " network parameter: \n",
      " [[ 2.66201104e-02  1.04232454e-01 -1.61904956e+00  5.98884773e-01\n",
      "  -2.82809595e-02 -8.96136472e-03]\n",
      " [ 4.01114366e-01  1.42266730e-01  1.58916732e+00  6.88026030e-01\n",
      "  -4.23682835e-02 -6.01395817e-02]\n",
      " [-7.30273413e-02  9.34775776e-02 -3.48701372e-01 -5.29307121e-01\n",
      "  -1.47240949e-02 -5.37945560e-05]]\n",
      "best candidates so far: ['ry0', 'cnot01', 'rz0']\n",
      "corresponding weights for each gate: [-1.6190495562985034, 0.0, -0.014724094945736954]\n",
      "----------epoch 490-----------\n",
      "batched average loss:  0.04976228840362948\n",
      "strcuture parameter: \n",
      " [[-3.87241212 -2.85077089  2.63506255  0.25896325  0.07995022  1.3465\n",
      "   1.80370686 -0.68681685]\n",
      " [-2.94364254 -1.72476469  1.48284202 -0.13112958  0.76233067 -0.30191621\n",
      "   5.26182547 -0.76243791]\n",
      " [ 1.5130618   1.0086917   2.13677828  2.4887622   4.79747567  4.56906223\n",
      "   1.50813331 -6.85905145]] \n",
      " network parameter: \n",
      " [[ 0.03313986  0.1034762  -1.6172524   0.59708835 -0.02657844 -0.00719578]\n",
      " [ 0.39208856  0.1630666   1.59170906  0.69656645 -0.04144895 -0.06259256]\n",
      " [-0.07305937  0.09493765 -0.34690445 -0.53109883 -0.01294077  0.00172629]]\n",
      "best candidates so far: ['ry0', 'cnot01', 'rz0']\n",
      "corresponding weights for each gate: [-1.617252401393024, 0.0, -0.012940768995802935]\n",
      "----------epoch 560-----------\n",
      "batched average loss:  0.05065420046477945\n",
      "strcuture parameter: \n",
      " [[-3.8739842  -2.85266121  2.6371075   0.25701232  0.0779132   1.34445655\n",
      "   1.80166269 -0.68888562]\n",
      " [-3.87431144 -1.73479083  1.48077359 -0.1180455   0.76187864 -0.30734926\n",
      "   5.26390441 -0.7996279 ]\n",
      " [ 1.52854329  1.02822466  2.13880714  2.48667702  4.79985267  4.57146409\n",
      "   1.50375206 -7.1518221 ]] \n",
      " network parameter: \n",
      " [[ 0.03952625  0.09952938 -1.6156793   0.59551608 -0.02510117 -0.0056556 ]\n",
      " [ 0.37625454  0.18679017  1.59450355  0.70332975 -0.04084047 -0.0651834 ]\n",
      " [-0.07341076  0.09584067 -0.3453315  -0.53266713 -0.01138222  0.00328135]]\n",
      "best candidates so far: ['ry0', 'cnot01', 'rz0']\n",
      "corresponding weights for each gate: [-1.6156792973551577, 0.0, -0.01138222012800648]\n",
      "----------epoch 630-----------\n",
      "batched average loss:  0.05169300354431061\n",
      "strcuture parameter: \n",
      " [[-3.87544608 -2.85438213  2.63892729  0.25527243  0.0760992   1.34263784\n",
      "   1.79984346 -0.69073143]\n",
      " [-4.81271904 -1.75890412  1.47892571 -0.1106167   0.75992863 -0.31691451\n",
      "   5.26576503 -0.83861308]\n",
      " [ 1.54468553  1.04577833  2.14061233  2.48481415  4.80202942  4.57367936\n",
      "   1.49248445 -7.36638329]] \n",
      " network parameter: \n",
      " [[ 0.04556425  0.08358514 -1.61433443  0.59417203 -0.02385289 -0.00434529]\n",
      " [ 0.35393466  0.2231968   1.59768193  0.70856295 -0.04056008 -0.06806482]\n",
      " [-0.07428592  0.0956737  -0.3439867  -0.53400786 -0.01005294  0.00460687]]\n",
      "best candidates so far: ['ry0', 'cnot01', 'rz0']\n",
      "corresponding weights for each gate: [-1.614334431628148, 0.0, -0.010052938889800701]\n",
      "----------epoch 699-----------\n",
      "batched average loss:  0.10256012079470897\n",
      "strcuture parameter: \n",
      " [[-3.787228   -2.7662122   2.55075055  0.34350642  0.16427612  1.43081451\n",
      "   1.88802021 -0.60259055]\n",
      " [-5.6073258  -2.03961941  1.56706312 -0.37366319  0.66712936 -0.41905248\n",
      "   5.17764218 -0.96685202]\n",
      " [ 1.46844149  0.80497766  2.05243854  2.57293907  4.71424946  4.48595762\n",
      "   1.38208153 -7.52953678]] \n",
      " network parameter: \n",
      " [[ 4.11010856e-02  9.05029846e-03 -1.60328084e+00  5.83119412e-01\n",
      "  -1.28978327e-02  6.67130794e-03]\n",
      " [ 3.37676048e-01  5.40067470e-01  1.61148200e+00  7.22459586e-01\n",
      "  -3.43741776e-02 -6.15320200e-02]\n",
      " [-6.44290250e-02  8.86161369e-02 -3.32933747e-01 -5.45057551e-01\n",
      "   9.83673124e-04  1.56395047e-02]]\n",
      "best candidates so far: ['ry0', 'cnot01', 'rz0']\n",
      "corresponding weights for each gate: [-1.6032808418929712, 0.0, 0.0009836731240883082]\n"
     ]
    }
   ],
   "source": [
    "verbose = True\n",
    "epochs = 700\n",
    "lr = tf.keras.optimizers.schedules.ExponentialDecay(0.05, 200, 0.5)\n",
    "structure_opt = tc.backend.optimizer(tf.keras.optimizers.Adam(0.04))\n",
    "network_opt = tc.backend.optimizer(tf.keras.optimizers.Adam(lr))\n",
    "nnp = K.implicit_randn(stddev=0.02, shape=[p, 6], dtype=rtype)\n",
    "stp = K.implicit_randn(stddev=0.02, shape=[p, 8], dtype=rtype)\n",
    "for epoch in range(epochs):\n",
    "\n",
    "    infd, (gnnp, gstp) = vag2(nnp, stp)\n",
    "\n",
    "    nnp = network_opt.update(gnnp, nnp)\n",
    "    stp = structure_opt.update(gstp, stp)\n",
    "    if epoch % 70 == 0 or epoch == epochs - 1:\n",
    "        print(\"----------epoch %s-----------\" % epoch)\n",
    "        print(\n",
    "            \"batched average loss: \",\n",
    "            np.mean(infd),\n",
    "        )\n",
    "        if verbose:\n",
    "            print(\n",
    "                \"strcuture parameter: \\n\",\n",
    "                stp.numpy(),\n",
    "                \"\\n network parameter: \\n\",\n",
    "                nnp.numpy(),\n",
    "            )\n",
    "\n",
    "        cand_preset = best_from_structure(stp)\n",
    "        print(\"best candidates so far:\", [ops_repr[i] for i in cand_preset])\n",
    "        print(\n",
    "            \"corresponding weights for each gate:\",\n",
    "            [K.numpy(nnp[j, i]) if i < 6 else 0.0 for j, i in enumerate(cand_preset)],\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fc3d1df",
   "metadata": {},
   "source": [
    "## Final Fine-tune\n",
    "\n",
    "For the obtained circuit layout we can further adjust the circuit weights to make the objective more close to zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "486aa70d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 3, 8), dtype=float32, numpy=\n",
       "array([[[0., 0., 1., 0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 0., 0., 1., 0.]]], dtype=float32)>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chosen_structure = K.onehot(np.array([2, 4, 6]), num=8)\n",
    "chosen_structure = K.reshape(chosen_structure, [1, p, ch])\n",
    "chosen_structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7d003077",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 loss:  0.9827084438054802\n",
      "60 loss:  0.9449745688150044\n",
      "120 loss:  0.8850948396917335\n",
      "180 loss:  0.8048454837720991\n",
      "240 loss:  0.706158632509899\n",
      "300 loss:  0.5901794549931197\n",
      "360 loss:  0.45808014651166296\n",
      "420 loss:  0.3113751664914397\n",
      "480 loss:  0.1520672098147883\n",
      "540 loss:  0.0014714944860031312\n",
      "599 loss:  0.0032763286672428237\n"
     ]
    }
   ],
   "source": [
    "network_opt = tc.backend.optimizer(tf.keras.optimizers.Adam(1e-3))\n",
    "nnp = K.implicit_randn(stddev=0.02, shape=[p, 6], dtype=rtype)\n",
    "verbose = True\n",
    "epochs = 600\n",
    "for epoch in range(epochs):\n",
    "    infd, gnnp = vag1(nnp, chosen_structure)\n",
    "    nnp = network_opt.update(gnnp, nnp)\n",
    "    if epoch % 60 == 0 or epoch == epochs - 1:\n",
    "        print(epoch, \"loss: \", K.numpy(infd[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58b1ddc3",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "1. https://arxiv.org/pdf/2010.08561.pdf"
   ]
  }
 ],
 "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}