Add ClassicalDataStore class to keep track of qubits measured

cirq-core/cirq/__init__.py

@@ -483,6 +483,9 @@
     canonicalize_half_turns,
     chosen_angle_to_canonical_half_turns,
     chosen_angle_to_half_turns,
+    ClassicalDataDictionaryStore,
+    ClassicalDataStore,
+    ClassicalDataStoreReader,
     Condition,
     Duration,
     DURATION_LIKE,
@@ -491,6 +494,7 @@
     LinearDict,
     MEASUREMENT_KEY_SEPARATOR,
     MeasurementKey,
+    MeasurementType,
     PeriodicValue,
     RANDOM_STATE_OR_SEED_LIKE,
     state_vector_to_probabilities,

cirq-core/cirq/contrib/quimb/mps_simulator.py

@@ -19,7 +19,7 @@
 
 import dataclasses
 import math
-from typing import Any, Dict, List, Optional, Sequence, Set, TYPE_CHECKING, Union
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
 import quimb.tensor as qtn
@@ -87,11 +87,11 @@ def __init__(
             seed=seed,
         )
 
-    def _create_partial_act_on_args(
+    def _create_partial_act_on_args_ex(
         self,
         initial_state: Union[int, 'MPSState'],
         qubits: Sequence['cirq.Qid'],
-        logs: Dict[str, Any],
+        classical_data: 'cirq.ClassicalDataStore',
     ) -> 'MPSState':
         """Creates MPSState args for simulating the Circuit.
 
@@ -101,7 +101,8 @@ def _create_partial_act_on_args(
             qubits: Determines the canonical ordering of the qubits. This
                 is often used in specifying the initial state, i.e. the
                 ordering of the computational basis states.
-            logs: A mutable object that measurements are recorded into.
+            classical_data: The shared classical data container for this
+                simulation.
 
         Returns:
             MPSState args for simulating the Circuit.
@@ -115,7 +116,7 @@ def _create_partial_act_on_args(
             simulation_options=self.simulation_options,
             grouping=self.grouping,
             initial_state=initial_state,
-            log_of_measurement_results=logs,
+            classical_data=classical_data,
         )
 
     def _create_step_result(
@@ -229,6 +230,7 @@ def __init__(
         grouping: Optional[Dict['cirq.Qid', int]] = None,
         initial_state: int = 0,
         log_of_measurement_results: Dict[str, Any] = None,
+        classical_data: 'cirq.ClassicalDataStore' = None,
     ):
         """Creates and MPSState
 
@@ -242,11 +244,13 @@ def __init__(
             initial_state: An integer representing the initial state.
             log_of_measurement_results: A mutable object that measurements are
                 being recorded into.
+            classical_data: The shared classical data container for this
+                simulation.
 
         Raises:
             ValueError: If the grouping does not cover the qubits.
         """
-        super().__init__(prng, qubits, log_of_measurement_results)
+        super().__init__(prng, qubits, log_of_measurement_results, classical_data)
         qubit_map = self.qubit_map
         self.grouping = qubit_map if grouping is None else grouping
         if self.grouping.keys() != self.qubit_map.keys():

cirq-core/cirq/contrib/quimb/mps_simulator_test.py

@@ -539,10 +539,10 @@ def test_state_act_on_args_initializer():
     s = ccq.mps_simulator.MPSState(
         qubits=(cirq.LineQubit(0),),
         prng=np.random.RandomState(0),
-        log_of_measurement_results={'test': 4},
+        log_of_measurement_results={'test': [4]},
     )
     assert s.qubits == (cirq.LineQubit(0),)
-    assert s.log_of_measurement_results == {'test': 4}
+    assert s.log_of_measurement_results == {'test': [4]}
 
 
 def test_act_on_gate():

cirq-core/cirq/json_resolver_cache.py

@@ -64,6 +64,7 @@ def _parallel_gate_op(gate, qubits):
         'Circuit': cirq.Circuit,
         'CircuitOperation': cirq.CircuitOperation,
         'ClassicallyControlledOperation': cirq.ClassicallyControlledOperation,
+        'ClassicalDataDictionaryStore': cirq.ClassicalDataDictionaryStore,
         'CliffordState': cirq.CliffordState,
         'CliffordTableau': cirq.CliffordTableau,
         'CNotPowGate': cirq.CNotPowGate,
@@ -104,6 +105,7 @@ def _parallel_gate_op(gate, qubits):
         'MixedUnitaryChannel': cirq.MixedUnitaryChannel,
         'MeasurementKey': cirq.MeasurementKey,
         'MeasurementGate': cirq.MeasurementGate,
+        'MeasurementType': cirq.MeasurementType,
         '_MeasurementSpec': cirq.work._MeasurementSpec,
         'Moment': cirq.Moment,
         'MutableDensePauliString': cirq.MutableDensePauliString,

cirq-core/cirq/ops/classically_controlled_operation.py

@@ -148,7 +148,6 @@ def _circuit_diagram_info_(
         sub_info = protocols.circuit_diagram_info(self._sub_operation, sub_args, None)
         if sub_info is None:
             return NotImplemented  # coverage: ignore
-
         control_count = len({k for c in self._conditions for k in c.keys})
         wire_symbols = sub_info.wire_symbols + ('^',) * control_count
         if any(not isinstance(c, value.KeyCondition) for c in self._conditions):
@@ -176,8 +175,8 @@ def _json_dict_(self) -> Dict[str, Any]:
             'sub_operation': self._sub_operation,
         }
 
-    def _act_on_(self, args: 'cirq.ActOnArgs') -> bool:
-        if all(c.resolve(args.log_of_measurement_results) for c in self._conditions):
+    def _act_on_(self, args: 'cirq.OperationTarget') -> bool:
+        if all(c.resolve(args.classical_data) for c in self._conditions):
             protocols.act_on(self._sub_operation, args)
         return True
 

cirq-core/cirq/ops/classically_controlled_operation_test.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+import numpy as np
 import pytest
 import sympy
 from sympy.parsing import sympy_parser
@@ -702,6 +704,40 @@ def test_sympy():
             assert result.measurements['m_result'][0][0] == (j > i)
 
 
+def test_sympy_qudits():
+    q0 = cirq.LineQid(0, 3)
+    q1 = cirq.LineQid(1, 5)
+    q_result = cirq.LineQubit(2)
+
+    class PlusGate(cirq.Gate):
+        def __init__(self, dimension, increment=1):
+            self.dimension = dimension
+            self.increment = increment % dimension
+
+        def _qid_shape_(self):
+            return (self.dimension,)
+
+        def _unitary_(self):
+            inc = (self.increment - 1) % self.dimension + 1
+            u = np.empty((self.dimension, self.dimension))
+            u[inc:] = np.eye(self.dimension)[:-inc]
+            u[:inc] = np.eye(self.dimension)[-inc:]
+            return u
+
+    for i in range(15):
+        digits = cirq.big_endian_int_to_digits(i, digit_count=2, base=(3, 5))
+        circuit = cirq.Circuit(
+            PlusGate(3, digits[0]).on(q0),
+            PlusGate(5, digits[1]).on(q1),
+            cirq.measure(q0, q1, key='m'),
+            cirq.X(q_result).with_classical_controls(sympy_parser.parse_expr('m % 4 <= 1')),
+            cirq.measure(q_result, key='m_result'),
+        )
+
+        result = cirq.Simulator().run(circuit)
+        assert result.measurements['m_result'][0][0] == (i % 4 <= 1)
+
+
 def test_sympy_path_prefix():
     q = cirq.LineQubit(0)
     op = cirq.X(q).with_classical_controls(sympy.Symbol('b'))

cirq-core/cirq/protocols/json_test_data/ClassicalDataDictionaryStore.json

@@ -0,0 +1,60 @@
+{
+  "cirq_type": "ClassicalDataDictionaryStore",
+  "measurements": [
+    [
+      {
+        "cirq_type": "MeasurementKey",
+        "name": "m",
+        "path": []
+      },
+      [0, 1]
+    ]
+  ],
+  "measured_qubits": [
+    [
+      {
+        "cirq_type": "MeasurementKey",
+        "name": "m",
+        "path": []
+      },
+      [
+        {
+          "cirq_type": "LineQubit",
+          "x": 0
+        },
+        {
+          "cirq_type": "LineQubit",
+          "x": 1
+        }
+      ]
+    ]
+  ],
+  "channel_measurements": [
+    [
+      {
+        "cirq_type": "MeasurementKey",
+        "name": "c",
+        "path": []
+      },
+      3
+    ]
+  ],
+  "measurement_types": [
+    [
+      {
+        "cirq_type": "MeasurementKey",
+        "name": "m",
+        "path": []
+      },
+      1
+    ],
+    [
+      {
+        "cirq_type": "MeasurementKey",
+        "name": "c",
+        "path": []
+      },
+      2
+    ]
+  ]
+}

cirq-core/cirq/protocols/json_test_data/ClassicalDataDictionaryStore.repr

@@ -0,0 +1 @@
+cirq.ClassicalDataDictionaryStore(_measurements={cirq.MeasurementKey('m'): [0, 1]}, _measured_qubits={cirq.MeasurementKey('m'): [cirq.LineQubit(0), cirq.LineQubit(1)]}, _channel_measurements={cirq.MeasurementKey('c'): 3}, _measurement_types={cirq.MeasurementKey('m'): cirq.MeasurementType.MEASUREMENT, cirq.MeasurementKey('c'): cirq.MeasurementType.CHANNEL})

cirq-core/cirq/protocols/json_test_data/MeasurementType.json

@@ -0,0 +1 @@
+[1, 2]

cirq-core/cirq/protocols/json_test_data/MeasurementType.repr

@@ -0,0 +1 @@
+[cirq.MeasurementType.MEASUREMENT, cirq.MeasurementType.CHANNEL]

cirq-core/cirq/protocols/measurement_key_protocol.py

@@ -20,9 +20,6 @@
 from cirq import value
 from cirq._doc import doc_private
 
-if TYPE_CHECKING:
-    import cirq
-
 if TYPE_CHECKING:
     import cirq
 

cirq-core/cirq/sim/act_on_args.py

@@ -29,7 +29,7 @@
 
 import numpy as np
 
-from cirq import protocols, ops
+from cirq import protocols, ops, value
 from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
 from cirq.sim.operation_target import OperationTarget
 
@@ -47,6 +47,7 @@ def __init__(
         prng: np.random.RandomState = None,
         qubits: Sequence['cirq.Qid'] = None,
         log_of_measurement_results: Dict[str, List[int]] = None,
+        classical_data: 'cirq.ClassicalDataStore' = None,
         ignore_measurement_results: bool = False,
     ):
         """Inits ActOnArgs.
@@ -59,6 +60,8 @@ def __init__(
                 ordering of the computational basis states.
             log_of_measurement_results: A mutable object that measurements are
                 being recorded into.
+            classical_data: The shared classical data container for this
+                simulation.
             ignore_measurement_results: If True, then the simulation
                 will treat measurement as dephasing instead of collapsing
                 process, and not log the result. This is only applicable to
@@ -68,11 +71,11 @@ def __init__(
             prng = cast(np.random.RandomState, np.random)
         if qubits is None:
             qubits = ()
-        if log_of_measurement_results is None:
-            log_of_measurement_results = {}
         self._set_qubits(qubits)
         self.prng = prng
-        self._log_of_measurement_results = log_of_measurement_results
+        # pylint: disable=line-too-long
+        self._classical_data = classical_data or value.ClassicalDataDictionaryStore(_measurements=log_of_measurement_results)  # type: ignore
+        # pylint: enable=line-too-long
         self._ignore_measurement_results = ignore_measurement_results
 
     def _set_qubits(self, qubits: Sequence['cirq.Qid']):
@@ -101,9 +104,9 @@ def measure(self, qubits: Sequence['cirq.Qid'], key: str, invert_mask: Sequence[
             return
         bits = self._perform_measurement(qubits)
         corrected = [bit ^ (bit < 2 and mask) for bit, mask in zip(bits, invert_mask)]
-        if key in self._log_of_measurement_results:
-            raise ValueError(f"Measurement already logged to key {key!r}")
-        self._log_of_measurement_results[key] = corrected
+        self._classical_data.record_measurement(
+            value.MeasurementKey.parse_serialized(key), corrected, qubits
+        )
 
     def get_axes(self, qubits: Sequence['cirq.Qid']) -> List[int]:
         return [self.qubit_map[q] for q in qubits]
@@ -117,7 +120,7 @@ def copy(self: TSelf) -> TSelf:
         """Creates a copy of the object."""
         args = copy.copy(self)
         self._on_copy(args)
-        args._log_of_measurement_results = self.log_of_measurement_results.copy()
+        args._classical_data = self._classical_data.copy()
         return args
 
     def _on_copy(self: TSelf, args: TSelf):
@@ -194,8 +197,8 @@ def _on_transpose_to_qubit_order(self: TSelf, qubits: Sequence['cirq.Qid'], targ
         functionality, if supported."""
 
     @property
-    def log_of_measurement_results(self) -> Dict[str, List[int]]:
-        return self._log_of_measurement_results
+    def classical_data(self) -> 'cirq.ClassicalDataStoreReader':
+        return self._classical_data
 
     @property
     def ignore_measurement_results(self) -> bool: