vQEC.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5ab00b00",
   "metadata": {},
   "source": [
    "# Variational Quantum Error Correction with QVECTOR\n",
    "\n",
    "In this notebook, we will optimize a Quantum Error Correction (QEC) problem. \n",
    "Here, our set of states $\\{\\psi_i\\}_{i=1}^N$ are sampled from the Haar distribution (or an approximation thereof) and we optimize the cost \n",
    "$$L(\\vec{p}, \\vec{q}) = \\sum_{\\vert \\psi_i \\rangle \\in \\mathcal{S}} \\frac{1}{\\vert\\mathcal{S}\\vert} E_{i}(\\vert \\psi_i \\rangle, \\vec{p}, \\vec{q})$$\n",
    "where $\\vec{p}, \\vec{q}$ are parameters of a variational circuit, and $E_{i}(\\vert \\psi_i \\rangle, \\vec{p}, \\vec{q}) = \\langle \\psi_i \\vert \\mathcal{V}(\\vec{p})^\\dagger \\mathcal{W}(\\vec{q}) \\mathcal{N} \\mathcal{V}(\\vec{p}) \\vert \\psi_i \\rangle$, with $\\mathcal{V}$ the encoding operator, $\\mathcal{N}$ the noise model, and $\\mathcal{W}$ the decoding operator. The noise model here will be a bitflip channel (with probability $0.1$) acting on all qubits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "646b11e9-44fe-4eb9-a1a1-6243a854b3de",
   "metadata": {},
   "outputs": [],
   "source": [
    "import typing as ty\n",
    "import pennylane as qml\n",
    "from pennylane import numpy as np\n",
    "from refoqus import Refoqus\n",
    "\n",
    "# Initialise the number of qubits that interest us.\n",
    "k, n, r = 1, 3, 0\n",
    "nbqbits = n + r\n",
    "# Number of layers to use in the V and W parametrised circuits.\n",
    "layers = 2\n",
    "# Parametrised template that will be used to construct the V and W circuits.\n",
    "ansatz = qml.BasicEntanglerLayers\n",
    "#ansatz = qml.StronglyEntanglingLayers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "153a6d9a-6649-4fb1-b2da-6fb8dcfc3db7",
   "metadata": {},
   "source": [
    "The dataset is composed of random unitary matrices sampled from a 2-design (or approximate 2-design here). For 1 qubit, the dataset can be simplified to be only the 6 one-qubit stabiliser states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f50d82b8-96cf-43c5-8dba-878ad58489f0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def random_angles(*shape: int) -> np.ndarray:\n",
    "    return 2 * np.pi * np.random.rand(*shape)\n",
    "\n",
    "if k == 1:\n",
    "    # Only consider the 6 one qubit stabilizer states:\n",
    "    dataset_size = 6\n",
    "    dataset = [\n",
    "        [qml.Identity(wires=[0])], [qml.PauliX(wires=[0])],  # |0>, |1>\n",
    "        [qml.Hadamard(wires=[0])], [qml.PauliX(wires=[0]), qml.Hadamard(wires=[0])],  # |+>, |->\n",
    "        [qml.Hadamard(wires=[0]), qml.S(wires=[0])], [qml.PauliX(wires=[0]), qml.Hadamard(wires=[0]), qml.S(wires=[0])]  # |+i>, |-i>\n",
    "    ]\n",
    "    \n",
    "elif k > 1:\n",
    "    # Build the dataset here with qml.SimplifiedTwoDesign\n",
    "    # See https://docs.pennylane.ai/en/stable/code/api/pennylane.SimplifiedTwoDesign.html\n",
    "    dataset_size = 200\n",
    "    dataset = [[qml.SimplifiedTwoDesign(initial_layer_weights=random_angles(k), weights=random_angles(layers, k - 1, 2), wires=range(k))] for _ in range(dataset_size)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b9b163d-2d20-495d-acac-a1e623698a1c",
   "metadata": {},
   "source": [
    "The cost function $E_{i}(\\vert \\psi_i \\rangle, \\vec{p}, \\vec{q})$ can be computed by applying $ S^\\dagger \\mathcal{V}(\\vec{p})^\\dagger \\mathcal{W}(\\vec{q}) \\mathcal{N} \\mathcal{V}(\\vec{p}) S$ to the $\\vert 0 \\rangle$ and measure the probability of getting the full zero outcome $\\vert 0 \\rangle^{\\otimes n}$. \n",
    "\n",
    "To avoid using a global cost function that would inevitably suffer from the Barren Plateau phenomenon, we are using the local cost function defined by the Hamiltonian\n",
    "$$\n",
    "H = \\mathbb{I} - \\frac{1}{k} \\sum_{i=0}^{k-1} \\mathbb{I}_{\\bar{i}} \\otimes \\vert 0 \\rangle \\langle 0 \\vert_i\n",
    "$$\n",
    "where $\\mathbb{I}_{\\bar{i}} \\otimes \\vert 0 \\rangle \\langle 0 \\vert_i$ means that we are projecting only the $i$-th qubit to the $\\vert 0 \\rangle$ state, leaving the other qubits.\n",
    "\n",
    "Note that our goal is to recover pure states, for which the $\\mathbb{I}$ part will evaluate to $1$, so we do not have to include it in the measured Hamiltonian and we can add it during post-processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9a855d60",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the Hamiltonian of interest\n",
    "coefficients_cost = [- 1.0 / k for _ in range(k)]\n",
    "projector = np.zeros((2, 2))\n",
    "projector[0, 0] = 1\n",
    "vqec_hamiltonian_term = [qml.Hermitian(projector, wires=i) for i in range(k)]\n",
    "hamiltonian_of_interest = qml.Hamiltonian(coefficients_cost, vqec_hamiltonian_term)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ebfe464",
   "metadata": {},
   "source": [
    "Next, we define functions to evaluate the true cost during optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2c0864ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def circuit_construction(\n",
    "    weights: np.ndarray,\n",
    "    data_circuit: ty.List[qml.operation.Operation],\n",
    "    parameterised_circuit = None,\n",
    "    noise_circuit = None,\n",
    "):\n",
    "    if parameterised_circuit is None:\n",
    "        parameterised_circuit = ansatz\n",
    "    if noise_circuit is None:\n",
    "        def noise_circuit():\n",
    "            for i in range(n):\n",
    "                qml.BitFlip(0.1, wires=i)\n",
    "\n",
    "    encoding_shape = parameterised_circuit.shape(layers, n)\n",
    "    decoding_shape = parameterised_circuit.shape(layers, n + r)\n",
    "    encoding_size = np.prod(encoding_shape)\n",
    "    decoding_size = np.prod(decoding_shape)\n",
    "    encoding_weights = weights[:encoding_size].reshape(encoding_shape)\n",
    "    decoding_weights = weights[encoding_size:encoding_size + decoding_size].reshape(decoding_shape)\n",
    "    \n",
    "    # Apply S         to range(k)     to prepare an approximate 2-design state\n",
    "    for op in data_circuit:\n",
    "        qml.apply(op)\n",
    "    \n",
    "    # Apply V(p)      to range(n)     to encode the state\n",
    "    parameterised_circuit(encoding_weights, wires=range(n))\n",
    "    # Apply potential noise\n",
    "    noise_circuit()\n",
    "    # Apply W(q)      to range(n + r) to correct potential errors\n",
    "    parameterised_circuit(decoding_weights, wires=range(n + r))\n",
    "    # Apply V^{-1}(p) to range(n)     to decode the state\n",
    "    qml.adjoint(parameterised_circuit(encoding_weights, wires=range(n)))\n",
    "    \n",
    "    # Apply S^{-1}    to range(k)     to un-prepare.\n",
    "    for op in reversed(data_circuit):\n",
    "        qml.apply(qml.adjoint(op))\n",
    "    \n",
    "def cost_function(\n",
    "    weights: np.ndarray,\n",
    "    hamiltonian_terms: qml.operation.Operator,\n",
    "    data_circuit: ty.List[qml.operation.Operation],\n",
    "    parameterised_circuit = None,\n",
    "):\n",
    "    # Here noise_circuit is left intentionally to default, allowing the default BitFlip noise\n",
    "    # to be applied.\n",
    "    circuit_construction(weights, data_circuit, parameterised_circuit)\n",
    "    return qml.sample(hamiltonian_terms)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b578f36d-3f2f-4b4b-b2fe-3a8a1ff46edb",
   "metadata": {},
   "source": [
    "In order to have access to the non-noisy cost during the optimisation process we define functions to estimate it using perfect, noise-free simulators."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f3cb65db-cb7d-4334-ac99-ec73d40343f8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "analytic_dev = qml.device(\"default.mixed\", wires=nbqbits, shots=None)\n",
    "\n",
    "@qml.qnode(analytic_dev)\n",
    "def cost_analytic_one_circuit(weights, index_datapoint, parameterised_circuit = None):\n",
    "    # Here, noise_circuit is explicitely set to a lambda that does nothing.\n",
    "    # This is to prevent any noise from being added to the computation because we want a noise-free estimation in this routine.\n",
    "    circuit_construction(weights, dataset[index_datapoint], parameterised_circuit, noise_circuit=lambda : None)\n",
    "    return qml.expval(hamiltonian_of_interest)\n",
    "\n",
    "def cost_analytic_alldataset(weights, parameterised_circuit = None):\n",
    "    cost = 0.0\n",
    "    for m in range(dataset_size):\n",
    "        cost += cost_analytic_one_circuit(weights, m, parameterised_circuit)\n",
    "    cost = 1.0 + cost / dataset_size\n",
    "    return cost"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18c1e86f",
   "metadata": {},
   "source": [
    "Now, the ansatz is defined as with StronglyEntanglingLayers. We also sample initial values and the corresponding cost."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba1d361f",
   "metadata": {},
   "source": [
    "Our adaptative optimizer will be Refoqus where we provide the necessary arguments as follows and we perform niter iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8f0f0d37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0: cost = 0.7885592318715495, shots_used = 48\n",
      "Step 1: cost = 0.6046781891568378, shots_used = 108\n",
      "Step 2: cost = 0.33305975485748573, shots_used = 162\n",
      "Step 3: cost = 0.12745506986684318, shots_used = 288\n",
      "Step 4: cost = 0.016385763262979025, shots_used = 416\n",
      "Step 5: cost = 0.006900780477326118, shots_used = 704\n",
      "Step 6: cost = 0.04513347218667174, shots_used = 1068\n",
      "Step 7: cost = 0.00441184534694683, shots_used = 1628\n",
      "Step 8: cost = 0.02682783170616554, shots_used = 2254\n",
      "Step 9: cost = 0.031992724753818536, shots_used = 2710\n",
      "Step 10: cost = 0.08351446141118402, shots_used = 3020\n",
      "Step 11: cost = 0.08187249083468018, shots_used = 3312\n",
      "Step 12: cost = 0.05533068625900561, shots_used = 3604\n",
      "Step 13: cost = 0.111717896296195, shots_used = 3764\n",
      "Step 14: cost = 0.021075139974568446, shots_used = 3988\n",
      "Step 15: cost = 0.020913140362164184, shots_used = 4262\n",
      "Step 16: cost = 0.03615837859137494, shots_used = 4572\n",
      "Step 17: cost = 0.027974805502202993, shots_used = 4844\n",
      "Step 18: cost = 0.016791096292577135, shots_used = 5108\n",
      "Step 19: cost = 0.05276632404903914, shots_used = 5382\n"
     ]
    }
   ],
   "source": [
    "parameter_size: int = np.prod(ansatz.shape(layers, n)) + np.prod(ansatz.shape(layers, n + r))\n",
    "params = random_angles(parameter_size)\n",
    "\n",
    "opt = Refoqus(nbqbits, dataset, vqec_hamiltonian_term, coefficients_cost, param_shape=(parameter_size,), function_cost_term_tosample=cost_function, min_shots=2, device_name=\"default.mixed\")\n",
    "niter = 20\n",
    "\n",
    "cost_refoqus = [cost_analytic_alldataset(params, ansatz)]\n",
    "shots_refoqus = [0]\n",
    "\n",
    "for i in range(niter):\n",
    "    params = opt.step(params)\n",
    "    cost_refoqus.append(cost_analytic_alldataset(params, ansatz))\n",
    "    shots_refoqus.append(opt.shots_used)\n",
    "    print(f\"Step {i}: cost = {cost_refoqus[-1]}, shots_used = {shots_refoqus[-1]}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}