{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial on creating a custom operator (e.g. gate)\n",
    "\n",
    "This tutorial demonstrates the process of creating your own gate operation.  One can view gate (or layer) operations in pyGSTi as simply parameterized process matrices: a mapping that associates with any given set of parameter values a process matrix.  This mapping is encapsulated by a `LinearOperator`-derived class in pyGSTi, and in addition to using the ones included with pyGSTi (e.g. `FullDenseOp`, see the [Operator tutorial](Operators.ipynb) for more examples) you're free to make your own.  That's exactly what we'll be doing here.\n",
    "\n",
    "There are lots of good reasons for doing this, the foremost being that you have a specific way you want to model a gate operation that is specific to your system's physics and not captured by pyGSTi's more generic built-in operation classes.  You also may want to make an operation whose parameters are exactly the \"knobs\" that you have access to in the lab.  Whatever the reason, pyGSTi has been designed to make the creation of custom operator types easy and straightforward.\n",
    "\n",
    "In this example, we'll be creating a custom 1-qubit gate operation.  It will be a $X(\\pi/2)$-rotation that may have some amount of depolarization and \"on-axis\" overrotation, but no other imperfections.  Thus, it will only have to parameters: the depolarization and the overrotation amounts.\n",
    "\n",
    "Here's a class which implements this operation.  The comments explain what different parts do."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygsti\n",
    "import numpy as np\n",
    "\n",
    "class MyXPi2Operator(pygsti.obj.DenseOperator):\n",
    "    def __init__(self):\n",
    "        #initialize with no noise\n",
    "        self.from_vector([0,0]) \n",
    "        super(MyXPi2Operator,self).__init__(self.base, \"densitymx\") # this is *super*-operator, so \"densitymx\"\n",
    "        \n",
    "    def num_params(self): \n",
    "        return 2 # we have two parameters\n",
    "    \n",
    "    def to_vector(self):\n",
    "        return np.array([self.depol_amt, self.over_rotation],'d') #our parameter vector\n",
    "        \n",
    "    def from_vector(self,v):\n",
    "        #initialize from parameter vector v\n",
    "        self.depol_amt = v[0]\n",
    "        self.over_rotation = v[1]\n",
    "        \n",
    "        theta = (np.pi/2 + self.over_rotation)/2\n",
    "        a = 1.0-self.depol_amt\n",
    "        b = a*2*np.cos(theta)*np.sin(theta)\n",
    "        c = a*(np.sin(theta)**2 - np.cos(theta)**2)\n",
    "        \n",
    "        # .base is a member of DenseOperator and is a numpy array that is \n",
    "        # the dense Pauli transfer matrix of this operator\n",
    "        self.base = np.array([[1,   0,   0,   0],\n",
    "                              [0,   a,   0,   0],\n",
    "                              [0,   0,   c,  -b],\n",
    "                              [0,   0,   b,   c]],'d')\n",
    "        \n",
    "    def transform(self, S):\n",
    "        # Update self with inverse(S) * self * S (used in gauge optimization)\n",
    "        raise NotImplementedError(\"MyXPi2Operator cannot be transformed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll add a `MyXPi2Operator` instance as the `\"Gx\"` gate in pyGSTi's standard {Idle, $X(\\pi/2)$, $Y(\\pi/2)$} model (see the [standard modules tutorial](StandardModules.ipynb) for more information on standard models)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rho0 = FullSPAMVec with dimension 4\n",
      " 0.71   0   0 0.71\n",
      "\n",
      "\n",
      "Mdefault = UnconstrainedPOVM with effect vectors:\n",
      "0: FullSPAMVec with dimension 4\n",
      " 0.71   0   0 0.71\n",
      "\n",
      "1: FullSPAMVec with dimension 4\n",
      " 0.71   0   0-0.71\n",
      "\n",
      "\n",
      "\n",
      "Gi = \n",
      "FullDenseOp with shape (4, 4)\n",
      " 1.00   0   0   0\n",
      "   0 1.00   0   0\n",
      "   0   0 1.00   0\n",
      "   0   0   0 1.00\n",
      "\n",
      "\n",
      "Gx = \n",
      "MyXPi2Operator with shape (4, 4)\n",
      " 1.00   0   0   0\n",
      "   0 1.00   0   0\n",
      "   0   0   0-1.00\n",
      "   0   0 1.00   0\n",
      "\n",
      "\n",
      "Gy = \n",
      "FullDenseOp with shape (4, 4)\n",
      " 1.00   0   0   0\n",
      "   0   0   0 1.00\n",
      "   0   0 1.00   0\n",
      "   0-1.00   0   0\n",
      "\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pygsti.construction import std1Q_XYI\n",
    "mdl = std1Q_XYI.target_model()\n",
    "mdl.operations['Gx'] = MyXPi2Operator()\n",
    "print(mdl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, to demonstrate everything is working like it should, we'll optimize this model using Gate Set tomography (see the [GST overview tutorial](../../algorithms/GST-Overview.ipynb) for the details on what all this stuff does).  GST by default attempts to gauge optimize its final estimate to look like the target model (see the [gauge optimization tutorial](../../algorithms/advanced/GaugeOpt.ipynb) for details).  This would requires all of the operators in our model to implement the (gauge) `transform` method.  Because `MyXPi2Operator` doesn't, we tell GST not to perform any gauge optimization by setting `gaugeOptParams=False` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Circuit Creation ---\n",
      "   1282 sequences created\n",
      "   Dataset has 1282 entries: 1282 utilized, 0 requested sequences were missing\n",
      "--- Iterative MLGST: Iter 1 of 5  92 operation sequences ---: \n",
      "  --- Minimum Chi^2 GST ---\n",
      "  Sum of Chi^2 = 77.8909 (92 data params - 42 model params = expected mean of 50; p-value = 0.00700319)\n",
      "  Completed in 0.5s\n",
      "  2*Delta(log(L)) = 78.6371\n",
      "  Iteration 1 took 0.5s\n",
      "  \n",
      "--- Iterative MLGST: Iter 2 of 5  168 operation sequences ---: \n",
      "  --- Minimum Chi^2 GST ---\n",
      "  Sum of Chi^2 = 143.522 (168 data params - 42 model params = expected mean of 126; p-value = 0.136116)\n",
      "  Completed in 0.4s\n",
      "  2*Delta(log(L)) = 146.001\n",
      "  Iteration 2 took 0.5s\n",
      "  \n",
      "--- Iterative MLGST: Iter 3 of 5  450 operation sequences ---: \n",
      "  --- Minimum Chi^2 GST ---\n",
      "  Sum of Chi^2 = 439.592 (450 data params - 42 model params = expected mean of 408; p-value = 0.135269)\n",
      "  Completed in 1.5s\n",
      "  2*Delta(log(L)) = 444.637\n",
      "  Iteration 3 took 1.6s\n",
      "  \n",
      "--- Iterative MLGST: Iter 4 of 5  862 operation sequences ---: \n",
      "  --- Minimum Chi^2 GST ---\n",
      "  Sum of Chi^2 = 879.21 (862 data params - 42 model params = expected mean of 820; p-value = 0.0742732)\n",
      "  Completed in 1.8s\n",
      "  2*Delta(log(L)) = 887.725\n",
      "  Iteration 4 took 1.9s\n",
      "  \n",
      "--- Iterative MLGST: Iter 5 of 5  1282 operation sequences ---: \n",
      "  --- Minimum Chi^2 GST ---\n",
      "  Sum of Chi^2 = 1308.01 (1282 data params - 42 model params = expected mean of 1240; p-value = 0.0877443)\n",
      "  Completed in 1.8s\n",
      "  2*Delta(log(L)) = 1317.04\n",
      "  Iteration 5 took 1.9s\n",
      "  \n",
      "  Switching to ML objective (last iteration)\n",
      "  --- MLGST ---\n",
      "    Maximum log(L) = 657.167 below upper bound of -1.96182e+06\n",
      "      2*Delta(log(L)) = 1314.33 (1282 data params - 42 model params = expected mean of 1240; p-value = 0.0698154)\n",
      "    Completed in 0.4s\n",
      "  2*Delta(log(L)) = 1314.33\n",
      "  Final MLGST took 0.4s\n",
      "  \n",
      "Iterative MLGST Total Time: 6.7s\n"
     ]
    }
   ],
   "source": [
    "# Generate \"fake\" data from a depolarized version of the target (ideal) model\n",
    "maxLengths = [1,2,4,8,16]\n",
    "mdl_datagen = std1Q_XYI.target_model().depolarize(op_noise=0.01, spam_noise=0.001)\n",
    "listOfExperiments = pygsti.construction.make_lsgst_experiment_list(\n",
    "    mdl_datagen, std1Q_XYI.prepStrs, std1Q_XYI.effectStrs, std1Q_XYI.germs, maxLengths)\n",
    "ds = pygsti.construction.generate_fake_data(mdl_datagen, listOfExperiments, nSamples=1000,\n",
    "                                            sampleError=\"binomial\", seed=1234)\n",
    "\n",
    "#Run GST *without* gauge optimization\n",
    "results = pygsti.do_long_sequence_gst(ds, mdl, std1Q_XYI.prepStrs, std1Q_XYI.effectStrs,\n",
    "                                      std1Q_XYI.germs, maxLengths, gaugeOptParams=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**That's it!  We just ran GST with a custom operation.**\n",
    "\n",
    "Our `MyXPi2Operator`-containing model fits the data pretty well (compare the actual and expected $2\\Delta \\log \\mathcal{L}$ values printed above).  This makes sense because the data was generated by a model containing only depolarization errors on the gates, and our custom gate class can model this type of noise.  We expect, since we know how the data was generated, that the `\"Gx\"` gate depolarizes with magnitude $0.01$ and has no (zero) over-rotation.  Indeed, this is what we find when we look at `\"Gx\"` of the estimated model: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyXPi2Operator with shape (4, 4)\n",
      " 1.00   0   0   0\n",
      "   0 0.99   0   0\n",
      "   0   0   0-0.99\n",
      "   0   0 0.99   0\n",
      "\n",
      "Estimated Gx depolarization = 0.00982963242432355\n",
      "Estimated Gx over-rotation = 0.0005135253334618614\n"
     ]
    }
   ],
   "source": [
    "mdl_estimate = results.estimates['default'].models['final iteration estimate']\n",
    "print(mdl_estimate['Gx'])\n",
    "est_depol, est_overrotation = mdl_estimate['Gx'].to_vector()\n",
    "print(\"Estimated Gx depolarization =\",est_depol)\n",
    "print(\"Estimated Gx over-rotation =\",est_overrotation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The reason these values aren't exactly $0.01$ and $0$ are due to the finite number of samples, and to a lesser extent  gauge degrees of freedom.\n",
    "\n",
    "## What's next?\n",
    "This tutorial showed you how to create a custom *dense* operation (a subclass of `DenseOperator`).  We'll be adding demonstrations of more complex custom operations in the future.  Here are some places you might want to go next:\n",
    "- The [operators tutorial](Operators.ipynb) explains and shows examples of pyGSTi's existing operations.\n",
    "- MORE TODO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}