Add option to store state of Model and Detectors

surface_sim/detectors/detectors.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+import pathlib
 from collections.abc import Callable, Collection, Mapping, Sequence
 from copy import deepcopy
 from typing import TypeVar
 
 import stim
+import yaml
 
 from ..layouts.layout import Layout
 
@@ -116,6 +118,44 @@ def new_circuit(self):
         self.gauge_detectors: set[str] = set()
         return
 
+    def store_state(self, filename: str | pathlib.Path) -> None:
+        """Stores the current state to the given YAML file.
+        This is useful in a conditioned circuit to not encode the first part
+        of the circuit for every realization."""
+        state = {
+            "anc_qubit_labels": deepcopy(self.anc_qubit_labels),
+            "frame": deepcopy(self.frame),
+            "anc_coords": deepcopy(self.anc_coords),
+            "include_gauge_dets": deepcopy(self.include_gauge_dets),
+            "detectors": deepcopy(self.detectors),
+            "num_rounds": deepcopy(self.num_rounds),
+            "total_num_rounds": deepcopy(self.total_num_rounds),
+            "update_dict_list": deepcopy(self.update_dict_list),
+            "gauge_detectors": deepcopy(self.gauge_detectors),
+        }
+        with open(filename, "w") as file:
+            yaml.dump(state, file)
+        return
+
+    @classmethod
+    def load_state(cls, filename: str | pathlib.Path) -> "Detectors":
+        """Loads the state inside the given YAML file.
+        See ``Model.store_state`` for more information."""
+        with open(filename, "r") as file:
+            state = yaml.safe_load(file)
+        detectors = cls(
+            anc_qubits=state["anc_qubit_labels"],
+            frame=state["frame"],
+            anc_coords=state["anc_coords"],
+            include_gauge_dets=state["include_gauge_dets"],
+        )
+        detectors.detectors = state["detectors"]
+        detectors.num_rounds = state["num_rounds"]
+        detectors.total_num_rounds = state["total_num_rounds"]
+        detectors.update_dict_list = state["update_dict_list"]
+        detectors.gauge_detectors = state["gauge_detectors"]
+        return detectors
+
     def activate_detectors(
         self, anc_qubits: Collection[str], gauge_dets: Collection[str] | None = None
     ):

surface_sim/models/model.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import pathlib
 from collections.abc import Collection, Sequence
 from copy import deepcopy
 
+import yaml
 from stim import Circuit, CircuitInstruction, GateTarget, target_rec
 from typing_extensions import override
 
@@ -181,6 +183,35 @@ def new_circuit(self) -> None:
         self._new_op: str = ""
         return
 
+    def store_state(self, filename: str | pathlib.Path) -> None:
+        """Stores the current state to the given YAML file.
+        This is useful in a conditioned circuit to not encode the first part
+        of the circuit for every realization."""
+        state = {
+            "meas_order": deepcopy(self._meas_order),
+            "num_meas": deepcopy(self._num_meas),
+            "last_op": deepcopy(self._last_op),
+            "new_op": deepcopy(self._new_op),
+            "setup": deepcopy(self._setup.to_dict()),
+            "qubit_inds": deepcopy(self._qubit_inds),
+        }
+        with open(filename, "w") as file:
+            yaml.dump(state, file)
+        return
+
+    @classmethod
+    def load_state(cls, filename: str | pathlib.Path) -> "Model":
+        """Loads the state inside the given YAML file.
+        See ``Model.store_state`` for more information."""
+        with open(filename, "r") as file:
+            state = yaml.safe_load(file)
+        model = cls(setup=Setup(state["setup"]), qubit_inds=state["qubit_inds"])
+        model._meas_order = state["meas_order"]
+        model._num_meas = state["num_meas"]
+        model._last_op = state["last_op"]
+        model._new_op = state["new_op"]
+        return model
+
     # annotation operations
     def tick(self) -> Circuit:
         if self._last_op != "tick":

surface_sim/setups/setup.py

@@ -136,14 +136,19 @@ def __init__(self, setup: SetupDict = dict(setup=[{}])) -> None:
         self._mode: str = "standard"
         self._qubit_params: dict[str | tuple[str, ...], dict[str, Param]] = dict()
         self._global_params: dict[str, Param] = dict()
-        self._var_params: dict[str, Param] = dict()
         self.uniform: bool = False
 
         _setup: SetupDict = deepcopy(setup)
         self.name: str | None = _setup.pop("name", None)
         self.description: str | None = _setup.pop("description", None)
         self._gate_durations: dict[str, float | int] = _setup.pop("gate_durations", {})
+
+        # self._load_setup requires self._var_params to be initialized.
+        # load var params after self._load_setup because it sets all of them to 'None'.
+        self._var_params: dict[str, Param] = {}
         self._load_setup(_setup)
+        self._var_params |= _setup.pop("var_params", {})
+
         if self._qubit_params == {}:
             self.uniform = True
 
@@ -183,6 +188,7 @@ def _load_setup(self, setup: SetupDict) -> None:
                 if isinstance(val, str):
                     for p in _get_var_params(val):
                         self._var_params[p] = None
+        return
 
     @property
     def free_params(self) -> list[str]:
@@ -225,6 +231,9 @@ def to_dict(self) -> SetupDict:
         setup["name"] = self.name
         setup["description"] = self.description
         setup["gate_durations"] = self._gate_durations
+        setup["var_params"] = {
+            k: v for k, v in self._var_params.items() if v is not None
+        }
 
         qubit_params: list[dict[str, Param]] = []
         if self._global_params:
@@ -356,7 +365,9 @@ def var_param(self, var_param: str) -> Param:
             return self.var_param(self.PARENTS[var_param])
 
         if val is None:
-            raise ValueError(f"Variable param {var_param} not in 'Setup.free_params'.")
+            raise ValueError(
+                f"Variable param {var_param} not specified or does not exist."
+            )
         return val
 
     def set_var_param(self, var_param: str, val: Param) -> None:

tests/detectors/test_detectors.py

@@ -349,3 +349,22 @@ def meas_rec(q, t):
     assert sorted([-110, -101, -201]) in detector_rec
 
     return
+
+
+def test_store_and_load_state(tmp_path_factory):
+    anc_qubits = ["X1", "Z1"]
+    detectors = Detectors(anc_qubits=anc_qubits, frame="pre-gate")
+    detectors.activate_detectors(anc_qubits, [])
+    new_stabs = {"X1": set(["X1", "Z1"]), "Z1": set(["Z1"])}
+    new_stabs_inv = {"X1": set(["X1", "Z1"]), "Z1": set(["Z1"])}
+    detectors.update(new_stabs, new_stabs_inv)
+
+    path = tmp_path_factory.mktemp("model_states")
+    detectors.store_state(path / "si1000.yaml")
+
+    new_detectors = Detectors.load_state(path / "si1000.yaml")
+
+    assert new_detectors.frame == detectors.frame
+    assert new_detectors.total_num_rounds == detectors.total_num_rounds
+
+    return

tests/models/test_model.py

@@ -1,4 +1,5 @@
 from surface_sim import Model, Setup
+from surface_sim.models import SI1000NoiseModel
 
 SETUP = {
     "gate_durations": {
@@ -61,3 +62,18 @@ def test_new_circuit():
     assert model._num_meas == 0
 
     return
+
+
+def test_store_and_load_state(tmp_path_factory):
+    model = SI1000NoiseModel({"D1": 300, "d3": 2})
+    model.setup.set_var_param("prob", 1e-3)
+
+    path = tmp_path_factory.mktemp("model_states")
+    model.store_state(path / "si1000.yaml")
+
+    new_model = SI1000NoiseModel.load_state(path / "si1000.yaml")
+
+    assert model.setup.var_param("prob") == new_model.setup.var_param("prob")
+    assert model._meas_order == new_model._meas_order
+
+    return

tests/setups/test_setup.py

@@ -34,6 +34,7 @@
     ],
     "name": "test",
     "description": "test description",
+    "var_params": {},
 }