cusp_demo_short.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/zapata.png\" width=\"100px\">\n",
    "\n",
    "# CUSP Code Demonstration with Cirq\n",
    "    \n",
    "Compressed Unsupervised State Preparation (CUSP) is a method for building more efficient quantum circuits by using a quantum autoencoder.  The protocol performs a kind of circuit synthesis that, if training is successful, results in a more compact circuit.  Since typically shorter-depth circuits are less prone to noise on a real quantum computer, this tool gives the opportunity to make a more accurate state preparation, resulting in better accuracy for a quantum computation.  In this demo, we will use the example of improving a circuit which computes the ground state energies of molecular hydrogen at various bond lengths.\n",
    "\n",
    "CUSP has 3 stages of learning (with an optional fourth stage that is not included in this demo).  The first stage is to simply train (e.g. using a circuit ansatz) the circuit that we would like to improve.  In the second stage, we will train a quantum autoencoder to compress the output of Stage 1 into a representation on a smaller number of qubits (called the \"latent space\").  In the last stage, we choose a small variational circuit in the latent space, apply the decoder that we learned from stage 2, and try to optimize the parameters to generate the best state we can.  If that sounds a bit abstract, don't worry!  We will step you more specifically through each of the stages below.\n",
    "\n",
    "Please note that this is the short version of the demo, which is a bit more elegant to look at and use, but does not get into many details of the code.  A new version is coming soon which will walk through the code in detail.  For a formal description of the CUSP protocol, please refer to `cusp_protocol.pdf`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by importing some necessary modules.  In this short demo, almost all of the code is hidden in these files so you don't have to see it.  The code that is implemented here utilizes OpenFermion for its chemistry integration and Google's quantum simulator Cirq for processing quantum circuits. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import modules\n",
    "import numpy as np\n",
    "from scipy.optimize import minimize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background\n",
    "\n",
    "For this demo, we will look to improving a circuit ansatz which is used for the variational quantum eigensolver (VQE) algorithm.  VQE is a variational algorithm that computes the approximate ground state energy of a molecule with a particular configuration.  Here, we will consider a very simple example of molecular hydrogen (H2) with a variable bond length (the bond length is the distance between the two hydrogen atoms).  Since the ground state energy of H2 changes as a function of the bond length, we will have to choose which bond lengths we would like to train on and compute.\n",
    "\n",
    "\n",
    "<img src=\"images/H2_curve.png\" width=\"500px\">\n",
    "\n",
    "\n",
    "<center>Bond dissociation curve for molecular hydrogen in the minimal basis (STO-3G)</center>\n",
    "\n",
    "\n",
    "For simplicity in this demo, the variational circuits have been chosen for you in each stage, which we will detail more later.  You will be left to choose which parameters of those circuits to fix, and which to optimize over."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings for CUSP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### This is the general environment in which CUSP will run, including strength of noise, number of statistical samples, and the set of training examples to consider.  In order to give a fair analysis of the algorithm, these settings will apply to all stages of the circuit.\n",
    "\n",
    "`num_trials` determines how many times we will run a circuit to gather our statistics. Optimization routines may struggle to converge if `num_trials` is small and noise is large.  However, the time to run the simulation is largely determined by the value of `num_trials` so we recommend `num_trials = 1000` as an upperbound if the noise is significant.  If the noise is subtle, `num_trials = 500` will likely be sufficient.\n",
    "\n",
    "Setting `include_gate_noise` to `True` allows you to include noise during the simulation.  Note that turning this on will increase the simulation time roughly by the scale of `num_trials`, because we will run that many instances of the circuit to get our measurement statistics for just one run.  To see a substantial difference between the original circuit and the CUSP optimized circuit, you will need to set `gate_noise = True`.  However, to simply see how the CUSP protocol operates, we recommend you set `gate_noise = False` since this will dramatically increase the speed of the simulation.\n",
    "\n",
    "`noise_level` determines the amount of (dephasing) noise to include during the simulations. This is a built-in Cirq function where the value of `noise_level` is associated with the probability of having an error after each operation. Since it is a probability, the parameter values range between 0 (no noise) and 1 (an error after every gate). By default, Cirq's noisy simulations run with a noise level of `0.001`.  Because the original circuit we are trying to optimize is already a bit short, we recommend `noise_level = 0.002` to see a marked difference between the original circuit and the CUSP-improved circuit.  Note that very high levels of noise may prevent the protocol from working since the autoencoder would rarely get reliable input to train on.\n",
    "\n",
    "`bond_lengths` determines what our training set will be throughout the CUSP protocol.  For the autoencoder to learn a good representation of the ground states, the training set should include a range of bond length examples.  Note that a larger training set will correspond to a longer simulation. We recommend `bond_lengths = [1.0, 1.5, 2.0, 2.5]` as a reasonable set.  The set of valid bond lengths to choose from are in the range of `0.2` to `4.1` in increments of `0.1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Settings\n",
    "num_trials = 500\n",
    "include_gate_noise = False\n",
    "noise_level = 0.002\n",
    "\n",
    "bond_lengths = [1.0, 1.5, 2.0, 2.5]\n",
    "\n",
    "user_settings = np.array([True, include_gate_noise, noise_level,\n",
    "                          num_trials, bond_lengths], dtype=object)\n",
    "np.save('data/user_settings', user_settings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now import some necessary files with the settings you have chosen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from config import CODE_DIRECTORY\n",
    "\n",
    "sys.path.append(CODE_DIRECTORY)\n",
    "\n",
    "# User settings for CUSP\n",
    "import settings\n",
    "from cusp_demo_utils import *\n",
    "\n",
    "# Module containing subroutines to set up cost functions at each stage\n",
    "import cusp_stage1\n",
    "\n",
    "# Module containing subroutines to set up optimization at each stage\n",
    "import stage1_opt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: If you change the settings above, you will need to restart your Jupyter notebook kernel and run the above cells again."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CUSP Stage 1: Preparation of Training Set\n",
    "\n",
    "<img src=\"images/stage1_alg.png\" width=\"500px\">\n",
    "\n",
    "Recall that in Stage 1 of this example, we will be trying to improve a circuit that prepares ground states of H2 for VQE.  The chosen circuit was generated from a circuit ansatz called \"unitary coupled cluster\", and has just a single variational parameter, `alpha`, which is the degree of rotation of a Z gate (circled below).  The circuit, as produced from Cirq's integrated circuit drawing function, looks like this:\n",
    "\n",
    "<img src=\"images/stage1.png\" width=\"500px\">\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code below will execute an optimization by trying to find the value of `alpha` which minimizes the ground state energy, and give a comparison to the actual ground state energy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#### STAGE 1 OF CUSP NOW RUNNING ####\n",
      "\n",
      "Optimizing for bond length 1.0 ... Optimal parameter setting is: [2.11250064]\n",
      "Exact ground state energy               : -1.1011503293035882\n",
      "VQE optimized energy                    : -1.1011502682427659\n",
      "Energy difference (absolute value)      : 6.106082239831778e-08\n",
      "\n",
      "Optimizing for bond length 1.5 ... Optimal parameter setting is: [2.23132465]\n",
      "Exact ground state energy               : -0.9981493524136991\n",
      "VQE optimized energy                    : -0.9981494323009096\n",
      "Energy difference (absolute value)      : 7.988721051344072e-08\n",
      "\n",
      "Optimizing for bond length 2.0 ... Optimal parameter setting is: [2.36051448]\n",
      "Exact ground state energy               : -0.9486411117424077\n",
      "VQE optimized energy                    : -0.9486411102426555\n",
      "Energy difference (absolute value)      : 1.4997522113091577e-09\n",
      "\n",
      "Optimizing for bond length 2.5 ... Optimal parameter setting is: [2.43916912]\n",
      "Exact ground state energy               : -0.9360549198442759\n",
      "VQE optimized energy                    : -0.9360548349623685\n",
      "Energy difference (absolute value)      : 8.488190739175394e-08\n",
      "\n",
      "CPU times: user 2.38 s, sys: 20 ms, total: 2.4 s\n",
      "Wall time: 2.39 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "print('#### STAGE 1 OF CUSP NOW RUNNING ####\\n')\n",
    "\n",
    "# Lists to store energies\n",
    "check_energies = []       # Energies of FCI/exact wavefunctions\n",
    "stage1_energies = []      # Energies of VQE wavefunctions\n",
    "\n",
    "# Run thru bond lengths (or the training set)\n",
    "for bond_length in bond_lengths:\n",
    "    \n",
    "    # Run VQE calculation for each training point/state\n",
    "    opt_stage1_params = stage1_opt.run_state_preparation_optimization(bond_length)\n",
    "    print('Optimizing for bond length {0} ... '\n",
    "          'Optimal parameter setting is: {1}'.format(bond_length, opt_stage1_params))\n",
    "\n",
    "    # Compute and store energies to check results\n",
    "    exact_energy = settings.fetch_ground_energy(bond_length)\n",
    "    check_energies.append(exact_energy)\n",
    "    opt_energy = cusp_stage1.compute_stage1_cost_function(opt_stage1_params,\n",
    "                                                          bond_length,\n",
    "                                                          n_repetitions=num_trials,\n",
    "                                                          exact=True,\n",
    "                                                          noisy=include_gate_noise)\n",
    "    stage1_energies.append(opt_energy)\n",
    "    \n",
    "    # Display stage 1 results\n",
    "    print('Exact ground state energy               : {}'.format(exact_energy))\n",
    "    print('VQE optimized energy                    : {}'.format(opt_energy))\n",
    "    print('Energy difference (absolute value)      : {}\\n'.format(\n",
    "            np.abs(opt_energy - exact_energy)))\n",
    "    \n",
    "    # Save these optimized VQE parameters into numpy arrays\n",
    "    np.save('data/stage1_param_{}'.format(bond_length), opt_stage1_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will likely observe that if `gate_noise = False`, these energies are very close to each other; if `gate_noise = True`, the energy difference typically gets larger as `noise_level` increases."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CUSP Stage 2: Training the Quantum Autoencoder (QAE)\n",
    "\n",
    "<img src=\"images/stage2_alg.png\" width=\"500px\">\n",
    "\n",
    "Now that we have trained an initial circuit, we wish to apply a quantum autoencoder to the output of the previous state.  Before we get ahead of ourselves, let's start by loading the parameters of the previous circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the optimized VQE parameters from stage 1 of CUSP\n",
    "stage1_param_list = []\n",
    "\n",
    "for bond_length in bond_lengths:\n",
    "    stage1_param = np.load('data/stage1_param_{}.npy'.format(bond_length))\n",
    "    stage1_param_list.append(stage1_param)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The autoencoder's job in Stage 2 is to compress the quantum information from the initial state preparation down to (in this example) just a single qubit.  We know this is in principle possible because we had only one parameter in the original state preparation.  The challenge at this stage is to find an autoencoder circuit that is general enough to accomplish this task, but is not as deep as the original circuit.  For simplicity in this demo, we have chosen the circuit for you---it is a set of three \"parameterized CNOT\" gates.  A regular CNOT, on Google's quantum hardware, is compiled in the following way:\n",
    "\n",
    "<img src=\"images/cnot.png\" width=\"400px\">\n",
    "\n",
    "We are going to replace these gates with the native parameterized gates that they are composed of on Google's hardware: we'll replace the Y gate with a variable W gate, the Z gate with a variable Z gate, and the CZ gate with a variable controlled-Z rotation.  The autoencoder circuit with the optimal parameters then looks like:\n",
    "\n",
    "<img src=\"images/stage2.png\" width=\"700px\">\n",
    "\n",
    "The entire Stage 2 circuit looks like:\n",
    "\n",
    "<img src=\"images/circuit_pic.png\" width=\"1100px\">\n",
    "\n",
    "Note that the bottom qubit will contain the latent space, while the top three are the reference qubits which we will use to train the autoencoder.  The training is performed by varying the parameters of the circuit and making repeated measurements of the reference qubits, attempting to maximize the frequency that $|000\\rangle$ is measured.  If the autoencoder is able to find a parameter setting where $|000\\rangle$ is achieved with high fidelity, it means that these qubits have been decoupled (or disentangled) from the remaining qubit.  As a result, all of the quantum information in the circuit must be possessed by the single qubit in the latent space, and so the autoencoder has succeeded in its task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, you have the choice of selecting which circuit variables to search over, and which to fix.  The optimal settings are:\n",
    "\n",
    "`ExpWGate(half_turns = 0.25, axis_half_turns = 0.5)`\n",
    "\n",
    "`ExpZGate(half_turns = 1)`\n",
    "\n",
    "`Exp11Gate(half_turns = 1)`\n",
    "\n",
    "We will call the variable parameters to search over `'w1'`, `'w2'`, `'z'`, and `'cz'` respectively.  In other words, the quantum autoencoder will use the settings:\n",
    "\n",
    "`ExpWGate(half_turns = w1, axis_half_turns = w2)`\n",
    "\n",
    "`ExpZGate(half_turns = z)`\n",
    "\n",
    "`Exp11Gate(half_turns = cz)`\n",
    "\n",
    "In the code block below, you may choose which parameters you wish to optimize in `search_parameters_stage2` (note that it is important to input them in the list as strings).  If you want to fix a specific value for any of the parameters, remove them from the `search_parameters_stage2` list and indicate their value in `fixed_<variable name>` (in other words, changing the fixed value while the parameter is still in the `search_parameters_stage2` list will have no effect).  Note that it is entirely possible to choose a set of parameters that prevents the autoencoder from ever completely decoupling the qubits.  Here is an example of the cost function (fidelity) landscape over `z` and `cz`, where `w1` and `w2` are fixed at `.25` and `.5` respectively:\n",
    "<img src=\"images/2d_plot.png\" width=\"600px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_parameters_stage2 = ['w1', 'w2', 'z', 'cz']\n",
    "fixed_w1 = .25\n",
    "fixed_w2 = .5\n",
    "fixed_z = 1\n",
    "fixed_cz = 1\n",
    "\n",
    "user_parameters_stage2 = np.array([search_parameters_stage2, fixed_w1, fixed_w2,\n",
    "                          fixed_z, fixed_cz], dtype=object)\n",
    "np.save('data/user_parameters_stage2', user_parameters_stage2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cusp_stage2\n",
    "import stage2_opt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, please refresh the notebook kernel when making any changes to the above parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code block will run the quantum autoencoder training.  Note a few options in the first few lines of the block, `threshold` and `n_qae_trials`.  Because the optimization routine can fail to give the global optimum if it gets caught in a local minimum, these settings allow the search to automatically restart up to `n_qae_trials` many times if the error of the autoencoder (computed as $1 - \\text{Fidelity}$) in each case is above `threshold`.  If `threshold` is set too large, the autoencoder will be allowed to terminate after finding a local optimimum, and so will not perform well in Stage 3.  As `noise_level` increases, the rate of successfully training the autoencoder tends to decrease slightly.  Note that as `noise_level` increases, the minimum error of the training will also increase, but this minimum will still usually correspond to an optimal setting for the parameters.  Keep this in mind if you are running the simulation with a high `noise_level`, since you may wish to increase `threshold` to ensure that you are not asking the autoencoder for too much accuracy during training.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stage 2 using the following bond lengths for training: [1.0, 1.5, 2.0, 2.5]\n",
      "\n",
      "#### STAGE 2 OF CUSP NOW RUNNING ####\n",
      "\n",
      "Trial 0: Quantum autoencoder learning had low fidelity. Trying again.\n",
      "Quantum autoencoder learning succeeded with error : -9.743289464125837e-08\n",
      "\n",
      "CPU times: user 30.3 s, sys: 20 ms, total: 30.3 s\n",
      "Wall time: 30.3 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "print('Stage 2 using the following bond lengths for training: {}\\n'.format(bond_lengths))\n",
    "\n",
    "# QAE settings\n",
    "threshold = 0.1\n",
    "n_qae_trials = 25\n",
    "\n",
    "print('#### STAGE 2 OF CUSP NOW RUNNING ####\\n')\n",
    "opt_qae_params = stage2_opt.run_qae_optimization(training_states=stage1_param_list,\n",
    "                                                 n_repetitions=num_trials,\n",
    "                                                 exact=True,\n",
    "                                                 noisy=include_gate_noise)\n",
    "\n",
    "# Repeat optimization of QAE circuit while error value is above threshold\n",
    "iter_count = 0\n",
    "while stage2_opt.compute_avg_fid_proxy(params=opt_qae_params,\n",
    "                                       training_states=stage1_param_list,\n",
    "                                       n_repetitions=num_trials,\n",
    "                                       exact=True,\n",
    "                                       noisy=include_gate_noise) > threshold:\n",
    "    if iter_count >= n_qae_trials:\n",
    "        print('Surpassed the QAE iteration limit. Exiting loop.')\n",
    "        break\n",
    "    \n",
    "    print('Trial {}: Quantum autoencoder learning had low fidelity. '\n",
    "          'Trying again.'.format(iter_count))\n",
    "    \n",
    "    opt_qae_params = stage2_opt.run_qae_optimization(training_states=stage1_param_list,\n",
    "                                                     n_repetitions=num_trials,\n",
    "                                                     exact=True,\n",
    "                                                     noisy=include_gate_noise)\n",
    "    iter_count += 1\n",
    "\n",
    "# Compute error of optimized QAE circuit\n",
    "err = stage2_opt.compute_avg_fid_proxy(opt_qae_params, training_states=stage1_param_list,\n",
    "                                       n_repetitions=num_trials, exact=True,\n",
    "                                       noisy=include_gate_noise)\n",
    "print('Quantum autoencoder learning succeeded with error : {}'.format(err))\n",
    "\n",
    "opt_qae_params = fix_list(opt_qae_params, stage2_opt.all_param,stage2_opt.var_param,\n",
    "                          stage2_opt.fixed_vals)\n",
    "# Save QAE results\n",
    "np.save('data/stage2_param', opt_qae_params)\n",
    "print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may note an interesting observation when training the autoencoder under noise.  It is tempting to identify the fidelity of the autoencoder training as the accuracy with which the autoencoder is learning to map to the latent space. Although the maximum fidelity of the autoencoder training will always decrease as noise increases, the autoencoder can often still find a set of parameters that is approximately correct.  This is because the optimal parameter setting still coincides with the maximum value of the fidelity.  You can see this trend in this single parameter (`z`) cost function landscape:  <img src=\"images/1d_plot.png\" width=\"600px\"> "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the training from Stage 2 is successful, the optimal parameters that are found are passed on to Stage 3.  These parameters are then used to construct the _decoder_ part of the quantum autoencoder, which is conveniently just the hermitian conjugate of the encoder.  In this example, for an optimal CNOT encoding, the decoder simply looks like:\n",
    "\n",
    "<img src=\"images/decoder.png\" width=\"150px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CUSP Stage 3: Generative Model Search\n",
    "\n",
    "<img src=\"images/stage3_alg.png\" width=\"500px\">\n",
    "\n",
    "We now address the final stage of CUSP wherein we construct the final circuit which will ideally produce a more accurate ground state energy than the original.  To do so, we need to search through the latent space to find a quantum state which approximates the compressed state at the end of Stage 2.  We will then use the autoencoder parameters found in the previous stage to construct a decoder which maps the state in the latent space back to the desired ground state of molecular hydrogen.  To optimize the latent space parameters, we will look to minimize the energy as we did in Stage 1.  In other words, we will perform VQE on our new circuit ansatz.\n",
    "\n",
    "Because our latent space is only one qubit in width, we have a conveniently small space to search. We will again choose the circuit for you: a two-parameter W gate followed by a parameterized Z gate.  The entire latent space search (circled in the figure) and decoder circuit looks like:\n",
    "\n",
    "<img src=\"images/stage3.png\" width=\"250px\">\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, you have the choice of which circuit variables to search over, and which to fix.  Unlike in Stage 2, because state preparation depends on which bond length you are trying to prepare, there are no a priori optimal settings.  However, there is some redundancy in the search space, so it is possible to fix certain gates and still find a good solution.  We will leave it up to you to discover exactly which redundancies those are.\n",
    "\n",
    "We will call the variable parameters to search over `'ht'`, `'aht'`, and `'zz'` respectively.  In other words, the latent space circuit will use the settings:\n",
    "\n",
    "`ExpWGate(half_turns = ht, axis_half_turns = aht)`\n",
    "\n",
    "`ExpZGate(half_turns = zz)`\n",
    "\n",
    "In the code block below, you may choose which parameters you wish to optimize in `search_parameters_stage3`   (note that it is important to input them in the list as strings).  As before, if you want to fix a specific value for any of the parameters, remove them from the `search_parameters_stage3` list and indicate their value in `fixed_<variable name>`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_parameters_stage3 = ['zz']\n",
    "fixed_aht = 0\n",
    "fixed_ht = 0\n",
    "fixed_zz = 0\n",
    "\n",
    "user_parameters_stage3 = np.array([search_parameters_stage3, fixed_aht,\n",
    "                                   fixed_ht, fixed_zz], dtype=object)\n",
    "np.save('data/user_parameters_stage3', user_parameters_stage3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cusp_stage3\n",
    "import stage3_opt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a reminder, please refresh the notebook kernel when making any changes to the above parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code block below now optimizes the latent space circuit, attempting to minimize the final state's energy.  The final result is compared to the true ground state energy, and to the energy that the original circuit found in Stage 1.  Because CUSP is designed to improve the gate depth, you will only see a significant improvement in the energy over the original circuit if `gate_noise=True`.  Otherwise, all of these energies are likely to be very close to oneanother."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameters used from Stage 2: [1.7500101294889823, 1.423778064710338, 1.0001189990659105, 0.9999606811201764]\n",
      "\n",
      "#### STAGE 3 OF CUSP NOW RUNNING ####\n",
      "\n",
      "Bond length                             : 1.0\n",
      "CUSP optimized energy                   : -1.0661088768899505\n",
      "Stage 1 energy                          : -1.1011502682427659\n",
      "Exact energy                            : -1.1011503293035882\n",
      "Energy difference (Stage 1 vs. exact)   : 6.106082239831778e-08\n",
      "Energy difference (CUSP    vs. exact)   : 0.035041452413637764\n",
      "\n",
      "Bond length                             : 1.5\n",
      "CUSP optimized energy                   : -0.9108737501492565\n",
      "Stage 1 energy                          : -0.9981494323009096\n",
      "Exact energy                            : -0.9981493524136991\n",
      "Energy difference (Stage 1 vs. exact)   : 7.988721051344072e-08\n",
      "Energy difference (CUSP    vs. exact)   : 0.08727560226444253\n",
      "\n",
      "Bond length                             : 2.0\n",
      "CUSP optimized energy                   : -0.7837928331430448\n",
      "Stage 1 energy                          : -0.9486411102426555\n",
      "Exact energy                            : -0.9486411117424077\n",
      "Energy difference (Stage 1 vs. exact)   : 1.4997522113091577e-09\n",
      "Energy difference (CUSP    vs. exact)   : 0.1648482785993629\n",
      "\n",
      "Bond length                             : 2.5\n",
      "CUSP optimized energy                   : -0.7029437250745418\n",
      "Stage 1 energy                          : -0.9360548349623685\n",
      "Exact energy                            : -0.9360549198442759\n",
      "Energy difference (Stage 1 vs. exact)   : 8.488190739175394e-08\n",
      "Energy difference (CUSP    vs. exact)   : 0.23311119476973408\n",
      "\n",
      "CPU times: user 430 ms, sys: 20 ms, total: 450 ms\n",
      "Wall time: 455 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# Optimal parameters from Stage 2\n",
    "print('Parameters used from Stage 2: {}\\n'.format(opt_qae_params))\n",
    "\n",
    "print('#### STAGE 3 OF CUSP NOW RUNNING ####\\n')\n",
    "\n",
    "stage3_energies = []\n",
    "cusp_params = {}\n",
    "\n",
    "for i, bond_length in enumerate(bond_lengths):\n",
    "    \n",
    "    # Initialize parameters\n",
    "    half_turn_min = 0\n",
    "    half_turn_max = 2\n",
    "    init_params = np.random.uniform(low=half_turn_min,\n",
    "                                    high=half_turn_max,\n",
    "                                    size=stage3_opt.num_param)\n",
    "\n",
    "    # Optimization using Nelder-Mead\n",
    "    stage3_fcn = lambda x: stage3_opt.stage3(x, bond_length=bond_length,\n",
    "                                             n_repetitions=num_trials)\n",
    "    res = minimize(stage3_fcn, init_params, args=(),\n",
    "                   method='Nelder-Mead', tol=None, \n",
    "                   options={'disp': False, 'maxiter': None, 'xatol': 0.001,\n",
    "                            'return_all': False, 'fatol': 0.001})\n",
    "    opt_cusp_param = res.x\n",
    "    opt_cusp_param = fix_list(opt_cusp_param, stage3_opt.all_param,stage3_opt.var_param,\n",
    "                              stage3_opt.fixed_vals)\n",
    "    cusp_params[bond_length] = opt_cusp_param\n",
    "    cusp_energy = cusp_stage3.run_sim_repetitions_stage3(*opt_cusp_param,\n",
    "                                                         bond_length=bond_length,\n",
    "                                                         n_repetitions=num_trials,\n",
    "                                                         exact=True,\n",
    "                                                         noisy=include_gate_noise)\n",
    "    stage3_energies.append(cusp_energy)\n",
    "    \n",
    "    print('Bond length                             : {}'.format(bond_length))\n",
    "    print('CUSP optimized energy                   : {}'.format(cusp_energy))\n",
    "    print('Stage 1 energy                          : {}'.format(stage1_energies[i]))\n",
    "    print('Exact energy                            : {}'.format(check_energies[i]))\n",
    "    print('Energy difference (Stage 1 vs. exact)   : {}'.format(\n",
    "            np.abs(stage1_energies[i] - check_energies[i])))\n",
    "    print('Energy difference (CUSP    vs. exact)   : {}\\n'.format(\n",
    "            np.abs(cusp_energy - check_energies[i])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final circuit parameters generated from the protocol can then be extracted:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Autoencoder parameters: [1.7500101294889823, 1.423778064710338, 1.0001189990659105, 0.9999606811201764]\n",
      "\n",
      "Latent circuit parameters: {1.0: [0, 0, 1.822578420639555], 2.0: [0, 0, 0.6538516336210352], 2.5: [0, 0, 1.840224600925812], 1.5: [0, 0, 0.7883810361965281]}\n"
     ]
    }
   ],
   "source": [
    "# Print the single-qubit circuit parameters\n",
    "print('Autoencoder parameters: {}\\n'.format(opt_qae_params))\n",
    "print('Latent circuit parameters: {}'.format(cusp_params))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You have reached the end of our tutorial!  Thank you for participating!  If you have questions on the operation of this algorithm, please feel free to forward them to jonny@zapatacomputing.com."
   ]
  },
  {
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}