Add support for converting OBSERVABLE_INCLUDE instructions with pauli targets to stimcirq (#952)

Author: AlexBourassa

glue/cirq/stimcirq/_obs_annotation.py

@@ -16,6 +16,7 @@ def __init__(
         *,
         parity_keys: Iterable[str] = (),
         relative_keys: Iterable[int] = (),
+        pauli_keys: Iterable[str] = (),
         observable_index: int,
     ):
         """
@@ -28,6 +29,7 @@ def __init__(
         """
         self.parity_keys = frozenset(parity_keys)
         self.relative_keys = frozenset(relative_keys)
+        self.pauli_keys = frozenset(pauli_keys)
         self.observable_index = observable_index
 
     @property
@@ -38,11 +40,12 @@ def with_qubits(self, *new_qubits) -> 'CumulativeObservableAnnotation':
         return self
 
     def _value_equality_values_(self) -> Any:
-        return self.parity_keys, self.relative_keys, self.observable_index
+        return self.parity_keys, self.relative_keys, self.pauli_keys, self.observable_index
 
     def _circuit_diagram_info_(self, args: Any) -> str:
         items: List[str] = [repr(e) for e in sorted(self.parity_keys)]
         items += [f'rec[{e}]' for e in sorted(self.relative_keys)]
+        items += sorted(self.pauli_keys)
         k = ",".join(str(e) for e in items)
         return f"Obs{self.observable_index}({k})"
 
@@ -51,6 +54,7 @@ def __repr__(self) -> str:
             f'stimcirq.CumulativeObservableAnnotation('
             f'parity_keys={sorted(self.parity_keys)}, '
             f'relative_keys={sorted(self.relative_keys)}, '
+            f'pauli_keys={sorted(self.pauli_keys)}, '
             f'observable_index={self.observable_index!r})'
         )
 
@@ -62,6 +66,7 @@ def _json_dict_(self) -> Dict[str, Any]:
         result = {
             'parity_keys': sorted(self.parity_keys),
             'observable_index': self.observable_index,
+            'pauli_keys': sorted(self.pauli_keys),
         }
         if self.relative_keys:
             result['relative_keys'] = sorted(self.relative_keys)
@@ -104,6 +109,12 @@ def _stim_conversion_(
                 rec_targets.append(stim.target_rec(-1 - offset))
                 if not remaining:
                     break
+        rec_targets.extend(
+            [
+                stim.target_pauli(qubit_index=int(k[1:]), pauli=k[0]) 
+                for k in sorted(self.pauli_keys)
+            ]
+        )
         if remaining:
             raise ValueError(
                 f"{self!r} was processed before measurements it referenced ({sorted(remaining)!r})."

glue/cirq/stimcirq/_obs_annotation_test.py

@@ -191,9 +191,30 @@ def test_json_serialization():
     c2 = cirq.read_json(json_text=json, resolvers=[*cirq.DEFAULT_RESOLVERS, stimcirq.JSON_RESOLVER])
     assert c == c2
 
+def test_json_serialization_with_pauli_keys():
+    c = cirq.Circuit(
+        stimcirq.CumulativeObservableAnnotation(parity_keys=["a", "b"], observable_index=5, pauli_keys=["X0", "Y1", "Z2"]),
+        stimcirq.CumulativeObservableAnnotation(
+            parity_keys=["a", "b"], relative_keys=[-1, -3], observable_index=5, pauli_keys=["X0", "Y1", "Z2"]
+        ),
+        stimcirq.CumulativeObservableAnnotation(observable_index=2, pauli_keys=["X0", "Y1", "Z2"]),
+        stimcirq.CumulativeObservableAnnotation(parity_keys=["d", "c"], observable_index=5, pauli_keys=["X0", "Y1", "Z2"]),
+    )
+    json = cirq.to_json(c)
+    c2 = cirq.read_json(json_text=json, resolvers=[*cirq.DEFAULT_RESOLVERS, stimcirq.JSON_RESOLVER])
+    assert c == c2
+
 
 def test_json_backwards_compat_exact():
     raw = stimcirq.CumulativeObservableAnnotation(parity_keys=['z'], relative_keys=[-2], observable_index=5)
-    packed = '{\n  "cirq_type": "CumulativeObservableAnnotation",\n  "parity_keys": [\n    "z"\n  ],\n  "observable_index": 5,\n  "relative_keys": [\n    -2\n  ]\n}'
-    assert cirq.read_json(json_text=packed, resolvers=[*cirq.DEFAULT_RESOLVERS, stimcirq.JSON_RESOLVER]) == raw
-    assert cirq.to_json(raw) == packed
+    packed_v1 = '{\n  "cirq_type": "CumulativeObservableAnnotation",\n  "parity_keys": [\n    "z"\n  ],\n  "observable_index": 5,\n  "relative_keys": [\n    -2\n  ]\n}'
+    packed_v2 ='{\n  "cirq_type": "CumulativeObservableAnnotation",\n  "parity_keys": [\n    "z"\n  ],\n  "observable_index": 5,\n  "pauli_keys": [],\n  "relative_keys": [\n    -2\n  ]\n}'
+    assert cirq.read_json(json_text=packed_v1, resolvers=[*cirq.DEFAULT_RESOLVERS, stimcirq.JSON_RESOLVER]) == raw
+    assert cirq.read_json(json_text=packed_v2, resolvers=[*cirq.DEFAULT_RESOLVERS, stimcirq.JSON_RESOLVER]) == raw
+    assert cirq.to_json(raw) == packed_v2
+
+    # With pauli_keys
+    raw = stimcirq.CumulativeObservableAnnotation(parity_keys=['z'], relative_keys=[-2], observable_index=5, pauli_keys=["X0", "Y1", "Z2"])
+    packed_v2 ='{\n  "cirq_type": "CumulativeObservableAnnotation",\n  "parity_keys": [\n    "z"\n  ],\n  "observable_index": 5,\n  "pauli_keys": [\n    "X0",\n    "Y1",\n    "Z2"\n  ],\n  "relative_keys": [\n    -2\n  ]\n}'
+    assert cirq.read_json(json_text=packed_v2, resolvers=[*cirq.DEFAULT_RESOLVERS, stimcirq.JSON_RESOLVER]) == raw
+    assert cirq.to_json(raw) == packed_v2

glue/cirq/stimcirq/_stim_to_cirq.py

@@ -340,19 +340,26 @@ def coords_after_offset(
 
     def resolve_measurement_record_keys(
         self, targets: Iterable[stim.GateTarget]
-    ) -> Tuple[List[str], List[int]]:
+    ) -> Tuple[List[str], List[int], List[str]]:
+        pauli_targets, meas_targets = [], []
+        for t in targets:
+            if t.is_measurement_record_target:
+                meas_targets.append(t)
+            else:
+                pauli_targets.append(f'{t.pauli_type}{t.value}')
+
         if self.have_seen_loop:
-            return [], [t.value for t in targets]
+            return [], [t.value for t in meas_targets], pauli_targets
         else:
-            return [str(self.num_measurements_seen + t.value) for t in targets], []
+            return [str(self.num_measurements_seen + t.value) for t in meas_targets], [], pauli_targets
 
     def process_detector(self, instruction: stim.CircuitInstruction) -> None:
         if instruction.tag:
             tags = [instruction.tag]
         else:
             tags = ()
         coords = self.coords_after_offset(instruction.gate_args_copy())
-        keys, rels = self.resolve_measurement_record_keys(instruction.targets_copy())
+        keys, rels, _ = self.resolve_measurement_record_keys(instruction.targets_copy())
         self.append_operation(
             DetAnnotation(parity_keys=keys, relative_keys=rels, coordinate_metadata=coords).with_tags(*tags)
         )
@@ -364,10 +371,10 @@ def process_observable_include(self, instruction: stim.CircuitInstruction) -> No
             tags = ()
         args = instruction.gate_args_copy()
         index = 0 if not args else int(args[0])
-        keys, rels = self.resolve_measurement_record_keys(instruction.targets_copy())
+        keys, rels, paulis = self.resolve_measurement_record_keys(instruction.targets_copy())
         self.append_operation(
             CumulativeObservableAnnotation(
-                parity_keys=keys, relative_keys=rels, observable_index=index
+                parity_keys=keys, relative_keys=rels, pauli_keys=paulis, observable_index=index
             ).with_tags(*tags)
         )
 

glue/cirq/stimcirq/_stim_to_cirq_test.py

@@ -763,3 +763,19 @@ def test_id_error_round_trip():
     cirq_circuit = stimcirq.stim_circuit_to_cirq_circuit(stim_circuit)
     restored_circuit = stimcirq.cirq_circuit_to_stim_circuit(cirq_circuit)
     assert restored_circuit == stim_circuit
+
+def test_round_trip_with_pauli_obs():
+    stim_circuit = stim.Circuit("""
+        QUBIT_COORDS(5, 5) 0
+        R 0
+        OBSERVABLE_INCLUDE(0) X0
+        TICK
+        H 0
+        TICK
+        M 0
+        OBSERVABLE_INCLUDE(0) rec[-1]
+        TICK
+    """)
+    cirq_circuit = stimcirq.stim_circuit_to_cirq_circuit(stim_circuit)
+    restored_circuit = stimcirq.cirq_circuit_to_stim_circuit(cirq_circuit)
+    assert restored_circuit == stim_circuit