Skip to content

Commit

Permalink
Add unit tests for checkpointing
Browse files Browse the repository at this point in the history 
This adds some unit tests for the checkpointing functionality. Decided to add this to the drivers test package for simplicity. Only test we may want to revisit is for the StandardGSTCheckpoint, which uses log-likelihoods temporarily due to unexpected sign behavior with CPTPLND parameterized models. May also want to add tests for when checkpoints are disabled and such, but then that is already largely covered already. Also refactored a few other tests to bring them in line with the recent deprecations of comma-separated strings for modes on StandardGST.
  • Loading branch information
Corey Ostrove committed Sep 6, 2023
1 parent 7d2e1e5 commit bf888fa
Show file tree
Hide file tree
Showing 5 changed files with 5,696 additions and 19 deletions.
2 changes: 1 addition & 1 deletion pygsti/protocols/gst.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1685,7 +1685,7 @@ def __init__(self, modes=('full TP','CPTPLND','Target'), gaugeopt_suite='stdgaug
+ " and may be removed in a future release. Please pass in a list or tuple"
+" (or other iterable) of strings")
else:
self.modes = list(modes) #Cast to a list for uniformity
self.modes = [modes] #Cast to a list for uniformity
else:
self.modes = modes
self.models_to_test = models_to_test
Expand Down
128 changes: 113 additions & 15 deletions test/test_packages/drivers/test_drivers.py
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,11 @@
import unittest

import numpy as np
from pygsti.forwardsims.mapforwardsim import MapForwardSimulator

from numpy.linalg import norm
import pygsti
from pygsti.modelpacks.legacy import std1Q_XYI as std
from ..testutils import BaseTestCase, compare_files, temp_files, regenerate_references


class DriversTestCase(BaseTestCase):

def setUp(self):
Expand Down Expand Up @@ -62,7 +61,8 @@ def test_longSequenceGST_fiducialPairReduction(self):
std.target_model(), std.fiducials, std.fiducials, std.germs, maxLens,
fid_pairs=fidPairs)

result = pygsti.run_long_sequence_gst_base(ds, std.target_model(), gfprStructs, verbosity=0)
result = pygsti.run_long_sequence_gst_base(ds, std.target_model(), gfprStructs, verbosity=0,
disable_checkpointing = True)
pygsti.report.create_standard_report(result, temp_files + "/full_report_GFPR",
"GFPR report", verbosity=2)

Expand All @@ -88,7 +88,8 @@ def test_longSequenceGST_fiducialPairReduction(self):
std.target_model(), std.fiducials, std.fiducials, std.germs, maxLens,
fid_pairs=fidPairsDict)

result = pygsti.run_long_sequence_gst_base(ds, std.target_model(), pfprStructs, verbosity=0)
result = pygsti.run_long_sequence_gst_base(ds, std.target_model(), pfprStructs, verbosity=0,
disable_checkpointing = True)
pygsti.report.create_standard_report(result, temp_files + "/full_report_PFPR",
"PFPR report", verbosity=2)

Expand All @@ -106,7 +107,8 @@ def test_longSequenceGST_randomReduction(self):
maxLens, fidPairs, ts, keep_fraction=0.5, keep_seed=1234)
result = self.runSilent(pygsti.run_long_sequence_gst_base,
ds, std.target_model(), reducedLists,
advanced_options={'truncScheme': ts})
advanced_options={'truncScheme': ts},
disable_checkpointing=True)

#create a report...
pygsti.report.create_standard_report(result, temp_files + "/full_report_RFPR",
Expand All @@ -120,7 +122,8 @@ def test_longSequenceGST_randomReduction(self):
maxLens, fidPairs, ts, keep_fraction=0.5, keep_seed=1234)
result2 = self.runSilent(pygsti.run_long_sequence_gst_base,
ds, std.target_model(), reducedLists,
advanced_options={'truncScheme': ts})
advanced_options={'truncScheme': ts},
disable_checkpointing=True)

#create a report...
pygsti.report.create_standard_report(result2, temp_files + "/full_report_RFPR2.html",
Expand All @@ -130,12 +133,12 @@ def test_longSequenceGST_CPTP(self):
ds = pygsti.data.DataSet(file_to_load_from=compare_files + "/drivers.dataset")

target_model = std.target_model()
target_model.set_all_parameterizations("CPTP")
target_model.set_all_parameterizations("CPTPLND")

maxLens = self.maxLens
result = self.runSilent(pygsti.run_long_sequence_gst,
ds, target_model, std.fiducials, std.fiducials,
std.germs, maxLens)
std.germs, maxLens, disable_checkpointing=True)

#create a report...
pygsti.report.create_standard_report(result, temp_files + "/full_report_CPTPGates",
Expand All @@ -151,7 +154,7 @@ def test_longSequenceGST_Sonly(self):
maxLens = self.maxLens
result = self.runSilent(pygsti.run_long_sequence_gst,
ds, target_model, std.fiducials, std.fiducials,
std.germs, maxLens)
std.germs, maxLens, disable_checkpointing=True)

#create a report...
pygsti.report.create_standard_report(result, temp_files + "/full_report_SGates.html",
Expand All @@ -174,7 +177,7 @@ def test_longSequenceGST_GLND(self):
maxLens = self.maxLens
result = self.runSilent(pygsti.run_long_sequence_gst,
ds, target_model, std.fiducials, std.fiducials,
std.germs, maxLens)
std.germs, maxLens, disable_checkpointing=True)

#create a report...
pygsti.report.create_standard_report(result, temp_files + "/full_report_SGates",
Expand All @@ -190,7 +193,7 @@ def test_longSequenceGST_HplusS(self):
maxLens = self.maxLens
result = self.runSilent(pygsti.run_long_sequence_gst,
ds, target_model, std.fiducials, std.fiducials,
std.germs, maxLens)
std.germs, maxLens, disable_checkpointing=True)

#create a report...
pygsti.report.create_standard_report(result, temp_files + "/full_report_HplusSGates",
Expand All @@ -205,7 +208,8 @@ def test_longSequenceGST_badfit(self):
maxLens = self.maxLens
result = self.runSilent(pygsti.run_long_sequence_gst,
ds, std.target_model(), std.fiducials, std.fiducials,
std.germs, maxLens, advanced_options={'bad_fit_threshold': -100})
std.germs, maxLens, advanced_options={'bad_fit_threshold': -100},
disable_checkpointing=True)

pygsti.report.create_standard_report(result, temp_files + "/full_report_badfit",
"badfit report", verbosity=2)
Expand All @@ -218,9 +222,10 @@ def test_stdpracticeGST(self):
maxLens = self.maxLens
result = self.runSilent(pygsti.run_stdpractice_gst,
ds, std.target_model().create_processor_spec(), std.fiducials, std.fiducials,
std.germs, maxLens, modes="full TP,CPTP,Test,Target",
std.germs, maxLens, modes=['full TP','CPTPLND','Test','Target'],
models_to_test = {"Test": mdl_guess},
comm=None, mem_limit=None, verbosity=5)
comm=None, mem_limit=None, verbosity=5,
disable_checkpointing=True)
pygsti.report.create_standard_report(result, temp_files + "/full_report_stdpractice",
"Std Practice Test Report", verbosity=2)

Expand All @@ -243,6 +248,99 @@ def test_bootstrap(self):
2, ds_defaultMaxLens, 'parametric', std.fiducials, std.fiducials,
std.germs, default_maxLens, input_model=mdl, target_model=tp_target,
return_data=False) #test when max_lengths == None ?

def test_GST_checkpointing(self):
ds = pygsti.data.DataSet(file_to_load_from=compare_files + "/drivers.dataset")
maxLens = self.maxLens

#Make list-of-lists of GST operation sequences
fullStructs = pygsti.circuits.make_lsgst_structs(
std.target_model(), std.fiducials, std.fiducials, std.germs, maxLens)

#Test GateSetTomographyCheckpoint:
#First run from scratch:
result_gst = pygsti.run_long_sequence_gst_base(ds, std.target_model(), fullStructs, verbosity=0,
checkpoint_path= temp_files + '/checkpoint_testing/GateSetTomography')

#double check that we can read in this checkpoint object correctly:
gst_checkpoint = pygsti.protocols.GateSetTomographyCheckpoint.read(temp_files + '/checkpoint_testing/GateSetTomography_iteration_0.json')

#run GST using this checkpoint
result_gst_warmstart = pygsti.run_long_sequence_gst_base(ds, std.target_model(), fullStructs, verbosity=0,
checkpoint = gst_checkpoint,
checkpoint_path= temp_files + '/checkpoint_testing/GateSetTomography')

diff = norm(result_gst.estimates['GateSetTomography'].models['final iteration estimate'].to_vector()-
result_gst_warmstart.estimates['GateSetTomography'].models['final iteration estimate'].to_vector())
print(f'{diff=}')
#Assert that this gives the same result as before:
self.assertTrue(diff<=1e-10)


def test_ModelTest_checkpointing(self):
ds = pygsti.data.DataSet(file_to_load_from=compare_files + "/drivers.dataset")
maxLens = self.maxLens

#Next test ModelTestCheckpoint
#First run from scratch:
result_modeltest = pygsti.run_model_test(std.target_model(), ds,std.target_model().create_processor_spec(),
std.fiducials, std.fiducials, std.germs,
maxLens, verbosity=0,
checkpoint_path= temp_files + '/checkpoint_testing/ModelTest')

#double check that we can read in this checkpoint object correctly:
model_test_checkpoint = pygsti.protocols.ModelTestCheckpoint.read(temp_files + '/checkpoint_testing/ModelTest_iteration_0.json')

#run GST using this checkpoint
result_modeltest_warmstart = pygsti.run_model_test(std.target_model(), ds,std.target_model().create_processor_spec(),
std.fiducials, std.fiducials, std.germs,
maxLens, verbosity=0,
checkpoint = model_test_checkpoint,
checkpoint_path= temp_files + '/checkpoint_testing/ModelTest')

diff = norm(np.array(result_modeltest.estimates['ModelTest'].parameters['model_test_values'])-
np.array(result_modeltest_warmstart.estimates['ModelTest'].parameters['model_test_values']))
#Assert that this gives the same result as before:
self.assertTrue(diff<=1e-10)



def test_StandardGST_checkpointing(self):
ds = pygsti.data.DataSet(file_to_load_from=compare_files + "/drivers.dataset")
maxLens = self.maxLens

#Finally test StandardGSTCheckpoint
#First run from scratch:
mdl_guess = std.target_model().depolarize(op_noise=0.01,spam_noise=0.01)

result_standardgst = pygsti.run_stdpractice_gst(ds, std.target_model().create_processor_spec(), std.fiducials, std.fiducials,
std.germs, maxLens, modes=['full TP','CPTPLND','Test','Target'],
models_to_test = {"Test": mdl_guess},
comm=None, mem_limit=None, verbosity=0,
checkpoint_path= temp_files + '/checkpoint_testing/StandardGST')

#double check that we can read in this checkpoint object correctly:
standardgst_checkpoint = pygsti.protocols.StandardGSTCheckpoint.read(temp_files + '/checkpoint_testing/StandardGST_CPTPLND_iteration_1.json')

#run GST using this checkpoint
result_standardgst_warmstart = pygsti.run_stdpractice_gst(ds, std.target_model().create_processor_spec(), std.fiducials, std.fiducials,
std.germs, maxLens, modes=['full TP','CPTPLND','Test','Target'],
models_to_test = {"Test": mdl_guess},
comm=None, mem_limit=None, verbosity=0,
checkpoint = standardgst_checkpoint,
checkpoint_path= temp_files + '/checkpoint_testing/StandardGST')

#Assert that this gives the same result as before:
#diff = norm(result_standardgst.estimates['CPTPLND'].models['final iteration estimate'].to_vector()-
# result_standardgst_warmstart.estimates['CPTPLND'].models['final iteration estimate'].to_vector())
diff = pygsti.tools.logl(result_standardgst.estimates['CPTPLND'].models['final iteration estimate'], ds)- \
pygsti.tools.logl(result_standardgst_warmstart.estimates['CPTPLND'].models['final iteration estimate'], ds)

diff1 = norm(result_standardgst.estimates['full TP'].models['final iteration estimate'].to_vector()-
result_standardgst_warmstart.estimates['full TP'].models['final iteration estimate'].to_vector())

self.assertTrue(abs(diff)<=1e-8)
self.assertTrue(diff1<=1e-10)


if __name__ == "__main__":
Expand Down
Loading

0 comments on commit bf888fa

Please sign in to comment.