{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some parts of `pygsti` are works-in-progress. Here, we investigate how to do the task of \"model selection\" within GST, essentially answering the question \"Can we do a better job of modeling the experiment by changing the assumptions within GST?\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing variable-gateset-dimension GST with model selection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###     Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import pygsti"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Load gateset and some string lists\n",
    "gs_target = pygsti.io.load_gateset(\"tutorial_files/Example_Gateset.txt\")\n",
    "fiducialList = pygsti.io.load_gatestring_list(\"tutorial_files/Example_FiducialList.txt\")\n",
    "germList = pygsti.io.load_gatestring_list(\"tutorial_files/Example_GermsList.txt\")\n",
    "specs = pygsti.construction.build_spam_specs(fiducialList)\n",
    "expList = [1,2,4]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Create some testing gate string lists\n",
    "lgstList = pygsti.construction.list_lgst_gatestrings(specs, gs_target.gates.keys())\n",
    "lsgstLists = [ lgstList[:] ]\n",
    "for exp in expList:\n",
    "    gsList = pygsti.construction.create_gatestring_list(\n",
    "                \"f0+germ*exp+f1\", f0=fiducialList, f1=fiducialList,\n",
    "                germ=germList, exp=exp, order=['germ','f0','f1'])\n",
    "    lsgstLists.append( lsgstLists[-1] +  gsList )\n",
    "    \n",
    "dsList = pygsti.remove_duplicates( lsgstLists[-1] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Test on fake data by depolarizing target set, increasing its dimension,\n",
    "# and adding leakage to the gates into the new dimension.\n",
    "\n",
    "gs_dataGen4 = gs_target.depolarize(gate_noise=0.1)\n",
    "gs_dataGen5 = gs_dataGen4.increase_dimension(5)\n",
    "leakGate = pygsti.construction.build_gate( [2,1],[('Q0',),('L0',)] , \"LX(pi/4.0,0,2)\",\"gm\") # X(pi,Q0)*LX(pi,0,2)\n",
    "\n",
    "gs_dataGen5['Gx'] = pygsti.objects.compose( gs_dataGen5['Gx'], leakGate)\n",
    "gs_dataGen5['Gy'] = pygsti.objects.compose( gs_dataGen5['Gy'], leakGate)\n",
    "print(gs_dataGen5.gates.keys())\n",
    "\n",
    "#Some debugging...\n",
    "#NOTE: with LX(pi,0,2) above, dim 5 test will choose a dimension 3 gateset, which may be sensible\n",
    "#       looking at the gate matrices in this case... but maybe LX(pi,...) is faulty?\n",
    "#print(gs_dataGen4)\n",
    "#print(gs_dataGen5)\n",
    "\n",
    "#Jmx = GST.JOps.jamiolkowski_iso(gs_dataGen4['Gx'])\n",
    "#Jmx = GST.JOps.jamiolkowski_iso(gs_dataGen5['Gx'],dimOrStateSpaceDims=[2,1])\n",
    "#print(\"J = \\n\",Jmx)\n",
    "#print(\"evals = \",eigvals(Jmx))\n",
    "\n",
    "dsFake4 = pygsti.construction.generate_fake_data(gs_dataGen4, dsList, nSamples=1000000, sampleError=\"binomial\", seed=1234)\n",
    "dsFake5 = pygsti.construction.generate_fake_data(gs_dataGen5, dsList, nSamples=1000000, sampleError=\"binomial\", seed=1234)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"Number of gates                        = \",len(gs_target.gates.keys()))\n",
    "print(\"Number of fiducials                    =\",len(fiducialList))\n",
    "print(\"Maximum length for a gate string in ds =\",max(map(len,dsList)))\n",
    "print(\"Number of LGST strings                 = \",len(lgstList))\n",
    "print(\"Number of LSGST strings                = \",map(len,lsgstLists))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test using dimension-4 fake data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#Run LGST to get an initial estimate for the gates in gs_target based on the data in ds\n",
    "# NOTE: with nSamples less than 1M (100K, 10K, 1K) this routine will choose a higher-than-4 dimensional gateset\n",
    "ds = dsFake4\n",
    "gs_lgst4 = pygsti.do_lgst(ds, specs, targetGateset=gs_target, svdTruncateTo=4, verbosity=3)\n",
    "gs_lgst6 = pygsti.do_lgst(ds, specs, targetGateset=gs_target, svdTruncateTo=6, verbosity=3)\n",
    "\n",
    "#Print chi^2 of 4-dim and 6-dim estimates\n",
    "chiSq4 = pygsti.chi2(ds, gs_lgst4, lgstList, minProbClipForWeighting=1e-4)\n",
    "chiSq6 = pygsti.chi2(ds, gs_lgst6, lgstList, minProbClipForWeighting=1e-4)\n",
    "print(\"LGST dim=4 chiSq = \", chiSq4)\n",
    "print(\"LGST dim=6 chiSq = \", chiSq6)\n",
    "\n",
    "# Least squares GST with model selection\n",
    "gs_lsgst = pygsti.do_iterative_mc2gst_with_model_selection(ds, gs_lgst4, 1, lsgstLists, verbosity=2,\n",
    "                                                           minProbClipForWeighting=1e-3, probClipInterval=(-1e5,1e5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(gs_lsgst)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test using dimension-5 fake data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Run LGST to get an initial estimate for the gates in gs_target based on the data in ds\n",
    "ds = dsFake5\n",
    "gs_lgst4 = pygsti.do_lgst(ds, specs, targetGateset=gs_target, svdTruncateTo=4, verbosity=3)\n",
    "gs_lgst6 = pygsti.do_lgst(ds, specs, targetGateset=gs_target, svdTruncateTo=6, verbosity=3)\n",
    "\n",
    "#Print chi^2 of 4-dim and 6-dim estimates\n",
    "chiSq4 = pygsti.chi2(ds, gs_lgst4, lgstList, minProbClipForWeighting=1e-2)\n",
    "chiSq6 = pygsti.chi2(ds, gs_lgst6, lgstList, minProbClipForWeighting=1e-2)\n",
    "print(\"LGST dim=4 chiSq = \", chiSq4)\n",
    "print(\"LGST dim=6 chiSq = \", chiSq6)\n",
    "\n",
    "# Least squares GST with model selection\n",
    "gs_lsgst = pygsti.do_iterative_mc2gst_with_model_selection(ds, gs_lgst4, 1, lsgstLists, verbosity=2, minProbClipForWeighting=1e-3, probClipInterval=(-1e5,1e5), useFreqWeightedChiSq=False, regularizeFactor=1.0, check=False, check_jacobian=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(gs_lsgst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.5.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}