{
 "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": null,
   "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",
    "        super(MyXPi2Operator,self).__init__(np.identity(4,'d'), \"densitymx\") # this is *super*-operator, so \"densitymx\"\n",
    "        self.from_vector([0,0]) \n",
    "    \n",
    "    @property\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, close=False, dirty_value=True):\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",
    "        # Technical note: use [:,:] instead of direct assignment so id of self.base doesn't change\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",
    "        self.dirty = dirty_value  # mark that parameter vector may have changed\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pygsti.modelpacks import smq1Q_XYI\n",
    "mdl = smq1Q_XYI.target_model()\n",
    "mdl.operations[('Gxpi2',0)] = 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 `gauge_opt_params=False` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate \"fake\" data from a depolarized version of the target (ideal) model\n",
    "maxLengths = [1,2,4,8,16]\n",
    "mdl_datagen = smq1Q_XYI.target_model().depolarize(op_noise=0.01, spam_noise=0.001)\n",
    "listOfExperiments = pygsti.construction.create_lsgst_circuits(\n",
    "    mdl_datagen, smq1Q_XYI.prep_fiducials(), smq1Q_XYI.meas_fiducials(), smq1Q_XYI.germs(), maxLengths)\n",
    "ds = pygsti.construction.simulate_data(mdl_datagen, listOfExperiments, num_samples=1000,\n",
    "                                            sample_error=\"binomial\", seed=1234)\n",
    "\n",
    "#Run GST *without* gauge optimization\n",
    "results = pygsti.run_long_sequence_gst(ds, mdl, smq1Q_XYI.prep_fiducials(), smq1Q_XYI.meas_fiducials(),\n",
    "                                      smq1Q_XYI.germs(), maxLengths, gauge_opt_params=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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mdl_estimate = results.estimates['GateSetTomography'].models['final iteration estimate']\n",
    "print(mdl_estimate[('Gxpi2',0)])\n",
    "est_depol, est_overrotation = mdl_estimate[('Gxpi2',0)].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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}