docs/tutorials/qcnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xLOXFOT5Q40E"
   },
   "source": [
    "##### Copyright 2020 The TensorFlow Authors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {},
    "colab_type": "code",
    "id": "iiQkM5ZgQ8r2",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uLeF5Nmdef0V"
   },
   "source": [
    "# Quantum Convolutional Neural Network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "i9Jcnb8bQQyd"
   },
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://www.tensorflow.org/quantum/tutorials/qcnn\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/quantum/blob/master/docs/tutorials/qcnn.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/tensorflow/quantum/blob/master/docs/tutorials/qcnn.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://storage.googleapis.com/tensorflow_docs/quantum/docs/tutorials/qcnn.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4D3xaWBHOIVg"
   },
   "source": [
    "This tutorial implements a simplified <a href=\"https://www.nature.com/articles/s41567-019-0648-8\" class=\"external\">Quantum Convolutional Neural Network</a> (QCNN), a proposed quantum analogue to a classical convolutional neural network that is also *translationally invariant*.\n",
    "\n",
    "This example demonstrates how to detect certain properties of a quantum data source, such as a quantum sensor or a complex simulation from a device. The quantum data source being a <a href=\"https://arxiv.org/pdf/quant-ph/0504097.pdf\" class=\"external\">cluster state</a> that may or may not have an excitation—what the QCNN will learn to detect (The dataset used in the paper was SPT phase classification)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FnjolLuz8o5C"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Aquwcz-0aHqz",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "!pip install tensorflow==2.15.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e_ZuLN_N8yhT"
   },
   "source": [
    "Install TensorFlow Quantum:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3Pl5PW-ACO9J",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "!pip install tensorflow-quantum==0.7.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4Ql5PW-ACO0J",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Update package resources to account for version changes.\n",
    "import importlib, pkg_resources\n",
    "\n",
    "importlib.reload(pkg_resources)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "TL_LvHXzPNjW"
   },
   "source": [
    "Now import TensorFlow and the module dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QytLEAtoejW5",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_quantum as tfq\n",
    "\n",
    "import cirq\n",
    "import sympy\n",
    "import numpy as np\n",
    "\n",
    "# visualization tools\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from cirq.contrib.svg import SVGCircuit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "j6331ZSsQGY3"
   },
   "source": [
    "## 1. Build a QCNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Qg85u3G--CGq"
   },
   "source": [
    "### 1.1 Assemble circuits in a TensorFlow graph\n",
    "\n",
    "TensorFlow Quantum (TFQ) provides layer classes designed for in-graph circuit construction. One example is the `tfq.layers.AddCircuit` layer that inherits from `tf.keras.Layer`. This layer can either prepend or append to the input batch of circuits, as shown in the following figure.\n",
    "\n",
    "<img src=\"./images/qcnn_1.png\" width=\"700\">\n",
    "\n",
    "The following snippet uses this layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "FhNf0G_OPLqZ",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "qubit = cirq.GridQubit(0, 0)\n",
    "\n",
    "# Define some circuits.\n",
    "circuit1 = cirq.Circuit(cirq.X(qubit))\n",
    "circuit2 = cirq.Circuit(cirq.H(qubit))\n",
    "\n",
    "# Convert to a tensor.\n",
    "input_circuit_tensor = tfq.convert_to_tensor([circuit1, circuit2])\n",
    "\n",
    "# Define a circuit that we want to append\n",
    "y_circuit = cirq.Circuit(cirq.Y(qubit))\n",
    "\n",
    "# Instantiate our layer\n",
    "y_appender = tfq.layers.AddCircuit()\n",
    "\n",
    "# Run our circuit tensor through the layer and save the output.\n",
    "output_circuit_tensor = y_appender(input_circuit_tensor, append=y_circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ShZbRZCXkvk5"
   },
   "source": [
    "Examine the input tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ImRynsUN4BSG",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print(tfq.from_tensor(input_circuit_tensor))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xkGU4ZTUk4gf"
   },
   "source": [
    "And examine the output tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tfff6dJp39Fg",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print(tfq.from_tensor(output_circuit_tensor))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "23JeZ7Ns5qy5"
   },
   "source": [
    "While it is possible to run the examples below without using `tfq.layers.AddCircuit`, it's a good opportunity to understand how complex functionality can be embedded into TensorFlow compute graphs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GcVplt9455Hi"
   },
   "source": [
    "### 1.2 Problem overview\n",
    "\n",
    "You will prepare a *cluster state* and train a quantum classifier to detect if it is \"excited\" or not. The cluster state is highly entangled but not necessarily difficult for a classical computer. For clarity, this is a simpler dataset than the one used in the paper.\n",
    "\n",
    "For this classification task you will implement a deep <a href=\"https://arxiv.org/pdf/quant-ph/0610099.pdf\" class=\"external\">MERA</a>-like QCNN architecture since:\n",
    "\n",
    "1.   Like the QCNN, the cluster state on a ring is translationally invariant.\n",
    "2.   The cluster state is highly entangled.\n",
    "\n",
    "This architecture should be effective at reducing entanglement, obtaining the classification by reading out a single qubit.\n",
    "\n",
    "<img src=\"./images/qcnn_2.png\" width=\"1000\">\n",
    "\n",
    "An \"excited\" cluster state is defined as a cluster state that had a `cirq.rx` gate applied to any of its qubits. Qconv and QPool are discussed later in this tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jpqtsGJH_I1d"
   },
   "source": [
    "### 1.3 Building blocks for TensorFlow\n",
    "\n",
    "<img src=\"./images/qcnn_3.png\" width=\"1000\">\n",
    "\n",
    "One way to solve this problem with TensorFlow Quantum is to implement the following:\n",
    "\n",
    "1. The input to the model is a circuit tensor—either an empty circuit or an X gate on a particular qubit indicating an excitation.\n",
    "2. The rest of the model's quantum components are constructed with `tfq.layers.AddCircuit` layers.\n",
    "3. For inference a `tfq.layers.PQC` layer is used. This reads $\\langle \\hat{Z} \\rangle$ and compares it to a label of 1 for an excited state, or -1 for a non-excited state."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oa7Q3m_ThDgO"
   },
   "source": [
    "### 1.4 Data\n",
    "Before building your model, you can generate your data. In this case it's going to be excitations to the cluster state (The original paper uses a more complicated dataset). Excitations are represented with `cirq.rx` gates. A large enough rotation is deemed an excitation and is labeled `1` and a rotation that isn't large enough is labeled `-1` and deemed not an excitation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iUrvTCU1hDgP",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def generate_data(qubits):\n",
    "    \"\"\"Generate training and testing data.\"\"\"\n",
    "    n_rounds = 20  # Produces n_rounds * n_qubits datapoints.\n",
    "    excitations = []\n",
    "    labels = []\n",
    "    for n in range(n_rounds):\n",
    "        for bit in qubits:\n",
    "            rng = np.random.uniform(-np.pi, np.pi)\n",
    "            excitations.append(cirq.Circuit(cirq.rx(rng)(bit)))\n",
    "            labels.append(1 if (-np.pi / 2) <= rng <= (np.pi / 2) else -1)\n",
    "\n",
    "    split_ind = int(len(excitations) * 0.7)\n",
    "    train_excitations = excitations[:split_ind]\n",
    "    test_excitations = excitations[split_ind:]\n",
    "\n",
    "    train_labels = labels[:split_ind]\n",
    "    test_labels = labels[split_ind:]\n",
    "\n",
    "    return tfq.convert_to_tensor(train_excitations), np.array(train_labels), \\\n",
    "        tfq.convert_to_tensor(test_excitations), np.array(test_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "wGsDkZnrhDgS"
   },
   "source": [
    "You can see that just like with regular machine learning you create a training and testing set to use to benchmark the model. You can quickly look at some datapoints with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eLJ-JHOihDgT",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "sample_points, sample_labels, _, __ = generate_data(cirq.GridQubit.rect(1, 4))\n",
    "print('Input:', tfq.from_tensor(sample_points)[0], 'Output:', sample_labels[0])\n",
    "print('Input:', tfq.from_tensor(sample_points)[1], 'Output:', sample_labels[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "sFiRlDt_0-DL"
   },
   "source": [
    "### 1.5 Define layers\n",
    "\n",
    "Now define the layers shown in the figure above in TensorFlow."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "s2B9geIqLWHK"
   },
   "source": [
    "#### 1.5.1 Cluster state\n",
    "\n",
    "The first step is to define the <a href=\"https://arxiv.org/pdf/quant-ph/0504097.pdf\" class=\"external\">cluster state</a> using <a href=\"https://github.com/quantumlib/Cirq\" class=\"external\">Cirq</a>, a Google-provided framework for programming quantum circuits. Since this is a static part of the model, embed it using the `tfq.layers.AddCircuit` functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qpQwVWKazU8g",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def cluster_state_circuit(bits):\n",
    "    \"\"\"Return a cluster state on the qubits in `bits`.\"\"\"\n",
    "    circuit = cirq.Circuit()\n",
    "    circuit.append(cirq.H.on_each(bits))\n",
    "    for this_bit, next_bit in zip(bits, bits[1:] + [bits[0]]):\n",
    "        circuit.append(cirq.CZ(this_bit, next_bit))\n",
    "    return circuit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e9qX1uN740vJ"
   },
   "source": [
    "Display a cluster state circuit for a rectangle of <a href=\"https://cirq.readthedocs.io/en/stable/generated/cirq.GridQubit.html\" class=\"external\"><code>cirq.GridQubit</code></a>s:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9tZt0aAO4r4F",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "SVGCircuit(cluster_state_circuit(cirq.GridQubit.rect(1, 4)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "4xElWnRf1ZC7"
   },
   "source": [
    "#### 1.5.2 QCNN layers\n",
    "\n",
    "Define the layers that make up the model using the <a href=\"https://arxiv.org/abs/1810.03787\" class=\"external\">Cong and Lukin QCNN paper</a>. There are a few prerequisites:\n",
    "\n",
    "* The one- and two-qubit parameterized unitary matrices from the <a href=\"https://arxiv.org/abs/quant-ph/0507171\" class=\"external\">Tucci paper</a>.\n",
    "* A general parameterized two-qubit pooling operation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oNRGOqky2exY",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def one_qubit_unitary(bit, symbols):\n",
    "    \"\"\"Make a Cirq circuit enacting a rotation of the bloch sphere about the X,\n",
    "    Y and Z axis, that depends on the values in `symbols`.\n",
    "    \"\"\"\n",
    "    return cirq.Circuit(\n",
    "        cirq.X(bit)**symbols[0],\n",
    "        cirq.Y(bit)**symbols[1],\n",
    "        cirq.Z(bit)**symbols[2])\n",
    "\n",
    "\n",
    "def two_qubit_unitary(bits, symbols):\n",
    "    \"\"\"Make a Cirq circuit that creates an arbitrary two qubit unitary.\"\"\"\n",
    "    circuit = cirq.Circuit()\n",
    "    circuit += one_qubit_unitary(bits[0], symbols[0:3])\n",
    "    circuit += one_qubit_unitary(bits[1], symbols[3:6])\n",
    "    circuit += [cirq.ZZ(*bits)**symbols[6]]\n",
    "    circuit += [cirq.YY(*bits)**symbols[7]]\n",
    "    circuit += [cirq.XX(*bits)**symbols[8]]\n",
    "    circuit += one_qubit_unitary(bits[0], symbols[9:12])\n",
    "    circuit += one_qubit_unitary(bits[1], symbols[12:])\n",
    "    return circuit\n",
    "\n",
    "\n",
    "def two_qubit_pool(source_qubit, sink_qubit, symbols):\n",
    "    \"\"\"Make a Cirq circuit to do a parameterized 'pooling' operation, which\n",
    "    attempts to reduce entanglement down from two qubits to just one.\"\"\"\n",
    "    pool_circuit = cirq.Circuit()\n",
    "    sink_basis_selector = one_qubit_unitary(sink_qubit, symbols[0:3])\n",
    "    source_basis_selector = one_qubit_unitary(source_qubit, symbols[3:6])\n",
    "    pool_circuit.append(sink_basis_selector)\n",
    "    pool_circuit.append(source_basis_selector)\n",
    "    pool_circuit.append(cirq.CNOT(source_qubit, sink_qubit))\n",
    "    pool_circuit.append(sink_basis_selector**-1)\n",
    "    return pool_circuit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "LoG0a3U_2qGA"
   },
   "source": [
    "To see what you created, print out the one-qubit unitary circuit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "T5uhvF-g2rpZ",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "SVGCircuit(one_qubit_unitary(cirq.GridQubit(0, 0), sympy.symbols('x0:3')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NWuMb_Us8ar2"
   },
   "source": [
    "And the two-qubit unitary circuit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aJTdRrfS2uIo",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "SVGCircuit(two_qubit_unitary(cirq.GridQubit.rect(1, 2), sympy.symbols('x0:15')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EXQD1R_V8jyk"
   },
   "source": [
    "And the two-qubit pooling circuit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DOHRbkvH2xGK",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "SVGCircuit(two_qubit_pool(*cirq.GridQubit.rect(1, 2), sympy.symbols('x0:6')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AzVauXWD3v8C"
   },
   "source": [
    "##### 1.5.2.1 Quantum convolution\n",
    "\n",
    "As in the <a href=\"https://arxiv.org/abs/1810.03787\" class=\"external\">Cong and Lukin</a> paper, define the 1D quantum convolution as the application of a two-qubit parameterized unitary to every pair of adjacent qubits with a stride of one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1Fa19Lzb3wnR",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def quantum_conv_circuit(bits, symbols):\n",
    "    \"\"\"Quantum Convolution Layer following the above diagram.\n",
    "    Return a Cirq circuit with the cascade of `two_qubit_unitary` applied\n",
    "    to all pairs of qubits in `bits` as in the diagram above.\n",
    "    \"\"\"\n",
    "    circuit = cirq.Circuit()\n",
    "    for first, second in zip(bits[0::2], bits[1::2]):\n",
    "        circuit += two_qubit_unitary([first, second], symbols)\n",
    "    for first, second in zip(bits[1::2], bits[2::2] + [bits[0]]):\n",
    "        circuit += two_qubit_unitary([first, second], symbols)\n",
    "    return circuit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "fTzOm_t394Gj"
   },
   "source": [
    "Display the (very horizontal) circuit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Bi6q2nmY3z_U",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "SVGCircuit(\n",
    "    quantum_conv_circuit(cirq.GridQubit.rect(1, 8), sympy.symbols('x0:15')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3svBAfap4xhP"
   },
   "source": [
    "##### 1.5.2.2 Quantum pooling\n",
    "\n",
    "A quantum pooling layer pools from $N$ qubits to $\\frac{N}{2}$ qubits using the two-qubit pool defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jD3fgcWO4yEU",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def quantum_pool_circuit(source_bits, sink_bits, symbols):\n",
    "    \"\"\"A layer that specifies a quantum pooling operation.\n",
    "    A Quantum pool tries to learn to pool the relevant information from two\n",
    "    qubits onto 1.\n",
    "    \"\"\"\n",
    "    circuit = cirq.Circuit()\n",
    "    for source, sink in zip(source_bits, sink_bits):\n",
    "        circuit += two_qubit_pool(source, sink, symbols)\n",
    "    return circuit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NX83NHDP_Q_Z"
   },
   "source": [
    "Examine a pooling component circuit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pFXow2OX47O5",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_bits = cirq.GridQubit.rect(1, 8)\n",
    "\n",
    "SVGCircuit(\n",
    "    quantum_pool_circuit(test_bits[:4], test_bits[4:], sympy.symbols('x0:6')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "23VcPLT45Lg7"
   },
   "source": [
    "### 1.6 Model definition\n",
    "\n",
    "Now use the defined layers to construct a purely quantum CNN. Start with eight qubits, pool down to one, then measure $\\langle \\hat{Z} \\rangle$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vzEsY6-n5NR0",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def create_model_circuit(qubits):\n",
    "    \"\"\"Create sequence of alternating convolution and pooling operators \n",
    "    which gradually shrink over time.\"\"\"\n",
    "    model_circuit = cirq.Circuit()\n",
    "    symbols = sympy.symbols('qconv0:63')\n",
    "    # Cirq uses sympy.Symbols to map learnable variables. TensorFlow Quantum\n",
    "    # scans incoming circuits and replaces these with TensorFlow variables.\n",
    "    model_circuit += quantum_conv_circuit(qubits, symbols[0:15])\n",
    "    model_circuit += quantum_pool_circuit(qubits[:4], qubits[4:],\n",
    "                                          symbols[15:21])\n",
    "    model_circuit += quantum_conv_circuit(qubits[4:], symbols[21:36])\n",
    "    model_circuit += quantum_pool_circuit(qubits[4:6], qubits[6:],\n",
    "                                          symbols[36:42])\n",
    "    model_circuit += quantum_conv_circuit(qubits[6:], symbols[42:57])\n",
    "    model_circuit += quantum_pool_circuit([qubits[6]], [qubits[7]],\n",
    "                                          symbols[57:63])\n",
    "    return model_circuit\n",
    "\n",
    "\n",
    "# Create our qubits and readout operators in Cirq.\n",
    "cluster_state_bits = cirq.GridQubit.rect(1, 8)\n",
    "readout_operators = cirq.Z(cluster_state_bits[-1])\n",
    "\n",
    "# Build a sequential model enacting the logic in 1.3 of this notebook.\n",
    "# Here you are making the static cluster state prep as a part of the AddCircuit and the\n",
    "# \"quantum datapoints\" are coming in the form of excitation\n",
    "excitation_input = tf.keras.Input(shape=(), dtype=tf.dtypes.string)\n",
    "cluster_state = tfq.layers.AddCircuit()(\n",
    "    excitation_input, prepend=cluster_state_circuit(cluster_state_bits))\n",
    "\n",
    "quantum_model = tfq.layers.PQC(create_model_circuit(cluster_state_bits),\n",
    "                               readout_operators)(cluster_state)\n",
    "\n",
    "qcnn_model = tf.keras.Model(inputs=[excitation_input], outputs=[quantum_model])\n",
    "\n",
    "# Show the keras plot of the model\n",
    "tf.keras.utils.plot_model(qcnn_model,\n",
    "                          show_shapes=True,\n",
    "                          show_layer_names=False,\n",
    "                          dpi=70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9jqTEe5VSbug"
   },
   "source": [
    "### 1.7 Train the model\n",
    "\n",
    "Train the model over the full batch to simplify this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_TFkAm1sQZEN",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Generate some training data.\n",
    "train_excitations, train_labels, test_excitations, test_labels = generate_data(\n",
    "    cluster_state_bits)\n",
    "\n",
    "\n",
    "# Custom accuracy metric.\n",
    "@tf.function\n",
    "def custom_accuracy(y_true, y_pred):\n",
    "    y_true = tf.squeeze(y_true)\n",
    "    y_pred = tf.map_fn(lambda x: 1.0 if x >= 0 else -1.0, y_pred)\n",
    "    return tf.keras.backend.mean(tf.keras.backend.equal(y_true, y_pred))\n",
    "\n",
    "\n",
    "qcnn_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),\n",
    "                   loss=tf.losses.mse,\n",
    "                   metrics=[custom_accuracy])\n",
    "\n",
    "history = qcnn_model.fit(x=train_excitations,\n",
    "                         y=train_labels,\n",
    "                         batch_size=16,\n",
    "                         epochs=25,\n",
    "                         verbose=1,\n",
    "                         validation_data=(test_excitations, test_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2tiCJOb5Qzcr",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "plt.plot(history.history['loss'][1:], label='Training')\n",
    "plt.plot(history.history['val_loss'][1:], label='Validation')\n",
    "plt.title('Training a Quantum CNN to Detect Excited Cluster States')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GyrkcEReQ5Bc"
   },
   "source": [
    "## 2. Hybrid models\n",
    "\n",
    "You don't have to go from eight qubits to one qubit using quantum convolution—you could have done one or two rounds of quantum convolution and fed the results into a classical neural network. This section explores quantum-classical hybrid models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "A2tOK22t7Kjm"
   },
   "source": [
    "### 2.1 Hybrid model with a single quantum filter\n",
    "\n",
    "Apply one layer of quantum convolution, reading out $\\langle \\hat{Z}_n \\rangle$ on all bits, followed by a densely-connected neural network.\n",
    "\n",
    "<img src=\"./images/qcnn_5.png\" width=\"1000\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lKXuOApgWYFa"
   },
   "source": [
    "#### 2.1.1 Model definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ut-U1hBkQ8Fs",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# 1-local operators to read out\n",
    "readouts = [cirq.Z(bit) for bit in cluster_state_bits[4:]]\n",
    "\n",
    "\n",
    "def multi_readout_model_circuit(qubits):\n",
    "    \"\"\"Make a model circuit with less quantum pool and conv operations.\"\"\"\n",
    "    model_circuit = cirq.Circuit()\n",
    "    symbols = sympy.symbols('qconv0:21')\n",
    "    model_circuit += quantum_conv_circuit(qubits, symbols[0:15])\n",
    "    model_circuit += quantum_pool_circuit(qubits[:4], qubits[4:],\n",
    "                                          symbols[15:21])\n",
    "    return model_circuit\n",
    "\n",
    "\n",
    "# Build a model enacting the logic in 2.1 of this notebook.\n",
    "excitation_input_dual = tf.keras.Input(shape=(), dtype=tf.dtypes.string)\n",
    "\n",
    "cluster_state_dual = tfq.layers.AddCircuit()(\n",
    "    excitation_input_dual, prepend=cluster_state_circuit(cluster_state_bits))\n",
    "\n",
    "quantum_model_dual = tfq.layers.PQC(\n",
    "    multi_readout_model_circuit(cluster_state_bits),\n",
    "    readouts)(cluster_state_dual)\n",
    "\n",
    "d1_dual = tf.keras.layers.Dense(8)(quantum_model_dual)\n",
    "\n",
    "d2_dual = tf.keras.layers.Dense(1)(d1_dual)\n",
    "\n",
    "hybrid_model = tf.keras.Model(inputs=[excitation_input_dual], outputs=[d2_dual])\n",
    "\n",
    "# Display the model architecture\n",
    "tf.keras.utils.plot_model(hybrid_model,\n",
    "                          show_shapes=True,\n",
    "                          show_layer_names=False,\n",
    "                          dpi=70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qDqoLZJuWcgH"
   },
   "source": [
    "#### 2.1.2 Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EyYw9kYIRCE7",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "hybrid_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),\n",
    "                     loss=tf.losses.mse,\n",
    "                     metrics=[custom_accuracy])\n",
    "\n",
    "hybrid_history = hybrid_model.fit(x=train_excitations,\n",
    "                                  y=train_labels,\n",
    "                                  batch_size=16,\n",
    "                                  epochs=25,\n",
    "                                  verbose=1,\n",
    "                                  validation_data=(test_excitations,\n",
    "                                                   test_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yL3jhGiBRJHt",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "plt.plot(history.history['val_custom_accuracy'], label='QCNN')\n",
    "plt.plot(hybrid_history.history['val_custom_accuracy'], label='Hybrid CNN')\n",
    "plt.title('Quantum vs Hybrid CNN performance')\n",
    "plt.xlabel('Epochs')\n",
    "plt.legend()\n",
    "plt.ylabel('Validation Accuracy')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NCNiNvheRNzq"
   },
   "source": [
    "As you can see, with very modest classical assistance, the hybrid model will usually converge faster than the purely quantum version."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nVUtWLZnRRDE"
   },
   "source": [
    "### 2.2 Hybrid convolution with multiple quantum filters\n",
    "\n",
    "Now let's try an architecture that uses multiple quantum convolutions and a classical neural network to combine them.\n",
    "\n",
    "<img src=\"./images/qcnn_6.png\" width=\"1000\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ldo_m5P3YBV7"
   },
   "source": [
    "#### 2.2.1 Model definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "W3TkNVm9RTBj",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "excitation_input_multi = tf.keras.Input(shape=(), dtype=tf.dtypes.string)\n",
    "\n",
    "cluster_state_multi = tfq.layers.AddCircuit()(\n",
    "    excitation_input_multi, prepend=cluster_state_circuit(cluster_state_bits))\n",
    "\n",
    "# apply 3 different filters and measure expectation values\n",
    "\n",
    "quantum_model_multi1 = tfq.layers.PQC(\n",
    "    multi_readout_model_circuit(cluster_state_bits),\n",
    "    readouts)(cluster_state_multi)\n",
    "\n",
    "quantum_model_multi2 = tfq.layers.PQC(\n",
    "    multi_readout_model_circuit(cluster_state_bits),\n",
    "    readouts)(cluster_state_multi)\n",
    "\n",
    "quantum_model_multi3 = tfq.layers.PQC(\n",
    "    multi_readout_model_circuit(cluster_state_bits),\n",
    "    readouts)(cluster_state_multi)\n",
    "\n",
    "# concatenate outputs and feed into a small classical NN\n",
    "concat_out = tf.keras.layers.concatenate(\n",
    "    [quantum_model_multi1, quantum_model_multi2, quantum_model_multi3])\n",
    "\n",
    "dense_1 = tf.keras.layers.Dense(8)(concat_out)\n",
    "\n",
    "dense_2 = tf.keras.layers.Dense(1)(dense_1)\n",
    "\n",
    "multi_qconv_model = tf.keras.Model(inputs=[excitation_input_multi],\n",
    "                                   outputs=[dense_2])\n",
    "\n",
    "# Display the model architecture\n",
    "tf.keras.utils.plot_model(multi_qconv_model,\n",
    "                          show_shapes=True,\n",
    "                          show_layer_names=True,\n",
    "                          dpi=70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2eNhDWwKY9N4"
   },
   "source": [
    "#### 2.2.2 Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "suRvxcAKRZK6",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "multi_qconv_model.compile(\n",
    "    optimizer=tf.keras.optimizers.Adam(learning_rate=0.02),\n",
    "    loss=tf.losses.mse,\n",
    "    metrics=[custom_accuracy])\n",
    "\n",
    "multi_qconv_history = multi_qconv_model.fit(x=train_excitations,\n",
    "                                            y=train_labels,\n",
    "                                            batch_size=16,\n",
    "                                            epochs=25,\n",
    "                                            verbose=1,\n",
    "                                            validation_data=(test_excitations,\n",
    "                                                             test_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-6NR7yAQRmOU",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "plt.plot(history.history['val_custom_accuracy'][:25], label='QCNN')\n",
    "plt.plot(hybrid_history.history['val_custom_accuracy'][:25], label='Hybrid CNN')\n",
    "plt.plot(multi_qconv_history.history['val_custom_accuracy'][:25],\n",
    "         label='Hybrid CNN \\n Multiple Quantum Filters')\n",
    "plt.title('Quantum vs Hybrid CNN performance')\n",
    "plt.xlabel('Epochs')\n",
    "plt.legend()\n",
    "plt.ylabel('Validation Accuracy')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "qcnn.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}