qgrnn_ising/qgrnn_ising.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "QGRNN_Ising_copy.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u3-uOLgO1R74",
        "colab_type": "text"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Quantum Authors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t8XVZWCF08e4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ls6iWjEfITx7",
        "colab_type": "text"
      },
      "source": [
        "# Quantum Graph Recurrent Neural Networks for Ising model\n",
        "\n",
        "Author : Jae H. Yoo, Google Research.\n",
        "\n",
        "Contributors : Guillaume Verdon (X company) Trevor McCourt, Antonio J. Martinez, Michael Broughton (Google Research)\n",
        "\n",
        "Created : 2020-Feb-01\n",
        "\n",
        "Last updated : 2020-Feb-02 "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6w4cSncf1Fcl",
        "colab_type": "text"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tensorflow/quantum/blob/research/qgrnn_ising/qgrnn_ising.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gJj9eXch1bv1",
        "colab_type": "text"
      },
      "source": [
        "## Introduction\n",
        "\n",
        "In this colab, we will learn how to train QGRNN, a variant of Quantum Graph Neural Networks ([Verdon et al.](https://arxiv.org/abs/1909.12264)) to learn the dynamics of the target Hamiltonian of given Ising model. Before going further, we will install related libraries."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WgaptKeDJgg1",
        "colab_type": "text"
      },
      "source": [
        "### Import & pip install libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JuXxC5fbaGAS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install --upgrade tensorflow==2.7.0"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "Xx3k3M39y0SW",
        "colab": {}
      },
      "source": [
        "!pip install tensorflow-quantum"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E-5MiNQHNpS3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import networkx as nx\n",
        "import numpy as np\n",
        "import sympy\n",
        "import cirq\n",
        "# TensorFlow\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras import Input, Model\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "import tensorflow.keras.backend as K\n",
        "# TensorFlow Quantum\n",
        "import tensorflow_quantum as tfq\n",
        "from tensorflow_quantum.python.layers import Expectation\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1ZtU7j3LJsGq",
        "colab_type": "text"
      },
      "source": [
        "### Overview\n",
        "\n",
        "Here are steps we will follow.\n",
        "\n",
        "- Preparation of Quantum Data with VQE on Ising model.\n",
        "- Construct a QGRNN model\n",
        "- Construct Fidelity with Swap Test\n",
        "- Calculate the average infidelity loss function\n",
        "- Train the QGRNN & get the final result\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kJctwuR0Od3W",
        "colab_type": "text"
      },
      "source": [
        "## Preparation of Quantum Data with VQE on Ising model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wAcIwSsNSS5t",
        "colab_type": "text"
      },
      "source": [
        "We will construct a target Hamiltonian $H_{target}$ by using a random ring graph $G$ with $N=6$ qubits. On top of it, we will find a low energy state near to ground state of $H_{target}$ by using Variational Quantum Eigensolver (VQE)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "di3PZEO9U9bN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Define the number of qubits of our target quantum system.\n",
        "N = 6"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qUVwr99DTkVB",
        "colab_type": "text"
      },
      "source": [
        "### Transverse field Ising model Hamiltonian\n",
        "\n",
        "Here are basic introduction of the target Hamiltonian of Ising model. Given\n",
        "* $J_{jk}$ for interacting spin pairs and\n",
        "* $B_{v}$ for site bias term of each spin,\n",
        "\n",
        "$H_{target} = \\sum_{j,k} J_{jk} Z_j Z_k + \\sum_{v} B_v Z_v + \\sum_{v} X_v$\n",
        "\n",
        "It is very easy to construct this Hamiltonian using networkx library in Python."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zd4zMjthU470",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "G = nx.cycle_graph(N)\n",
        "weights = [4*(np.random.random()-0.5) for _ in G.edges]\n",
        "biases = [4*(np.random.random()-0.5) for _ in G.nodes]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YEdTzOq1rK4h",
        "colab_type": "text"
      },
      "source": [
        "Also, networkx provides graphic APIs to draw graph data. You can see that nodes have Brown to Blue Green colormap, and edges have Red to Blue one."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8XJlkeCTrSYf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def draw(graph, pos, weights, biases, title):\n",
        "    label = {i:'{}'.format(i) for i in graph.nodes}\n",
        "    edge_options = {\n",
        "        \"edge_color\": weights,\n",
        "        \"width\": 4,\n",
        "        \"edge_cmap\": plt.cm.RdBu,\n",
        "        \"edge_vmin\" : -2,\n",
        "        \"edge_vmax\" : 2,\n",
        "    }\n",
        "    node_options = {\n",
        "        \"node_color\": biases,\n",
        "        \"cmap\": plt.cm.BrBG,\n",
        "        \"vmin\" : -2,\n",
        "        \"vmax\" : 2,\n",
        "    }\n",
        "    nx.draw_networkx_labels(graph, pos, label, font_color=\"w\")\n",
        "    nodes = nx.draw_networkx_nodes(graph, pos, **node_options)\n",
        "    edges = nx.draw_networkx_edges(graph, pos, **edge_options)\n",
        "    edges.set_cmap(plt.cm.RdBu)\n",
        "    edges.set_clim(-2, 2)\n",
        "\n",
        "    plt.title(title)\n",
        "    plt.colorbar(nodes)\n",
        "    plt.colorbar(edges)\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "pos = nx.circular_layout(G)\n",
        "draw(G, pos, weights, biases, 'Target Ising model')\n",
        "print(*zip(G.edges, weights))\n",
        "print(*zip(G.nodes, biases))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ifqGlfp0AQeB",
        "colab_type": "text"
      },
      "source": [
        "Now that we have the graph structure, weights of edges & nodes, we can construct `cirq` based Hamiltonian which can be directly calculated in `cirq` and `tfq`. To create Hamiltonian by using `cirq.PauliSum`'s or `cirq.PauliString`'s we need to assign appropriate qubits on them. We can bring qubits by using `cirq.GridQubit`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2AamI46wAgkN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def Hamiltonian(graph, weights, biases, qubits):\n",
        "    H_cost = [w * cirq.Z(qubits[i]) * cirq.Z(qubits[j]) \\\n",
        "            for (i, j), w in zip(graph.edges, weights)]\n",
        "    H_cost += [b * cirq.Z(qubits[v]) for v, b in enumerate(biases)]\n",
        "    H_mixer = [cirq.X(q) for q in qubits]\n",
        "    return H_cost, H_mixer\n",
        "\n",
        "qubits = cirq.GridQubit.rect(1, N)\n",
        "true_H_cost, true_H_mixer = Hamiltonian(G, weights, biases, qubits) \n",
        "for cost in true_H_cost:\n",
        "    print(cost)\n",
        "for mixer in true_H_mixer:\n",
        "    print(mixer)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iGGHlPZPT0T1",
        "colab_type": "text"
      },
      "source": [
        "### Variational Quantum Eigensolver (VQE)\n",
        "\n",
        "Here, we will construct a Variational Quantum Eigensolver (VQE) to find out a low energy state $|\\psi_0\\rangle$ near to the ground state."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-S_H9EcwCqlR",
        "colab_type": "text"
      },
      "source": [
        "#### Variational method\n",
        "\n",
        "[Variational method](https://en.wikipedia.org/wiki/Variational_method_(quantum_mechanics)) provides the way to find approximated ground energy states. Let $|\\psi(\\vec\\theta)\\rangle$ be a variational ansatz. We can control parameters $\\vec\\theta$ to change the quantum state. Then, the following inequality is guaranteed for any given Hamiltonian $H$.\n",
        "\n",
        "- $\\langle\\psi(\\vec\\theta)|H|\\psi(\\vec\\theta)\\rangle \\ge E_0$, where $E_0$ is the energy of ground state.\n",
        "- We can find out $\\vec\\theta$ giving us approximated ground state by minimizing the above expectation value as a loss function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AlKAegEBE7zO",
        "colab_type": "text"
      },
      "source": [
        "#### Variational Ansatz\n",
        "\n",
        "We can make any quantum state by using $X$ rotation & $Z$ rotation in Bloch sphere, which could be a simple and useful variational ansatz for our problem. Because our ansatz will be used for Ising model & QGRNN later, we should not assign qubits to the rotational gates because qubits are immutable in `cirq.Circuit`. That's why we use `VQE.gates` and `VQE.circuit` separately.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vNCnp0jXFPA8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class VQE:\n",
        "    \"\"\"Variational Quantum Eigensolver\"\"\"\n",
        "    def __init__(self, qubits):\n",
        "        self.qubits = qubits\n",
        "        # Parameters\n",
        "        t_x = ['t_x{}'.format(i) for i, _ in enumerate(qubits)]\n",
        "        t_z = ['t_z{}'.format(i) for i, _ in enumerate(qubits)]\n",
        "        self.symbols = t_x + t_z\n",
        "        # Parameterized quantum gates without Qubits\n",
        "        gates = []\n",
        "        gates.append([cirq.XPowGate(exponent=sympy.Symbol(x)) for x in t_x])\n",
        "        gates.append([cirq.ZPowGate(exponent=sympy.Symbol(z)) for z in t_z])\n",
        "        self._gates = gates\n",
        "        self.circuit = self.get_state(qubits)\n",
        "\n",
        "    def get_state(self, qubits, params=None):\n",
        "        \"\"\"Outputs quantum data with given qubits.\"\"\"\n",
        "        circuit = cirq.Circuit(\n",
        "            [g(i) for gates in self._gates for i, g in zip(qubits, gates)])\n",
        "        if params is None:\n",
        "            return circuit\n",
        "        resolver = {k: v for k, v in zip(self.symbols, params)}\n",
        "        return cirq.resolve_parameters(circuit, resolver)\n",
        "\n",
        "vqe = VQE(qubits)\n",
        "print(vqe.circuit)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xzz4vipJKUfP",
        "colab_type": "text"
      },
      "source": [
        "#### Construct $\\langle \\psi (\\vec\\theta) | H | \\psi (\\vec\\theta)\\rangle$\n",
        "\n",
        "TensorFlow Quantum provides `tfq.layers.Expectation` Keras layer to provide easy interface to calculate the expectation value of given ansatz & Hamiltonian."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BmUn4JHHKpLb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def vqe_keras_model(vqe, H_target):\n",
        "    # Construct measurement layers.\n",
        "    circuit_input = Input(shape=(), dtype=tf.string)\n",
        "    output = Expectation()(\n",
        "            circuit_input,\n",
        "            symbol_names=vqe.symbols,\n",
        "            operators=tfq.convert_to_tensor([H_target]),\n",
        "            initializer=tf.keras.initializers.RandomNormal())\n",
        "    # Each term in H_target is calculated respectively.\n",
        "    # Here, we sum them up to get the final <H>.\n",
        "    output = tf.math.reduce_sum(output, axis=-1, keepdims=True)\n",
        "\n",
        "    # Model compile\n",
        "    model = Model(inputs=circuit_input, outputs=output)\n",
        "    adam = Adam(learning_rate=0.05)\n",
        "    model.compile(optimizer=adam, loss='mse')\n",
        "    return model\n",
        "\n",
        "H_target = true_H_cost + true_H_mixer\n",
        "model = vqe_keras_model(vqe, H_target)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SsfJPODEJzKx",
        "colab_type": "text"
      },
      "source": [
        "#### Minimizing $\\langle \\psi (\\vec\\theta) | H | \\psi (\\vec\\theta)\\rangle$\n",
        "\n",
        "Keras model is used for training VQE. To feed quantum data in the form of `cirq.Circuit`, `tfq.convert_to_tensor()` will be frequently used to make them `tf.Tensor`.\n",
        "\n",
        "We have quantum input data. What's the output true value? Isn't it just a minimization problem?\n",
        "\n",
        "Here is some tip for training. Setting the output true value to theoretical lower bound, we can minimize our expectation value in the Keras model fit framework. The how can we calculate the lower bound? By the fact that the expectation value of any PauliString is bounded in [-1, 1], we can easily find the lower bound.\n",
        "\n",
        "- $\\langle \\psi (\\vec\\theta) | H | \\psi (\\vec\\theta)\\rangle = \\sum_{jk}J_{jk}\\langle Z_jZ_k\\rangle + \\sum_{v}B_{v}\\langle Z_v\\rangle + \\sum_{v}\\langle X_v\\rangle \\ge \\sum_{jk}(-)|J_{jk}| -\\sum_{v}|B_{v}| - N $"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zIbP9w8AJ8_I",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "lower_bound = -np.sum(np.abs(weights + biases)) - N\n",
        "vqe_input = tfq.convert_to_tensor([vqe.circuit])\n",
        "vqe_output = tf.convert_to_tensor([[lower_bound]])\n",
        "print('Before training : <H>={}'.format(model.predict(x=vqe_input)))\n",
        "history = model.fit(x=vqe_input, y=vqe_output, batch_size=1, epochs=100,\n",
        "                    verbose=0)\n",
        "plt.plot(history.history['loss'])\n",
        "print('After training : <H>={} >= {}'.format(model.predict(x=vqe_input),\n",
        "                                             lower_bound))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u6UG2-XSlzxE",
        "colab_type": "text"
      },
      "source": [
        "In the next sections, we will use this low energy state as initial states of Ising model & QGRNN. Since they have different qubits, we need to create both."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NbdHQpzcmAJW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "vqe_params = model.get_weights()[0]\n",
        "low_energy_ising = vqe.get_state(qubits, vqe_params)\n",
        "# For QGRNN, get new qubit indices (0, N)~(0,2*N-1)\n",
        "qubits_qgrnn = cirq.GridQubit.rect(1, N, 0, N)\n",
        "low_energy_qgrnn = vqe.get_state(qubits_qgrnn, vqe_params)\n",
        "print(low_energy_ising)\n",
        "print(low_energy_qgrnn)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QiQzwRN0oy8a",
        "colab_type": "text"
      },
      "source": [
        "### Construct Ising model\n",
        "\n",
        "In VQE, we can just use Hamiltonian $H$ directly. However, we need to exponentiate Hamiltonians to construct Ising model. Moreover, it is proved that exponentiation of two non-commutable observables is not equal to exponentiation of the sum of them. That is, $e^{A+B} \\neq e^Ae^B$ if $[A, B]\\neq0$.\n",
        "\n",
        "There are two options to deal with this problem.\n",
        "1. [Baker-Campbell-Hausdorff (BCH) formula](https://en.wikipedia.org/wiki/Baker%E2%80%93Campbell%E2%80%93Hausdorff_formula)\n",
        "\n",
        "  - $e^Ae^B=e^Ae^Be^{\\frac{1}{2}[A, B]}e^{\\frac{1}{12}[A, [A, B]]}e^{-\\frac{1}{12}[B, [A, B]]} ...$\n",
        "  - Analytic solution, but it is hard to be calculated.\n",
        "\n",
        "2. [Suzuki-Trotter expansion](https://en.wikipedia.org/wiki/Time-evolving_block_decimation#The_Suzuki-Trotter_expansion)\n",
        "\n",
        "  - $e^Ae^B \\simeq \\prod e^{\\delta t A}e^{\\delta t B}$\n",
        "  - Tractable, but it is an approximated solution\n",
        "\n",
        "We will use the second option. Let $P=\\frac{T_j}{\\delta t}$.\n",
        "\n",
        "  - $|\\psi_{T_j}\\rangle = U^{j}_{Ising}|\\psi_0\\rangle = e^{-i T_j H_{target}}|\\psi_0\\rangle\\sim [\\prod^{P}e^{-i {\\delta t}H_{mixer}} e^{-i {\\delta t}H_{cost}}]|\\psi_0\\rangle$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JoqmPfBMo56q",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class IsingModel:\n",
        "\n",
        "    def __init__(self, qubits, graph, weights, biases, eta=0.01):\n",
        "        self.qubits = qubits\n",
        "        self.graph = graph\n",
        "        self.weights = weights\n",
        "        self.biases = biases\n",
        "        self.eta = eta\n",
        "\n",
        "        # Construct Hamiltonian\n",
        "        _weights = [eta * w for w in weights]\n",
        "        _biases = [eta * b for b in biases]\n",
        "        H_cost, H_mixer = Hamiltonian(graph, _weights, _biases, qubits)\n",
        "        # Exponentiate each Hamiltonian\n",
        "        # The reason why we split cost & mixer is they are not commutable\n",
        "        self._cost_step = tfq.util.exponential(operators=H_cost)\n",
        "        self._mix_step = tfq.util.exponential(operators=H_mixer)\n",
        "\n",
        "    def __call__(self, input_state, depth):\n",
        "        \"\"\"Trotterizaiton\"\"\"\n",
        "        add = tfq.layers.AddCircuit()\n",
        "        output_state = add(cirq.Circuit(), append=input_state)\n",
        "        for _ in range(depth):\n",
        "            output_state = add(output_state, append=self._cost_step)\n",
        "            output_state = add(output_state, append=self._mix_step)\n",
        "        return output_state\n",
        "\n",
        "ising = IsingModel(qubits, G, weights, biases)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JQW032AZT7Jv",
        "colab_type": "text"
      },
      "source": [
        "### Time evolution of Ising model\n",
        "Let's construct an Ising model and evolve $|\\psi_0\\rangle$ with randomly sampled timesteps $T_j\\in [0, T_{max}]$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VTza-tRUe6Rg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def random_timestep_sample(batch_size, T_max=0.1):\n",
        "    return [T_max * np.random.uniform() for _ in range(batch_size)]\n",
        "\n",
        "batch_size = 15\n",
        "T = random_timestep_sample(batch_size)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hxk4j9Zgm_dE",
        "colab_type": "text"
      },
      "source": [
        "We can evolve the above `low_energy_ising` and `low_energy_qgrnn` by using exponentiation of Hamiltonian & Trotterization. If we say $P$ is a Trotterization depth with time unit ${\\delta t}$, we can get the depth $P=\\frac{T}{\\delta t}$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ua-G3kgGnXdk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "dt = 0.01\n",
        "depth = [int(t/dt)+1 for t in T] # Circuit depth for Ising & QGRNN model\n",
        "print(depth)\n",
        "\n",
        "true_final_states = []\n",
        "for P in depth:\n",
        "    true_final_states.append(ising(low_energy_ising, depth=P))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oyCwkeK8T-zZ",
        "colab_type": "text"
      },
      "source": [
        "Now that we have a set $\\{(|\\psi_0\\rangle, |\\psi_{T_j}\\rangle) | j = 1..M\\}$ where $M$ is the number of data, or batch size, we finished to generate quantum data for QGRNN."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VwKKQMcUOhCK",
        "colab_type": "text"
      },
      "source": [
        "## Construct QGRNN model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZDsRvEVLpuWq",
        "colab_type": "text"
      },
      "source": [
        "In this section, we will implement QGRNN model with \n",
        "\n",
        "- Trainable parameters : $\\theta_{jk}$, $\\phi_{v}$, $\\alpha_{v}$ \n",
        "  - $\\alpha_{v}$ can be set to a constant 1.\n",
        "- $H_{QGRNN} = \\sum_{j,k} \\theta_{jk} Z_j Z_k + \\sum_{v} \\phi_v Z_v + \\sum_{v} \\alpha_v X_v$\n",
        "\n",
        "Because the target Hamiltonian is unknown to QGRNN, we need to initialize a new random graph inside our QGRNN. In this example we will use 4-regular graph (each node has 4 edges. For 6 nodes, graph has 12 edges.) But, for simplicity, we add missing edges in the ring structure."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sJu51fYwtVYL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class QuantumGraphRNN:\n",
        "\n",
        "    def __init__(self, qubits, G, eta=0.01, train_mixer=False):\n",
        "        self.qubits = qubits\n",
        "        self.eta = eta\n",
        "        self.graph = G\n",
        "        # Set parameters for cost\n",
        "        self._theta = ['theta{}'.format(e) for e in G.edges]\n",
        "        self._phi = ['phi{}'.format(v) for v in G.nodes]\n",
        "        self.symbols = self._theta + self._phi \n",
        "        # Set parameters for mixer\n",
        "        self.train_mixer = train_mixer\n",
        "        if train_mixer:\n",
        "            self._alpha = ['alpha{}'.format(v) for v in G.nodes]\n",
        "            self.symbols += self._alpha\n",
        "        else:\n",
        "            self._alpha = [1.0 for _ in G.nodes]\n",
        "\n",
        "        # Construct Hamiltonian\n",
        "        weights = [eta] * len(G.edges)\n",
        "        biases = [eta] * len(G.nodes)\n",
        "        H_cost, H_mixer = Hamiltonian(G, weights, biases, qubits)\n",
        "        # Exponentiate each Hamiltonian with parameters\n",
        "        self._cost_step = tfq.util.exponential(\n",
        "            operators=H_cost, coefficients=self._theta + self._phi)\n",
        "        self._mix_step = tfq.util.exponential(\n",
        "            operators=H_mixer, coefficients=self._alpha)\n",
        "\n",
        "    def __call__(self, input_state, depth):\n",
        "        add = tfq.layers.AddCircuit()\n",
        "        output_state = add(cirq.Circuit(), append=input_state)\n",
        "        for _ in range(depth):\n",
        "            output_state = add(output_state, append=self._cost_step)\n",
        "            output_state = add(output_state, append=self._mix_step)\n",
        "        return output_state\n",
        "\n",
        "# The true graph of Ising model is unknown to QGRNN.\n",
        "# Think a new 4-regular random graph with at least one cycle\n",
        "# 0->1->2->3->4->5->0\n",
        "G_qgrnn = nx.random_regular_graph(n=N, d=4)\n",
        "node_color = [2.0 for _ in G_qgrnn.nodes]\n",
        "missing = []\n",
        "add_cycle = True\n",
        "if add_cycle:\n",
        "    for i in range(N):\n",
        "        j = (i+1) % N\n",
        "        if (i, j) in G_qgrnn.edges or (j, i) in G_qgrnn.edges:\n",
        "            continue\n",
        "        print('add : ', i, j)\n",
        "        G_qgrnn.add_edge(i, j)\n",
        "        missing.extend([(i, j), (j, i)])\n",
        "edge_color = [-2.0 if e in missing else 2.0 for e in G_qgrnn.edges]\n",
        "draw(G_qgrnn, pos, edge_color, node_color, '4-regular random graph with cycle')\n",
        "qgrnn = QuantumGraphRNN(qubits_qgrnn, G_qgrnn)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XyeCeJOpqmfj",
        "colab_type": "text"
      },
      "source": [
        "Let's evolve quantum input data $|\\psi_0\\rangle$ according to QGRNN ansatz, and get $U^{j}_{QGRNN}(\\theta, \\phi)|\\psi_0\\rangle$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NCtsOE79vCkR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pred_final_states = []\n",
        "for P in depth:\n",
        "    pred_final_states.append(qgrnn(low_energy_qgrnn, depth=P))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4fZ-p83KSUlH",
        "colab_type": "text"
      },
      "source": [
        "## Construct Fidelity with Swap Test"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VYeYcPa0TEe3",
        "colab_type": "text"
      },
      "source": [
        "Now we have both (1) quantum data from true Hamiltonian of Ising model and (2) predicted quantum data from QGRNN. Because they are all quantum data, we can't compare them each other as we do in the calculation of loss function at the classical deep neural networks. Measurement on a qubit will destory other qubit informations. That's why we introduce Fidelity with [Swap Test](https://en.wikipedia.org/wiki/Swap_test) here.\n",
        "\n",
        "- The expectation value of swap test of two quantum states $|\\psi\\rangle$ and $|\\phi\\rangle$ is the square of the inner product of them.\n",
        "  - $\\operatorname{Prob}(0)=\\operatorname{Prob}(Z=+1)=\\frac{1}{2}+\\frac{1}{2}|\\langle\\phi|\\psi\\rangle|^2$\n",
        "  - $\\langle Z_{test} \\rangle = 1 \\times \\operatorname{Prob}(Z=+1) + (-1) \\times \\operatorname{Prob}(Z=-1)=2\\operatorname{Prob}(Z=+1) - 1$\n",
        "  - $\\therefore \\langle Z_{test} \\rangle = |\\langle\\phi|\\psi\\rangle|^2$\n",
        "- We have ground truth Ising model state $|\\psi_{T_j}\\rangle$ and predicted state $U^{j}_{QGRNN}(\\theta, \\phi)|\\psi_0\\rangle$\n",
        "    - That is, $\\langle Z_{test} \\rangle_j = |\\langle \\psi_{T_j} | U^{j}_{QGRNN}(\\theta, \\phi)|\\psi_0\\rangle|^2 \\ge 0$\n",
        "\n",
        "As you can see, the expectation has lower bound 0.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4OuggTHHxnVZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class SwapTestFidelity:\n",
        "\n",
        "    def __init__(self, qubits1, qubits2, batch_size):\n",
        "        circuit = cirq.Circuit()\n",
        "        test_bit = set(cirq.GridQubit.rect(1, len(qubits1) + len(qubits2) + 1))\n",
        "        test_bit -= set(qubits1 + qubits2)\n",
        "        test_bit = list(test_bit)[0]\n",
        "        circuit.append(cirq.H(test_bit))\n",
        "        for a, b in zip(qubits1, qubits2):\n",
        "            circuit.append(cirq.decompose(cirq.CSwapGate().on(test_bit, a, b)))\n",
        "        circuit.append(cirq.H(test_bit))\n",
        "        self.circuit = tfq.convert_to_tensor([circuit] * batch_size)\n",
        "        self.op = tfq.convert_to_tensor([[cirq.Z(test_bit)]] * batch_size)\n",
        "\n",
        "    def __call__(self, input_state1, input_state2):\n",
        "        add = tfq.layers.AddCircuit()\n",
        "        return add(add(input_state1, append=input_state2), append=self.circuit)\n",
        "        \n",
        "fidelity = SwapTestFidelity(qubits, qubits_qgrnn, batch_size)\n",
        "\n",
        "# Construct measurement layers.\n",
        "state_true = Input(shape=(), dtype=tf.string)\n",
        "state_pred = Input(shape=(), dtype=tf.string)\n",
        "fid_output = fidelity(state_true, state_pred)\n",
        "fid_output = Expectation()(fid_output,\n",
        "                           symbol_names=qgrnn.symbols,\n",
        "                           operators=fidelity.op)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JC5gg6zP2iEB",
        "colab_type": "text"
      },
      "source": [
        "## Calculate the average infidelity"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3tmY4N9dTKUm",
        "colab_type": "text"
      },
      "source": [
        "In fact, QGRNN is also a recurrent neural network that learns various time step evolution from the same input. We need to aggregate the results among the different timestep evolutions to train the QGRNN model. Here we introduce the average fidelity.\n",
        "\n",
        "  - $L(\\theta, \\phi) = 1 - \\frac{1}{B} \\sum^{B}_{j=1} |\\langle \\psi_{T_j} | U^{j}_{QGRNN}(\\theta, \\phi)|\\psi_0\\rangle|^2 = 1 - \\frac{1}{B} \\sum^{B}_{j=1} \\langle Z_{test} \\rangle_j $\n",
        "\n",
        "We can implement this custom keras loss function like this."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p9rvt2NGyUde",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def average_fidelity(y_true, y_pred):\n",
        "    return 1 - K.mean(y_pred)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ckA4FbKoS7Lu",
        "colab_type": "text"
      },
      "source": [
        "## Train the QGRNN & get the final result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2LU0bfW8mBzR",
        "colab_type": "text"
      },
      "source": [
        "Again, we can use Keras model fit. To feed a batch of quantum data, we can use `tf.concat` because the quantum circuits are already in `tf.Tensor`. In this case, we know that the lower bound of fidelity is 0, but the true `model_output` is not used in our custom loss function `average_fidelity`. For the purpose of comparison, `initial_params` contains the initial weights & biases of QGRNN. We set learning rate of Adam optimizer to $0.05$ as our paper described."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TW5R4fCIysbs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Model compile\n",
        "model = Model(inputs=[state_true, state_pred], outputs=fid_output)\n",
        "adam = Adam(learning_rate=0.05)\n",
        "model.compile(optimizer=adam, loss=average_fidelity)\n",
        "initial_params = model.get_weights()\n",
        "\n",
        "# Data preparation\n",
        "y_true = tf.concat(true_final_states, axis=0)\n",
        "y_pred = tf.concat(pred_final_states, axis=0)\n",
        "model_input = [y_true, y_pred]\n",
        "# Lower bound of average fidelity = 0, but not used in average_fidelity\n",
        "model_output = tf.convert_to_tensor([[0]] * batch_size)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e7_rGM8QG53c",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%%time\n",
        "history = model.fit(x=model_input,\n",
        "                    y=model_output,\n",
        "                    batch_size=batch_size,\n",
        "                    epochs=500)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zQTZALG2Ybfq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Learning curve\n",
        "plt.plot(history.history['loss'])\n",
        "plt.title('Learning curve (min loss={:.3e})'.format(\n",
        "    np.min(history.history['loss'])))\n",
        "plt.xlabel('iteration')\n",
        "plt.ylabel('log fidelity')\n",
        "plt.yscale('log')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r94jsV_VkU_y",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Target Ising model graph\n",
        "draw(G, pos, weights, biases, 'Target Ising model')\n",
        "# Initial QGRNN graph.\n",
        "init_edge_color = initial_params[0][:len(G_qgrnn.edges)]\n",
        "init_node_color = initial_params[0][len(G_qgrnn.edges):]\n",
        "draw(G_qgrnn, pos, init_edge_color, init_node_color, 'Initial QGRNN model')\n",
        "# Trained QGRNN graph.\n",
        "final_params = model.get_weights()\n",
        "final_edge_color = final_params[0][:len(G_qgrnn.edges)]\n",
        "final_node_color = final_params[0][len(G_qgrnn.edges):]\n",
        "draw(G_qgrnn, pos, final_edge_color, final_node_color, 'Trained QGRNN model')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vgS2kGNmmy_d",
        "colab_type": "text"
      },
      "source": [
        "# Conclusion\n",
        "\n",
        "In this colab, we implemented a variant of Quantum Graph Neural Networks called Quantum Graph `Recurrent` Neural Network. As the name suggests, this model learns time evolution dynamics of given target Hamiltonian from its quantum data. You can modify and train your own configurations of the model. Here are some suggestions for you.\n",
        "\n",
        "- Set `train_mixer=True` in class `QuantumGraphRNN` and retrain the model. Then the mixer Hamiltonian in QGRNN becomes also trainable. What can you observe after training?\n",
        "\n",
        "- We added missing edges in the cycle `0->1->...->5->0`. What if we remove the cycle? It means the QGRNN doesn't know any clue of the true Hamiltonian at all except the number of qubits. Set `add_cycle=False`. Is the loss still going to the order of `1e-4`?\n",
        "\n"
      ]
    }
  ]
}