{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyGSTi CHP Object Test\n",
    "\n",
    "<font color='red'>This notebook is under construction and will have more description in the near future.</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function #python 2 & 3 compatibility\n",
    "import pygsti\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from pygsti.modelmembers.operations import LinearOperator, StaticStandardOp, StochasticNoiseOp, DepolarizeOp, ComposedOp, EmbeddedOp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LinearOperator and StaticStandardOp\n",
    "\n",
    "Now with 'chp' evotype."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "StaticStandardOp with name Gxpi and evotype chp\n",
      "\n",
      "['h 0', 'p 0', 'p 0', 'h 0']\n"
     ]
    }
   ],
   "source": [
    "Gx = StaticStandardOp('Gxpi', evotype='chp')\n",
    "print(Gx)\n",
    "print(Gx._rep._chp_ops())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'chp'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Gx.evotype.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Can also make custom CHP operations\n",
    "# Here I'm making a (deterministic) Hadamard on qubit 0 and CNOT on qubits 1 and 2\n",
    "rep = pygsti.evotypes.chp.opreps.OpRep(['h 0', 'c 1 2'], state_space=3)\n",
    "c = LinearOperator(rep, 'chp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<pygsti.modelmembers.operations.linearop.LinearOperator object at 0x000001A47A83BFD0>\n",
      "['h 0', 'c 1 2']\n"
     ]
    }
   ],
   "source": [
    "print(c)\n",
    "print(c._rep._chp_ops())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "StaticStandardOp with name Gc20 and evotype chp\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(StaticStandardOp('Gc20', evotype='chp'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## StochasticNoiseOp and DepolarizeOp\n",
    "\n",
    "Now with 'chp' evotype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stochastic noise operation map with state space = QubitSpace((0,)), num params = 3\n",
      "Rates: [0.5 0.1 0.1]\n",
      "\n",
      "['h 0', 'p 0', 'p 0', 'h 0']\n",
      "['h 0', 'p 0', 'p 0', 'h 0']\n",
      "[]\n",
      "['h 0', 'p 0', 'p 0', 'h 0']\n"
     ]
    }
   ],
   "source": [
    "nqubits = 1\n",
    "scop = StochasticNoiseOp(nqubits, basis='pp', evotype='chp', initial_rates=[0.5, 0.1, 0.1], seed_or_state=2021)\n",
    "print(scop)\n",
    "for _ in range(4):\n",
    "    print(scop._rep._chp_ops())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Depolarize noise operation map with dim = 4, num params = 1\n",
      "Strength: [0.7]\n",
      "\n",
      "['p 0', 'h 0', 'p 0', 'p 0', 'h 0', 'p 0', 'p 0', 'p 0']\n",
      "['p 0', 'h 0', 'p 0', 'p 0', 'h 0', 'p 0', 'p 0', 'p 0']\n",
      "[]\n",
      "['h 0', 'p 0', 'p 0', 'h 0']\n"
     ]
    }
   ],
   "source": [
    "nqubits = 1\n",
    "dop = DepolarizeOp(nqubits, basis='pp', evotype='chp', initial_rate=0.7, seed_or_state=2021)\n",
    "print(dop)\n",
    "for _ in range(4): # With seed 2021, pulls Z, I (no output), X, Y\n",
    "    print(dop._rep._chp_ops())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ComposedOp + EmbeddedOp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "StaticStandardOp with name Gzpi and evotype chp\n",
      "Factor 1:\n",
      "StaticStandardOp with name Gxpi and evotype chp\n",
      "\n",
      "['p 0', 'p 0', 'h 0', 'p 0', 'p 0', 'h 0']\n"
     ]
    }
   ],
   "source": [
    "# ComposedOp\n",
    "Gzx_composed = ComposedOp([StaticStandardOp('Gzpi', evotype='chp'), StaticStandardOp('Gxpi', evotype='chp')])\n",
    "print(Gzx_composed)\n",
    "print(Gzx_composed._rep._chp_ops())\n",
    "#print(Gzx_composed.get_chp_str([2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q0',) space\n",
      "StaticStandardOp with name Gxpi and evotype chp\n",
      "\n",
      "['h 0', 'p 0', 'p 0', 'h 0']\n"
     ]
    }
   ],
   "source": [
    "# EmbeddedOp\n",
    "Gxi_embedded = EmbeddedOp(['Q0', 'Q1'], ['Q0'], StaticStandardOp('Gxpi', evotype='chp'))\n",
    "print(Gxi_embedded)\n",
    "print(Gxi_embedded._rep._chp_ops())\n",
    "#print(Gxi_embedded.get_chp_str([5,7]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q1',) space\n",
      "StaticStandardOp with name Gxpi and evotype chp\n",
      "\n",
      "['h 1', 'p 1', 'p 1', 'h 1']\n"
     ]
    }
   ],
   "source": [
    "Gix_embedded = EmbeddedOp(['Q0', 'Q1'], ['Q1'], StaticStandardOp('Gxpi', evotype='chp'))\n",
    "print(Gix_embedded)\n",
    "print(Gix_embedded._rep._chp_ops())\n",
    "#print(Gix_embedded.get_chp_str([5,7]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedded operation with full dimension 256 and state space QubitSpace(('Q0', 'Q1', 'Q2', 'Q3'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q1',) space\n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "StaticStandardOp with name Gzpi and evotype chp\n",
      "Factor 1:\n",
      "StaticStandardOp with name Gxpi and evotype chp\n",
      "\n",
      "['p 1', 'p 1', 'h 1', 'p 1', 'p 1', 'h 1']\n"
     ]
    }
   ],
   "source": [
    "# EmbeddedOp made of ComposedOps\n",
    "Gzx_comp_embed = EmbeddedOp(['Q0', 'Q1', 'Q2', 'Q3'], ['Q1'], Gzx_composed)\n",
    "print(Gzx_comp_embed)\n",
    "print(Gzx_comp_embed._rep._chp_ops())\n",
    "#print(Gzx_comp_embed.get_chp_str([5, 6, 7, 8]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CHPForwardSimulator + Explicit Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "#This is the directory where the chp directory lives during github testing, replace this with the \n",
    "#correct directory for your personal installation\n",
    "pygsti.evotypes.chp.chpexe = 'chp'\n",
    "sim = pygsti.forwardsims.WeakForwardSimulator(shots=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rho0 = Computational Z-basis state vec for 2 qubits w/z-values: [0 0]\n",
      "\n",
      "Mdefault = Computational(Z)-basis POVM on 2 qubits and filter None\n",
      "\n",
      "\n",
      "Gii = \n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q0',) space\n",
      "StaticStandardOp with name Gi and evotype chp\n",
      "Factor 1:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q1',) space\n",
      "StaticStandardOp with name Gi and evotype chp\n",
      "\n",
      "\n",
      "Gxi = \n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q0',) space\n",
      "StaticStandardOp with name Gxpi and evotype chp\n",
      "Factor 1:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q1',) space\n",
      "StaticStandardOp with name Gi and evotype chp\n",
      "\n",
      "\n",
      "Gix = \n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q0',) space\n",
      "StaticStandardOp with name Gi and evotype chp\n",
      "Factor 1:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q1',) space\n",
      "StaticStandardOp with name Gxpi and evotype chp\n",
      "\n",
      "\n",
      "Gxx = \n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q0',) space\n",
      "StaticStandardOp with name Gxpi and evotype chp\n",
      "Factor 1:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q1',) space\n",
      "StaticStandardOp with name Gxpi and evotype chp\n",
      "\n",
      "\n",
      "Gyi = \n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q0',) space\n",
      "StaticStandardOp with name Gypi and evotype chp\n",
      "Factor 1:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q1',) space\n",
      "StaticStandardOp with name Gi and evotype chp\n",
      "\n",
      "\n",
      "Giy = \n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q0',) space\n",
      "StaticStandardOp with name Gi and evotype chp\n",
      "Factor 1:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q1',) space\n",
      "StaticStandardOp with name Gypi and evotype chp\n",
      "\n",
      "\n",
      "Gyy = \n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q0',) space\n",
      "StaticStandardOp with name Gypi and evotype chp\n",
      "Factor 1:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace(('Q0', 'Q1'))\n",
      " that embeds the following 4-dimensional operation into acting on the ('Q1',) space\n",
      "StaticStandardOp with name Gypi and evotype chp\n",
      "\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#Initialize an empty Model object\n",
    "model = pygsti.models.ExplicitOpModel(['Q0', 'Q1'], simulator=sim, evotype='chp')\n",
    "\n",
    "def make_2Q_op(name0, name1):\n",
    "    return ComposedOp([\n",
    "        EmbeddedOp(['Q0', 'Q1'], ['Q0'], StaticStandardOp(name0, evotype='chp')),\n",
    "        EmbeddedOp(['Q0', 'Q1'], ['Q1'], StaticStandardOp(name1, evotype='chp')),\n",
    "    ])\n",
    "\n",
    "#Populate the Model object with states, effects, gates\n",
    "# For CHP, prep must be all-zero ComputationalSPAMVec\n",
    "# and povm must be ComputationalBasisPOVM\n",
    "model['rho0'] = pygsti.modelmembers.states.ComputationalBasisState([0, 0], evotype='chp')\n",
    "model['Mdefault'] = pygsti.modelmembers.povms.ComputationalBasisPOVM(2, evotype='chp')\n",
    "\n",
    "model['Gii'] = make_2Q_op('Gi', 'Gi')\n",
    "model['Gxi'] = make_2Q_op('Gxpi', 'Gi')\n",
    "model['Gix'] = make_2Q_op('Gi', 'Gxpi')\n",
    "model['Gxx'] = make_2Q_op('Gxpi', 'Gxpi')\n",
    "model['Gyi'] = make_2Q_op('Gypi', 'Gi')\n",
    "model['Giy'] = make_2Q_op('Gi', 'Gypi')\n",
    "model['Gyy'] = make_2Q_op('Gypi', 'Gypi')\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OutcomeLabelDict([(('01',), 1.0000000000000007)])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "circ = pygsti.circuits.Circuit(['Gix'])\n",
    "model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OutcomeLabelDict([(('11',), 1.0000000000000007)])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "circ = pygsti.circuits.Circuit(['Gix', 'Gxi'])\n",
    "model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OutcomeLabelDict([(('11',), 1.0000000000000007)])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "circ = pygsti.circuits.Circuit(['rho0', 'Gxx', 'Mdefault'])\n",
    "model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced State Prep and Measurement\n",
    "\n",
    "<font color='red'>TODO: This section does not work due to non-CHP related issues. Come back to this once other issues are fixed.</font>\n",
    "\n",
    "### State Prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OutcomeLabelDict([(('01',), 1.0000000000000007)])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Initialize an empty Model object\n",
    "prep01_model = pygsti.models.ExplicitOpModel(['Q0', 'Q1'], simulator=sim, evotype='chp')\n",
    "\n",
    "# Make a ComputationalSPAMVec with one bit in 1 state\n",
    "prep01_model.preps['rho0'] = pygsti.modelmembers.states.ComputationalBasisState([0, 1], evotype='chp')\n",
    "prep01_model.povms['Mdefault'] = pygsti.modelmembers.povms.ComputationalBasisPOVM(2, evotype='chp')\n",
    "\n",
    "circ = pygsti.circuits.Circuit([])\n",
    "prep01_model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OutcomeLabelDict([(('01',), 0.11999999999999998),\n",
       "                  (('00',), 0.8800000000000006)])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Initialize an empty Model object\n",
    "prep00noise_model = pygsti.models.ExplicitOpModel(['Q0', 'Q1'], simulator=sim, evotype='chp')\n",
    "\n",
    "# Make a ComposedSPAMVec where second qubit has X error\n",
    "rho0 = pygsti.modelmembers.states.ComposedState(\n",
    "    pygsti.modelmembers.states.ComputationalBasisState([0, 0], evotype='chp'), # Pure SPAM vec is 00 state\n",
    "    make_2Q_op('Gi', 'Gxpi2')) # Second qubit has X(pi/2) error (partial flip on qubit 1)\n",
    "\n",
    "prep00noise_model['rho0'] = rho0\n",
    "prep00noise_model['Mdefault'] = pygsti.modelmembers.povms.ComputationalBasisPOVM(2, 'chp')\n",
    "\n",
    "circ = pygsti.circuits.Circuit([])\n",
    "prep00noise_model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OutcomeLabelDict([(('10',), 0.46000000000000024),\n",
       "                  (('11',), 0.5400000000000003)])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Initialize an empty Model object\n",
    "prep11noise_model = pygsti.models.ExplicitOpModel(['Q0', 'Q1'], simulator=sim, evotype='chp')\n",
    "\n",
    "# Make a ComposedSPAMVec where second qubit has X error AND is initialized to 1 state\n",
    "rho0 = pygsti.modelmembers.states.ComposedState(\n",
    "    pygsti.modelmembers.states.ComputationalBasisState([1, 1], evotype='chp'), # Pure SPAM vec is 00 state\n",
    "    make_2Q_op('Gi', 'Gxpi2')) # Second qubit has X(pi/2) error (partial flip on qubit 1)\n",
    "\n",
    "prep11noise_model['rho0'] = rho0\n",
    "prep11noise_model['Mdefault'] = pygsti.modelmembers.povms.ComputationalBasisPOVM(2, 'chp')\n",
    "\n",
    "circ = pygsti.circuits.Circuit([])\n",
    "prep11noise_model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Measurement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pygsti.evotypes.chp.opreps.OpRepComposed at 0x132840040>"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "make_2Q_op('Gi', 'Gxpi2')._rep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "##Initialize an empty Model object\n",
    "#povm01_model = pygsti.models.ExplicitOpModel(['Q0', 'Q1'], simulator=sim, evotype='chp')\n",
    "#\n",
    "## Make a measurement with a bitflip error on qubit 1\n",
    "#povm01_model.preps['rho0'] = pygsti.modelmembers.states.ComputationalBasisState([0, 1], evotype='chp')\n",
    "#povm01_model.povms['Mdefault'] = pygsti.modelmembers.povms.ComposedPOVM(\n",
    "#    make_2Q_op('Gi', 'Gxpi2'),\n",
    "#    pygsti.modelmembers.povms.ComputationalBasisPOVM(2, evotype='chp'),\n",
    "#    mx_basis='pp')\n",
    "#\n",
    "#povm01_model._primitive_povm_label_dict['Mdefault'] = povm01_model['Mdefault']\n",
    "#\n",
    "#circ = pygsti.circuits.Circuit([])\n",
    "#povm01_model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CHPForwardSimulator + LocalNoiseModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Define stochastic Pauli noise operators\n",
    "# Note that the probabilities here are the \"error rates\" that would be model parameters (currently just static)\n",
    "noise_1q = StochasticNoiseOp(1, basis='pp', evotype='chp', initial_rates=[0.1, 0.01, 0.01], seed_or_state=2021)\n",
    "\n",
    "# Also need two-qubit version\n",
    "# Here we just make it independent stochastic Pauli noise\n",
    "noise_2q = ComposedOp([EmbeddedOp([0, 1], [0], noise_1q), EmbeddedOp([0, 1], [1], noise_1q)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 2: Define gate dict of noisy gates\n",
    "# Using equivalent of XYICNOT modelpack\n",
    "gatedict = {}\n",
    "gatedict['Gi'] = noise_1q\n",
    "gatedict['Gx'] = ComposedOp([StaticStandardOp('Gxpi', evotype='chp'), noise_1q])\n",
    "gatedict['Gy'] = ComposedOp([StaticStandardOp('Gypi', evotype='chp'), noise_1q])\n",
    "# Note that first Gcnot is now key in model, whereas second Gcnot is a standard gatename known to CHPOp constructor\n",
    "gatedict['Gcnot'] = ComposedOp([StaticStandardOp('Gcnot', evotype='chp'), noise_2q])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pygsti.models.localnoisemodel import LocalNoiseModel\n",
    "from pygsti.modelmembers.states import ComputationalBasisState\n",
    "from pygsti.modelmembers.povms import ComputationalBasisPOVM\n",
    "from pygsti.processors import QubitProcessorSpec\n",
    "\n",
    "pspec = QubitProcessorSpec(4, list(gatedict.keys()), geometry='line',\n",
    "                           availability={'Gcnot': [(0,1),(1,2),(2,3)]})\n",
    "\n",
    "rho0 = ComputationalBasisState([0,]*4, evotype='chp')\n",
    "Mdefault = ComputationalBasisPOVM(4, evotype='chp')\n",
    "\n",
    "ln_model = LocalNoiseModel(pspec, gatedict=gatedict, prep_layers=[rho0], povm_layers=[Mdefault],\n",
    "                           simulator=sim, evotype='chp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "State prep building blocks (.prep_blks):\n",
      " layers :  rho0\n",
      "\n",
      "POVM building blocks (.povm_blks):\n",
      " layers :  M0\n",
      "\n",
      "Operation building blocks (.operation_blks):\n",
      " gates :  Gi, Gx, Gy, Gcnot\n",
      " layers :  {auto_global_idle}, Gi:0, Gi:1, Gi:2, Gi:3, Gx:0, Gx:1, Gx:2, Gx:3, Gy:0, Gy:1, Gy:2, Gy:3, Gcnot:0:1, Gcnot:1:2, Gcnot:2:3\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Step 4: Profit?? Worked way too quickly...\n",
    "def print_implicit_model_blocks(mdl, showSPAM=False):\n",
    "    if showSPAM:\n",
    "        print('State prep building blocks (.prep_blks):')\n",
    "        for blk_lbl,blk in mdl.prep_blks.items():\n",
    "            print(\" \" + blk_lbl, \": \", ', '.join(map(str,blk.keys())))\n",
    "        print()\n",
    "\n",
    "        print('POVM building blocks (.povm_blks):')\n",
    "        for blk_lbl,blk in mdl.povm_blks.items():\n",
    "            print(\" \"  + blk_lbl, \": \", ', '.join(map(str,blk.keys())))\n",
    "        print()\n",
    "    \n",
    "    print('Operation building blocks (.operation_blks):')\n",
    "    for blk_lbl,blk in mdl.operation_blks.items():\n",
    "        print(\" \" + blk_lbl, \": \", ', '.join(map(str,blk.keys())))\n",
    "    print()\n",
    "\n",
    "print_implicit_model_blocks(ln_model, showSPAM=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Computational Z-basis state vec for 4 qubits w/z-values: [0 0 0 0]\n"
     ]
    }
   ],
   "source": [
    "print(ln_model.prep_blks['layers']['rho0'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "StaticStandardOp with name Gxpi and evotype chp\n",
      "Factor 1:\n",
      "Stochastic noise operation map with state space = QubitSpace((0,)), num params = 3\n",
      "Rates: [0.1  0.01 0.01]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(ln_model.operation_blks['gates']['Gx'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedded operation with full dimension 256 and state space QubitSpace((0, 1, 2, 3))\n",
      " that embeds the following 16-dimensional operation into acting on the (1, 2) space\n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "StaticStandardOp with name Gcnot and evotype chp\n",
      "Factor 1:\n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace((0, 1))\n",
      " that embeds the following 4-dimensional operation into acting on the (0,) space\n",
      "Stochastic noise operation map with state space = QubitSpace((0,)), num params = 3\n",
      "Rates: [0.1  0.01 0.01]\n",
      "Factor 1:\n",
      "Embedded operation with full dimension 16 and state space QubitSpace((0, 1))\n",
      " that embeds the following 4-dimensional operation into acting on the (1,) space\n",
      "Stochastic noise operation map with state space = QubitSpace((0,)), num params = 3\n",
      "Rates: [0.1  0.01 0.01]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "Gcnot_layer_op = ln_model.operation_blks['layers']['Gcnot', 1, 2]\n",
    "print(ln_model.operation_blks['layers']['Gcnot', 1, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OutcomeLabelDict([(('0100',), 0.9000000000000006),\n",
       "                  (('0000',), 0.09999999999999999)])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Step 5: Actually run circuits with local noise model\n",
    "circ = pygsti.circuits.Circuit([('Gx', 1)], num_lines=4)\n",
    "ln_model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OutcomeLabelDict([(('0110',), 0.7100000000000004),\n",
       "                  (('0000',), 0.13999999999999999),\n",
       "                  (('0100',), 0.07),\n",
       "                  (('0010',), 0.08)])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "circ = pygsti.circuits.Circuit([('Gx', 1), ('Gcnot', 1, 2)], num_lines=4)\n",
    "ln_model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Could also define correlated noise for 2-qubit error?\n",
    "pp = pygsti.baseobjs.Basis.cast('pp', 16)\n",
    "rates_2q = [0.01,]*15\n",
    "rates_2q[pp.labels.index('XX')] = 0.1 # Set XX to much higher\n",
    "\n",
    "noise_2q_correlated = StochasticNoiseOp(2, basis='pp', evotype='chp', initial_rates=rates_2q, seed_or_state=2021)\n",
    "\n",
    "gatedict = {}\n",
    "gatedict['Gi'] = noise_1q\n",
    "gatedict['Gx'] = ComposedOp([StaticStandardOp('Gxpi', evotype='chp'), noise_1q])\n",
    "gatedict['Gy'] = ComposedOp([StaticStandardOp('Gypi', evotype='chp'), noise_1q])\n",
    "# Note that first Gcnot is now key in model, whereas second Gcnot is a standard gatename known to CHPOp constructor\n",
    "gatedict['Gcnot'] = ComposedOp([StaticStandardOp('Gcnot', evotype='chp'), noise_2q_correlated])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho0 = ComputationalBasisState([0,]*4, evotype='chp')\n",
    "Mdefault = ComputationalBasisPOVM(4, evotype='chp')\n",
    "\n",
    "sim = pygsti.forwardsims.WeakForwardSimulator(shots=100)\n",
    "\n",
    "ln_model_corr = LocalNoiseModel(pspec, gatedict=gatedict, prep_layers=[rho0], povm_layers=[Mdefault],\n",
    "                                simulator=sim, evotype='chp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedded operation with full dimension 256 and state space QubitSpace((0, 1, 2, 3))\n",
      " that embeds the following 16-dimensional operation into acting on the (1, 2) space\n",
      "Composed operation of 2 factors:\n",
      "Factor 0:\n",
      "StaticStandardOp with name Gcnot and evotype chp\n",
      "Factor 1:\n",
      "Stochastic noise operation map with state space = QubitSpace((0, 1)), num params = 15\n",
      "Rates: [0.01 0.01 0.01 0.01 0.01 0.1  0.01 0.01 0.01 0.01 0.01 0.01 0.01 0.01\n",
      " 0.01]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Now the CNOT gates have a 2-qubit stochastic gate instead of independent 1-qubit ones\n",
    "print(ln_model_corr.operation_blks['layers']['Gcnot', 1, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OutcomeLabelDict([(('0000',), 0.10999999999999999),\n",
       "                  (('0100',), 0.8900000000000006)])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "circ = pygsti.circuits.Circuit([('Gx', 1)], num_lines=4)\n",
    "ln_model_corr.probabilities(circ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OutcomeLabelDict([(('0110',), 0.7000000000000004),\n",
       "                  (('0000',), 0.21000000000000005),\n",
       "                  (('0010',), 0.05),\n",
       "                  (('0100',), 0.04)])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "circ = pygsti.circuits.Circuit([('Gx', 1), ('Gcnot', 1, 2)], num_lines=4)\n",
    "ln_model_corr.probabilities(circ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Crosstalk-Free Model Construction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import pygsti.models.modelconstruction as mc\n",
    "#\n",
    "#sim = pygsti.forwardsims.WeakForwardSimulator(shots=100, base_seed=2021)\n",
    "#\n",
    "#pspec = QubitProcessorSpec(4, ['Gi', 'Gxpi', 'Gypi', 'Gcnot'], availability={'Gcnot': [(0,1),(1,2),(2,3)]})\n",
    "#\n",
    "## Use the same 2-qubit stochastic noise for CNOT as above\n",
    "#ctf_model = mc.create_crosstalk_free_model(pspec,\n",
    "#    depolarization_strengths={'Gi': 0.1, 'Gxpi': 0.1},\n",
    "#    stochastic_error_probs={'Gypi': [0.1, 0.1, 0.1], 'Gcnot': rates_2q},\n",
    "#    simulator=sim, evotype='chp')\n",
    "#\n",
    "#print_implicit_model_blocks(ctf_model, showSPAM=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for name, gate in ctf_model.operation_blks['gates'].items():\n",
    "#    print(f'Gate {name}')\n",
    "#    print(gate)\n",
    "#    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "#circ = pygsti.circuits.Circuit([('Gxpi', 1)], num_lines=4)\n",
    "#ctf_model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "#circ = pygsti.circuits.Circuit([('Gxpi', 1), ('Gcnot', 1, 2)], num_lines=4)\n",
    "#ctf_model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Marginalized POVMs now work!\n",
    "#circ = pygsti.circuits.Circuit([('Gxpi', 1), ('Gcnot', 1, 2)])\n",
    "#ctf_model.probabilities(circ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's try a model with only readout error\n",
    "#sim = pygsti.forwardsims.CHPForwardSimulator(chpexe, shots=1000) # Bump up shots for better noise resolution\n",
    "#\n",
    "#ctf_povm_model = mc.create_crosstalk_free_model(pspec,\n",
    "#    stochastic_error_probs={'povm': [0.05, 0.0, 0.0]}, # 5% X error on prep\n",
    "#    simulator=sim, evotype='chp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "#circ = pygsti.circuits.Circuit([])\n",
    "#ctf_povm_model.probabilities(circ) # Expect about 80% all 0, 5% on weight one errors, 0.25% on weight 2, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}