{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Partial state tomography\n",
    "\n",
    "In this tutorial we look at the convergence of partial state tomography, and how this can be optimized by grouping.\n",
    "\n",
    "1. An example algorithm\n",
    "1. Sampling noise convergence\n",
    "1. Calculating and propagating error bounds on sampling noise\n",
    "1. Setting up parallel measurements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cirq\n",
    "import sympy\n",
    "import openfermioncirq\n",
    "from openfermion import MolecularData, jordan_wigner, eigenspectrum\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# An example algorithm\n",
    "\n",
    "Let's take our VQE for the H2 molecule from last week; here's a simplified form.\n",
    "First, code for the Hamiltonian:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "diatomic_bond_length = .7414\n",
    "\n",
    "geometry = [('H', (0., 0., 0.)), \n",
    "            ('H', (0., 0., diatomic_bond_length))]\n",
    "\n",
    "basis = 'sto-3g'\n",
    "multiplicity = 1\n",
    "charge = 0\n",
    "description = format(diatomic_bond_length)\n",
    "\n",
    "molecule = MolecularData(\n",
    "    geometry,\n",
    "    basis,\n",
    "    multiplicity,\n",
    "    description=description)\n",
    "molecule.load()\n",
    "\n",
    "hamiltonian = jordan_wigner(molecule.get_molecular_hamiltonian())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0, 0): ───Ry(π)──────Rx(0.5π)───@───────────────────────────────────────@───────────Rx(-0.5π)───\n",
      "                                 │                                       │\n",
      "(0, 1): ───Ry(π)──────Ry(0.5π)───X───@───────────────────────@───────────X───────────Ry(-0.5π)───\n",
      "                                     │                       │\n",
      "(1, 0): ───Ry(0.5π)──────────────────X───@───────────────@───X───────────Ry(-0.5π)───────────────\n",
      "                                         │               │\n",
      "(1, 1): ───Ry(0.5π)──────────────────────X───Rz(theta)───X───Ry(-0.5π)───────────────────────────\n"
     ]
    }
   ],
   "source": [
    "theta = sympy.Symbol('theta')\n",
    "qubits = [cirq.GridQubit(i,j) for i in range(2) for j in range(2)]\n",
    "\n",
    "class H2Ansatz(openfermioncirq.VariationalAnsatz):\n",
    "    \n",
    "    def params(self):\n",
    "        return [theta]\n",
    "\n",
    "    def operations(self, qubits):\n",
    "        yield [cirq.ry(np.pi).on(qubits[0]),\n",
    "               cirq.ry(np.pi).on(qubits[1])]\n",
    "        yield [cirq.rx(np.pi/2).on(qubits[0]),\n",
    "               cirq.ry(np.pi/2).on(qubits[1]),\n",
    "               cirq.ry(np.pi/2).on(qubits[2]),\n",
    "               cirq.ry(np.pi/2).on(qubits[3])]\n",
    "        yield [cirq.CNOT(qubits[0], qubits[1]),\n",
    "               cirq.CNOT(qubits[1], qubits[2]),\n",
    "               cirq.CNOT(qubits[2], qubits[3])]\n",
    "        yield cirq.rz(theta).on(qubits[3])\n",
    "        yield [cirq.CNOT(qubits[2], qubits[3]),\n",
    "               cirq.CNOT(qubits[1], qubits[2]),\n",
    "               cirq.CNOT(qubits[0], qubits[1])]\n",
    "        yield [cirq.rx(-np.pi/2).on(qubits[0]),\n",
    "               cirq.ry(-np.pi/2).on(qubits[1]),\n",
    "               cirq.ry(-np.pi/2).on(qubits[2]),\n",
    "               cirq.ry(-np.pi/2).on(qubits[3])]\n",
    "        \n",
    "    def _generate_qubits(self):\n",
    "        return qubits\n",
    "\n",
    "ansatz = H2Ansatz()\n",
    "\n",
    "objective = openfermioncirq.HamiltonianObjective(hamiltonian)\n",
    "study = openfermioncirq.VariationalStudy(\n",
    "    name='UCC_single_term',\n",
    "    ansatz=ansatz,\n",
    "    objective=objective)\n",
    "\n",
    "print(study.circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's find the optimal angle choice for our ground state, and fix it in our VQE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimized VQE result: -1.1372701737048119\n",
      "Target Hamiltonian eigenvalue: -1.137270174625328\n",
      "Optimal angle: [-0.22618398]\n",
      "(0, 0): ───Ry(π)──────Rx(0.5π)───@─────────────────────────────────────────@───────────Rx(-0.5π)───\n",
      "                                 │                                         │\n",
      "(0, 1): ───Ry(π)──────Ry(0.5π)───X───@─────────────────────────@───────────X───────────Ry(-0.5π)───\n",
      "                                     │                         │\n",
      "(1, 0): ───Ry(0.5π)──────────────────X───@─────────────────@───X───────────Ry(-0.5π)───────────────\n",
      "                                         │                 │\n",
      "(1, 1): ───Ry(0.5π)──────────────────────X───Rz(-0.072π)───X───Ry(-0.5π)───────────────────────────\n"
     ]
    }
   ],
   "source": [
    "from openfermioncirq.optimization import COBYLA, OptimizationParams\n",
    "optimization_params = OptimizationParams(\n",
    "    algorithm=COBYLA,\n",
    "    initial_guess=[0.01])\n",
    "result = study.optimize(optimization_params)\n",
    "print('Optimized VQE result: {}'.format(result.optimal_value))\n",
    "print('Target Hamiltonian eigenvalue: {}'.format(eigenspectrum(hamiltonian)[0]))\n",
    "print('Optimal angle: {}'.format(result.optimal_parameters))\n",
    "\n",
    "from cirq import ParamResolver, resolve_parameters\n",
    "resolver = ParamResolver({'theta': result.optimal_parameters})\n",
    "resolved_circuit = resolve_parameters(study.circuit, resolver)\n",
    "\n",
    "print(resolved_circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the rest of the tutorial we'll just worry about this circuit, and in particular how we estimate the energy of the prepared state.\n",
    "\n",
    "Let's start by investigating the Hamiltonian above a bit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.09886397351781592 [] +\n",
       "-0.04532220209856541 [X0 X1 Y2 Y3] +\n",
       "0.04532220209856541 [X0 Y1 Y2 X3] +\n",
       "0.04532220209856541 [Y0 X1 X2 Y3] +\n",
       "-0.04532220209856541 [Y0 Y1 X2 X3] +\n",
       "0.17119774853325856 [Z0] +\n",
       "0.16862219143347554 [Z0 Z1] +\n",
       "0.12054482186554413 [Z0 Z2] +\n",
       "0.16586702396410954 [Z0 Z3] +\n",
       "0.17119774853325856 [Z1] +\n",
       "0.16586702396410954 [Z1 Z2] +\n",
       "0.12054482186554413 [Z1 Z3] +\n",
       "-0.22278592890107016 [Z2] +\n",
       "0.17434844170557132 [Z2 Z3] +\n",
       "-0.22278592890107016 [Z3]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hamiltonian"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to estimate $\\langle \\Psi|H|\\Psi\\rangle$, we want to split it into the different terms as above, and sum the expectation values of the individual terms.\n",
    "\n",
    "Let's first take a single term - $Z_0$:\n",
    "\n",
    "**Problem 1:** Extract the wavefunction generated by the above circuit acting on $|0000\\rangle$, and calculate the expectation value of $Z_0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Insert your code here!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately, on a quantum computer we never have direct access to the wavefunction. Instead, we have to estimate $\\langle\\Psi|Z_0|\\Psi\\rangle$ by inference from repeated preparation of $|\\Psi\\rangle$ and measurement in the $Z_0$ basis.\n",
    "\n",
    "**Problem 2:** Copy the resolved circuit, and add a measurement of qubit $0$ in the $Z$ basis.\n",
    "\n",
    "**Problem 3:** Extract an estimation of $\\langle\\Psi|Z_0|\\Psi\\rangle$ from the output of $10^4$ repetitions of the above circuit, and compare to your answer to Problem 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Insert your code here!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This works, but it's not perfect, is it? The result can be improved by turning up the number of repetitions, let's try it:\n",
    "\n",
    "**Problem 4:** Repeat the above using $100$, $500$, $1000$, $5000$, and $10^4$ repetitions. For each choice of the number of repetitions, average the error in your approximation over $10$ experiments.\n",
    "\n",
    "**Problem 5:** Create a log-log plot of the above results and determine the scaling of the error as a function of the number of repetitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Insert your code here!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This scaling is called sampling noise scaling, and is fundamental to any quantum device. In particular, this implies that extracting any expectation value to exponentially small precision is generally exponentially difficult (though exceptions do exist). Luckily, we can still achieve a speedup over classical computers with only polynomially small error.\n",
    "\n",
    "This circuitry above works for the single operator we focused on, but we need to get the rest. Luckily, we can parallelise this significantly. In particular, all products of $Z$ operators are amenable --- they share the same tensor factors on each qubit --- and can be read out simultaneously in a single experiment.\n",
    "\n",
    "**Problem 6**: write code to estimate expectation values for all tensor products of $Z$ operators in the above Hamiltonian using the output from a single circuit, and execute it using $10^4$ shots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Insert your code here!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This leaves four expectation values to calculate. Unfortunately none of the corresponding operators are amenable, and so the expectation values cannot be estimated simultaneously with single-qubit rotations. The simplest (and lowest-depth) method to estimate the remaining four expectation values is in parallel.\n",
    "\n",
    "**Problem 7**: write circuits and code to estimate the expectation values of the remaining four operators, and evaluate this using $10^4$ repetitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Insert your code here!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It remains to combine the values that you have to estimate the expectation value of the Hamiltonian itself.\n",
    "\n",
    "**Problem 8:** Take an appropriate linear combination of the values calculated above to estimate the expectation value of the Hamiltonian, and compare to the value originally found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Insert your code here!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Extension problem 1:** Write down analytic formulae for the expectation values of single Pauli operators, and use this to write a function to predict the error in the estimation of the Hamiltonian. Compare your results with simulation.\n",
    "\n",
    "**Extension problem 2:** In the above, we chose a fixed number of samples for each estimation, regardless of how much this contributed to the final Hamiltonian. How could you optimize this choice?\n",
    "\n",
    "**Extension problem 3:** Although amenability is sufficient for simultaneous measurement of operators, it is not necessary; the necessary condition is just that the operators commute. Which of the operators you studied in Problem 7 commute? Can you design a circuit to perform simultaneous measurement of them?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}