{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VQE example with `pytket-qujax`\n",
    "\n",
    "**Download this notebook - {nb-download}`pytket-qujax_heisenberg_vqe.ipynb`**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See the docs for [qujax](https://cqcl.github.io/qujax/) and [pytket-qujax](https://cqcl.github.io/pytket-qujax/api/index.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import numpy as jnp, random, value_and_grad, jit\n",
    "from pytket import Circuit\n",
    "from pytket.circuit.display import render_circuit_jupyter\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's start with a TKET circuit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qujax\n",
    "from pytket.extensions.qujax.qujax_convert import tk_to_qujax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We place barriers to stop tket automatically rearranging gates and we also store the number of circuit parameters as we'll need this later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_circuit(n_qubits, depth):\n",
    "    n_params = 2 * n_qubits * (depth + 1)\n",
    "    param = jnp.zeros((n_params,))\n",
    "    circuit = Circuit(n_qubits)\n",
    "    k = 0\n",
    "    for i in range(n_qubits):\n",
    "        circuit.H(i)\n",
    "    for i in range(n_qubits):\n",
    "        circuit.Rx(param[k], i)\n",
    "        k += 1\n",
    "    for i in range(n_qubits):\n",
    "        circuit.Ry(param[k], i)\n",
    "        k += 1\n",
    "    for _ in range(depth):\n",
    "        for i in range(0, n_qubits - 1):\n",
    "            circuit.CZ(i, i + 1)\n",
    "        circuit.add_barrier(range(0, n_qubits))\n",
    "        for i in range(n_qubits):\n",
    "            circuit.Rx(param[k], i)\n",
    "            k += 1\n",
    "        for i in range(n_qubits):\n",
    "            circuit.Ry(param[k], i)\n",
    "            k += 1\n",
    "    return circuit, n_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_qubits = 4\n",
    "depth = 2\n",
    "circuit, n_params = get_circuit(n_qubits, depth)\n",
    "render_circuit_jupyter(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now let's invoke qujax\n",
    "The `pytket.extensions.qujax.tk_to_qujax` function will generate a parameters -> statetensor function for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_to_st = tk_to_qujax(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try it out on some random parameters values. Be aware that's JAX's random number generator requires a `jax.random.PRNGkey` every time it's called - more info on that [here](https://jax.readthedocs.io/en/latest/jax.random.html).\n",
    "Be aware that we still have convention where parameters are specified as multiples of $\\pi$ - that is in [0,2]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = random.uniform(random.PRNGKey(0), shape=(n_params,), minval=0.0, maxval=2.0)\n",
    "statetensor = param_to_st(params)\n",
    "print(statetensor)\n",
    "print(statetensor.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that this function also has an optional second argument where an initiating `statetensor_in` can be provided. If it is not provided it will default to the all 0s state (as we use here)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can obtain statevector by simply calling `.flatten()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "statevector = statetensor.flatten()\n",
    "statevector.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And sampling probabilities by squaring the absolute value of the statevector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_probs = jnp.square(jnp.abs(statevector))\n",
    "plt.bar(jnp.arange(statevector.size), sample_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cost function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have our `param_to_st` function we are free to define a cost function that acts on bitstrings (e.g. maxcut) or integers by directly wrapping a function around `param_to_st`. However, cost functions defined via quantum Hamiltonians are a bit more involved.\n",
    "Fortunately, we can encode an Hamiltonian in JAX via the `qujax.get_statetensor_to_expectation_func` function which generates a statetensor -> expected value function for us.\n",
    "It takes three arguments as input\n",
    "- `gate_seq_seq`: A list of string (or array) lists encoding the gates in each term of the Hamiltonian. I.e. `[['X','X'], ['Y','Y'], ['Z','Z']]` corresponds to $H = aX_iX_j +  bY_kY_l + cZ_mZ_n$ with qubit indices $i,j,k,l,m,n$ specified in the second argument and coefficients $a,b,c$ specified in the third argument\n",
    "- `qubit_inds_seq`: A list of integer lists encoding which qubit indices to apply the aforementioned gates. I.e. `[[0, 1],[0,1],[0,1]]`. Must have the same structure as `gate_seq_seq` above.\n",
    "- `coefficients`: A list of floats encoding any coefficients in the Hamiltonian. I.e. `[2.3, 0.8, 1.2]` corresponds to $a=2.3,b=0.8,c=1.2$ above. Must have the same length as the two above arguments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More specifically let's consider the problem of finding the ground state of the quantum Heisenberg Hamiltonian"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{equation}\n",
    "H = \\sum_{i=1}^{n_\\text{qubits}-1} X_i X_{i+1} + Y_i Y_{i+1} + Z_i Z_{i+1}.\n",
    "\\end{equation}\n",
    "$$\n",
    "\n",
    "As described, we define the Hamiltonian via its gate strings, qubit indices and coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hamiltonian_gates = [[\"X\", \"X\"], [\"Y\", \"Y\"], [\"Z\", \"Z\"]] * (n_qubits - 1)\n",
    "hamiltonian_qubit_inds = [\n",
    "    [int(i), int(i) + 1] for i in jnp.repeat(jnp.arange(n_qubits), 3)\n",
    "]\n",
    "coefficients = [1.0] * len(hamiltonian_qubit_inds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Gates:\\t\", hamiltonian_gates)\n",
    "print(\"Qubits:\\t\", hamiltonian_qubit_inds)\n",
    "print(\"Coefficients:\\t\", coefficients)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's get the Hamiltonian as a pure JAX function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st_to_expectation = qujax.get_statetensor_to_expectation_func(\n",
    "    hamiltonian_gates, hamiltonian_qubit_inds, coefficients\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check it works on the statetensor we've already generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expected_val = st_to_expectation(statetensor)\n",
    "expected_val"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's wrap the `param_to_st` and `st_to_expectation` together to give us an all in one `param_to_expectation` cost function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_to_expectation = lambda param: st_to_expectation(param_to_st(param))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_to_expectation(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sanity check that a different, randomly generated set of parameters gives us a new expected value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_params = random.uniform(\n",
    "    random.PRNGKey(1), shape=(n_params,), minval=0.0, maxval=2.0\n",
    ")\n",
    "param_to_expectation(new_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exact gradients within a VQE algorithm\n",
    "The `param_to_expectation` function we created is a pure JAX function and outputs a scalar. This means we can pass it to `jax.grad` (or even better `jax.value_and_grad`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost_and_grad = value_and_grad(param_to_expectation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `cost_and_grad` function returns a tuple with the exact cost value and exact gradient evaluated at the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost_and_grad(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now we have all the tools we need to design our VQE!\n",
    "We'll just use vanilla gradient descent with a constant stepsize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vqe(init_param, n_steps, stepsize):\n",
    "    params = jnp.zeros((n_steps, n_params))\n",
    "    params = params.at[0].set(init_param)\n",
    "    cost_vals = jnp.zeros(n_steps)\n",
    "    cost_vals = cost_vals.at[0].set(param_to_expectation(init_param))\n",
    "    for step in range(1, n_steps):\n",
    "        cost_val, cost_grad = cost_and_grad(params[step - 1])\n",
    "        cost_vals = cost_vals.at[step].set(cost_val)\n",
    "        new_param = params[step - 1] - stepsize * cost_grad\n",
    "        params = params.at[step].set(new_param)\n",
    "        print(\"Iteration:\", step, \"\\tCost:\", cost_val, end=\"\\r\")\n",
    "    print(\"\\n\")\n",
    "    return params, cost_vals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ok enough talking, let's run (and whilst we're at it we'll time it too)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "%time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vqe_params, vqe_cost_vals = vqe(params, n_steps=250, stepsize=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the results..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(vqe_cost_vals)\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.ylabel(\"Cost\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pretty good!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `jax.jit` speedup\n",
    "One last thing... We can significantly speed up the VQE above via the `jax.jit`. In our current implementation, the expensive `cost_and_grad` function is compiled to [XLA](https://www.tensorflow.org/xla)  and then executed at each call. By invoking `jax.jit` we ensure that the function is compiled only once (on the first call) and then simply executed at each future call - this is much faster!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost_and_grad = jit(cost_and_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll demonstrate this using the second set of initial parameters we randomly generated (to be sure of no caching)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "%time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_vqe_params, new_vqe_cost_vals = vqe(new_params, n_steps=250, stepsize=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's some speedup!\n",
    "But let's also plot the training to be sure it converged correctly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(new_vqe_cost_vals)\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.ylabel(\"Cost\")"
   ]
  }
 ],
 "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}