{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Leakage characterization using GST\n",
    "This tutorial demonstrates how to perform GST on a \"leaky-qubit\" described by a 3-level (instead of the desired 2-level) system. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygsti\n",
    "import pygsti.modelpacks.smq1Q_XYI as smq1Q\n",
    "from pygsti.baseobjs import Label\n",
    "from pygsti.circuits import Circuit\n",
    "import numpy as np\n",
    "import scipy.linalg as sla\n",
    "#import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_3level_unitary(U_2level):\n",
    "    U_3level = np.zeros((3,3),complex)\n",
    "    U_3level[0:2,0:2] = U_2level\n",
    "    U_3level[2,2] = 1.0\n",
    "    return U_3level\n",
    "\n",
    "def unitary_to_gmgate(U):\n",
    "    return pygsti.tools.change_basis( \n",
    "        pygsti.tools.unitary_to_std_process_mx(U), 'std','gm')\n",
    "\n",
    "def state_to_gmvec(state):\n",
    "    pygsti.tools.stdmx_to_gmvec\n",
    "\n",
    "Us = pygsti.tools.internalgates.standard_gatename_unitaries()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mdl_2level_ideal = smq1Q.target_model(qubit_labels=[\"Qubit\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho0 = np.array( [[1,0,0],\n",
    "                  [0,0,0],\n",
    "                  [0,0,0]], complex)\n",
    "E0 = rho0\n",
    "E1 = np.array( [[0,0,0],\n",
    "                [0,1,0],\n",
    "                [0,0,1]], complex)\n",
    "\n",
    "sslbls = pygsti.baseobjs.ExplicitStateSpace(['Qubit_leakage'],[3])\n",
    "mdl_3level_ideal = pygsti.models.ExplicitOpModel(sslbls, 'gm', simulator='matrix')\n",
    "mdl_3level_ideal['rho0'] = pygsti.tools.stdmx_to_gmvec(rho0)\n",
    "mdl_3level_ideal['Mdefault'] = pygsti.modelmembers.povms.TPPOVM([('0',pygsti.tools.stdmx_to_gmvec(E0)),\n",
    "                                                                 ('1',pygsti.tools.stdmx_to_gmvec(E1))],\n",
    "                                                                evotype='default')\n",
    "\n",
    "mdl_3level_ideal[tuple()] = unitary_to_gmgate( to_3level_unitary(Us['Gi']))\n",
    "mdl_3level_ideal['Gxpi2', 'Qubit_leakage'] = unitary_to_gmgate( to_3level_unitary(Us['Gxpi2']))\n",
    "mdl_3level_ideal['Gypi2', 'Qubit_leakage'] = unitary_to_gmgate( to_3level_unitary(Us['Gypi2']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigmaX = np.array([[0,1],[1,0]],complex)\n",
    "rot = sla.expm(1j * 0.1 * sigmaX)\n",
    "Uleakage = np.identity(3,complex)\n",
    "Uleakage[1:3,1:3] = rot\n",
    "leakageOp = unitary_to_gmgate(Uleakage)\n",
    "#print(Uleakage)\n",
    "\n",
    "#Guess of a model w/just unitary leakage\n",
    "mdl_3level_guess = mdl_3level_ideal.copy()\n",
    "mdl_3level_guess[tuple()] = np.dot(leakageOp, mdl_3level_guess[tuple()])\n",
    "#mdl_3level_guess['Gxpi2', 'Qubit_leakage'] = np.dot(leakageOp, mdl_3level_guess['Gxpi2', 'Qubit_leakage'])\n",
    "#mdl_3level_guess['Gypi2', 'Qubit_leakage'] = np.dot(leakageOp, mdl_3level_guess['Gypi2', 'Qubit_leakage'])\n",
    "\n",
    "#Actual model used for data generation (some depolarization too)\n",
    "mdl_3level_noisy = mdl_3level_ideal.depolarize(op_noise=0.005, spam_noise=0.01)\n",
    "mdl_3level_noisy[tuple()] = np.dot(leakageOp, mdl_3level_noisy[tuple()])\n",
    "#mdl_3level_noisy['Gxpi2', 'Qubit_leakage'] = np.dot(leakageOp, mdl_3level_noisy['Gxpi2', 'Qubit_leakage'])\n",
    "#mdl_3level_noisy['Gypi2', 'Qubit_leakage'] = np.dot(leakageOp, mdl_3level_noisy['Gypi2', 'Qubit_leakage'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(mdl_3level_guess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get sequences using expected model\n",
    "find_fiducials = True\n",
    "\n",
    "if find_fiducials:\n",
    "    prepfids, measfids = pygsti.algorithms.find_fiducials(\n",
    "        mdl_3level_guess, omit_identity=False, candidate_fid_counts={4: \"all upto\"}, verbosity=4)\n",
    "    pygsti.io.write_circuit_list(\"example_files/leakage_prepfids.txt\", prepfids)\n",
    "    pygsti.io.write_circuit_list(\"example_files/leakage_measfids.txt\", measfids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If files missing, run previous cell at least once with find_fiducials = True\n",
    "prepfids = pygsti.io.read_circuit_list(\"example_files/leakage_prepfids.txt\")\n",
    "measfids = pygsti.io.read_circuit_list(\"example_files/leakage_measfids.txt\")\n",
    "germs = smq1Q.germs(qubit_labels=[\"Qubit_leakage\"])\n",
    "maxLengths = [1,]\n",
    "expList = pygsti.circuits.create_lsgst_circuits(mdl_3level_noisy, prepfids, measfids, germs, maxLengths)\n",
    "ds = pygsti.data.simulate_data(mdl_3level_noisy, expList, 1000, 'binomial', seed=1234)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# We have found out prep fids, meas fids, and germs, as well as simulated noisy data, for the 3 level model\n",
    "# If we want to run GST on another model, we need to get versions of the circuits will the correct state space labels\n",
    "\n",
    "def map_2level_sslbls(circuit):\n",
    "    sslbl_map = {'Qubit_leakage': 'Qubit'}\n",
    "    return circuit.map_state_space_labels(sslbl_map)\n",
    "\n",
    "prepfids_2level = [map_2level_sslbls(c) for c in prepfids]\n",
    "measfids_2level = [map_2level_sslbls(c) for c in measfids]\n",
    "germs_2level = [map_2level_sslbls(c) for c in germs]\n",
    "ds_2level = ds.process_circuits(map_2level_sslbls)\n",
    "\n",
    "results_2level = pygsti.run_stdpractice_gst(ds_2level, mdl_2level_ideal, prepfids_2level, measfids_2level,\n",
    "                                           germs_2level, maxLengths, modes=\"CPTPLND\", verbosity=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pygsti.report.construct_standard_report(results_2level, \"2-level Leakage Example Report\").write_html('example_files/leakage_report_2level')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open the report [here](example_files/leakage_report_2level/main.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "results_3level = pygsti.run_stdpractice_gst(ds, mdl_3level_ideal, prepfids, measfids,\n",
    "                                           germs, maxLengths, modes=[\"CPTPLND\",\"True\"],\n",
    "                                           models_to_test={'True': mdl_3level_noisy}, \n",
    "                                           verbosity=4, advanced_options={'all': {'tolerance': 1e-2}})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pygsti.report.construct_standard_report(results_3level, \"3-level Leakage Example Report\").write_html('example_files/leakage_report')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open the report [here](example_files/leakage_report/main.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#try a different basis:\n",
    "gm_basis = pygsti.baseobjs.Basis.cast('gm',9)\n",
    "   \n",
    "leakage_basis_mxs = [ np.sqrt(2)/3*(np.sqrt(3)*gm_basis[0] + 0.5*np.sqrt(6)*gm_basis[8]),\n",
    "                      gm_basis[1], gm_basis[4], gm_basis[7],\n",
    "                     gm_basis[2], gm_basis[3], gm_basis[5], gm_basis[6],\n",
    "                     1/3*(np.sqrt(3)*gm_basis[0] - np.sqrt(6)*gm_basis[8]) ]\n",
    "#for mx in leakage_basis_mxs:\n",
    "#    pygsti.tools.print_mx(mx)\n",
    "\n",
    "check = np.zeros( (9,9), complex)\n",
    "for i,m1 in enumerate(leakage_basis_mxs):\n",
    "    for j,m2 in enumerate(leakage_basis_mxs):\n",
    "        check[i,j] = np.trace(np.dot(m1,m2))\n",
    "assert(np.allclose(check, np.identity(9,complex)))\n",
    "\n",
    "leakage_basis = pygsti.baseobjs.ExplicitBasis(leakage_basis_mxs, name=\"LeakageBasis\",  \n",
    "                                        longname=\"2+1 level leakage basis\", real=True,\n",
    "                                        labels=['I','X','Y','Z','LX0','LX1','LY0','LY1','L'])\n",
    "\n",
    "def changebasis_3level_model(mdl):\n",
    "    new_mdl = mdl.copy()\n",
    "    new_mdl.preps['rho0'] = pygsti.modelmembers.states.FullState(\n",
    "        pygsti.tools.change_basis(mdl.preps['rho0'].to_dense(), gm_basis, leakage_basis))\n",
    "    new_mdl.povms['Mdefault'] = pygsti.modelmembers.povms.UnconstrainedPOVM(\n",
    "        [('0', pygsti.tools.change_basis(mdl.povms['Mdefault']['0'].to_dense(), gm_basis, leakage_basis)),\n",
    "         ('1', pygsti.tools.change_basis(mdl.povms['Mdefault']['1'].to_dense(), gm_basis, leakage_basis))],\n",
    "        evotype='default')\n",
    "    \n",
    "    for lbl,op in mdl.operations.items():\n",
    "        new_mdl.operations[lbl] = pygsti.modelmembers.operations.FullArbitraryOp(\n",
    "            pygsti.tools.change_basis(op.to_dense(), gm_basis, leakage_basis))\n",
    "    new_mdl.basis = leakage_basis\n",
    "    return new_mdl\n",
    "\n",
    "def changebasis_3level_results(results):\n",
    "    new_results = results.copy()\n",
    "    for estlbl,est in results.estimates.items():\n",
    "        for mlbl,mdl in est.models.items():\n",
    "            if isinstance(mdl,(list,tuple)): #assume a list/tuple of models\n",
    "                new_results.estimates[estlbl].models[mlbl] = \\\n",
    "                    [ changebasis_3level_model(m) for m in mdl ]\n",
    "            else:\n",
    "                new_results.estimates[estlbl].models[mlbl] = changebasis_3level_model(mdl)\n",
    "    return new_results\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_3level_leakage_basis = changebasis_3level_results( results_3level )    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pygsti.report.construct_standard_report(results_3level_leakage_basis, \"3-level with Basis Change Leakage Example Report\"\n",
    "                                        ).write_html('example_files/leakage_report_basis')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open the report [here](example_files/leakage_report_basis/main.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use \"kite\" density-matrix structure\n",
    "def to_2plus1_superop(superop_2level):\n",
    "    ret = np.zeros((5,5),'d')\n",
    "    ret[0:4,0:4] = superop_2level.to_dense()\n",
    "    ret[4,4] = 1.0 #leave leakage population where it is\n",
    "    return ret\n",
    "\n",
    "#Tack on a single extra \"0\" for the 5-th dimension corresponding\n",
    "# to the classical leakage level population.\n",
    "eps = 0.01 # ideally zero, a smallish number to seed the GST optimiation away from 0-leakage so it doesn't get stuck there.\n",
    "rho0 = np.concatenate( (mdl_2level_ideal.preps['rho0'].to_dense(),[eps]), axis=0)\n",
    "E0 = np.concatenate( (mdl_2level_ideal.povms['Mdefault']['0'].to_dense(),[eps]), axis=0)\n",
    "E1 = np.concatenate( (mdl_2level_ideal.povms['Mdefault']['1'].to_dense(),[eps]), axis=0)\n",
    "\n",
    "\n",
    "statespace = pygsti.baseobjs.ExplicitStateSpace([('Qubit',),('Leakage',)], [(2,), (1,)])\n",
    "mdl_2plus1_ideal = pygsti.models.ExplicitOpModel(statespace, 'gm', simulator='matrix')\n",
    "mdl_2plus1_ideal['rho0'] = rho0\n",
    "mdl_2plus1_ideal['Mdefault'] = pygsti.modelmembers.povms.UnconstrainedPOVM([('0',E0),('1',E1)],\n",
    "                                                                           evotype='default', state_space=statespace)\n",
    "\n",
    "mdl_2plus1_ideal[tuple()] = to_2plus1_superop(mdl_2level_ideal[tuple()])\n",
    "mdl_2plus1_ideal['Gxpi2'] = to_2plus1_superop(mdl_2level_ideal['Gxpi2', 'Qubit'])\n",
    "mdl_2plus1_ideal['Gypi2'] = to_2plus1_superop(mdl_2level_ideal['Gypi2', 'Qubit'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We have found out prep fids, meas fids, and germs, as well as simulated noisy data, for the 3 level model\n",
    "# If we want to run GST on another model, we need to get versions of the circuits will the correct state space labels\n",
    "\n",
    "# We do this in a slightly different/awkward way here for this case since our state space labels are not a single entry\n",
    "# This would not be necessary if we were rebuilding the circuits/dataset from scratch, only hacky since we are reusing the 3-level information\n",
    "def map_2plus1_circuit_linelabels(circuit):\n",
    "    return Circuit([Label(l.name) if l.name != \"COMPOUND\" else tuple() for l in circuit.layertup],\n",
    "                   \"*\", None, not circuit._static)\n",
    "\n",
    "prepfids_2plus1 = [map_2plus1_circuit_linelabels(c) for c in prepfids]\n",
    "measfids_2plus1 = [map_2plus1_circuit_linelabels(c) for c in measfids]\n",
    "germs_2plus1 = [map_2plus1_circuit_linelabels(c) for c in germs]\n",
    "ds_2plus1 = ds.process_circuits(map_2plus1_circuit_linelabels)\n",
    "\n",
    "results_2plus1 = pygsti.run_long_sequence_gst(ds_2plus1, mdl_2plus1_ideal, prepfids_2plus1, measfids_2plus1,\n",
    "                                             germs_2plus1, maxLengths, verbosity=2,\n",
    "                                             advanced_options={\"starting_point\": \"target\",\n",
    "                                                               \"tolerance\": 1e-8,  # (lowering tolerance from 1e-6 gave a better fit)\n",
    "                                                               \"estimate_label\": \"kite\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbval-skip"
    ]
   },
   "outputs": [],
   "source": [
    "# TODO: This is currently broken\n",
    "pygsti.report.construct_standard_report(results_2plus1,\"2+1 Leakage Example Report\"\n",
    ").write_html('example_files/leakage_report_2plus1', autosize='none')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open the report [here](example_files/leakage_report/main.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pygsti",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
