{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pygsti\n",
    "from pygsti.extras import rb\n",
    "from pygsti.baseobjs import Label\n",
    "\n",
    "n_1 = 4\n",
    "glist = ['Gx','Gy','Gcnot']\n",
    "pspec_1 = pygsti.obj.ProcessorSpec(n_1,glist,verbosity=0)\n",
    "\n",
    "n_2 = 3\n",
    "glist = ['Gxpi','Gypi','Gzpi','Gh','Gp','Gcphase']\n",
    "availability = {'Gcphase':[(0,1),(1,2)]}\n",
    "pspec_2 = pygsti.obj.ProcessorSpec(n_2,glist,availability=availability,verbosity=0)\n",
    "\n",
    "#Fix processor specs to work with n=1 qubits.\n",
    "#n = 1\n",
    "#glist = ['Gh','Gp']\n",
    "#pspec_3 = pygsti.obj.ProcessorSpec(n=1,glist,verbosity=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#pspec_1.models['clifford'].gates.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# --------------------------------------- #\n",
    "# --- Test the circuit layer samplers --- #\n",
    "# --------------------------------------- #\n",
    "\n",
    "# Tests for the sampling by pairs function\n",
    "layer = rb.sample.circuit_layer_by_pairings(pspec_1, twoQprob=0.0, oneQgatenames='all', twoQgatenames='all', \n",
    "                                  gatesetname = 'clifford')\n",
    "assert(len(layer) == n_1)\n",
    "layer = rb.sample.circuit_layer_by_pairings(pspec_1, twoQprob=1.0, oneQgatenames='all', twoQgatenames='all', \n",
    "                                  gatesetname = 'clifford')\n",
    "assert(len(layer) == n_1//2)\n",
    "layer = rb.sample.circuit_layer_by_pairings(pspec_1, twoQprob=0.0, oneQgatenames=['Gx',], twoQgatenames='all', \n",
    "                                  gatesetname = 'target')\n",
    "assert(layer[0].name == 'Gx')\n",
    "\n",
    "layer = rb.sample.circuit_layer_by_Qelimination(pspec_2, twoQprob=0.0, oneQgates='all', twoQgates='all',\n",
    "                                                gatesetname='clifford')\n",
    "assert(len(layer) == n_2)\n",
    "layer = rb.sample.circuit_layer_by_Qelimination(pspec_2, twoQprob=1.0, oneQgates='all', twoQgates='all',\n",
    "                                                gatesetname='clifford')\n",
    "assert(len(layer) == (n_2 % 2) + n_2//2)\n",
    "layer = rb.sample.circuit_layer_by_pairings(pspec_1, twoQprob=0.0, oneQgatenames=['Gxpi'], twoQgatenames='all', \n",
    "                                  gatesetname = 'target')\n",
    "assert(layer[0].name == 'Gxpi')\n",
    "\n",
    "# Tests for the sampling by sectors function\n",
    "C01 = Label('Gcnot',(0,1))\n",
    "C23 = Label('Gcnot',(2,3))\n",
    "sectors = [[],[C01,C23]]\n",
    "layer = rb.sample.circuit_layer_by_sectors(pspec_1, sectors, sectorsprob='uniform', twoQprob=1.0, \n",
    "                                           oneQgatenames='all', gatesetname='clifford')\n",
    "assert(len(layer) == n_1 or len(layer) == n_1//2)\n",
    "layer = rb.sample.circuit_layer_by_sectors(pspec_1, sectors, sectorsprob=[0.,1.], twoQprob=1.0, \n",
    "                                           oneQgatenames='all', gatesetname='clifford')\n",
    "assert(len(layer) == n_1//2)\n",
    "layer = rb.sample.circuit_layer_by_sectors(pspec_1, sectors, sectorsprob=[1.,0.], twoQprob=1.0, \n",
    "                                           oneQgatenames=['Gx',], gatesetname='clifford')\n",
    "assert(len(layer) == n_1)\n",
    "assert(layer[0].name == 'Gx')\n",
    "layer = rb.sample.circuit_layer_by_sectors(pspec_1, sectors, sectorsprob=[0.25,0.75], twoQprob=0.5, \n",
    "                                           oneQgatenames='all', gatesetname='clifford')\n",
    "\n",
    "# Tests for the sampling a layer of 1Q gates.\n",
    "layer = rb.sample.circuit_layer_of_1Q_gates(pspec_1, oneQgatenames='all', pdist='uniform',\n",
    "                                            gatesetname='clifford')\n",
    "assert(len(layer) == n_1)\n",
    "layer = rb.sample.circuit_layer_of_1Q_gates(pspec_1, oneQgatenames=['Gx','Gy'], pdist=[1.,0.],\n",
    "                                            gatesetname='clifford')\n",
    "assert(layer[0].name == 'Gx')\n",
    "layer = rb.sample.circuit_layer_of_1Q_gates(pspec_1, oneQgatenames=['Gx'], pdist=[3.,],\n",
    "                                            gatesetname='clifford')\n",
    "assert(layer[0].name == 'Gx')\n",
    "layer = rb.sample.circuit_layer_of_1Q_gates(pspec_1, oneQgatenames=['Gx'], pdist='uniform',\n",
    "                                            gatesetname='clifford')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "circuit = rb.sample.circuit(pspec_1, length=100, sampler='Qelimination')\n",
    "assert(circuit.depth() == 100)\n",
    "circuit = rb.sample.circuit(pspec_2, length=100, sampler='Qelimination', samplerargs=[0.1,], addlocal = True)\n",
    "assert(circuit.depth() == 201)\n",
    "assert(len(circuit.get_circuit_layer(0)) == n_2)\n",
    "circuit = rb.sample.circuit(pspec_1, length=100, sampler='pairings')\n",
    "circuit = rb.sample.circuit(pspec_1, length=10, sampler='pairings', samplerargs=[0.1,['Gx',]])\n",
    "\n",
    "circuit = rb.sample.circuit(pspec_1, length=100, sampler='sectors', samplerargs=[sectors])\n",
    "circuit = rb.sample.circuit(pspec_1, length=100, sampler='sectors', samplerargs=[sectors,[0.1,0.2],0.1], \n",
    "                            addlocal = True, lsargs=[['Gx',]])\n",
    "assert(circuit.depth() == 201)\n",
    "circuit = rb.sample.circuit(pspec_1, length=5, sampler='local')\n",
    "assert(circuit.size() == n_1*5)\n",
    "circuit = rb.sample.circuit(pspec_1, length=5, sampler='local',samplerargs=[['Gx']])\n",
    "assert(circuit.line_items[0][0].name == 'Gx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "circuit = rb.sample.circuit(pspec_1, length=5, sampler='local')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(circuit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "order = [1,0,3,2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "circuit.relabel_qubits(order)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(circuit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}