{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5b2ac00e-c7fc-45cc-888e-8c26dabf80d2",
   "metadata": {},
   "source": [
    "# Readout Error Mitigation\n",
    "\n",
    "Readout errors are caused by imperfect qubit measurement and are a common source of error in quantum computing. Properly modelling these errors in simulation can give the user tools to understand better what these errors are and how to mitigate them when running on actual quantum devices.\n",
    "\n",
    "Readout errors can be mitigated with a confusion matrix $A$. It is a square matrix of size $2^n \\times 2^n$ and tells about the probability of observing the state $|y\\rangle$ given the true state $|x\\rangle$. The confusion matrix characterizes the readout error of the device and is circuit-independent. Once $A$ is estimated, we can compute its pseudoinverse $A^+$ which can be applied to the noisy probability distribution $p_{\\text{noisy}}$ to obtain an adjusted quasi-probability distribution \n",
    "\n",
    "$$\n",
    "p_{\\text{mitigated}} = A^+ p_{\\text{noisy}}\n",
    "$$\n",
    "\n",
    "In this tutorial, we show how to build a confusion matrix with the following approaches.\n",
    "\n",
    "- Using a  single qubit model\n",
    "- Using $k$ local confusion matrices\n",
    "- A full confusion matrix for each $2^n$ combination\n",
    "\n",
    "The last method works well for correcting correlated errors (which affect multiple qubits together). However, it becomes impractical for large numbers of qubits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "04b58854-7dab-4f1b-8c59-d8b9506bce44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the relevant packages.\n",
    "\n",
    "!pip install matplotlib==3.8.4 pandas==2.2.2 scipy==1.13.1 seaborn==0.13.2 -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "710a2a46-4e55-4508-a0b7-21969aac6a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cudaq\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "import itertools\n",
    "from functools import reduce\n",
    "from scipy.optimize import minimize\n",
    "from typing import Union\n",
    "\n",
    "sns.set_style(\"dark\")\n",
    "\n",
    "if cudaq.num_available_gpus() > 0 and cudaq.has_target(\"nvidia\"):\n",
    "    cudaq.set_target(\"nvidia\")\n",
    "else:\n",
    "    print(\"CUDA or GPU support is unavailable. Running with CPU simulator. Performance may be significantly reduced.\")\n",
    "    cudaq.set_target(\"density-matrix-cpu\")\n",
    "\n",
    "seed = 42\n",
    "\n",
    "cudaq.set_random_seed(seed)\n",
    "np.random.seed(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eab761a-0a82-4e20-8cf8-54602cd05d6a",
   "metadata": {},
   "source": [
    "To model the readout error, we apply a bitflip channel on every qubit at the end of the circuit using an Identity gate. The probability of bitflip `probs` can be the same or different for all the qubits.\n",
    "\n",
    "<div class=\"alert alert-block alert-success\"><b>Note: </b>\n",
    "In principle, readout error is applied to the measurement gate but we use the Identity gate because adding a quantum error on the measurement gate is not yet supported. Also, the Identity gate is not present in the <kbd>add_channel</kbd> method, so we model the Identity gate using a rotation gate with an angle of 0.0.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d84bb420-4fab-4c26-919a-7efa8d3a219b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_noise(n_qubits: int, probs: Union[float, list[float]]):\n",
    "    \"\"\"\n",
    "    Creates a noise model with specified probabilities of bit flip errors on each qubit.\n",
    "\n",
    "    Args:\n",
    "        n_qubits (int): The number of qubits in the quantum circuit.\n",
    "        probs (Union[float, list[float]]): Probabilities of a bit flip error occurring on each qubit.\n",
    "        If a single probability value is provided, it will be applied to all qubits. If a list of\n",
    "        probabilities is provided, each qubit will have its own probability.\n",
    "\n",
    "    Returns:\n",
    "        cudaq noise model\n",
    "    \"\"\"\n",
    "    noise = cudaq.NoiseModel()\n",
    "\n",
    "    if not isinstance(probs, list):\n",
    "        probs = [probs] * n_qubits\n",
    "\n",
    "    for i, p in enumerate(probs):\n",
    "        bit_flip = cudaq.BitFlipChannel(p)\n",
    "        noise.add_channel(\"rx\", [i], bit_flip)\n",
    "\n",
    "    return noise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb4c7a14-1c0c-4fab-a6aa-7174c23baa7f",
   "metadata": {},
   "source": [
    "We define a cudaq kernel which will help create a set of quantum circuits to take measurements of the basis states for $n$ qubits. The kernel takes the number of qubits and the basis state as arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "53e3f450-000f-49b8-ad5c-57aa74461694",
   "metadata": {},
   "outputs": [],
   "source": [
    "@cudaq.kernel\n",
    "def kernel(n_qubits: int, state_label: list[int]):\n",
    "    qvector = cudaq.qvector(n_qubits)\n",
    "    for i, val in enumerate(state_label):\n",
    "        if val == 1:\n",
    "            x(qvector[i])\n",
    "        rx(0.0, qvector[i]\n",
    "          )  # Identity gate at the end on each qubit to model readout error\n",
    "    mz(qvector)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95411463-cd1e-499f-baef-d624acc0122c",
   "metadata": {},
   "source": [
    "For the tutorial, we will use the GHZ state on three qubits to test readout error mitigation. The GHZ state on 3 qubits is described as:\n",
    "\n",
    "$$\n",
    "\\frac{|000\\rangle + |111\\rangle}{\\sqrt{2}}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c1476028-923d-4f2b-b5fc-39ebcc8cbab9",
   "metadata": {},
   "outputs": [],
   "source": [
    "@cudaq.kernel\n",
    "def ghz_kernel(n_qubits: int):\n",
    "    qvector = cudaq.qvector(n_qubits)\n",
    "\n",
    "    h(qvector[0])\n",
    "    for i in range(n_qubits - 1):\n",
    "        cx(qvector[i], qvector[i + 1])\n",
    "\n",
    "    # Apply identity gates for readout error mitigation\n",
    "    for i in range(n_qubits):\n",
    "        rx(0.0, qvector[i])\n",
    "\n",
    "    mz(qvector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "900d2110-d613-426a-8aa9-6c12b1238536",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to draw the confusion matrix\n",
    "def plot_cmat(mat):\n",
    "    fig, ax = plt.subplots()\n",
    "    n = len(mat)\n",
    "    im2 = ax.matshow(mat, cmap=plt.cm.Reds, vmin=0, vmax=1.0)\n",
    "    ax.set_yticks(np.arange(n))\n",
    "    ax.set_xticks(np.arange(n))\n",
    "    ax.set_yticklabels(n * [\"\"])\n",
    "    ax.set_xticklabels(n * [\"\"])\n",
    "    ax.set_title(r\"Confusion Matrix\", fontsize=16)\n",
    "    ax.set_xlabel(\"Prepared State\")\n",
    "    ax.xaxis.set_label_position(\"top\")\n",
    "    ax.set_ylabel(\"Measured State\")\n",
    "    fig.colorbar(im2, ax=ax)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "032f6898-c953-4d30-a2bf-e8b46ee77c62",
   "metadata": {},
   "source": [
    "It is possible that the adjusted quasi-probability distribution $p_{\\text{mitigated}}$ obtained by application of $A^+$ to $p_{\\text{noisy}}$ may be non-positive. We need to find the closest positive probability distribution in such a case.\n",
    "\n",
    " $$ p'' = \\min_{p_{\\rm positive}} \\|p_{\\rm noisy} - p_{\\rm positive}\\|_1$$\n",
    "\n",
    "Then, we can draw samples from $p''$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9b355dcb-bf80-42a1-ae11-dcc9f11ba1e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_closest_distribution(empirical_dist):\n",
    "    \"\"\"\n",
    "    Find the closest distribution to an empirical distribution by minimizing the L1 norm.\n",
    "\n",
    "    Args:\n",
    "        empirical_dist: Empirical distribution that you want to find the closest distribution to.\n",
    "\n",
    "    Returns:\n",
    "        Closest distribution to `empirical_dist`\n",
    "    \"\"\"\n",
    "\n",
    "    def objective(x):\n",
    "        return np.linalg.norm(empirical_dist - x, ord=1)\n",
    "\n",
    "    # Constraint: all elements of p must be positive, and the distribution must sum to 1\n",
    "    cons = (\n",
    "        {\n",
    "            \"type\": \"ineq\",\n",
    "            \"fun\": lambda p: p\n",
    "        },\n",
    "        {\n",
    "            \"type\": \"eq\",\n",
    "            \"fun\": lambda p: np.sum(p) - 1\n",
    "        },\n",
    "    )\n",
    "    bnds = [(0, 1) for _ in range(len(empirical_dist))]\n",
    "    initial_value = np.random.uniform(size=len(empirical_dist))\n",
    "\n",
    "    res = minimize(\n",
    "        objective,\n",
    "        initial_value,\n",
    "        method=\"SLSQP\",\n",
    "        options={\"maxiter\": 1000},\n",
    "        bounds=bnds,\n",
    "        constraints=cons,\n",
    "    )\n",
    "\n",
    "    return res.x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0ceb5a15-cecf-47ed-b1cc-cea79fa936f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_counts_from_distribution(n_qubits, size, dist):\n",
    "    \"\"\"\n",
    "    Generates samples based on a given distribution and returns the counts of each sample value.\n",
    "\n",
    "    Args:\n",
    "        n_qubits: The number of qubits in the quantum circuit.\n",
    "        dist: The probability distribution from which samples are drawn.\n",
    "\n",
    "    Returns:\n",
    "        An array of counts for each possible value in the distribution. The array has a length of 2^n_qubits.\n",
    "    \"\"\"\n",
    "    samples = np.random.choice(np.arange(2**n_qubits), size=size, p=dist)\n",
    "    values, counts = np.unique(samples, return_counts=True)\n",
    "    res = np.zeros(2**n_qubits, dtype=int)\n",
    "    res[values] = counts\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1225d4de-fcc8-4a56-846b-44d6a52a272f",
   "metadata": {},
   "source": [
    "Let's run the `ghz_kernel` with the noise model. We will perform two experiments. First, we keep the bitflip probability the same for all the qubits. In the second experiment, we use different bitflip probabilities for qubits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "53453484-a8d5-416c-8aad-7c982cad7191",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_qubits = 3\n",
    "shots = 1024\n",
    "\n",
    "labels = list(map(list, itertools.product([0, 1], repeat=n_qubits)))\n",
    "states = list(map(lambda label: \"\".join(map(str, label)), labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2d147fe7-ce5e-475c-9014-59d4ff63a644",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bitflip probability: 0.1\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'000': 352,\n",
       " '010': 43,\n",
       " '111': 367,\n",
       " '001': 47,\n",
       " '011': 47,\n",
       " '100': 54,\n",
       " '101': 61,\n",
       " '110': 53}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p = 0.1\n",
    "print(f\"bitflip probability: {p}\\n\")\n",
    "\n",
    "noise_1 = get_noise(n_qubits, p)\n",
    "\n",
    "ghz_result = cudaq.sample(ghz_kernel,\n",
    "                          n_qubits,\n",
    "                          shots_count=shots,\n",
    "                          noise_model=noise_1)\n",
    "noisy_dict_1 = dict(list(ghz_result.items()))\n",
    "\n",
    "noisy_res_1 = np.array(\n",
    "    [noisy_dict_1.get(state, 0) for i, state in enumerate(states)])\n",
    "\n",
    "noisy_dict_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "976f13da-277b-4754-9a58-277a5988f739",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bitflip probability: [0.149816047538945, 0.3802857225639665, 0.292797576724562]\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'000': 212,\n",
       " '010': 122,\n",
       " '111': 189,\n",
       " '001': 111,\n",
       " '011': 87,\n",
       " '100': 78,\n",
       " '101': 131,\n",
       " '110': 94}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p = list(np.random.uniform(low=0.0, high=0.4, size=n_qubits))\n",
    "print(f\"bitflip probability: {p}\\n\")\n",
    "\n",
    "noise_2 = get_noise(n_qubits, p)\n",
    "\n",
    "ghz_result = cudaq.sample(ghz_kernel,\n",
    "                          n_qubits,\n",
    "                          shots_count=shots,\n",
    "                          noise_model=noise_2)\n",
    "noisy_dict_2 = dict(list(ghz_result.items()))\n",
    "\n",
    "noisy_res_2 = np.array(\n",
    "    [noisy_dict_2.get(state, 0) for i, state in enumerate(states)])\n",
    "\n",
    "noisy_dict_2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1889f50a-f051-4f0e-bace-18477122eafe",
   "metadata": {},
   "source": [
    "## Inverse confusion matrix from single-qubit noise model\n",
    "\n",
    "Here we assume that the readout error of each qubit is independent with the same confusion probabilities. Then, the confusion matrix $A_1$ for a single qubit can be defined as\n",
    "\n",
    "$$\n",
    "A_1 = \\begin{bmatrix}\n",
    "P(0|0) & P(0|1)\\\\\n",
    "P(1|0) & P(1|1)\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "where $P(y|x)$ is the probability of observing state $|y\\rangle$ when measuring the true state $|x\\rangle$. Using $A_1$, the full $2^n \\times 2^n$ confusion matrix $A$ can be written as \n",
    "\n",
    "$$\n",
    "A = A_1 \\otimes \\dots \\otimes A_1\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e6ceb0a9-c2ec-4fcd-bb23-457890a8b57e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 becomes {'0': 924, '1': 100}\n",
      "1 becomes {'0': 104, '1': 920}\n"
     ]
    }
   ],
   "source": [
    "single_qubit_labels = [[0], [1]]\n",
    "single_qubit_states = [\"0\", \"1\"]  # Used for dictionary keys\n",
    "\n",
    "# we use noise_1 which has same bitflip probability for all qubits\n",
    "results = [\n",
    "    cudaq.sample(kernel, 1, label, shots_count=shots, noise_model=noise_1)\n",
    "    for label in single_qubit_labels\n",
    "]\n",
    "\n",
    "for i, state in enumerate(single_qubit_states):\n",
    "    res = dict(list(results[i].items()))\n",
    "    print(f\"{state} becomes {res}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9f030f99-0ad9-4f89-8d99-a570b3c5e745",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.90234375 0.1015625 ]\n",
      " [0.09765625 0.8984375 ]]\n"
     ]
    }
   ],
   "source": [
    "# Construct confusion matrix for single qubit\n",
    "\n",
    "A_1 = np.zeros((2, 2))\n",
    "for i, state in enumerate(single_qubit_states):\n",
    "    true_state = int(state, 2)\n",
    "    for key, val in results[i].items():\n",
    "        observed_state = int(key, 2)\n",
    "        A_1[observed_state][true_state] = val / shots\n",
    "\n",
    "print(A_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "cbc4a228-935f-413c-9ad3-b6e26f22e802",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mitigated counts:\n",
      "[476   3   0   1  12  19   9 501]\n"
     ]
    }
   ],
   "source": [
    "A = reduce(np.kron, [A_1] * n_qubits)  # joint confusion matrix\n",
    "A_pinv = np.linalg.pinv(A)  # Generate pseudoinverse confusion matrix.\n",
    "mitigated = np.array(np.dot(A_pinv, noisy_res_1),\n",
    "                     dtype=int)  # Obtain mitigated counts\n",
    "print(f\"Mitigated counts:\\n{mitigated}\")\n",
    "\n",
    "if not np.all(mitigated >= 0):\n",
    "    positive_dist = find_closest_distribution(mitigated / shots)\n",
    "    mitigated = get_counts_from_distribution(n_qubits, shots, positive_dist)\n",
    "    print(f\"\\nCorrected for negative counts:\\n{mitigated}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b65c4657-87f7-44df-b205-7fc6195b95a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31620 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df = pd.DataFrame({\n",
    "    \"states\": states,\n",
    "    \"noisy\": np.around(noisy_res_1 / sum(noisy_res_1), 3),\n",
    "    \"mitigated_sg\": np.around(mitigated / sum(mitigated), 3),\n",
    "})\n",
    "\n",
    "ax = df.plot(x=\"states\",\n",
    "             y=[\"noisy\", \"mitigated_sg\"],\n",
    "             kind=\"bar\",\n",
    "             figsize=(8, 5))\n",
    "ax.bar_label(ax.containers[0], labels=df[\"noisy\"])\n",
    "ax.bar_label(ax.containers[1], labels=df[\"mitigated_sg\"])\n",
    "ax.set_ylabel(\"probabilities\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2927e72-98bb-41cf-8e77-35ea9572474a",
   "metadata": {},
   "source": [
    "## Inverse confusion matrix from k local confusion matrices\n",
    "\n",
    "This method works under the assumption that the readout errors for different qubits are independent of each other and the confusion matrix $A$ can be factorized into the product of $k$ small local confusion matrices. We consider the special case when $k = n$ for $n$ qubits and different $2 \\times 2$ confusion matrices for each qubit\n",
    "\n",
    "$$\n",
    "A = A_1 \\otimes \\dots \\otimes A_n\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1e8fd13a-e0cf-4139-84ca-d81f7a71baa6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "000 becomes {'000': 388, '010': 235, '111': 18, '001': 174, '011': 79, '100': 58, '101': 28, '110': 44}\n",
      "111 becomes {'000': 23, '010': 38, '111': 370, '001': 40, '011': 60, '100': 107, '101': 224, '110': 162}\n"
     ]
    }
   ],
   "source": [
    "local_labels = [[0] * n_qubits, [1] * n_qubits]\n",
    "local_states = [\"0\" * n_qubits, \"1\" * n_qubits]\n",
    "\n",
    "results = [\n",
    "    cudaq.sample(kernel,\n",
    "                 n_qubits,\n",
    "                 label,\n",
    "                 shots_count=shots,\n",
    "                 noise_model=noise_2) for label in local_labels\n",
    "]\n",
    "\n",
    "for i, state in enumerate(local_states):\n",
    "    res = dict(list(results[i].items()))\n",
    "    print(f\"{state} becomes {res}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f592fb8-2ded-498e-b4fe-eff6eac74750",
   "metadata": {},
   "source": [
    "The local confusion matrices are generated using the marginal probability distribution for each qubit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7fabcc1f-8d9c-42c2-b02a-3ce42c1e34c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = [\n",
    "    dict(list(results[i].items())) for i, state in enumerate(local_states)\n",
    "]\n",
    "matrices = []\n",
    "\n",
    "for k in range(n_qubits):\n",
    "    matrix = np.zeros([2, 2], dtype=float)\n",
    "    marginalized_counts = []\n",
    "    total_shots = []\n",
    "    for i in range(2):\n",
    "        marginal_cts = dict(results[i].get_marginal_counts([k]).items())\n",
    "        marginalized_counts.append(marginal_cts)\n",
    "        total_shots.append(sum(marginal_cts.values()))\n",
    "\n",
    "    # matrix[i][j] is the probability of counting i for expected j\n",
    "    for i in range(2):\n",
    "        for j in range(2):\n",
    "            matrix[i][j] = marginalized_counts[j].get(str(i),\n",
    "                                                      0) / total_shots[j]\n",
    "    matrices.append(matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7f016f9b-3243-42b3-ab59-4d03f7765b82",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0.85546875 0.15722656]\n",
      " [0.14453125 0.84277344]]\n",
      "\n",
      "[[0.6328125  0.38476562]\n",
      " [0.3671875  0.61523438]]\n",
      "\n",
      "[[0.70800781 0.32226562]\n",
      " [0.29199219 0.67773438]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for m in matrices:\n",
    "    print(m)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "51734299-307f-4081-b8cc-1968ebfc7d52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mitigated counts:\n",
      "[559  -1 -58  32 -28  26 -16 510]\n",
      "\n",
      "Corrected for negative counts:\n",
      "[521   0   0  35   0  22   0 446]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 21132 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# To find the joint pseudoinverse confusion matrix, we can first perform pseudoinverse of the\n",
    "# local confusion matrices and then do the tensor product\n",
    "\n",
    "pinv_confusion_matrices = [np.linalg.pinv(m) for m in matrices]\n",
    "A_pinv = reduce(np.kron, pinv_confusion_matrices)\n",
    "\n",
    "mitigated = np.array(np.dot(A_pinv, noisy_res_2), dtype=int)\n",
    "print(f\"Mitigated counts:\\n{mitigated}\")\n",
    "\n",
    "if not np.all(mitigated >= 0):\n",
    "    positive_dist = find_closest_distribution(mitigated / shots)\n",
    "    mitigated = get_counts_from_distribution(n_qubits, shots, positive_dist)\n",
    "    print(f\"\\nCorrected for negative counts:\\n{mitigated}\")\n",
    "\n",
    "A_joint = reduce(np.kron, matrices)\n",
    "plot_cmat(A_joint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4261d8a3-b701-4939-8219-c0d823d42747",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33832 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df = pd.DataFrame({\n",
    "    \"states\": states,\n",
    "    \"noisy\": np.around(noisy_res_2 / sum(noisy_res_2), 3),\n",
    "    \"mitigated_k_local\": np.around(mitigated / sum(mitigated), 3),\n",
    "})\n",
    "\n",
    "ax = df.plot(x=\"states\",\n",
    "             y=[\"noisy\", \"mitigated_k_local\"],\n",
    "             kind=\"bar\",\n",
    "             figsize=(8, 5))\n",
    "ax.bar_label(ax.containers[0], labels=df[\"noisy\"])\n",
    "ax.bar_label(ax.containers[1], labels=df[\"mitigated_k_local\"])\n",
    "ax.set_ylabel(\"probabilities\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9329a30d-9224-42b1-9a66-1e87ba7c2674",
   "metadata": {},
   "source": [
    "## Inverse of full confusion matrix\n",
    "\n",
    "Here we generate a quantum circuit for each of the basis states of $n$ qubits (i.e., $2^n$ combinations). This gives more accurate readout error mitigation. Since it scales exponentially with the number of qubits, it is only practical for small devices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "cf6025e8-1755-4c8f-96fa-1b67830d448d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "000 becomes {'000': 367, '010': 248, '111': 17, '001': 159, '011': 83, '100': 68, '101': 35, '110': 47}\n",
      "001 becomes {'000': 161, '010': 93, '111': 46, '001': 363, '011': 253, '100': 17, '101': 76, '110': 15}\n",
      "010 becomes {'000': 234, '010': 378, '111': 27, '001': 100, '011': 150, '100': 36, '101': 25, '110': 74}\n",
      "011 becomes {'000': 106, '010': 145, '111': 65, '001': 251, '011': 386, '100': 11, '101': 30, '110': 30}\n",
      "100 becomes {'000': 83, '010': 34, '111': 99, '001': 29, '011': 18, '100': 388, '101': 144, '110': 229}\n",
      "101 becomes {'000': 22, '010': 13, '111': 247, '001': 70, '011': 43, '100': 163, '101': 371, '110': 95}\n",
      "110 becomes {'000': 32, '010': 77, '111': 157, '001': 14, '011': 30, '100': 233, '101': 85, '110': 396}\n",
      "111 becomes {'000': 14, '010': 31, '111': 393, '001': 38, '011': 60, '100': 95, '101': 233, '110': 160}\n"
     ]
    }
   ],
   "source": [
    "results = [\n",
    "    cudaq.sample(kernel,\n",
    "                 n_qubits,\n",
    "                 label,\n",
    "                 shots_count=shots,\n",
    "                 noise_model=noise_2) for label in labels\n",
    "]\n",
    "\n",
    "for i, state in enumerate(states):\n",
    "    res = dict(list(results[i].items()))\n",
    "    print(f\"{state} becomes {res}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "28f2b2cb-13d3-4806-852d-06ea91e58c1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate full confusion matrix\n",
    "A_full = np.zeros((2**n_qubits, 2**n_qubits))\n",
    "for i, state in enumerate(states):\n",
    "    true_state = int(state, 2)\n",
    "    for key, val in results[i].items():\n",
    "        observed_state = int(key, 2)\n",
    "        A_full[observed_state][true_state] = val / shots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "c3b0b80b-63ac-438a-9f3d-a6a3687da167",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21116 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_cmat(A_full)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "507852a8-6715-4802-b6a0-10967d09ed94",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.09734490298929"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Distance between full confusion matrix and tensored confusion matrix\n",
    "np.linalg.norm(A_joint - A_full)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "9e6c3874-2905-4c25-9bf0-22ab27756fed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mitigated counts:\n",
      "[ 714 -100 -211  152  -49   63   32  422]\n",
      "\n",
      "Corrected for negative counts:\n",
      "[683   0   0  39   0   0   0 302]\n"
     ]
    }
   ],
   "source": [
    "A_pinv = np.linalg.pinv(A_full)\n",
    "mitigated = np.array(np.dot(A_pinv, noisy_res_2), dtype=int)\n",
    "print(f\"Mitigated counts:\\n{mitigated}\")\n",
    "\n",
    "if not np.all(mitigated >= 0):\n",
    "    positive_dist = find_closest_distribution(mitigated / shots)\n",
    "    mitigated = get_counts_from_distribution(n_qubits, shots, positive_dist)\n",
    "    print(f\"\\nCorrected for negative counts:\\n{mitigated}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "cb109820-5478-4179-84b8-66e6af1008fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 43124 characters]",
      "text/plain": [
       "<Figure size 1200x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df[\"mitigated_full\"] = np.around(mitigated / sum(mitigated), 3)\n",
    "\n",
    "ax = df.plot(\n",
    "    x=\"states\",\n",
    "    y=[\"noisy\", \"mitigated_k_local\", \"mitigated_full\"],\n",
    "    kind=\"bar\",\n",
    "    figsize=(12, 5),\n",
    ")\n",
    "ax.bar_label(ax.containers[0], labels=df[\"noisy\"])\n",
    "ax.bar_label(ax.containers[1], labels=df[\"mitigated_k_local\"])\n",
    "ax.bar_label(ax.containers[2], labels=df[\"mitigated_full\"])\n",
    "ax.set_ylabel(\"probabilities\")\n",
    "plt.show()"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
