Enhancements and bugfixes in Report.write_notebook

optional-requirements.txt

@@ -5,9 +5,11 @@ psutil
 matplotlib
 ipython
 notebook
+nbformat
 msgpack
 cython
 cvxopt
 cvxpy
 seaborn
 packaging
+pytest

pygsti/baseobjs/label.py

@@ -201,7 +201,6 @@ def is_simple(self):
         return self.IS_SIMPLE
 
 
-
 class LabelTup(Label, tuple):
     """
     A label consisting of a string along with a tuple of integers or state-space-names.

pygsti/circuits/circuit.py

@@ -25,7 +25,6 @@
 
 import numpy as _np
 from pygsti.baseobjs.label import Label as _Label, CircuitLabel as _CircuitLabel
-
 from pygsti.baseobjs import outcomelabeldict as _ld, _compatibility as _compat
 from pygsti.tools import internalgates as _itgs
 from pygsti.tools import slicetools as _slct
@@ -512,7 +511,6 @@ def __init__(self, layer_labels=(), line_labels='auto', num_lines=None, editable
         self._bare_init(labels, my_line_labels, editable, name, stringrep, 
                         occurrence, compilable_layer_indices_tup)
 
-
     @classmethod
     def _fastinit(cls, labels, line_labels, editable, name='', stringrep=None, occurrence=None,
                   compilable_layer_indices_tup=()):
@@ -4942,6 +4940,7 @@ def done_editing(self):
         self._hashable_tup = self.tup
         self._hash = hash(self._hashable_tup)
 
+
 class CompressedCircuit(object):
     """
     A "compressed" Circuit that requires less disk space.

pygsti/protocols/estimate.py

@@ -24,11 +24,12 @@
 from pygsti.protocols.confidenceregionfactory import ConfidenceRegionFactory as _ConfidenceRegionFactory
 from pygsti.models.explicitmodel import ExplicitOpModel as _ExplicitOpModel
 from pygsti.objectivefns import objectivefns as _objfns
-from pygsti.circuits.circuitlist import CircuitList as _CircuitList
+from pygsti.circuits import CircuitList as _CircuitList, Circuit as _Circuit
 from pygsti.circuits.circuitstructure import PlaquetteGridCircuitStructure as _PlaquetteGridCircuitStructure
 from pygsti.baseobjs.verbosityprinter import VerbosityPrinter as _VerbosityPrinter
 from pygsti.baseobjs.mongoserializable import MongoSerializable as _MongoSerializable
 
+
 #Class for holding confidence region factory keys
 CRFkey = _collections.namedtuple('CRFkey', ['model', 'circuit_list'])
 
@@ -80,9 +81,18 @@ def from_dir(cls, dirname, quick_load=False):
         """
         ret = cls.__new__(cls)
         _MongoSerializable.__init__(ret)
-        ret.__dict__.update(_io.load_meta_based_dir(_pathlib.Path(dirname), 'auxfile_types', quick_load=quick_load))
+        state = _io.load_meta_based_dir(_pathlib.Path(dirname), 'auxfile_types', quick_load=quick_load)
+        ret.__dict__.update(state)
         for crf in ret.confidence_region_factories.values():
             crf.set_parent(ret)  # re-link confidence_region_factories
+        if ret.circuit_weights is not None:
+            from pygsti.circuits.circuitparser import parse_circuit
+            cws : dict[_Circuit, float] = dict()
+            for cstr, w in ret.circuit_weights.items():
+                lbls = parse_circuit(cstr, True, True)[0]
+                ckt = _Circuit(lbls)
+                cws[ckt] = w
+            ret.circuit_weights = cws
         return ret
 
     @classmethod
@@ -236,7 +246,17 @@ def write(self, dirname):
         -------
         None
         """
+        old_cw = self.circuit_weights
+        if isinstance(old_cw, dict):
+            new_cw : dict[str, float] = dict()
+            for c, w in old_cw.items():
+                if not isinstance(c, _Circuit):
+                    raise ValueError()
+                new_cw[c.str] = w
+            self.circuit_weights = new_cw
         _io.write_obj_to_meta_based_dir(self, dirname, 'auxfile_types')
+        self.circuit_weights = old_cw
+        return
 
     def _add_auxiliary_write_ops_and_update_doc(self, doc, write_ops, mongodb, collection_name, overwrite_existing):
         _io.add_obj_auxtree_write_ops_and_update_doc(self, doc, write_ops, mongodb, collection_name,

pygsti/report/report.py

@@ -12,13 +12,15 @@
 
 import pickle as _pickle
 import time as _time
+import os as _os
 from collections import defaultdict as _defaultdict
 from pathlib import Path as _Path
 
 from pygsti.report import merge_helpers as _merge
 from pygsti.report import workspace as _ws
 from pygsti.report.notebook import Notebook as _Notebook
 from pygsti.baseobjs import VerbosityPrinter as _VerbosityPrinter
+from pygsti.protocols import ModelEstimateResults as _ModelEstimateResults
 
 
 # TODO this whole thing needs to be rewritten with different reports as derived classes
@@ -196,7 +198,7 @@ def write_html(self, path, auto_open=False, link_to=None,
                 verbosity=verbosity
             )
 
-    def write_notebook(self, path, auto_open=False, connected=False, verbosity=0):
+    def write_notebook(self, path, auto_open=False, connected=False, verbosity=0, use_pickle=False):
         """
         Write this report to the disk as an IPython notebook
 
@@ -234,22 +236,24 @@ def write_notebook(self, path, auto_open=False, connected=False, verbosity=0):
         title = self._global_qtys['title']
         confidenceLevel = self._report_params['confidence_level']
 
+        if not path.endswith('.ipynb'):
+            raise ValueError(f'path={path} must have the `.ipynb` suffix for a Jupyter Notebook.')
+
         path = _Path(path)
         printer = _VerbosityPrinter.create_printer(verbosity)
-        templatePath = _Path(__file__).parent / 'templates' / self._templates['notebook']
-        outputDir = path.parent
+        notebook_templates_path = _Path(__file__).parent / 'templates' / self._templates['notebook']
 
         #Copy offline directory into position
         if not connected:
-            _merge.rsync_offline_dir(outputDir)
+            _merge.rsync_offline_dir(path.parent)
 
         #Save results to file
-        # basename = _os.path.splitext(_os.path.basename(filename))[0]
-        basename = path.stem
-        results_file_base = basename + '_results.pkl'
-        results_file = outputDir / results_file_base
-        with open(str(results_file), 'wb') as f:
-            _pickle.dump(self._results, f)
+        results_path = str(path.parent / (path.stem + '_results'))
+        if use_pickle:
+            results_path = results_path + '.pkl'
+            self.write_results_dict(results_path)
+        else:
+            self.write_results_dict(results_path)
 
         nb = _Notebook()
         nb.add_markdown('# {title}\n(Created on {date})'.format(
@@ -261,20 +265,20 @@ def write_notebook(self, path, auto_open=False, connected=False, verbosity=0):
         "classic Jupyter notebooks for PyGSTi report notebooks. To track this issue, " +
         "see https://github.com/pyGSTio/pyGSTi/issues/205.</font>")
 
-        nb.add_code("""\
-            import pickle
-            import pygsti""")
+        nb.add_code("""
+        import pygsti
+        from pygsti.report import Report
+        """)
 
         dsKeys = list(self._results.keys())
         results = self._results[dsKeys[0]]
         #Note: `results` is always a single Results obj from here down
 
-        nb.add_code("""\
+        nb.add_code(f"""
         #Load results dictionary
-        with open('{infile}', 'rb') as infile:
-            results_dict = pickle.load(infile)
-        print("Available dataset keys: ", ', '.join(results_dict.keys()))\
-        """.format(infile=results_file_base))
+        results_path = '{results_path}'
+        results_dict = Report.results_dict_from_dir(results_path)
+        """)
 
         nb.add_code("""\
         #Set which dataset should be used below
@@ -334,11 +338,15 @@ def write_notebook(self, path, auto_open=False, connected=False, verbosity=0):
             ws.init_notebook_mode(connected={conn}, autodisplay=True)\
             """.format(conn=str(connected)))
 
+        # The line below injects a whole BUNCH of cell definitions into 
+        # the notebook. Relative to the top-level of the pyGSTi repo,
+        # these files should be located in the folder
+        #       pygsti/report/templates/report_notebook/
         nb.add_notebook_text_files([
-            templatePath / 'summary.txt',
-            templatePath / 'goodness.txt',
-            templatePath / 'gauge_invariant.txt',
-            templatePath / 'gauge_variant.txt'])
+            notebook_templates_path / 'summary.txt',
+            notebook_templates_path / 'goodness.txt',
+            notebook_templates_path / 'gauge_invariant.txt',
+            notebook_templates_path / 'gauge_variant.txt'])
 
         #Insert multi-dataset specific analysis
         if len(dsKeys) > 1:
@@ -352,15 +360,16 @@ def write_notebook(self, path, auto_open=False, connected=False, verbosity=0):
             dscmp_circuits = results_dict[dslbl1].circuit_lists['final']
             ds1 = results_dict[dslbl1].dataset
             ds2 = results_dict[dslbl2].dataset
-            dscmp = pygsti.baseobjs.DataComparator([ds1, ds2], ds_names=[dslbl1, dslbl2])
+            dscmp = pygsti.data.DataComparator([ds1, ds2], ds_names=[dslbl1, dslbl2])
+            dscmp.run()
             """.format(dsLbl1=dsKeys[0], dsLbl2=dsKeys[1]))
             nb.add_notebook_text_files([
-                templatePath / 'data_comparison.txt'])
+                notebook_templates_path / 'data_comparison.txt'])
 
         #Add reference material
         nb.add_notebook_text_files([
-            templatePath / 'input.txt',
-            templatePath / 'meta.txt'])
+            notebook_templates_path / 'input.txt',
+            notebook_templates_path / 'meta.txt'])
 
         printer.log("Report Notebook created as %s" % path)
 
@@ -457,3 +466,43 @@ def write_pdf(self, path, latex_cmd='pdflatex', latex_flags=None,
 
         printer.log("Compiling with `{} {}`".format(latex_cmd, ' '.join(latex_flags)))
         _merge.compile_latex_report(str(path.parent / path.stem), [latex_cmd] + latex_flags, printer, auto_open)
+
+    def write_results_dict(self, path_name: str) -> None:
+        if path_name.endswith('.pkl'):
+            with open(path_name, 'wb') as f:
+                _pickle.dump(self._results, f)
+        else:
+            path : _Path = _Path(path_name)
+            if not path.parent.exists():
+                raise ValueError(f'Parent folder of path_name={path_name} does not exist.')
+            if path.is_file():
+                raise ValueError(f'path_name={path_name} points to a file, but we require a folder.')
+            if not path.exists():
+                path.mkdir()
+            for dskey, mer in self._results.items():
+                mer.write(path / dskey)
+        return
+
+    @staticmethod
+    def results_dict_from_dir(path_name: str) -> dict[str, _ModelEstimateResults]:
+        if path_name.endswith('.pkl'):
+            with open(path_name, 'rb') as infile:
+                results_dict = _pickle.load(infile)
+        else:
+            path = _Path(path_name)
+            if path.is_file():
+                raise ValueError(f'path_name={path_name} points to a file, but we require a folder.')
+            results_dict = dict()
+            for child in path.iterdir():
+                if child.is_file():
+                    continue
+                inner_dict = dict()
+                for estname in _os.listdir(str(child / 'results')):
+                    res = _ModelEstimateResults.from_dir(str(child), name=estname)
+                    inner_dict[estname] = res
+                _, res = inner_dict.popitem()
+                for en, e in inner_dict.items():
+                    res.add_estimate(en, e.estimates[en])
+                results_dict[child.stem + child.suffix] = res
+        return results_dict
+

pygsti/report/templates/report_notebook/gauge_variant.txt

@@ -74,4 +74,4 @@ ws.ChoiTable(mdl, None, cri, display=("barplot",))
 A heat map of the Error Generator for each gate, which is the Lindbladian $\mathbb{L}$ that describes *how* the gate is failing to match the target, along with the result of projecting each generator onto some subspaces of the error generator space.
 @@code
 errgen_type = "logTiG" # or "logGTi" or "logG-logT"
-ws.ErrgenTable(mdl, target_model, cri, ("errgen","H","S","A"), "boxes", errgen_type)
+ws.ErrgenTable(mdl, target_model, cri, ("errgen","H","S","CA"), "boxes", errgen_type)

pygsti/report/workspacetables.py

@@ -1542,7 +1542,11 @@ def _create(self, model, target_model,
                 colHeadings.append('%sStochastic Projections' % basisPrefix)
             elif disp == "CA":
                 colHeadings.append('%sActive\\Correlation Projections' % basisPrefix)
-            else: raise ValueError("Invalid display element: %s" % disp)
+            else: 
+                msg = "Invalid display element: %s" % disp
+                if disp in {'C','A'}:
+                    msg += f'\nYou probably meant to use "CA" instead of {disp}.'
+                raise ValueError(msg)
 
         assert(display_as == "boxes" or display_as == "numbers")
         table = _ReportTable(colHeadings, (None,) * len(colHeadings),

pyproject.toml

@@ -84,6 +84,7 @@ testing_no_cython_mpi = [
     'pytest-xdist',
     'pytest-cov',
     'nbval',
+    'nbformat',
     'packaging',
     'zmq',
     'seaborn',

requirements.txt

@@ -5,4 +5,4 @@ plotly
 pandas
 networkx
 stim
-tqdm
+tqdm

rtd-requirements.txt

@@ -7,4 +7,4 @@ networkx
 numpydoc
 sphinx==6.2.1
 sphinx_rtd_theme>=1.2.2
-sphinx-autoapi
+sphinx-autoapi