{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial: Data Analysis\n",
    "\n",
    "This is the follow up to [Tutorial: Data Collection](./Readout-Data-Collection.ipynb). We have measured bitstrings for the single-qubit circuit $R_y(\\theta)$ for various `theta`s. In this analysis, we compute $\\langle Z \\rangle (\\theta)$, compare to the anayltically expected true value, and fit to a depolarizing noise model with T1 decay during readout."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading data\n",
    "\n",
    "We can use utilities in ReCirq to query the filesystem and load in a dataset. Please recall that all tasks have an associated `EXPERIMENT_NAME` and a `dataset_id` which define the top two hierarchies in the filesystem. We import these values from the data collection script to ensure consistency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cirq\n",
    "import recirq\n",
    "\n",
    "from recirq.readout_scan.tasks import EXPERIMENT_NAME, DEFAULT_BASE_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`recirq.iterload_records` uses these two bits of information to iterate over records saved using `recirq.save` (in the data collection script.\n",
    "\n",
    "This also gives you a chance to do post-processing on the data. In general, you should do some massaging of the data and put the results into a pandas DataFrame. DataFrames are great for doing statistics and visualizations across tabular data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "records = []\n",
    "# Load all data, do some light processing\n",
    "for record in recirq.iterload_records(dataset_id='2020-02-tutorial', base_dir=DEFAULT_BASE_DIR):\n",
    "    # Expand task dataclass into columns\n",
    "    recirq.flatten_dataclass_into_record(record, 'task')\n",
    "    \n",
    "    # Unwrap BitArray into np.ndarray\n",
    "    all_bitstrings = [ba.bits for ba in record['all_bitstrings']]\n",
    "    \n",
    "    # Compute <Z>\n",
    "    record['z_vals'] = [np.mean((-1)**bitstrings, axis=0).item() for bitstrings in all_bitstrings]\n",
    "    \n",
    "    # Don't need to carry around the full array of bits anymore\n",
    "    del record['all_bitstrings']\n",
    "    records.append(record)\n",
    "    \n",
    "df = pd.DataFrame(records)\n",
    "print(len(df))\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the data\n",
    "\n",
    "A good first step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "entry = df.iloc[0] # Pick the first qubit\n",
    "\n",
    "plt.plot([], []) # advance color cycle in anticipation of future analysis\n",
    "plt.plot(entry['thetas'], entry['z_vals'], 'o-')\n",
    "plt.xlabel('Theta', fontsize=14)\n",
    "plt.ylabel(r'$\\langle Z \\rangle$', fontsize=14)\n",
    "plt.title(\"Qubit {}\".format(entry['qubit']), fontsize=14)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How does it compare to analytical results?\n",
    "\n",
    "You could imagine setting up a separate task for computing and saving analytic results. For this single qubit example, we'll just compute it on the fly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qubit = cirq.LineQubit(0)\n",
    "thetas = df.iloc[0]['thetas']\n",
    "\n",
    "class _DummyMeasurementGate(cirq.IdentityGate):\n",
    "    \"\"\"A dummy measurement used to trick simulators into applying\n",
    "    readout error when using PauliString.expectation_from_xxx.\"\"\"\n",
    "\n",
    "    def _measurement_key_(self):\n",
    "        return 'dummy!'\n",
    "\n",
    "    def __repr__(self):\n",
    "        if self.num_qubits() == 1:\n",
    "            return '_DummyMeasurementGate'\n",
    "        return '_DummyMeasurementGate({!r})'.format(self.num_qubits())\n",
    "\n",
    "    def __str__(self):\n",
    "        if (self.num_qubits() == 1):\n",
    "            return 'dummyM'\n",
    "        else:\n",
    "            return 'dummyM({})'.format(self.num_qubits())\n",
    "\n",
    "    def _circuit_diagram_info_(self, args):\n",
    "        from cirq import protocols\n",
    "        return protocols.CircuitDiagramInfo(\n",
    "            wire_symbols=('dM',) * self.num_qubits(), connected=True)\n",
    "\n",
    "def dummy_measure(qubits):\n",
    "    return _DummyMeasurementGate(num_qubits=len(qubits)).on(*qubits)\n",
    "\n",
    "def get_circuit(theta):\n",
    "    return cirq.Circuit([\n",
    "        cirq.ry(theta).on(qubit),\n",
    "        dummy_measure([qubit])\n",
    "    ])\n",
    "\n",
    "true_z_vals = []\n",
    "for theta in thetas:\n",
    "    wf = cirq.final_wavefunction(get_circuit(theta))\n",
    "    op = cirq.Z(qubit) * 1.\n",
    "    true_z_val = op.expectation_from_wavefunction(wf, qubit_map={qubit:0}, check_preconditions=False)\n",
    "    true_z_vals.append(np.real_if_close(true_z_val).item())\n",
    "\n",
    "true_z_vals = np.array(true_z_vals)\n",
    "true_z_vals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4))\n",
    "ax1.plot(thetas, true_z_vals, '-', label='True')\n",
    "ax1.plot(entry['thetas'], entry['z_vals'], 'o-', label='Data')\n",
    "\n",
    "ax2.plot([], []) # advance color cycle\n",
    "ax2.plot(entry['thetas'], np.abs(true_z_vals - entry['z_vals']), 'o-', label='|Data - True|')\n",
    "\n",
    "ax1.legend(loc='best', frameon=False)\n",
    "ax2.legend(loc='best', frameon=False)\n",
    "ax1.set_xlabel('Theta', fontsize=14)\n",
    "ax2.set_xlabel('Theta', fontsize=14)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learn a model\n",
    "\n",
    "Our experimental data has some wiggles in it, but it also has a clear pattern of deviation from the true values. We can hypothesize a (parameterized) noise model and then use function minimization to fit the noise model parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.optimize\n",
    "import cirq.contrib.noise_models as ccn\n",
    "\n",
    "def get_obj_func(data_expectations):\n",
    "    all_results = []\n",
    "    def obj_func(x):\n",
    "        depol_prob, decay_prob, readout_prob = x\n",
    "        \n",
    "        if depol_prob < 0 or decay_prob < 0 or readout_prob < 0:\n",
    "            # emulate constraints by returning a high cost if we\n",
    "            # stray into invalid territory\n",
    "            return 1000\n",
    "\n",
    "        sim = cirq.DensityMatrixSimulator(\n",
    "            noise=ccn.DepolarizingWithDampedReadoutNoiseModel(\n",
    "                depol_prob=depol_prob, decay_prob=decay_prob, bitflip_prob=readout_prob))\n",
    "        \n",
    "        results = []\n",
    "        for theta in thetas:            \n",
    "            density_result = sim.simulate(get_circuit(theta))\n",
    "            op = cirq.Z(qubit) * 1.\n",
    "            true_z_val = op.expectation_from_wavefunction(density_result.final_density_matrix, qubit_map=density_result.qubit_map, check_preconditions=False)\n",
    "            results.append(np.real_if_close(true_z_val).item())\n",
    "\n",
    "        results = np.array(results)\n",
    "        all_results.append(results)\n",
    "        cost = np.sum(np.abs(results - data_expectations))\n",
    "        return cost\n",
    "    \n",
    "    return obj_func, all_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_result(x):\n",
    "        depol_prob, decay_prob, readout_prob = x\n",
    "        print(f'depol   = {depol_prob:.2%}')\n",
    "        print(f'decay   = {decay_prob:.2%}')\n",
    "        print(f'readout = {readout_prob:.2%}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfb = df\n",
    "dfb = dfb.head(5) # Remove this to do all qubits\n",
    "len(dfb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Initial values\n",
    "depol_prob = 0.01\n",
    "decay_prob = 0.01\n",
    "readout_prob = 0.01\n",
    "\n",
    "opt_results = []\n",
    "for i, entry in dfb.iterrows():\n",
    "    ofunc, results = get_obj_func(entry['z_vals'])    \n",
    "    opt_result = scipy.optimize.minimize(ofunc, \n",
    "                                         [depol_prob, decay_prob, readout_prob],\n",
    "                                         method='nelder-mead',\n",
    "                                         options={'disp': True})\n",
    "    label = f\"{entry['qubit'].row}, {entry['qubit'].col}\"\n",
    "    print(\"Qubit\", label)\n",
    "    print_result(opt_result.x)\n",
    "    opt_results.append(opt_result)\n",
    "    \n",
    "    data_expectations = entry['z_vals']\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4))\n",
    "    ax1.plot(thetas, true_z_vals, label='True')\n",
    "    ax1.plot(thetas, data_expectations, 'o-', label=f'{label} Data')\n",
    "    ax1.plot(thetas, results[-1], '.-', label='Fit')\n",
    "    \n",
    "    ax2.plot([], []) # advance color cycle\n",
    "    ax2.plot(thetas, np.abs(true_z_vals - data_expectations), 'o-', label='|Data - True|')\n",
    "    ax2.plot(thetas, np.abs(true_z_vals - results[-1]), '-', label='|Fit - True|')\n",
    "    ax1.legend(loc='best')\n",
    "    ax2.legend(loc='best')\n",
    "    fig.tight_layout()\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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}