{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example for qutrit GST\n",
    "This notebook demonstrates how to construct the operation sequences and perform the analysis for qutrit GST when the model consists of symmetric $\\pi/2$-rotations on each single qubit separately, `X`, `Y` and a 2-qubit Molmer-Sorenson gate which rotates around the `XX` axis by $\\pi/2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygsti\n",
    "from pygsti.models import qutrit\n",
    "from pygsti.algorithms.fiducialselection import find_fiducials\n",
    "from pygsti.algorithms.germselection import find_germs\n",
    "\n",
    "from numpy import pi, array\n",
    "import pickle\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we construct the target model.  This functionality is built into pyGSTi, so we just need to specify the single-qubit and M-S angles.\n",
    "Note there are alternative approaches for building a qutrit model in pygsti using processor specification objects, but for this particular class of qutrit models in this example notebook there exist helper functions for creating the relevant models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "target_model = qutrit.create_qutrit_model(error_scale=0, x_angle=pi/2, y_angle=pi/2, ms_global=pi/2, ms_local=0, basis=\"qt\")\n",
    "#change the forward simulator for the purposes of experiment design code\n",
    "target_model.sim = 'matrix'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now construct the operation sequences needed by GST. Then we construct an empty dataset containing all of the necessary experimental sequences which can serve as a template for the actual experimental results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fiducialPrep, fiducialMeasure = find_fiducials(target_model, candidate_fid_counts={4: 'all upto'}, algorithm= 'greedy')\n",
    "germs = find_germs(target_model, randomize=False, candidate_germ_counts={4: 'all upto'}, mode= 'compactEVD', assume_real=True, float_type=np.double)\n",
    "maxLengths = [1,2,4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"%d prep fiducials\" % len(fiducialPrep))\n",
    "print(\"%d meas fiducials\" % len(fiducialMeasure))\n",
    "print(\"%d germs\" % len(germs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#generate data template\n",
    "expList = pygsti.circuits.create_lsgst_circuits(target_model.operations.keys(), fiducialPrep, fiducialMeasure, germs,  maxLengths)\n",
    "pygsti.io.write_empty_dataset(\"example_files/dataTemplate_qutrit_maxL=4.txt\", expList, \"## Columns = 0bright count, 1bright count, 2bright count\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point **STOP** and create/fill a dataset file using the template written in the above cell.  Then proceed with the lines below to run GST on the data and create (hopefully useful) reports telling you about your gates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mdl_datagen = target_model.depolarize(op_noise=0.05, spam_noise = .01)\n",
    "DS = pygsti.data.simulate_data(mdl_datagen, expList, 1000, sample_error='multinomial', seed=2018)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#DS = pygsti.io.load_dataset('PATH_TO_YOUR_DATASET',cache=True) # (cache=True speeds up future loads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#Run qutrit GST... which could take a while on a single CPU.  Please adjust memLimit to machine specs \n",
    "# (now 3GB; usually set to slightly less than the total machine memory)\n",
    "#Setting max_iterations lower than default for the sake of the example running faster. \n",
    "target_model.sim = \"matrix\"\n",
    "result = pygsti.run_stdpractice_gst(DS, target_model, fiducialPrep, fiducialMeasure, germs, maxLengths,\n",
    "                                    verbosity=3, comm=None, mem_limit=3*(1024)**3, modes=\"CPTPLND\",\n",
    "                                    advanced_options= {'max_iterations':50})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create a report\n",
    "ws = pygsti.report.construct_standard_report(\n",
    "    result, \"Example Qutrit Report\", verbosity=3\n",
    ").write_html('example_files/sampleQutritReport', auto_open=False, verbosity=3)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "api_updates",
   "language": "python",
   "name": "api_updates"
  },
  "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
}