docs/tutorials/hello_many_worlds.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xLOXFOT5Q40E"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "iiQkM5ZgQ8r2"
      },
      "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": "j6331ZSsQGY3"
      },
      "source": [
        "# Hello, many worlds"
      ]
    },
    {
      "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/hello_many_worlds\"><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/hello_many_worlds.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/hello_many_worlds.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/hello_many_worlds.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": "6tYn2HaAUgH0"
      },
      "source": [
        "This tutorial shows how a classical neural network can learn to correct qubit calibration errors. It introduces <a target=\"_blank\" href=\"https://github.com/quantumlib/Cirq\" class=\"external\">Cirq</a>, a Python framework to create, edit, and invoke Noisy Intermediate Scale Quantum (NISQ) circuits, and demonstrates how Cirq interfaces with TensorFlow Quantum."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "sPZoNKvpUaqa"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "TorxE5tnkvb2"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow==2.4.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FxkQA6oblNqI"
      },
      "source": [
        "Install TensorFlow Quantum:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "saFHsRDpkvkH"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow-quantum"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "F1L8h1YKUvIO"
      },
      "source": [
        "Now import TensorFlow and the module dependencies:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "enZ300Bflq80"
      },
      "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": "b08Mmbs8lr81"
      },
      "source": [
        "## 1. The Basics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "y31qSRCczI-L"
      },
      "source": [
        "### 1.1 Cirq and parameterized quantum circuits\n",
        "\n",
        "Before exploring TensorFlow Quantum (TFQ), let's look at some <a target=\"_blank\" href=\"https://github.com/quantumlib/Cirq\" class=\"external\">Cirq</a> basics. Cirq is a Python library for quantum computing from Google. You use it to define circuits, including static and parameterized gates.\n",
        "\n",
        "Cirq uses <a target=\"_blank\" href=\"https://www.sympy.org\" class=\"external\">SymPy</a> symbols to represent free parameters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "2yQdmhQLCrzQ"
      },
      "outputs": [],
      "source": [
        "a, b = sympy.symbols('a b')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "itUlpbKmDYNW"
      },
      "source": [
        "The following code creates a two-qubit circuit using your parameters:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Ps-pd2mndXs7"
      },
      "outputs": [],
      "source": [
        "# Create two qubits\n",
        "q0, q1 = cirq.GridQubit.rect(1, 2)\n",
        "\n",
        "# Create a circuit on these qubits using the parameters you created above.\n",
        "circuit = cirq.Circuit(\n",
        "    cirq.rx(a).on(q0),\n",
        "    cirq.ry(b).on(q1), cirq.CNOT(control=q0, target=q1))\n",
        "\n",
        "SVGCircuit(circuit)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zcCX109cJUaz"
      },
      "source": [
        "To evaluate circuits, you can use the `cirq.Simulator` interface. You replace free parameters in a circuit with specific numbers by passing in a `cirq.ParamResolver` object. The following code calculates the raw state vector output of your parameterized circuit:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VMq7EayNRyQb"
      },
      "outputs": [],
      "source": [
        "# Calculate a state vector with a=0.5 and b=-0.5.\n",
        "resolver = cirq.ParamResolver({a: 0.5, b: -0.5})\n",
        "output_state_vector = cirq.Simulator().simulate(circuit, resolver).final_state_vector\n",
        "output_state_vector"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-SUlLpXBeicF"
      },
      "source": [
        "State vectors are not directly accessible outside of simulation (notice the complex numbers in the output above). To be physically realistic, you must specify a measurement, which converts a state vector into a real number that classical computers can understand. Cirq specifies measurements using combinations of the <a target=\"_blank\" href=\"https://en.wikipedia.org/wiki/Pauli_matrices\" class=\"external\">Pauli operators</a> $\\hat{X}$, $\\hat{Y}$, and $\\hat{Z}$. As illustration, the following code measures $\\hat{Z}_0$ and $\\frac{1}{2}\\hat{Z}_0 + \\hat{X}_1$ on the state vector you just simulated:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "hrSnOCi3ehr_",
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "z0 = cirq.Z(q0)\n",
        "\n",
        "qubit_map={q0: 0, q1: 1}\n",
        "\n",
        "z0.expectation_from_state_vector(output_state_vector, qubit_map).real"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "OZ0lWFXv6pII"
      },
      "outputs": [],
      "source": [
        "z0x1 = 0.5 * z0 + cirq.X(q1)\n",
        "\n",
        "z0x1.expectation_from_state_vector(output_state_vector, qubit_map).real"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "bkC-yjIolDNr"
      },
      "source": [
        "### 1.2 Quantum circuits as tensors\n",
        "\n",
        "TensorFlow Quantum (TFQ) provides `tfq.convert_to_tensor`, a function that converts Cirq objects into tensors. This allows you to send Cirq objects to our <a target=\"_blank\" href=\"https://www.tensorflow.org/quantum/api_docs/python/tfq/layers\">quantum layers</a> and <a target=\"_blank\" href=\"https://www.tensorflow.org/quantum/api_docs/python/tfq/get_expectation_op\">quantum ops</a>. The function can be called on lists or arrays of Cirq Circuits and Cirq Paulis:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1gLQjA02mIyy",
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "# Rank 1 tensor containing 1 circuit.\n",
        "circuit_tensor = tfq.convert_to_tensor([circuit])\n",
        "\n",
        "print(circuit_tensor.shape)\n",
        "print(circuit_tensor.dtype)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SJy6AkbU6pIP"
      },
      "source": [
        "This encodes the Cirq objects as `tf.string` tensors that `tfq` operations decode as needed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aX_vEmCKmpQS"
      },
      "outputs": [],
      "source": [
        "# Rank 1 tensor containing 2 Pauli operators.\n",
        "pauli_tensor = tfq.convert_to_tensor([z0, z0x1])\n",
        "pauli_tensor.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FI1JLWe6m8JF"
      },
      "source": [
        "### 1.3 Batching circuit simulation\n",
        "\n",
        "TFQ provides methods for computing expectation values, samples, and state vectors. For now, let's focus on *expectation values*.\n",
        "\n",
        "The highest-level interface for calculating expectation values is the `tfq.layers.Expectation` layer, which is a `tf.keras.Layer`. In its simplest form, this layer is equivalent to simulating a parameterized circuit over many `cirq.ParamResolvers`; however, TFQ allows batching following TensorFlow semantics, and circuits are simulated using efficient C++ code.\n",
        "\n",
        "Create a batch of values to substitute for our `a` and `b` parameters:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1fsVZhF5lIXp"
      },
      "outputs": [],
      "source": [
        "batch_vals = np.array(np.random.uniform(0, 2 * np.pi, (5, 2)), dtype=np.float32)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ip7jlGXIf22u"
      },
      "source": [
        "Batching circuit execution over parameter values in Cirq requires a loop:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RsfF53UCJtr9"
      },
      "outputs": [],
      "source": [
        "cirq_results = []\n",
        "cirq_simulator = cirq.Simulator()\n",
        "\n",
        "for vals in batch_vals:\n",
        "    resolver = cirq.ParamResolver({a: vals[0], b: vals[1]})\n",
        "    final_state_vector = cirq_simulator.simulate(circuit, resolver).final_state_vector\n",
        "    cirq_results.append(\n",
        "        [z0.expectation_from_state_vector(final_state_vector, {\n",
        "            q0: 0,\n",
        "            q1: 1\n",
        "        }).real])\n",
        "\n",
        "print('cirq batch results: \\n {}'.format(np.array(cirq_results)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "W0JlZEu-f9Ac"
      },
      "source": [
        "The same operation is simplified in TFQ:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kGZVdcZ6y9lC"
      },
      "outputs": [],
      "source": [
        "tfq.layers.Expectation()(circuit,\n",
        "                         symbol_names=[a, b],\n",
        "                         symbol_values=batch_vals,\n",
        "                         operators=z0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "wppQ3TJ23mWC"
      },
      "source": [
        "## 2. Hybrid quantum-classical optimization\n",
        "\n",
        "Now that you've seen the basics, let's use TensorFlow Quantum to construct a *hybrid quantum-classical neural net*. You will train a classical neural net to control a single qubit. The control will be optimized to correctly prepare the qubit in the `0` or `1` state, overcoming a simulated systematic calibration error. This figure shows the architecture:\n",
        "\n",
        "<img src=\"./images/nn_control1.png\" width=\"1000\">\n",
        "\n",
        "Even without a neural network this is a straightforward problem to solve, but the theme is similar to the real quantum control problems you might solve using TFQ. It demonstrates an end-to-end example of a quantum-classical computation using the `tfq.layers.ControlledPQC` (Parametrized Quantum Circuit) layer inside of a `tf.keras.Model`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NlyxF3Q-6pIe"
      },
      "source": [
        "For the implementation of this tutorial, this architecture is split into 3 parts:\n",
        "\n",
        "- The *input circuit* or *datapoint circuit*: The first three $R$ gates.\n",
        "- The *controlled circuit*: The other three $R$ gates.\n",
        "- The *controller*: The classical neural-network setting the parameters of the controlled circuit."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VjDf-nTM6ZSs"
      },
      "source": [
        "### 2.1 The controlled circuit definition\n",
        "\n",
        "Define a learnable single bit rotation, as indicated in the figure above. This will correspond to our controlled circuit."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "N-j7SCl-51-q"
      },
      "outputs": [],
      "source": [
        "# Parameters that the classical NN will feed values into.\n",
        "control_params = sympy.symbols('theta_1 theta_2 theta_3')\n",
        "\n",
        "# Create the parameterized circuit.\n",
        "qubit = cirq.GridQubit(0, 0)\n",
        "model_circuit = cirq.Circuit(\n",
        "    cirq.rz(control_params[0])(qubit),\n",
        "    cirq.ry(control_params[1])(qubit),\n",
        "    cirq.rx(control_params[2])(qubit))\n",
        "\n",
        "SVGCircuit(model_circuit)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "wfjSbsvb7g9f"
      },
      "source": [
        "### 2.2 The controller\n",
        "\n",
        "Now define controller network: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1v4CK2jD6pIj"
      },
      "outputs": [],
      "source": [
        "# The classical neural network layers.\n",
        "controller = tf.keras.Sequential([\n",
        "    tf.keras.layers.Dense(10, activation='elu'),\n",
        "    tf.keras.layers.Dense(3)\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "QNimbsAt6pIm"
      },
      "source": [
        "Given a batch of commands, the controller outputs a batch of control signals for the controlled circuit. \n",
        "\n",
        "The controller is randomly initialized so these outputs are not useful, yet."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kZbYRTe16pIm"
      },
      "outputs": [],
      "source": [
        "controller(tf.constant([[0.0],[1.0]])).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "XizLExg56pIp"
      },
      "source": [
        "### 2.3 Connect the controller to the circuit"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "I5Pmy5-V6pIq"
      },
      "source": [
        "Use `tfq` to connect the controller to the controlled circuit, as a single `keras.Model`. \n",
        "\n",
        "See the [Keras Functional API guide](https://www.tensorflow.org/guide/keras/functional) for more about this style of model definition.\n",
        "\n",
        "First define the inputs to the model:  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "UfHF8NNE6pIr"
      },
      "outputs": [],
      "source": [
        "# This input is the simulated miscalibration that the model will learn to correct.\n",
        "circuits_input = tf.keras.Input(shape=(),\n",
        "                                # The circuit-tensor has dtype `tf.string` \n",
        "                                dtype=tf.string,\n",
        "                                name='circuits_input')\n",
        "\n",
        "# Commands will be either `0` or `1`, specifying the state to set the qubit to.\n",
        "commands_input = tf.keras.Input(shape=(1,),\n",
        "                                dtype=tf.dtypes.float32,\n",
        "                                name='commands_input')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "y9xN2mNl6pIu"
      },
      "source": [
        "Next apply operations to those inputs, to define the computation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Zvt2YGmZ6pIu"
      },
      "outputs": [],
      "source": [
        "dense_2 = controller(commands_input)\n",
        "\n",
        "# TFQ layer for classically controlled circuits.\n",
        "expectation_layer = tfq.layers.ControlledPQC(model_circuit,\n",
        "                                             # Observe Z\n",
        "                                             operators = cirq.Z(qubit))\n",
        "expectation = expectation_layer([circuits_input, dense_2])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ip2jNA9h6pIy"
      },
      "source": [
        "Now package this computation as a `tf.keras.Model`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Xs6EMhah6pIz"
      },
      "outputs": [],
      "source": [
        "# The full Keras model is built from our layers.\n",
        "model = tf.keras.Model(inputs=[circuits_input, commands_input],\n",
        "                       outputs=expectation)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "w7kgqm3t6pI3"
      },
      "source": [
        "The network architecture is indicated by the plot of the model below.\n",
        "Compare this model plot to the architecture diagram to verify correctness.\n",
        "\n",
        "Note: May require a system install of the `graphviz` package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ERXNPe4F6pI4"
      },
      "outputs": [],
      "source": [
        "tf.keras.utils.plot_model(model, show_shapes=True, dpi=70)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-Pbemgww6pI7"
      },
      "source": [
        "This model takes two inputs: The commands for the controller, and the input-circuit whose output the controller is attempting to correct. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hpnIBK916pI8"
      },
      "source": [
        "### 2.4 The dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yJSC9qH76pJA"
      },
      "source": [
        "The model attempts to output the correct correct measurement value of $\\hat{Z}$ for each command. The commands and correct values are defined below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ciMIJAuH6pJA"
      },
      "outputs": [],
      "source": [
        "# The command input values to the classical NN.\n",
        "commands = np.array([[0], [1]], dtype=np.float32)\n",
        "\n",
        "# The desired Z expectation value at output of quantum circuit.\n",
        "expected_outputs = np.array([[1], [-1]], dtype=np.float32)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kV1LM_hZ6pJD"
      },
      "source": [
        "This is not the entire training dataset for this task. \n",
        "Each datapoint in the dataset also needs an input circuit."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "bbiVHvSYVW4H"
      },
      "source": [
        "### 2.4 Input circuit definition\n",
        "\n",
        "The input-circuit below defines the random miscalibration the model will learn to correct."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_VYfzHffWo7n"
      },
      "outputs": [],
      "source": [
        "random_rotations = np.random.uniform(0, 2 * np.pi, 3)\n",
        "noisy_preparation = cirq.Circuit(\n",
        "  cirq.rx(random_rotations[0])(qubit),\n",
        "  cirq.ry(random_rotations[1])(qubit),\n",
        "  cirq.rz(random_rotations[2])(qubit)\n",
        ")\n",
        "datapoint_circuits = tfq.convert_to_tensor([\n",
        "  noisy_preparation\n",
        "] * 2)  # Make two copied of this circuit"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FvOkMyKI6pJI"
      },
      "source": [
        "There are two copies of the circuit, one for each datapoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6nk2Yr3e6pJJ"
      },
      "outputs": [],
      "source": [
        "datapoint_circuits.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gB--UhZZYgVY"
      },
      "source": [
        "### 2.5 Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jATjqUIv6pJM"
      },
      "source": [
        "With the inputs defined you can test-run the `tfq` model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Lwphqvs96pJO"
      },
      "outputs": [],
      "source": [
        "model([datapoint_circuits, commands]).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9gyg5qSL6pJR"
      },
      "source": [
        "Now run a standard training process to adjust these values towards the `expected_outputs`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "dtPYqbNi8zeZ"
      },
      "outputs": [],
      "source": [
        "optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)\n",
        "loss = tf.keras.losses.MeanSquaredError()\n",
        "model.compile(optimizer=optimizer, loss=loss)\n",
        "history = model.fit(x=[datapoint_circuits, commands],\n",
        "                    y=expected_outputs,\n",
        "                    epochs=30,\n",
        "                    verbose=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "azE-qV0OaC1o"
      },
      "outputs": [],
      "source": [
        "plt.plot(history.history['loss'])\n",
        "plt.title(\"Learning to Control a Qubit\")\n",
        "plt.xlabel(\"Iterations\")\n",
        "plt.ylabel(\"Error in Control\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "GTd5DGcRmmOK"
      },
      "source": [
        "From this plot you can see that the neural network has learned to overcome the systematic miscalibration."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "C2RfWismj66S"
      },
      "source": [
        "### 2.6 Verify outputs\n",
        "Now use the trained model, to correct the qubit calibration errors. With Cirq:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RoIlb7r7j5SY"
      },
      "outputs": [],
      "source": [
        "def check_error(command_values, desired_values):\n",
        "  \"\"\"Based on the value in `command_value` see how well you could prepare\n",
        "  the full circuit to have `desired_value` when taking expectation w.r.t. Z.\"\"\"\n",
        "  params_to_prepare_output = controller(command_values).numpy()\n",
        "  full_circuit = noisy_preparation + model_circuit\n",
        "\n",
        "  # Test how well you can prepare a state to get expectation the expectation\n",
        "  # value in `desired_values`\n",
        "  for index in [0, 1]:\n",
        "    state = cirq_simulator.simulate(\n",
        "        full_circuit,\n",
        "        {s:v for (s,v) in zip(control_params, params_to_prepare_output[index])}\n",
        "    ).final_state_vector\n",
        "    expt = cirq.Z(qubit).expectation_from_state_vector(state, {qubit: 0}).real\n",
        "    print(f'For a desired output (expectation) of {desired_values[index]} with'\n",
        "          f' noisy preparation, the controller\\nnetwork found the following '\n",
        "          f'values for theta: {params_to_prepare_output[index]}\\nWhich gives an'\n",
        "          f' actual expectation of: {expt}\\n')\n",
        "\n",
        "\n",
        "check_error(commands, expected_outputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "wvW_ZDwmsws6"
      },
      "source": [
        "The value of the loss function during training provides a rough idea of how well the model is learning. The lower the loss, the closer the expectation values in the above cell is to `desired_values`. If you aren't as concerned with the parameter values, you can always check the outputs from above using `tfq`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aYskLTacs8Ku"
      },
      "outputs": [],
      "source": [
        "model([datapoint_circuits, commands])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jNrW0NXR-lDC"
      },
      "source": [
        "## 3 Learning to prepare eigenstates of different operators\n",
        "\n",
        "The choice of the $\\pm \\hat{Z}$ eigenstates corresponding to 1 and 0 was arbitrary. You could have just as easily wanted 1 to correspond to the $+ \\hat{Z}$ eigenstate and 0 to correspond to the $-\\hat{X}$ eigenstate. One way to accomplish this is by specifying a different measurement operator for each command, as indicated in the figure below:\n",
        "\n",
        "<img src=\"./images/nn_control2.png\" width=\"1000\">\n",
        "\n",
        "This requires use of <code>tfq.layers.Expectation</code>. Now your input has grown to include three objects: circuit, command, and operator. The output is still the expectation value."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ci3WMZ9CjEM1"
      },
      "source": [
        "### 3.1 New model definition\n",
        "\n",
        "Lets take a look at the model to accomplish this task:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "hta0G3Nc6pJY"
      },
      "outputs": [],
      "source": [
        "# Define inputs.\n",
        "commands_input = tf.keras.layers.Input(shape=(1),\n",
        "                                       dtype=tf.dtypes.float32,\n",
        "                                       name='commands_input')\n",
        "circuits_input = tf.keras.Input(shape=(),\n",
        "                                # The circuit-tensor has dtype `tf.string` \n",
        "                                dtype=tf.dtypes.string,\n",
        "                                name='circuits_input')\n",
        "operators_input = tf.keras.Input(shape=(1,),\n",
        "                                 dtype=tf.dtypes.string,\n",
        "                                 name='operators_input')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dtdnkrZm6pJb"
      },
      "source": [
        "Here is the controller network:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "n_aTG4g3-y0F"
      },
      "outputs": [],
      "source": [
        "# Define classical NN.\n",
        "controller = tf.keras.Sequential([\n",
        "    tf.keras.layers.Dense(10, activation='elu'),\n",
        "    tf.keras.layers.Dense(3)\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "q9aN2ciy6pJf"
      },
      "source": [
        "Combine the circuit and the controller into a single `keras.Model` using `tfq`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "IMHjiKit6pJg"
      },
      "outputs": [],
      "source": [
        "dense_2 = controller(commands_input)\n",
        "\n",
        "# Since you aren't using a PQC or ControlledPQC you must append\n",
        "# your model circuit onto the datapoint circuit tensor manually.\n",
        "full_circuit = tfq.layers.AddCircuit()(circuits_input, append=model_circuit)\n",
        "expectation_output = tfq.layers.Expectation()(full_circuit,\n",
        "                                              symbol_names=control_params,\n",
        "                                              symbol_values=dense_2,\n",
        "                                              operators=operators_input)\n",
        "\n",
        "# Contruct your Keras model.\n",
        "two_axis_control_model = tf.keras.Model(\n",
        "    inputs=[circuits_input, commands_input, operators_input],\n",
        "    outputs=[expectation_output])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VQTM6CCiD4gU"
      },
      "source": [
        "### 3.2 The dataset\n",
        "\n",
        "Now you will also include the operators you wish to measure for each datapoint you supply for `model_circuit`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "4gw_L3JG0_G0"
      },
      "outputs": [],
      "source": [
        "# The operators to measure, for each command.\n",
        "operator_data = tfq.convert_to_tensor([[cirq.X(qubit)], [cirq.Z(qubit)]])\n",
        "\n",
        "# The command input values to the classical NN.\n",
        "commands = np.array([[0], [1]], dtype=np.float32)\n",
        "\n",
        "# The desired expectation value at output of quantum circuit.\n",
        "expected_outputs = np.array([[1], [-1]], dtype=np.float32)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ALCKSvwh0_G2"
      },
      "source": [
        "### 3.3 Training\n",
        "\n",
        "Now that you have your new inputs and outputs you can train once again using keras."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nFuGA73MAA4p"
      },
      "outputs": [],
      "source": [
        "optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)\n",
        "loss = tf.keras.losses.MeanSquaredError()\n",
        "\n",
        "two_axis_control_model.compile(optimizer=optimizer, loss=loss)\n",
        "\n",
        "history = two_axis_control_model.fit(\n",
        "    x=[datapoint_circuits, commands, operator_data],\n",
        "    y=expected_outputs,\n",
        "    epochs=30,\n",
        "    verbose=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Cf_G-GdturLL"
      },
      "outputs": [],
      "source": [
        "plt.plot(history.history['loss'])\n",
        "plt.title(\"Learning to Control a Qubit\")\n",
        "plt.xlabel(\"Iterations\")\n",
        "plt.ylabel(\"Error in Control\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "sdCPDH9NlJBl"
      },
      "source": [
        "The loss function has dropped to zero."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NzY8eSVm6pJs"
      },
      "source": [
        "The `controller` is available as a stand-alone model. Call the controller, and check its response to each command signal. It would take some work to correctly compare these outputs to the contents of `random_rotations`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "uXmH0TQ76pJt"
      },
      "outputs": [],
      "source": [
        "controller.predict(np.array([0,1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n2WtXnsxubD2"
      },
      "source": [
        "Success: See if you can adapt the `check_error` function from your first model to work with this new model architecture."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "hello_many_worlds.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}