{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: yacc table file version is out of date\n"
     ]
    }
   ],
   "source": [
    "from __future__ import division, print_function, absolute_import, unicode_literals\n",
    "import numpy as np\n",
    "from scipy.linalg import expm\n",
    "\n",
    "import pygsti\n",
    "from pygsti.construction import std1Q_XYI as std\n",
    "from pygsti.construction import std2Q_XYICNOT as std2Q\n",
    "from pygsti.extras import newrb as prb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# A list of standard hard-coded gate labels\n",
    "gllist = ['Gi','Gh','Gp','Gcnot'] \n",
    "\n",
    "# The number of qubits\n",
    "n = 4\n",
    "\n",
    "# The availiability of gates that are not available to all qubit / qubit pairs\n",
    "availability = {}\n",
    "\n",
    "a = [(i,i+1) for i in range(0,n-1)] + [(n-1,0)]\n",
    "availability={'Gcnot':a}\n",
    "\n",
    "pspec = pygsti.obj.ProcessorSpec(n, gllist, availability=availability, verbosity=0)\n",
    "\n",
    "# Add a target dm model\n",
    "pspec.add_std_model('dm_target', parameterization=\"static\", sim_type=\"dmmap\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Basic tests for perfect gates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Qubit 0 ---|Gp |-|Gi |-|  |---\n",
      "Qubit 1 ---|  |-|Gh |-Gcnot:1:2|-|--\n",
      "Qubit 2 ---|  |-|  |-Gcnot:1:2|-|--\n",
      "Qubit 3 ---|  |-|Gp |-|  |---\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create a fairly trivial circuit\n",
    "c = pygsti.objects.Circuit(gatestring=[('Gp',0),('Gi',0),('Gh',1),('Gp',3),('Gcnot',1,2)],num_lines=n)\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "vs_target_out = c.simulate(pspec.models['target'])\n",
    "dm_target_out = c.simulate(pspec.models['dm_target'])\n",
    "for key in vs_target_out:\n",
    "    assert(abs(vs_target_out[key] - dm_target_out[key]) < 1e-10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic tests for imperfect gates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pspec.add_std_model('dmsim', parameterization=\"static\", sim_type=\"dmmap\")\n",
    "pspec.add_std_model('svsim', parameterization=\"static\", sim_type=\"svmap\")\n",
    "\n",
    "# Nothing should be different from above, but let's check.\n",
    "vs_target_out = c.simulate(pspec.models['target'])\n",
    "dm_target_out = c.simulate(pspec.models['dm_target'])\n",
    "for key in vs_target_out:\n",
    "    assert(abs(vs_target_out[key] - dm_target_out[key]) < 1e-10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Change the Gi gates on qubit 0 to be small X rotations\n",
    "theta = 0.01*np.pi\n",
    "for glabel in list(pspec.models['svsim'].gates.keys()):\n",
    "    if glabel.name == 'Gi':\n",
    "        if glabel.qubits == (0,):\n",
    "            unitary = expm(-1j*theta*np.array([[0.,1.],[1.,0.]])/2) \n",
    "            pspec.models['svsim'][glabel] = unitary\n",
    "            pspec.models['dmsim'][glabel] = pygsti.unitary_to_pauligate(unitary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Qubit 0 ---|Gi |-|Gi |---\n",
      "Qubit 1 ---|  |-|  |---\n",
      "Qubit 2 ---|  |-|  |---\n",
      "Qubit 3 ---|  |-|  |---\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create a particularly simple circuit\n",
    "c2 = pygsti.objects.Circuit(gatestring=[('Gi',0),('Gi',0)],num_lines=n)\n",
    "print(c2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sv_out = c2.simulate(pspec.models['svsim'])\n",
    "dm_out = c2.simulate(pspec.models['dmsim'])\n",
    "for key in sv_out:\n",
    "    assert(abs(sv_out[key] - dm_out[key]) < 1e-10)\n",
    "assert(abs(np.cos(2*theta/2)**2 - sv_out['0'*n]) < 1e-10) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.999013364214\n",
      "0.9990133642141359\n",
      "0.999013364214136\n"
     ]
    }
   ],
   "source": [
    "print(np.cos(2*theta/2)**2)\n",
    "print(sv_out['0'*n])\n",
    "print(dm_out['0'*n])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Let's make the special Idle be the Gi gate.\n",
    "c2.replace_gatename('I','Gi')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Nothing should be different from above, but let's check.\n",
    "sv_out = c2.simulate(pspec.models['svsim'])\n",
    "dm_out = c2.simulate(pspec.models['dmsim'])\n",
    "for key in sv_out:\n",
    "    assert(abs(sv_out[key] - dm_out[key]) < 1e-10)\n",
    "assert(abs(np.cos(2*theta/2)**2 - sv_out['0'*n]) < 1e-10) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Change the Gi gates on all qubits to be small X rotations\n",
    "for glabel in list(pspec.models['svsim'].gates.keys()):\n",
    "    if glabel.name == 'Gi':\n",
    "        unitary = expm(-1j*theta*np.array([[0.,1.],[1.,0.]])/2) \n",
    "        pspec.models['svsim'][glabel] = unitary\n",
    "        pspec.models['dmsim'][glabel] = pygsti.unitary_to_pauligate(unitary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# The probability of an error-free circuit is now the nth power of the previous probability. Check that is the case\n",
    "sv_out = c2.simulate(pspec.models['svsim'])\n",
    "dm_out = c2.simulate(pspec.models['dmsim'])\n",
    "for key in sv_out:\n",
    "    assert(abs(sv_out[key] - dm_out[key]) < 1e-10)\n",
    "assert(abs(np.cos(2*theta/2)**(2*n) - sv_out['0'*n]) < 1e-10) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test the simulators using RB circuits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Let's get some RB circuits to test the simulators on\n",
    "sectors = [[pygsti.objects.Label('Gcnot',(0,1))],[pygsti.objects.Label('Gcnot',(1,2))]]\n",
    "cprb = prb.sample_prb_circuit(pspec,10,sampler='sectors',sampler_args={'sectors':sectors,'two_qubit_prob':0.5},return_partitioned=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Let's test the the target output probability is 1. for the '0000...' state\n",
    "prbout_sv_target = cprb.simulate(pspec.models['target'])\n",
    "prbout_dm_target = cprb.simulate(pspec.models['dm_target'])\n",
    "for key in prbout_sv_target:\n",
    "    assert(abs(prbout_sv_target[key] - prbout_dm_target[key]) < 1e-10)\n",
    "assert(abs(prbout_sv_target ['0'*n] - 1.) < 1e-10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Let's create a more complicated unitary error model\n",
    "\n",
    "thetarange = np.linspace(0,theta,n)\n",
    "cnottheta = theta\n",
    "for glabel in list(pspec.models['svsim'].gates.keys()):\n",
    "    # Change the Gi and Gh gates on all qubits to have a small X rotation error, increasing with qubit number\n",
    "    if glabel.name == 'Gi':\n",
    "        unitary = expm(-1j*thetarange[glabel.qubits[0]]*np.array([[0.,1.],[1.,0.]])/2) \n",
    "        pspec.models['svsim'][glabel] = unitary\n",
    "        pspec.models['dmsim'][glabel] = pygsti.unitary_to_pauligate(unitary)\n",
    "    if glabel.name == 'Gh':\n",
    "        unitary = np.dot(expm(-1j*thetarange[glabel.qubits[0]]*np.array([[0.,1.],[1.,0.]])/2),np.array([[1.,1.],[1.,-1.]],complex)/np.sqrt(2))\n",
    "        pspec.models['svsim'][glabel] = unitary\n",
    "        pspec.models['dmsim'][glabel] = pygsti.unitary_to_pauligate(unitary)\n",
    "    #Change the Gp to have a small Z rotation error, decreasing with qubit number\n",
    "    if glabel.name == 'Gp':\n",
    "        unitary = expm(-1j*(thetarange[n-1-glabel.qubits[0]]+np.pi/2)*np.array([[1.,0.],[0.,-1.]])/2) \n",
    "        pspec.models['svsim'][glabel] = unitary\n",
    "        pspec.models['dmsim'][glabel] = pygsti.unitary_to_pauligate(unitary)\n",
    "    # CNOT gates are followd by a control Z rotation with a small rotation angle.\n",
    "    if glabel.name == 'Gnot':\n",
    "        unitary = np.dot(np.array([[1.,0.,0.,0.],[0.,1.,0.,0.],[0.,0.,1.,0.],[0.,0.,0.,exp(cnottheta*1j)]],complex),np.array([[1.,0.,0.,0.],[0.,1.,0.,0.],[0.,0.,0.,1.],[0.,0.,1.,0.]],complex))\n",
    "        pspec.models['svsim'][glabel] = unitary\n",
    "        pspec.models['dmsim'][glabel] = pygsti.unitary_to_pauligate(unitary)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Let's replace the 'I' gates in the circuit with the 'Gi' gate\n",
    "cprb.replace_gatename('I','Gi')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The actual survival probability: 0.8949392394072598\n",
      "An order-of-magnitude estimate of the survival probabilitiy: 0.942499338239\n"
     ]
    }
   ],
   "source": [
    "# Let's test it. Let's check the simulators agree and the the survival probability is plausable\n",
    "\n",
    "prbout_sv = cprb.simulate(pspec.models['svsim'])\n",
    "prbout_dm = cprb.simulate(pspec.models['dmsim'])\n",
    "for key in prbout_sv:\n",
    "    assert(abs(prbout_sv[key] - prbout_dm[key]) < 1e-10)\n",
    "\n",
    "# These numbers are *not* guaranteed to always be even close (because, coherent addition/cancelation can happen\n",
    "# depending on the sampled circuit). But in a typical example they will likely be of the same order of magnitude.\n",
    "\n",
    "print('The actual survival probability:', prbout_sv['0'*n])\n",
    "print('An order-of-magnitude estimate of the survival probabilitiy:', np.cos(theta/2)**(2*n*cprb.depth()))"
   ]
  }
 ],
 "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}