{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from pprint import pprint\n",
    "np.set_printoptions(precision=4, linewidth=120, floatmode='maxprec_equal')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygsti\n",
    "import pygsti.extras.interpygate as interp\n",
    "from pygsti.tools.basistools import change_basis\n",
    "from pygsti.modelpacks import smq1Q_XY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "working_dir = Path.cwd()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build model gate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigI = np.array([[1.,0],[0, 1]], dtype='complex')\n",
    "sigX = np.array([[0, 1],[1, 0]], dtype='complex')\n",
    "sigY = np.array([[0,-1],[1, 0]], dtype='complex') * 1.j\n",
    "sigZ = np.array([[1, 0],[0,-1]], dtype='complex')\n",
    "sigM = (sigX - 1.j*sigY)/2.\n",
    "sigP = (sigX + 1.j*sigY)/2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SingleQubitTargetOp(pygsti.modelmembers.operations.OpFactory):\n",
    "\n",
    "    def __init__(self):\n",
    "        self.process = self.create_target_gate\n",
    "        pygsti.modelmembers.operations.OpFactory.__init__(self, 1, evotype=\"densitymx\")\n",
    "        self.dim = 4\n",
    "\n",
    "    def create_target_gate(self, v):\n",
    "        \n",
    "        phi, theta = v\n",
    "        target_unitary = (np.cos(theta/2) * sigI + \n",
    "                          1.j * np.sin(theta/2) * (np.cos(phi) * sigX + np.sin(phi) * sigY))\n",
    "        superop = change_basis(np.kron(target_unitary.conj(), target_unitary), 'col', 'pp')\n",
    "\n",
    "        return superop\n",
    "    \n",
    "    def create_object(self, args=None, sslbls=None):\n",
    "        assert(sslbls is None)\n",
    "        mx = self.process([*args])\n",
    "        return pygsti.modelmembers.operations.StaticArbitraryOp(mx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SingleQubitGate(interp.PhysicalProcess):\n",
    "    def __init__(self, \n",
    "                 verbose=False,\n",
    "                 cont_param_gate = False,\n",
    "                 num_params = None,\n",
    "#                  process_shape = (4, 4),\n",
    "                 item_shape = (4,4),\n",
    "                 aux_shape = None,\n",
    "                 num_params_evaluated_as_group = 0,\n",
    "                 ):\n",
    "\n",
    "        self.verbose = verbose\n",
    "\n",
    "        self.cont_param_gate = cont_param_gate\n",
    "\n",
    "        self.num_params = num_params\n",
    "        self.item_shape = item_shape\n",
    "\n",
    "        self.aux_shape = aux_shape\n",
    "        self.num_params_evaluated_as_group = num_params_evaluated_as_group\n",
    "\n",
    "   \n",
    "    def create_process_matrix(self, v, comm=None, return_generator=False):                                                                                                                                                                                             \n",
    "\n",
    "        processes = []\n",
    "        phi, theta, t = v\n",
    "        theta = theta * t\n",
    "        target_unitary = (np.cos(theta/2) * sigI + \n",
    "                          1.j * np.sin(theta/2) * (np.cos(phi) * sigX + np.sin(phi) * sigY))\n",
    "        superop = change_basis(np.kron(target_unitary.conj(), target_unitary), 'col', 'pp')\n",
    "        processes += [superop]\n",
    "        return np.array(processes) if (processes is not None) else None\n",
    "\n",
    "    def create_aux_info(self, v, comm=None):\n",
    "        return []  # matches aux_shape=() above\n",
    "    \n",
    "    def create_process_matrices(self, v, grouped_v, comm=None):\n",
    "        assert(len(grouped_v) == 1)  # we expect a single \"grouped\" parameter\n",
    "\n",
    "        processes = []\n",
    "        times = grouped_v[0]\n",
    "        phi_in, theta_in = v\n",
    "        for t in times:\n",
    "            phi = phi_in\n",
    "            theta = theta_in * t\n",
    "            target_unitary = (np.cos(theta/2) * sigI + \n",
    "                              1.j * np.sin(theta/2) * (np.cos(phi) * sigX + np.sin(phi) * sigY))\n",
    "            superop = change_basis(np.kron(target_unitary.conj(), target_unitary), 'col', 'pp')\n",
    "            processes += [superop]\n",
    "        return np.array(processes) if (processes is not None) else None\n",
    "\n",
    "    def create_aux_infos(self, v, grouped_v, comm=None):\n",
    "        import numpy as np\n",
    "        times = grouped_v[0]\n",
    "        return [ [] for t in times] # list elements must match aux_shape=() above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_ranges = [(0.9,1.1,3)]\n",
    "\n",
    "arg_ranges = [2*np.pi*(1+np.cos(np.linspace(np.pi,0, 7)))/2,\n",
    "              (0, np.pi, 3)] \n",
    "arg_indices = [0,1]\n",
    "\n",
    "\n",
    "target_op = SingleQubitTargetOp()\n",
    "gate_process = SingleQubitGate(num_params = 3,num_params_evaluated_as_group = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opfactory_linear = interp.InterpolatedOpFactory.create_by_interpolating_physical_process(\n",
    "    target_op, gate_process, argument_ranges=arg_ranges, parameter_ranges=param_ranges, \n",
    "    argument_indices=arg_indices, interpolator_and_args='linear')\n",
    "\n",
    "opfactory_spline = interp.InterpolatedOpFactory.create_by_interpolating_physical_process(\n",
    "    target_op, gate_process, argument_ranges=arg_ranges, parameter_ranges=param_ranges, \n",
    "    argument_indices=arg_indices, interpolator_and_args='spline')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check that the interpolator is working"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if False:\n",
    "    indices = (2,3)\n",
    "    nparams = 30\n",
    "\n",
    "    x = np.linspace(0,2*np.pi, nparams)\n",
    "    y = np.linspace(0, np.pi, nparams)\n",
    "    for p in np.linspace(.9,1.1,5):\n",
    "\n",
    "        def interp_linear(x, y):    \n",
    "            op = opfactory_linear.create_op([x, y])\n",
    "            return op.base_interpolator([x,y,p])[indices]\n",
    "\n",
    "        def interp_spline(x, y):    \n",
    "            op = opfactory_spline.create_op([x, y])\n",
    "            return op.base_interpolator([x,y,p])[indices]\n",
    "\n",
    "        def truth(x, y):\n",
    "            tmp_gate = gate_process.create_process_matrix([x,y,p])[0]\n",
    "            tar_gate = target_op.create_target_gate([x,y])\n",
    "            return pygsti.error_generator(tmp_gate, tar_gate, 'pp', 'logGTi')[indices]\n",
    "\n",
    "\n",
    "        X, Y = np.meshgrid(x, y, indexing='ij')\n",
    "        Z_linear = np.zeros([nparams, nparams])\n",
    "        Z_spline = np.zeros([nparams, nparams])\n",
    "        Z_truth  = np.zeros([nparams, nparams])\n",
    "        for idx, xx in enumerate(x):\n",
    "            for idy, yy in enumerate(y):\n",
    "                Z_linear[idx,idy] = interp_linear(xx,yy)\n",
    "                Z_spline[idx,idy] = interp_spline(xx,yy)\n",
    "                Z_truth[idx,idy]  = truth(xx,yy)\n",
    "\n",
    "        fig = plt.figure(figsize=(10,10))\n",
    "        ax = plt.axes(projection='3d')\n",
    "\n",
    "        ax.plot_surface(X, Y, Z_linear-Z_truth, rstride=1, cstride=1,\n",
    "                        edgecolor='none', alpha=.8)\n",
    "        ax.plot_surface(X, Y, Z_spline-Z_truth, rstride=1, cstride=1,\n",
    "                        cmap='viridis', edgecolor='none', alpha=.8)\n",
    "\n",
    "        ax.set_xlabel('x')\n",
    "        ax.set_ylabel('y')\n",
    "        ax.set_zlabel('z')\n",
    "    #     ax.set_zlim(-1,1)\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a model from this gate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_gate = opfactory_spline.create_op([0,np.pi/4])\n",
    "y_gate = opfactory_spline.create_op([np.pi/2,np.pi/4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_gate.from_vector([1.03])\n",
    "y_gate.from_vector([1.0])\n",
    "print(np.round_(x_gate,4))\n",
    "print()\n",
    "print(np.round_(y_gate,4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_gate.parameter_bounds = np.array([[0.91, 1.09]])\n",
    "y_gate.parameter_bounds = np.array([[0.91, 1.09]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make a fake dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Model only has Gx and Gy gates.  Let's rename them.\n",
    "\n",
    "model = pygsti.models.ExplicitOpModel([0],'pp')\n",
    "\n",
    "model['rho0'] = [ 1/np.sqrt(2), 0, 0, 1/np.sqrt(2) ] # density matrix [[1, 0], [0, 0]] in Pauli basis\n",
    "model['Mdefault'] = pygsti.modelmembers.povms.UnconstrainedPOVM(\n",
    "    {'0': [ 1/np.sqrt(2), 0, 0, 1/np.sqrt(2) ],   # projector onto [[1, 0], [0, 0]] in Pauli basis\n",
    "     '1': [ 1/np.sqrt(2), 0, 0, -1/np.sqrt(2) ] }, evotype=\"default\") # projector onto [[0, 0], [0, 1]] in Pauli basis\n",
    "model['Gxpi2',0] = x_gate\n",
    "model['Gypi2',0] = y_gate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.num_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the error model used to generate data\n",
    "datagen_model = model.copy()\n",
    "datagen_params = datagen_model.to_vector()\n",
    "datagen_params[-2:] = [1.03,1.00]\n",
    "datagen_model.from_vector(datagen_params)\n",
    "datagen_model.probabilities( (('Gxpi2',0),('Gypi2',0),))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.parameter_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Link the over-rotation errors on Gx and Gy\n",
    "#model.collect_parameters(model.parameter_labels[-2:], 'Shared Gx/Gy physical parameter')\n",
    "#print(model.parameter_labels)\n",
    "print(model.num_params)\n",
    "\n",
    "# # Define the error model used to generate data\n",
    "# datagen_model = model.copy()\n",
    "# datagen_params = datagen_model.to_vector()\n",
    "# datagen_params[-1:] = [1.02]\n",
    "# datagen_model.from_vector(datagen_params)\n",
    "# datagen_model.probabilities( (('Gxpi2',0),('Gypi2',0),))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the perfect target model\n",
    "target_model = model.copy()\n",
    "target_params = target_model.to_vector()\n",
    "target_params[-2:] = [1,1]\n",
    "# target_model.from_vector(target_params)\n",
    "target_model.probabilities( (('Gxpi2',0),('Gypi2',0),))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Germ and fiducial selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "final_germs = pygsti.algorithms.germselection.find_germs(\n",
    "                target_model, randomize=False, force=None, algorithm='greedy', \n",
    "                verbosity=4, num_nongauge_params=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fiducial_pairs = pygsti.algorithms.fiducialpairreduction.find_sufficient_fiducial_pairs_per_germ(\n",
    "                                model, \n",
    "                                smq1Q_XY.prep_fiducials(),\n",
    "                                smq1Q_XY.meas_fiducials(), \n",
    "                                final_germs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Reduce the number of fiducial pairs by hand, if you want\n",
    "\n",
    "# fiducial_pairs2 = fiducial_pairs.copy()\n",
    "# for key in fiducial_pairs2.keys():\n",
    "#     fiducial_pairs2[key] = fiducial_pairs2[key][0:2]\n",
    "# fiducial_pairs = fiducial_pairs2\n",
    "\n",
    "# print(fiducial_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use fiducial pair reductions\n",
    "exp_design = pygsti.protocols.StandardGSTDesign(model, \n",
    "                                                smq1Q_XY.prep_fiducials(), \n",
    "                                                smq1Q_XY.meas_fiducials(), \n",
    "                                                final_germs, \n",
    "                                                max_lengths=[1,2,4,8,16,32,64,128,256], \n",
    "                                                fiducial_pairs=fiducial_pairs,\n",
    "                                                include_lgst=False)\n",
    "\n",
    "dataset = pygsti.data.simulate_data(datagen_model, exp_design.all_circuits_needing_data,\n",
    "                                    num_samples=1000, seed=1234)\n",
    "\n",
    "data = pygsti.protocols.ProtocolData(exp_design, dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(data.dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fisher information matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fim = pygsti.tools.edesigntools.calculate_fisher_information_matrix(model,\n",
    "                                                                    exp_design.all_circuits_needing_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.log(np.linalg.inv(fim))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.matshow(np.linalg.inv(fim))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run GST on the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proto = pygsti.protocols.GateSetTomography(model, gaugeopt_suite=None)\n",
    "results = proto.run(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# What is the estimated value of the error parameter?\n",
    "\n",
    "final_model = results.estimates['GateSetTomography'].models['final iteration estimate']\n",
    "print('Actual: ', datagen_model.to_vector()[-2:])\n",
    "print('Estimated: ', final_model.to_vector()[-2:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pprint(np.sqrt(2)*final_model.to_vector()[0:4])\n",
    "pprint(np.sqrt(2)*final_model.to_vector()[4:8])\n",
    "pprint(np.sqrt(2)*final_model.to_vector()[8:12])"
   ]
  },
  {
   "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
}