{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to cirq 2: variational methods\n",
    "\n",
    "In this tutorial we will apply the quantum alternating operator ansatz (QAOA) to find the ground state of a random-bond Ising model.\n",
    "\n",
    "1. VQE construction\n",
    "    1. VQE ansatz construction\n",
    "        1. Initial state preparation\n",
    "        1. Circuit generators\n",
    "    1. Cost function construction\n",
    "1. VQE simulation and optimization\n",
    "    1. Simulating a single VQE instance\n",
    "    1. Symbolic parameters\n",
    "    1. Grid search with sweeps\n",
    "    1. Choosing a classical optimization strategy\n",
    "    \n",
    "(This tutorial is adapted from the Cirq tutorial at https://cirq.readthedocs.io/en/stable/tutorial.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VQE construction\n",
    "## VQE ansatz construction\n",
    "Recall that a VQE consists of\n",
    "1. A parameterized quantum circuit\n",
    "2. A starting state\n",
    "3. A cost function and a method for measuring it\n",
    "4. A classical optimization outer loop\n",
    "\n",
    "The QAOA is a popular variational circuit that consists of repeating two layers of operators - one defined by the starting state, one by the target state.\n",
    "In particular, in each layer we evolve by $e^{iH_at}$, where $H_0$ is chosen so that our starting state is an eigenstate, and $H_1$ is a Hamiltonian that defines the problem we wish to solve.\n",
    "This is a popular strategy for many applications; optimization and quantum simulation to name a few.\n",
    "\n",
    "In this tutorial we will focus on constructing the QAOA ansatz for a native target problem - the random bond Ising model. This is defined by a Hamiltonian\n",
    "$$\n",
    "H_1=\\sum_{\\langle i,j\\rangle}t_{i,j}Z_iZ_j,\n",
    "$$\n",
    "where $\\langle i,j\\rangle$ denotes nearest neighbours on some lattice, and $t_{i,j}=\\pm 1$.\n",
    "\n",
    "We note that this is a terrible problem to try to solve on a quantum device, as the ground state is known to be QMA-hard to prepare. But it still makes a good example problem :).\n",
    "\n",
    "For a starting state, let us take the uniform superposition across all ground states - $\\otimes_i|+\\rangle$. The corresponding Hamiltonian is then\n",
    "$$\n",
    "H_0=\\sum_{i}X_i.\n",
    "$$\n",
    "\n",
    "Let us implement the algorithm on a square lattice of size $n_x\\times n_y$. If we set $n_x=n_y=3$ this should be quite quick to simulate.\n",
    "\n",
    "### Initial state preparation\n",
    "We first need to construct a circuit to prepare our initial state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cirq\n",
    "nx = 3\n",
    "ny = 3\n",
    "def init_state_circuit(nx, ny):\n",
    "     yield [cirq.H(cirq.GridQubit(i,j)) for i in range(nx) for j in range(ny)]\n",
    "\n",
    "circuit = cirq.Circuit(init_state_circuit(nx, ny))\n",
    "print(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Circuit construction\n",
    "\n",
    "Now we need to make the two blocks of our circuit. Let us first set our random Hamiltonian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "rng = np.random.RandomState(seed=42)\n",
    "\n",
    "# H1_right[i,j] defines the link between qubit i,j and qubit i+1,j\n",
    "H1_right = np.array([[rng.choice([+1, -1]) for j in range(ny)] for i in range(nx-1)])\n",
    "# H1_down[i,j] defines the link between qubit i,j and qubit i,j+1\n",
    "H1_down = np.array([[rng.choice([+1, -1]) for j in range(ny-1)] for i in range(nx)])\n",
    "\n",
    "# To check that we have the right dimensions and deal with the standard\n",
    "# x, y\n",
    "print(H1_right.T)\n",
    "print(H1_down.T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defining our circuit, we need to be careful; if we arrange the gates for H1 at random it will not be constant depth. The trick here is to use a clockwork format - pick one qubit, and do gates to the right, then down, then left, then up around this qubit (and all gates facing in the same direction that fit with this one).\n",
    "\n",
    "**Problem 1:** Write code to implement evolution by $e^{iH_1\\theta}$ (I've started it off for you)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def H0_block(alpha, nx, ny):\n",
    "    xgate = cirq.XPowGate(exponent=alpha)\n",
    "    yield [xgate(cirq.GridQubit(i,j)) for i in range(nx) for j in range(ny)]\n",
    "\n",
    "def H1_block(H1_right, H1_down, beta, nx, ny):\n",
    "    zzgate = cirq.ZZPowGate(exponent=beta)\n",
    "    # Right gates\n",
    "    yield [zzgate(cirq.GridQubit(i, j), cirq.GridQubit(i+1, j))**H1_right[i,j]\n",
    "           for i in range(0,nx-1,2) for j in range(ny)]\n",
    "    \n",
    "    ### Add your code here!\n",
    "    \n",
    "circuit = cirq.Circuit(H1_block(H1_right, H1_down, 0.5, nx, ny))\n",
    "# Check that this is equal to 4!\n",
    "print('Number of blocks in circuit = {}'.format(len(circuit)))\n",
    "\n",
    "#Print the circuit to check\n",
    "print(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now need to handle measurement of the system - or rather how the measurement will be stored classically. In order for ease of access, we can read out all the qubits in a single measurement step, but we need to be worried about qubit order here --- or at least, that we know what order we have chosen for future access."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def msmt_circuit(nx, ny):\n",
    "    msmt_order = [cirq.GridQubit(i,j) for i in range(nx) for j in range(ny)]\n",
    "    yield cirq.measure(*msmt_order, key='Zmeasurements')\n",
    "    \n",
    "# Just to show how to go between the two values    \n",
    "test = [cirq.GridQubit(i,j) for i in range(nx) for j in range(ny)]\n",
    "print(test)\n",
    "array_reset = np.array([[test[i*ny+j] for j in range(ny)] for i in range(nx)])\n",
    "print(array_reset)\n",
    "print(array_reset[1,2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The full circuit now consists of initialization, $p$ repetitions of the two parts of the ansatz, and finally measurement. Note that we should do the evolution by $H_1$ before the evolution by $H_0$, and that each repetition should be controlled by a different parameter.\n",
    "\n",
    "**Problem 2:** Write a function to construct the full QAOA circuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def QAOA_circuit(H1_right, H1_down, p, alpha_vec, beta_vec, nx, ny):\n",
    "    '''\n",
    "    Args:\n",
    "        H1_right: Matrix defining the horizontal couplings in the random bond Ising model\n",
    "            H1_right[i,j] couples qubits (i,j) and (i+1,j)\n",
    "        H1_down: Matrix defining the vertical couplings in the random bond Ising model\n",
    "            H1_down[i,j] couples qubits (i,j) and (i,j+1)\n",
    "        p: number of layers in the ansatz\n",
    "        alpha_vec(list of p floats): list of alpha values for each QAOA layer\n",
    "        beta_vec(list of p floats): list of beta values for each QAOA layer\n",
    "        nx, ny(integers): dimensions of model\n",
    "    '''\n",
    "    ### Insert your code here!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cost function construction\n",
    "We now need to write code to calculate our cost function on the system. In this case we will use $\\langle H_1\\rangle$. As we are measuring our wavefunction in the $Z$ basis, we may calculate the cost function immediately for each instance of the system and then just average over this.\n",
    "\n",
    "**Problem 3:** write a function to calculate the energy of the random bond ising model given a set of measurements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def energy_func(H1_right, H1_down, nx, ny):\n",
    "    def energy(measurements):\n",
    "        # Rearrange to match the qubit array + change from 0 and 1 to 1 and -1\n",
    "        msmt_array = 1 - 2 * np.array([[measurements[i*ny+j] for j in range(ny)]\n",
    "                                        for i in range(nx)]).astype(int)\n",
    "        ### Insert your code here!\n",
    "        \n",
    "    return energy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our cost function is then the average over the results of multiple measurements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cost_func(result, H1_right, H1_down, nx, ny):\n",
    "    energy_hist = result.histogram(key='Zmeasurements',\n",
    "                                   fold_func=energy_func(H1_right, H1_down, nx, ny))\n",
    "    return np.sum([k * v for k,v in energy_hist.items()]) / result.repetitions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VQE simulation and optimization\n",
    "## Simulating a single VQE instance\n",
    "Before we go any further, let's use what we learnt last week to simulate an instance of our VQE using some fixed parameters and measuring the output.\n",
    "\n",
    "**Problem 4:** Simulate the above VQE for $p=1$, and $\\beta,\\alpha=0.25$. If you average over 100000 instances you should get an energy somewhere between $4.6-4.7$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Insert your code here!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Symbolic parameters\n",
    "\n",
    "We want the parameters $\\alpha$ and $\\beta$ to be optimized during the VQE. This is most easily done by passing around sympy.Symbol values, which work well with cirq."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy\n",
    "beta_vec = [sympy.Symbol('beta{}'.format(j)) for j in range(p)]\n",
    "alpha_vec = [sympy.Symbol('alpha{}'.format(j)) for j in range(p)]\n",
    "circuit = cirq.Circuit(QAOA_circuit(H1_right, H1_down, p, alpha_vec, beta_vec, nx, ny))\n",
    "print(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, some of the circuits now have gates that are parameterized. These can be resolved in cirq, most usefully by a 'sweep' over the parameters. This may be done via the cirq.Linspace object, which is a set of instructions to resolve a certain parameter that may be chained together, and the simulator.run_sweep function, which can interpret this. Let's try to run this with our simple $p=1$ algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sweep_size = 10\n",
    "sweep = (cirq.Linspace(key='alpha0', start=0, stop=1, length=sweep_size)\n",
    "        * cirq.Linspace(key='beta0', start=0, stop=1, length=sweep_size))\n",
    "results = simulator.run_sweep(circuit, params=sweep, repetitions=10000)\n",
    "for result in results:\n",
    "    print(result.params.param_dict, cost_func(result, H1_right, H1_down, nx, ny))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you have the ingredients to run the VQE, all you need to do is minimize and find the result.\n",
    "\n",
    "**Problem 5:** Find the minimum parameters and corresponding energy in the $p=1$ VQE. Use whatever method you want.\n",
    "**Problem 6:** Find the minimum parameters and corresponding energy in the QAOA with $p=2$. Compare the time taken to find the minimum and the result to the $p=1$ result."
   ]
  },
  {
   "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.5rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}