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": 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": "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TorxE5tnkvb2",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "!pip install tensorflow==2.15.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FxkQA6oblNqI"
   },
   "source": [
    "Install TensorFlow Quantum:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "saFHsRDpkvkH",
    "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": "F1L8h1YKUvIO"
   },
   "source": [
    "Now import TensorFlow and the module dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "enZ300Bflq80",
    "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": "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2yQdmhQLCrzQ",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ps-pd2mndXs7",
    "vscode": {
     "languageId": "python"
    }
   },
   "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(cirq.rx(a).on(q0), cirq.ry(b).on(q1), cirq.CNOT(q0, 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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VMq7EayNRyQb",
    "vscode": {
     "languageId": "python"
    }
   },
   "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,\n",
    "                                                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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hrSnOCi3ehr_",
    "scrolled": true,
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OZ0lWFXv6pII",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1gLQjA02mIyy",
    "scrolled": true,
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aX_vEmCKmpQS",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1fsVZhF5lIXp",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "batch_vals = np.array(np.random.uniform(0, 2 * np.pi, (5, 2)), dtype=float)"
   ]
  },
  {
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RsfF53UCJtr9",
    "vscode": {
     "languageId": "python"
    }
   },
   "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,\n",
    "                                                 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",
    "\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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kGZVdcZ6y9lC",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "N-j7SCl-51-q",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1v4CK2jD6pIj",
    "vscode": {
     "languageId": "python"
    }
   },
   "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)])"
   ]
  },
  {
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kZbYRTe16pIm",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UfHF8NNE6pIr",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# This input is the simulated miscalibration that the model will learn to correct.\n",
    "circuits_input = tf.keras.Input(\n",
    "    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')"
   ]
  },
  {
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Zvt2YGmZ6pIu",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dense_2 = controller(commands_input)\n",
    "\n",
    "# TFQ layer for classically controlled circuits.\n",
    "expectation_layer = tfq.layers.ControlledPQC(\n",
    "    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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Xs6EMhah6pIz",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ERXNPe4F6pI4",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ciMIJAuH6pJA",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_VYfzHffWo7n",
    "vscode": {
     "languageId": "python"
    }
   },
   "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",
    "datapoint_circuits = tfq.convert_to_tensor([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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6nk2Yr3e6pJJ",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Lwphqvs96pJO",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dtPYqbNi8zeZ",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "azE-qV0OaC1o",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RoIlb7r7j5SY",
    "vscode": {
     "languageId": "python"
    }
   },
   "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(full_circuit, {\n",
    "            s: v\n",
    "            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, {\n",
    "            qubit: 0\n",
    "        }).real\n",
    "        print(\n",
    "            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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aYskLTacs8Ku",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hta0G3Nc6pJY",
    "vscode": {
     "languageId": "python"
    }
   },
   "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(\n",
    "    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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "n_aTG4g3-y0F",
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Define classical NN.\n",
    "controller = tf.keras.Sequential(\n",
    "    [tf.keras.layers.Dense(10, activation='elu'),\n",
    "     tf.keras.layers.Dense(3)])"
   ]
  },
  {
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IMHjiKit6pJg",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4gw_L3JG0_G0",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nFuGA73MAA4p",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Cf_G-GdturLL",
    "vscode": {
     "languageId": "python"
    }
   },
   "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": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uXmH0TQ76pJt",
    "vscode": {
     "languageId": "python"
    }
   },
   "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
}