Add checkpointing to tutorials

Add sections on checkpointing an its usage to the tutorials on GST and model testing protocols. Also makes a minor change to one of the unit tests that is inexplicably failing on older python versions.
sandialabs · Sep 7, 2023 · a58d1fc · a58d1fc
1 parent 57b9a9f
commit a58d1fc
Show file tree

Hide file tree

Showing 4 changed files with 255 additions and 203 deletions.
diff --git a/.gitignore b/.gitignore
@@ -47,6 +47,8 @@ jupyter_notebooks/Tutorials/tutorial_files/exampleMultiDataSetReport
 jupyter_notebooks/Tutorials/tutorial_files/exampleBriefReport
 jupyter_notebooks/Tutorials/tutorial_files/*.ipynb
 jupyter_notebooks/Tutorials/tutorial_files/tempTest
+jupyter_notebooks/Tutorials/tutorial_files/*checkpoints
+
 
 
 # Compiled source #

diff --git a/jupyter_notebooks/Tutorials/algorithms/GST-Protocols.ipynb b/jupyter_notebooks/Tutorials/algorithms/GST-Protocols.ipynb
@@ -28,7 +28,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from __future__ import print_function\n",
     "import pygsti"
    ]
   },
@@ -63,7 +62,7 @@
     "## `GateSetTomography`\n",
     "This protocol performs a single model optimization, and so computes a **single GST estimate** given a `DataSet`, a target `Model`, and other parameters.  (The returned `ModelEstimateResults` object may sometimes contain multiple related estimates in certain cases, but in these cases all the estimates are closely related.)  The experiment design provides all of the information about the GST circuits, in this case a *standard*  (*prep_fiducial + germ^power + meas_fiducial*) set, so the only thing needed by the protocol is an initial `Model` to optimize.  Thus, the `GateSetTomography` protocol is essentially just a model optimizer that you give an initial point.  Importantly, this initial point (a `Model`) also specifies the *parameterization*, i.e. the space of parameters that are optimized over.\n",
     "\n",
-    "Minimally, when using `GateSetTomography` you should set the parameterization of the initial model.  This can be viewed as setting the constraints on the optimization.  For instance, when the gates in the model are parameterized as trace-preserving (TP) maps, the optimization will be constrained to trying gate sets with TP gates (because every set of parameters corresponds to a set of TP gates).  In the cell below, we constrain the optimization to TP gate sets by using `.target_model(\"TP\")`, which returns a version of the target model where all the gates are TP-parameterized, the state preparation has trace = 1, and the POVM effects always add to the identity.  This could also be done by calling `set_all_parameterizations(\"TP\")` on the fully-parameterized target model returned by `.target_model()`.  See the [tutorial on explicit models](../objects/ExplicitModel.ipynb) for more information on setting a model's parameterization."
+    "Minimally, when using `GateSetTomography` you should set the parameterization of the initial model.  This can be viewed as setting the constraints on the optimization.  For instance, when the gates in the model are parameterized as trace-preserving (TP) maps, the optimization will be constrained to trying gate sets with TP gates (because every set of parameters corresponds to a set of TP gates).  In the cell below, we constrain the optimization to TP gate sets by using `.target_model(\"full TP\")`, which returns a version of the target model where all the gates are TP-parameterized, the state preparation has trace = 1, and the POVM effects always add to the identity.  This could also be done by calling `set_all_parameterizations(\"TP\")` on the fully-parameterized target model returned by `.target_model()`.  See the [tutorial on explicit models](../objects/ExplicitModel.ipynb) for more information on setting a model's parameterization."
    ]
   },
   {
@@ -77,7 +76,7 @@
     "from pygsti.modelpacks import smq1Q_XYI\n",
     "target_model_TP = smq1Q_XYI.target_model(\"full TP\")\n",
     "proto = pygsti.protocols.GateSetTomography(target_model_TP)\n",
-    "results_TP = proto.run(data)"
+    "results_TP = proto.run(data, disable_checkpointing=True)"
    ]
   },
   {
@@ -126,7 +125,7 @@
     "    target_model_TP2, name=\"GSTwithMyGO\",\n",
     "    gaugeopt_suite={'my_gauge_opt': {'item_weights': {'gates': 1.0, 'spam': 0.001}}}\n",
     "    )\n",
-    "results_TP2 = proto.run(data)"
+    "results_TP2 = proto.run(data, disable_checkpointing=True)"
    ]
   },
   {
@@ -168,7 +167,7 @@
     "\n",
     "\n",
     "proto = pygsti.protocols.GateSetTomography(target_model_TP2, name=\"GSTwithReducedData\")\n",
-    "results_reduced = proto.run(reduced_data)"
+    "results_reduced = proto.run(reduced_data, disable_checkpointing=True)"
    ]
   },
   {
@@ -200,7 +199,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "results_stdprac = pygsti.protocols.StandardGST().run(data)"
+    "results_stdprac = pygsti.protocols.StandardGST().run(data, disable_checkpointing=True)"
    ]
   },
   {
@@ -227,7 +226,7 @@
    "outputs": [],
    "source": [
     "proto = pygsti.protocols.StandardGST(gaugeopt_suite=\"varySpam\", name=\"StdGST_varySpam\")\n",
-    "results_stdprac_nondefaultgo = proto.run(data)"
+    "results_stdprac_nondefaultgo = proto.run(data, disable_checkpointing = True)"
    ]
   },
   {
@@ -261,7 +260,7 @@
     "\n",
     "proto = pygsti.protocols.StandardGST(gaugeopt_suite=my_gaugeopt_suite,\n",
     "                                     name=\"StdGST_myGO\")\n",
-    "results_stdprac_nondefaultgo = proto.run(data)"
+    "results_stdprac_nondefaultgo = proto.run(data, disable_checkpointing=True)"
    ]
   },
   {
@@ -321,6 +320,163 @@
     "# pickle.dump(results_stdprac, open('../tutorial_files/exampleResults_stdprac.pkl',\"wb\"))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Checkpointing/Warmstarting\n",
+    "\n",
+    "The `GateSetTomography` and `StandardGST` protocols both support checkpointing to enable resuming GST analysis after an unexpected failure, such as an out-of-memory error, or an unexpected timeout in resource limited compute environments (clusters etc.), or for whatever other reason. Checkpointing is enabled by default, so no additional changes are needed in order to have these generated. \n",
+    "\n",
+    "Each protocol has a corresponding checkpoint object, `GateSetTomographyCheckpoint` and `StandardGSTCheckpoint`, which are saved to disk over the course of an iterative fit in serialized json format. By default checkpoint files associated with a `GateSetTomographyCheckpoint` object are saved to a new directory located in whichever current working directory the protocol is being run from named 'gst_checkpoints'. A new file is written to disk after each iteration with default naming of the form `GateSetTomography_iteration_{i}.json` where i is the index of the completed GST iteration associated with that checkpoint. Similarly, for a `StandardGSTCheckpoint` object the checkpoints are by default saved to a directory named 'standard_gst_checkpoints' with default file names of the form `StandardGST_{mode}_iteration_{i}` where mode corresponds to the current parameterized fit or model test associated with that file (including checkpoint information for all previously completed modes prior to the currently running one) and i is the index of the completed iteration within that current mode.\n",
+    "\n",
+    "Below we repeat our first example of the notebook, but this time with checkpointing enabled (as is the default)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pygsti.modelpacks import smq1Q_XYI\n",
+    "target_model_TP = smq1Q_XYI.target_model(\"full TP\")\n",
+    "proto = pygsti.protocols.GateSetTomography(target_model_TP)\n",
+    "results_TP = proto.run(data, checkpoint_path = '../tutorial_files/gst_checkpoints/GateSetTomography')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note that in the example above we have specified a value for an additional kwarg called `checkpoint_path`. This allows for overriding the default behavior for the save location and naming of checkpoint files. The expected format is `{path}/{name}` where path is the directory to save the checkpoint files to (with that directory being created is required) and where name is the stem of the checkpoint file names `{name}_iteration_{i}.json`. Inspecting the contents of the directory we just specified, we can see that it is now populated by 8 new checkpoint files."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.listdir('../tutorial_files/gst_checkpoints/')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Suppose hypothetically that a GST fit had failed at iteration 5 and we wanted to restart from that point without redoing all of the previous iterations from scratch again. We'll call this warmstarting. We can do so by reading in the appropriate serialized checkpoint object using the `read` class method of `GateSetTomographyCheckpoint` and passing that now loaded checkpoint object in for the `checkpoint` kwarg of `run`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pygsti.protocols import GateSetTomographyCheckpoint\n",
+    "gst_iter_5_checkpoint = GateSetTomographyCheckpoint.read('../tutorial_files/gst_checkpoints/GateSetTomography_iteration_5.json')\n",
+    "results_TP_from_iter_5= proto.run(data, checkpoint= gst_iter_5_checkpoint, checkpoint_path = '../tutorial_files/gst_checkpoints/GateSetTomography')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can see from the output that we indeed started from iteration 6 (note the output log indexes from 1 instead of 0). Moreover we can see that we've indeed produced the same output as before without warmstarting, as we would expect/hope:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "all(results_TP.estimates['GateSetTomography'].models['final iteration estimate'].to_vector() == \\\n",
+    "results_TP_from_iter_5.estimates['GateSetTomography'].models['final iteration estimate'].to_vector())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The checkpoint object itself contains information that could be useful for diagnostics or debugging, including the current list of models associated each iterative fit, the last completed iteration it is associated with, and the list of circuits for the last completed iteration it is associated with."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Checkpointing with the StandardGST protocol works similarly:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "proto_standard_gst = pygsti.protocols.StandardGST(modes=['full TP', 'CPTPLND', 'Target'], verbosity=3)\n",
+    "results_stdprac = proto_standard_gst.run(data, checkpoint_path = '../tutorial_files/standard_gst_checkpoints/StandardGST')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Except this time we have significantly more files saved, as during the course of the StandardGST protocol we're actually running three subprotocols:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.listdir('../tutorial_files/standard_gst_checkpoints/')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note that the StandardGST protocol runs the subprotocols in the order listed in the `modes` argument, and checkpoint objects labeled with a given model label additionally contain the checkpointing information for the final iterations of any preceding modes which have been completed. i.e. the CPTPLND checkpoint objects contain the information required for full TP. Likewise, checkpoints for Target contain the information required for the full TP and CPTPLND modes. As before, imagine that our fitting failed for whatever reason during iteration 5 of CPTPLND, we can warmstart the protocol by loading in the checkpoint object associated with iteration 4 as below:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "from pygsti.protocols import StandardGSTCheckpoint\n",
+    "standard_gst_checkpoint = StandardGSTCheckpoint.read('../tutorial_files/standard_gst_checkpoints/StandardGST_CPTPLND_iteration_4.json')\n",
+    "results_stdprac_warmstart= proto_standard_gst.run(data, checkpoint= standard_gst_checkpoint, checkpoint_path = '../tutorial_files/standard_gst_checkpoints/StandardGST')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Notice that we've indeed skipped past the previously completed full TP mode and jumped straight to the 6th iteration of the CPTPLND fit as expected. \n",
+    "\n",
+    "As for the GateSetTomographyCheckpoint object described above, the `StandardGSTCheckpoint` can often be useful to inspect as a debugging/diagnostic tool. `StandardGSTCheckpoints` are essentially structured as container object that hold a set of child `GateSetTomographyCheckpoint` and `ModelTestCheckpoint` (more on these in the ModelTest tutorial) objects for each of the modes being run (and potentially more types of chile checkpoints in the future as we add additional functionality). These children can be accessed using the `children` attribute of a `StandardGSTCheckpoint` instance which is a dictionary with keys given by the mode names contained therein."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(standard_gst_checkpoint.children['CPTPLND'])"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -331,9 +487,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "gst_checkpointing",
    "language": "python",
-   "name": "python3"
+   "name": "gst_checkpointing"
   },
   "language_info": {
    "codemirror_mode": {