<commit_before>{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "sensing.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.5rc1"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xLOXFOT5Q40E"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "iiQkM5ZgQ8r2",
        "colab": {}
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "j6331ZSsQGY3"
      },
      "source": [
        "# Quantum sensing"
      ]
    },
    {
      "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/sensing\"><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/sensing.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/sensing.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_quantum/docs/tutorials/sensing.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sb8xt8c9habl",
        "colab_type": "text"
      },
      "source": [
        "Quantum sensing applications measure the quantum properties or quantum phenomena of physical systems. This includes motion, gravity, electric and magnetic fields, and imaging.\n",
        "\n",
        "This tutorial shows a simple sensing problem and uses TensorFlow Quantum to learn to amplify a weak quantum signal. You will use a <a href=\"https://wikipedia.org/wiki/Greenberger%E2%80%93Horne%E2%80%93Zeilinger_state\" class=\"external\">Greenberger–Horne–Zeilinger (GHZ) state</a> that interacts with a \"signal cavity\" to amplify the signal."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bK3mhgcgg56v",
        "colab_type": "text"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Download and install the required packages:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "fAVlg7rxkvUw",
        "colab": {}
      },
      "source": [
        "%%capture\n",
        "!pip install --upgrade pip\n",
        "!pip install cirq==0.7.0"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "TorxE5tnkvb2",
        "colab": {}
      },
      "source": [
        "%%capture\n",
        "!pip install --upgrade tensorflow==2.1.0"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rn6Bj5sVhBck",
        "colab_type": "text"
      },
      "source": [
        "Note: If the following code cell fails, execute the first code cells and then restart the Colab runtime (*Runtime > Restart Runtime*)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "saFHsRDpkvkH",
        "colab": {}
      },
      "source": [
        "%%capture\n",
        "h = \"2dfcfceb9726fa73c40381c037dc01facd3d061e\"\n",
        "!cd ~/\n",
        "!rm -r -f TFQuantum/\n",
        "!git clone https://{h}:{h}@github.com/quantumlib/TFQuantum.git;cd TFQuantum/\n",
        "!pip install --upgrade ./TFQuantum/wheels/tfquantum-0.2.0-cp36-cp36m-linux_x86_64.whl"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AiqayQeNhNDq",
        "colab_type": "text"
      },
      "source": [
        "Now import TensorFlow and the module dependencies:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "2NGNsLzqN4Lb",
        "colab": {}
      },
      "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"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "b08Mmbs8lr81"
      },
      "source": [
        "## Amplify signals for sensing\n",
        "\n",
        "Start with the following circuit:\n",
        "\n",
        "<img src=\"./images/sensing_1.png\" width=\"1000\">\n",
        "\n",
        "If you follow the state of the `(0,0)` qubit to the end of this circuit, you see it is rotated on the x-axis away from the zero state by an angle of $N \\theta$. This effectively \"amplifies\" the input rotation by a factor of the number of qubits involved in the entangled state.\n",
        "\n",
        "It turns out this amplification effect is sensitive to the signal injection of exactly one rotation on the z-axis. Can a parametric circuit be trained to perform the same task for an arbitrary axis of rotation?\n",
        "\n",
        "For the correct values of $\\vec{\\phi}_n$, measuring $ \\langle \\hat{Z} \\rangle$ should produce 1 for any $\\theta$:\n",
        "\n",
        "<img src=\"./images/sensing_2.png\" width=\"1000\">\n",
        "\n",
        "But can $\\vec{\\phi}_n$ be found in practice? Let's explore this using TensorFlow Quantum."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "u7hmqfTR9VuD"
      },
      "source": [
        "## 1. Create layers for in-graph circuit construction\n",
        "\n",
        "Use the above diagrams as reference to define the static parts of the circuit. First, build a circuit that prepares a GHZ state:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "b6ir7z3Y-GFQ",
        "colab": {}
      },
      "source": [
        "def get_ghz_circuit(qubits):\n",
        "    # This method will return a Cirq circuit representing the GHZ\n",
        "    # state preparations circuit to prepend to the input of this layer\n",
        "    # at runtime.\n",
        "    circuit = cirq.Circuit()\n",
        "    circuit.append(cirq.H(qubits[0]))\n",
        "\n",
        "    for control, target in zip(qubits[:-1], qubits[1:]):\n",
        "        circuit.append(cirq.CNOT(control=control, target=target))\n",
        "    return circuit"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "zTZvLBCZQPlN",
        "colab": {}
      },
      "source": [
        "SVGCircuit(get_ghz_circuit(cirq.GridQubit.rect(1, 4)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "q3ZgFfaTRJRh"
      },
      "source": [
        "Create a layer that applies an arbitrary single qubit rotation to each qubit:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "AJM-MuvGRcrn",
        "colab": {}
      },
      "source": [
        "def one_qubit_rotation(bit, parameters):\n",
        "    return [\n",
        "        cirq.X(bit)**parameters[0],\n",
        "        cirq.Y(bit)**parameters[1],\n",
        "        cirq.Z(bit)**parameters[2]\n",
        "    ]\n",
        "\n",
        "\n",
        "def get_single_qubit_rotation_wall(qubits, params):\n",
        "    circuit = cirq.Circuit()\n",
        "    for bit, param in zip(qubits, params):\n",
        "        circuit.append(one_qubit_rotation(bit, param))\n",
        "    return circuit"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "9YN4sLasSgIj",
        "scrolled": false,
        "colab": {}
      },
      "source": [
        "SVGCircuit(\n",
        "    get_single_qubit_rotation_wall(cirq.GridQubit.rect(1, 4),\n",
        "                                   np.random.uniform(0, 2 * np.pi,\n",
        "                                                     (4, 3))))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gG1TSEJWk8mp"
      },
      "source": [
        "The GHZ circuit previously defined can be inverted to collect the entanglement:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "1maK-MOElnut",
        "scrolled": false,
        "colab": {}
      },
      "source": [
        "SVGCircuit(get_ghz_circuit(cirq.GridQubit.rect(1, 4)) ** -1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "3W5vCsgTlyc-"
      },
      "source": [
        "## 2. Model definition\n",
        "\n",
        "Now build a model. Signals are simulated with one qubit rotations."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "5stYOW-ll3Tf",
        "colab": {}
      },
      "source": [
        "# Qubits for this problem.\n",
        "sensing_qubits = cirq.GridQubit.rect(1, 6)\n",
        "\n",
        "# Some random signal injection angle.\n",
        "signal_injection_angles = np.random.uniform(0, 2 * np.pi,\n",
        "                                            (len(sensing_qubits), 3))\n",
        "\n",
        "# Phi parameters that need to be learned.\n",
        "phis = [sympy.symbols('x_{}_0:3'.format(i)) for i in range(len(sensing_qubits))]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "RDvejUwLofa-",
        "scrolled": false,
        "colab": {}
      },
      "source": [
        "# Input for the wall of Rz gates that inject the simulated signal.\n",
        "signal_input = tf.keras.layers.Input(shape=(), dtype=tf.dtypes.string)\n",
        "\n",
        "# Input for the Rx gate that should undo the signal injection during training.\n",
        "expected_unrotation_input = tf.keras.layers.Input(shape=(),\n",
        "                                                  dtype=tf.dtypes.string)\n",
        "\n",
        "# Wall of random 1 qubit rotations that randomly select a signal injection axis.\n",
        "injection_angle_randomizer = tfq.layers.AddCircuit()(\n",
        "    signal_input, prepend=get_single_qubit_rotation_wall(\n",
        "        sensing_qubits, signal_injection_angles))\n",
        "\n",
        "# Wall of parameterized 1 qubit rotations that will be trained.\n",
        "injection_angle_selector = tfq.layers.AddCircuit()(\n",
        "    injection_angle_randomizer, prepend=get_single_qubit_rotation_wall(\n",
        "        sensing_qubits, phis))\n",
        "\n",
        "# GHZ prep and unprep.\n",
        "ghz_prep = tfq.layers.AddCircuit()(\n",
        "    injection_angle_selector, prepend=get_ghz_circuit(sensing_qubits))\n",
        "\n",
        "ghz_unprep = tfq.layers.AddCircuit()(\n",
        "    ghz_prep, append=get_ghz_circuit(sensing_qubits)**-1)\n",
        "\n",
        "# Add the \"unrotation\" to each input using lower level append_circuit op.\n",
        "expected_unrotation = tfq.append_circuit(ghz_unprep, expected_unrotation_input)\n",
        "\n",
        "# In this case keras can train all of the circuit parameters, so you can pass\n",
        "# circuit parameters to trainable_params instead of feed_in_params.\n",
        "expectation_output = tfq.layers.Expectation()(\n",
        "    expected_unrotation, symbol_names=list(np.array(phis).flatten()),\n",
        "    operators=[cirq.Z(sensing_qubits[0])])\n",
        "\n",
        "sensing_model = tf.keras.Model(inputs=[signal_input, expected_unrotation_input],\n",
        "                               outputs=[expectation_output])\n",
        "\n",
        "# Display model architecture\n",
        "tf.keras.utils.plot_model(sensing_model,\n",
        "                          show_shapes=True,\n",
        "                          show_layer_names=False,\n",
        "                          dpi=70)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ipx77hn8uU9S"
      },
      "source": [
        "## 3. Data definition\n",
        "The signal input is $Rz(\\theta)$ on each qubit. This leads to a state rotated on the x-axis by $N \\theta$ at the output of the GHZ un-preparation. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "ht1r21O_vZCm",
        "scrolled": false,
        "colab": {}
      },
      "source": [
        "inputs = []\n",
        "un_rotations = []\n",
        "thetas = []\n",
        "\n",
        "for theta in np.arange(0, 2 * np.pi, step=0.05):\n",
        "    thetas.append(theta)\n",
        "    # Signal injection is an Rz on each qubit.\n",
        "    inputs.append(cirq.Circuit(\n",
        "        cirq.Rz(theta).on(bit) for bit in sensing_qubits))\n",
        "\n",
        "    # During training you want to undo the x rotation at the end of the circuit to send\n",
        "    # the qubit back to the zero state.\n",
        "    un_rotations.append(\n",
        "        cirq.Circuit(\n",
        "            cirq.Rx(-1 * len(sensing_qubits) * theta).on(sensing_qubits[0])))\n",
        "\n",
        "signal_injection_tensor = tfq.convert_to_tensor(inputs)\n",
        "un_rotation_tensor = tfq.convert_to_tensor(un_rotations)\n",
        "\n",
        "SVGCircuit(inputs[5])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "cDF4EsxI48ZG"
      },
      "source": [
        "## 4. Untrained performance\n",
        "\n",
        "Let's see how the *untrained* circuit performs at amplifying the input signal:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "3g1U-GYd9Dci",
        "colab": {}
      },
      "source": [
        "# when evaluating the performance of the model you don't want to un-rotate\n",
        "# at the end, so send some empty circuits into that input of the model.\n",
        "empty_circuits = tfq.convert_to_tensor([cirq.Circuit()] *\n",
        "                                       signal_injection_tensor.shape[0])\n",
        "\n",
        "untrained_outputs = sensing_model.predict(\n",
        "    x=[signal_injection_tensor, empty_circuits])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "N_J_QNhc5elN",
        "scrolled": false,
        "colab": {}
      },
      "source": [
        "plt.plot(thetas, np.cos(thetas), label='Un-amplified signal')\n",
        "plt.plot(thetas, untrained_outputs, label='Untrained Amplified Output')\n",
        "plt.title(\"Untrained Amplification Performance\")\n",
        "plt.legend(loc='lower right')\n",
        "plt.xlabel(r\"\\theta\")\n",
        "plt.ylabel(r\"$-\\langle \\hat{Z} \\rangle$\")\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vLE6ucT8zqnI"
      },
      "source": [
        "## 5. Train the model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "DaINoOSuzsAE",
        "colab": {}
      },
      "source": [
        "optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)\n",
        "loss = lambda x, y: -1 * y\n",
        "\n",
        "sensing_model.compile(optimizer=optimizer, loss=loss)\n",
        "\n",
        "history = sensing_model.fit(x=[signal_injection_tensor, un_rotation_tensor],\n",
        "                            y=tf.zeros_like(signal_injection_tensor,\n",
        "                                            dtype=tf.float32),\n",
        "                            epochs=10,\n",
        "                            batch_size=20,\n",
        "                            verbose=1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "hDqB6EXs1TEM",
        "scrolled": true,
        "colab": {}
      },
      "source": [
        "plt.plot(history.history['loss'])\n",
        "plt.title(\"Learning to Sense A Randomized Feild\")\n",
        "plt.xlabel(\"Epochs\")\n",
        "plt.ylabel(r\"$-\\langle \\hat{Z} \\rangle$\")\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "IwlCaAgvEvSC"
      },
      "source": [
        "## 6. Trained performance\n",
        "\n",
        "And here's how the *trained* sensing model performs:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "qXzCn2_B2Y3U",
        "colab": {}
      },
      "source": [
        "empty_circuits = tfq.convert_to_tensor([cirq.Circuit()] *\n",
        "                                       signal_injection_tensor.shape[0])\n",
        "\n",
        "trained_outputs = sensing_model.predict(\n",
        "    x=[signal_injection_tensor, empty_circuits])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "oduATtbj3SAp",
        "scrolled": false,
        "colab": {}
      },
      "source": [
        "plt.plot(thetas, np.cos(thetas), label='Un-amplified signal')\n",
        "plt.plot(thetas, trained_outputs, label='Amplified Output')\n",
        "plt.title(\"Trained Amplification Performance\")\n",
        "plt.legend(loc='lower right')\n",
        "plt.xlabel(r\"\\theta\")\n",
        "plt.ylabel(r\"$-\\langle \\hat{Z} \\rangle$\")\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8z58CjVejOCV"
      },
      "source": [
        "From the plot you can see that the trained model has amplified the signal from the cavity using the GHZ state. You can measure this signal (which is easier now since it's amplified) or use it as a component in another quantum algorithm."
      ]
    }
  ]
}<commit_msg>Re-ordered sensing demo. (#45)<commit_after>{
 "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"
   },
   "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": [
    "# Quantum sensing"
   ]
  },
  {
   "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/sensing\"><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/sensing.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/sensing.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_quantum/docs/tutorials/sensing.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": "Sb8xt8c9habl"
   },
   "source": [
    "Quantum sensing applications measure the quantum properties or quantum phenomena of physical systems. This includes motion, gravity, electric and magnetic fields, and imaging.\n",
    "\n",
    "This tutorial shows a simple sensing problem and uses TensorFlow Quantum to learn to amplify a weak quantum signal. You will use a <a href=\"https://wikipedia.org/wiki/Greenberger%E2%80%93Horne%E2%80%93Zeilinger_state\" class=\"external\">Greenberger–Horne–Zeilinger (GHZ) state</a> that interacts with a \"signal cavity\" to amplify the signal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bK3mhgcgg56v"
   },
   "source": [
    "## Setup\n",
    "\n",
    "Download and install the required packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fAVlg7rxkvUw"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install --upgrade pip\n",
    "!pip install cirq==0.7.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TorxE5tnkvb2"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install --upgrade tensorflow==2.1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Rn6Bj5sVhBck"
   },
   "source": [
    "Note: If the following code cell fails, execute the first code cells and then restart the Colab runtime (*Runtime > Restart Runtime*)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "saFHsRDpkvkH"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "h = \"2dfcfceb9726fa73c40381c037dc01facd3d061e\"\n",
    "!cd ~/\n",
    "!rm -r -f TFQuantum/\n",
    "!git clone https://{h}:{h}@github.com/quantumlib/TFQuantum.git;cd TFQuantum/\n",
    "!pip install --upgrade ./TFQuantum/wheels/tfquantum-0.2.0-cp36-cp36m-linux_x86_64.whl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AiqayQeNhNDq"
   },
   "source": [
    "Now import TensorFlow and the module dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2NGNsLzqN4Lb"
   },
   "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": [
    "## Amplify signals for sensing\n",
    "\n",
    "Start with the following circuit:\n",
    "\n",
    "<img src=\"./images/sensing_1.png\" width=\"1000\">\n",
    "\n",
    "If you follow the state of the `(0,0)` qubit to the end of this circuit, you see it is rotated on the x-axis away from the zero state by an angle of $N \\theta$. This effectively \"amplifies\" the input rotation by a factor of the number of qubits involved in the entangled state.\n",
    "\n",
    "It turns out this amplification effect is sensitive to the signal injection of exactly one rotation on the z-axis. Can a parametric circuit be trained to perform the same task for an arbitrary axis of rotation?\n",
    "\n",
    "For the correct values of $\\vec{\\phi}_n$, measuring $ \\langle \\hat{Z} \\rangle$ should produce 1 for any $\\theta$:\n",
    "\n",
    "<img src=\"./images/sensing_2.png\" width=\"1000\">\n",
    "\n",
    "But can $\\vec{\\phi}_n$ be found in practice? Let's explore this using TensorFlow Quantum."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "u7hmqfTR9VuD"
   },
   "source": [
    "## 1. Create layers for in-graph circuit construction\n",
    "\n",
    "Use the above diagrams as reference to define the static parts of the circuit. First, build a circuit that prepares a GHZ state:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "b6ir7z3Y-GFQ"
   },
   "outputs": [],
   "source": [
    "def get_ghz_circuit(qubits):\n",
    "    # This method will return a Cirq circuit representing the GHZ\n",
    "    # state preparations circuit to prepend to the input of this layer\n",
    "    # at runtime.\n",
    "    circuit = cirq.Circuit()\n",
    "    circuit.append(cirq.H(qubits[0]))\n",
    "\n",
    "    for control, target in zip(qubits[:-1], qubits[1:]):\n",
    "        circuit.append(cirq.CNOT(control=control, target=target))\n",
    "    return circuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zTZvLBCZQPlN"
   },
   "outputs": [],
   "source": [
    "SVGCircuit(get_ghz_circuit(cirq.GridQubit.rect(1, 4)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "q3ZgFfaTRJRh"
   },
   "source": [
    "Create a layer that applies an arbitrary single qubit rotation to each qubit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "AJM-MuvGRcrn"
   },
   "outputs": [],
   "source": [
    "def one_qubit_rotation(bit, parameters):\n",
    "    return [\n",
    "        cirq.X(bit)**parameters[0],\n",
    "        cirq.Y(bit)**parameters[1],\n",
    "        cirq.Z(bit)**parameters[2]\n",
    "    ]\n",
    "\n",
    "\n",
    "def get_single_qubit_rotation_wall(qubits, params):\n",
    "    circuit = cirq.Circuit()\n",
    "    for bit, param in zip(qubits, params):\n",
    "        circuit.append(one_qubit_rotation(bit, param))\n",
    "    return circuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9YN4sLasSgIj",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "SVGCircuit(\n",
    "    get_single_qubit_rotation_wall(cirq.GridQubit.rect(1, 4),\n",
    "                                   np.random.uniform(0, 2 * np.pi, (4, 3))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gG1TSEJWk8mp"
   },
   "source": [
    "The GHZ circuit previously defined can be inverted to collect the entanglement:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1maK-MOElnut",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "SVGCircuit(get_ghz_circuit(cirq.GridQubit.rect(1, 4))**-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Define the data\n",
    "The signal input is $Rz(\\theta)$ on each qubit. This leads to a state rotated on the x-axis by $N \\theta$ at the output of the GHZ un-preparation. During the model building stage you will see how to input these datapoints in the middle of your model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Qubits for this problem.\n",
    "sensing_qubits = cirq.GridQubit.rect(1, 6)\n",
    "\n",
    "inputs = []\n",
    "un_rotations = []\n",
    "thetas = []\n",
    "\n",
    "for theta in np.arange(0, 2 * np.pi, step=0.05):\n",
    "    thetas.append(theta)\n",
    "    # Signal injection is an Rz on each qubit.\n",
    "    inputs.append(cirq.Circuit(\n",
    "        cirq.Rz(theta).on(bit) for bit in sensing_qubits))\n",
    "\n",
    "    # During training you want to undo the x rotation at the end of the circuit to send\n",
    "    # the qubit back to the zero state.\n",
    "    un_rotations.append(\n",
    "        cirq.Circuit(\n",
    "            cirq.Rx(-1 * len(sensing_qubits) * theta).on(sensing_qubits[0])))\n",
    "\n",
    "signal_injection_tensor = tfq.convert_to_tensor(inputs)\n",
    "un_rotation_tensor = tfq.convert_to_tensor(un_rotations)\n",
    "\n",
    "SVGCircuit(inputs[5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3W5vCsgTlyc-"
   },
   "source": [
    "## 3. Define the model\n",
    "\n",
    "Now build a model. Signals are simulated with one qubit rotations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5stYOW-ll3Tf"
   },
   "outputs": [],
   "source": [
    "# Some random signal injection angle (this will be the random rotation to undo).\n",
    "signal_injection_angles = np.random.uniform(0, 2 * np.pi,\n",
    "                                            (len(sensing_qubits), 3))\n",
    "\n",
    "# Phi parameters that need to be learned.\n",
    "phis = [sympy.symbols('x_{}_0:3'.format(i)) for i in range(len(sensing_qubits))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RDvejUwLofa-",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Input for the wall of Rz gates that inject the simulated signal.\n",
    "signal_input = tf.keras.layers.Input(shape=(), dtype=tf.dtypes.string)\n",
    "\n",
    "# Input for the Rx gate that should undo the signal injection during training.\n",
    "expected_unrotation_input = tf.keras.layers.Input(shape=(),\n",
    "                                                  dtype=tf.dtypes.string)\n",
    "\n",
    "# Wall of random 1 qubit rotations that randomly select a signal injection axis.\n",
    "injection_angle_randomizer = tfq.layers.AddCircuit()(\n",
    "    signal_input,\n",
    "    prepend=get_single_qubit_rotation_wall(sensing_qubits,\n",
    "                                           signal_injection_angles))\n",
    "\n",
    "# Wall of parameterized 1 qubit rotations that will be trained.\n",
    "injection_angle_selector = tfq.layers.AddCircuit()(\n",
    "    injection_angle_randomizer,\n",
    "    prepend=get_single_qubit_rotation_wall(sensing_qubits, phis))\n",
    "\n",
    "# GHZ prep and unprep.\n",
    "ghz_prep = tfq.layers.AddCircuit()(injection_angle_selector,\n",
    "                                   prepend=get_ghz_circuit(sensing_qubits))\n",
    "\n",
    "ghz_unprep = tfq.layers.AddCircuit()(ghz_prep,\n",
    "                                     append=get_ghz_circuit(sensing_qubits)**-1)\n",
    "\n",
    "# Add the \"unrotation\" to each input using lower level append_circuit op.\n",
    "expected_unrotation = tfq.append_circuit(ghz_unprep, expected_unrotation_input)\n",
    "\n",
    "# In this case keras can train all of the circuit parameters, so you can pass\n",
    "# circuit parameters to trainable_params instead of feed_in_params.\n",
    "expectation_output = tfq.layers.Expectation()(\n",
    "    expected_unrotation,\n",
    "    symbol_names=list(np.array(phis).flatten()),\n",
    "    operators=[cirq.Z(sensing_qubits[0])])\n",
    "\n",
    "sensing_model = tf.keras.Model(inputs=[signal_input, expected_unrotation_input],\n",
    "                               outputs=[expectation_output])\n",
    "\n",
    "# Display model architecture\n",
    "tf.keras.utils.plot_model(sensing_model,\n",
    "                          show_shapes=True,\n",
    "                          show_layer_names=False,\n",
    "                          dpi=70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cDF4EsxI48ZG"
   },
   "source": [
    "## 4. Untrained performance\n",
    "\n",
    "Let's see how the *untrained* circuit performs at amplifying the input signal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3g1U-GYd9Dci"
   },
   "outputs": [],
   "source": [
    "# when evaluating the performance of the model you don't want to un-rotate\n",
    "# at the end, so send some empty circuits into that input of the model.\n",
    "empty_circuits = tfq.convert_to_tensor([cirq.Circuit()] *\n",
    "                                       signal_injection_tensor.shape[0])\n",
    "\n",
    "untrained_outputs = sensing_model.predict(\n",
    "    x=[signal_injection_tensor, empty_circuits])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "N_J_QNhc5elN",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.plot(thetas, np.cos(thetas), label='Un-amplified signal')\n",
    "plt.plot(thetas, untrained_outputs, label='Untrained Amplified Output')\n",
    "plt.title(\"Untrained Amplification Performance\")\n",
    "plt.legend(loc='lower right')\n",
    "plt.xlabel(r\"\\theta\")\n",
    "plt.ylabel(r\"$-\\langle \\hat{Z} \\rangle$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vLE6ucT8zqnI"
   },
   "source": [
    "## 5. Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DaINoOSuzsAE"
   },
   "outputs": [],
   "source": [
    "optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)\n",
    "loss = lambda x, y: -1 * y\n",
    "\n",
    "sensing_model.compile(optimizer=optimizer, loss=loss)\n",
    "\n",
    "history = sensing_model.fit(x=[signal_injection_tensor, un_rotation_tensor],\n",
    "                            y=tf.zeros_like(signal_injection_tensor,\n",
    "                                            dtype=tf.float32),\n",
    "                            epochs=10,\n",
    "                            batch_size=20,\n",
    "                            verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hDqB6EXs1TEM",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plt.plot(history.history['loss'])\n",
    "plt.title(\"Learning to Sense A Randomized Feild\")\n",
    "plt.xlabel(\"Epochs\")\n",
    "plt.ylabel(r\"$-\\langle \\hat{Z} \\rangle$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IwlCaAgvEvSC"
   },
   "source": [
    "## 6. Trained performance\n",
    "\n",
    "And here's how the *trained* sensing model performs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qXzCn2_B2Y3U"
   },
   "outputs": [],
   "source": [
    "empty_circuits = tfq.convert_to_tensor([cirq.Circuit()] *\n",
    "                                       signal_injection_tensor.shape[0])\n",
    "\n",
    "trained_outputs = sensing_model.predict(\n",
    "    x=[signal_injection_tensor, empty_circuits])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oduATtbj3SAp",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.plot(thetas, np.cos(thetas), label='Un-amplified signal')\n",
    "plt.plot(thetas, trained_outputs, label='Amplified Output')\n",
    "plt.title(\"Trained Amplification Performance\")\n",
    "plt.legend(loc='lower right')\n",
    "plt.xlabel(r\"\\theta\")\n",
    "plt.ylabel(r\"$-\\langle \\hat{Z} \\rangle$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8z58CjVejOCV"
   },
   "source": [
    "From the plot you can see that the trained model has amplified the signal from the cavity using the GHZ state. You can measure this signal (which is easier now since it's amplified) or use it as a component in another quantum algorithm."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "sensing.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
{'fraction_non_alphanumeric': 0.17799328792814134, 'fraction_numerical': 0.010956470239857862, 'mean_word_length': 2.128135854882285, 'pattern_counts': {'":': 531, '<': 55, '<?xml version=': 0, '>': 57, 'https://': 22, 'lorem ipsum': 0, 'www.': 12, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 45, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3599438', 'n_tokens_mistral': 14655, 'n_tokens_neox': 14033, 'n_words': 2920}