Classical sympy conditions with qudits

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
@@ -92,6 +92,7 @@ def _create_partial_act_on_args(
         initial_state: Union[int, 'MPSState'],
         qubits: Sequence['cirq.Qid'],
         logs: Dict[str, Any],
+        measured_qubits: Dict[str, Tuple['cirq.Qid', ...]],
     ) -> 'MPSState':
         """Creates MPSState args for simulating the Circuit.
 
@@ -102,6 +103,8 @@ def _create_partial_act_on_args(
                 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.
+            measured_qubits: A dictionary that contains the qubits that were
+                measured in each measurement.
 
         Returns:
             MPSState args for simulating the Circuit.
@@ -116,6 +119,7 @@ def _create_partial_act_on_args(
             grouping=self.grouping,
             initial_state=initial_state,
             log_of_measurement_results=logs,
+            measured_qubits=measured_qubits,
         )
 
     def _create_step_result(
@@ -229,6 +233,7 @@ def __init__(
         grouping: Optional[Dict['cirq.Qid', int]] = None,
         initial_state: int = 0,
         log_of_measurement_results: Dict[str, Any] = None,
+        measured_qubits: Dict[str, Tuple['cirq.Qid', ...]] = None,
     ):
         """Creates and MPSState
 
@@ -242,11 +247,13 @@ def __init__(
             initial_state: An integer representing the initial state.
             log_of_measurement_results: A mutable object that measurements are
                 being recorded into.
+            measured_qubits: A dictionary that contains the qubits that were
+                measured in each measurement.
 
         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, measured_qubits)
         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/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,9 @@ 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:
+        measurements, qubits = args.log_of_measurement_results, args.measured_qubits
+        if all(c.resolve(measurements, qubits) 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, 3)
+    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(9):
+        digits = cirq.big_endian_int_to_digits(i, digit_count=2, base=3)
+        circuit = cirq.Circuit(
+            PlusGate(3, digits[0]).on(q0),
+            PlusGate(3, digits[1]).on(q1),
+            cirq.measure(q0, q1, key='m'),
+            cirq.X(q_result).with_classical_controls(sympy_parser.parse_expr('m > 4')),
+            cirq.measure(q_result, key='m_result'),
+        )
+
+        result = cirq.Simulator().run(circuit)
+        assert result.measurements['m_result'][0][0] == (i > 4)
+
+
 def test_sympy_path_prefix():
     q = cirq.LineQubit(0)
     op = cirq.X(q).with_classical_controls(sympy.Symbol('b'))

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

@@ -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,
+        measured_qubits: Dict[str, Tuple['cirq.Qid', ...]] = 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.
+            measured_qubits: A dictionary that contains the qubits that were
+                measured in each measurement.
             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
@@ -70,9 +73,12 @@ def __init__(
             qubits = ()
         if log_of_measurement_results is None:
             log_of_measurement_results = {}
+        if measured_qubits is None:
+            measured_qubits = {}
         self._set_qubits(qubits)
         self.prng = prng
         self._log_of_measurement_results = log_of_measurement_results
+        self._measured_qubits = measured_qubits
         self._ignore_measurement_results = ignore_measurement_results
 
     def _set_qubits(self, qubits: Sequence['cirq.Qid']):
@@ -104,6 +110,7 @@ def measure(self, qubits: Sequence['cirq.Qid'], key: str, invert_mask: Sequence[
         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._measured_qubits[key] = tuple(qubits)
 
     def get_axes(self, qubits: Sequence['cirq.Qid']) -> List[int]:
         return [self.qubit_map[q] for q in qubits]
@@ -118,6 +125,7 @@ def copy(self: TSelf) -> TSelf:
         args = copy.copy(self)
         self._on_copy(args)
         args._log_of_measurement_results = self.log_of_measurement_results.copy()
+        args._measured_qubits = self.measured_qubits.copy()
         return args
 
     def _on_copy(self: TSelf, args: TSelf):
@@ -197,6 +205,10 @@ def _on_transpose_to_qubit_order(self: TSelf, qubits: Sequence['cirq.Qid'], targ
     def log_of_measurement_results(self) -> Dict[str, List[int]]:
         return self._log_of_measurement_results
 
+    @property
+    def measured_qubits(self) -> Dict[str, Tuple['cirq.Qid', ...]]:
+        return self._measured_qubits
+
     @property
     def ignore_measurement_results(self) -> bool:
         return self._ignore_measurement_results

cirq-core/cirq/sim/act_on_args_container.py

@@ -14,15 +14,16 @@
 
 from collections import abc
 from typing import (
+    Any,
     Dict,
-    TYPE_CHECKING,
     Generic,
-    Sequence,
-    Optional,
     Iterator,
-    Any,
-    Tuple,
     List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    TYPE_CHECKING,
     Union,
 )
 
@@ -51,6 +52,7 @@ def __init__(
         qubits: Sequence['cirq.Qid'],
         split_untangled_states: bool,
         log_of_measurement_results: Dict[str, Any],
+        measured_qubits: Dict[str, Tuple['cirq.Qid', ...]] = None,
     ):
         """Initializes the class.
 
@@ -63,11 +65,14 @@ def __init__(
                 at the end.
             log_of_measurement_results: A mutable object that measurements are
                 being recorded into.
+            measured_qubits: A dictionary that contains the qubits that were
+                measured in each measurement.
         """
         self.args = args
         self._qubits = tuple(qubits)
         self.split_untangled_states = split_untangled_states
         self._log_of_measurement_results = log_of_measurement_results
+        self._measured_qubits = measured_qubits if measured_qubits is not None else {}
 
     def create_merged_state(self) -> TActOnArgs:
         if not self.split_untangled_states:
@@ -132,9 +137,11 @@ def _act_on_fallback_(
 
     def copy(self) -> 'cirq.ActOnArgsContainer[TActOnArgs]':
         logs = self.log_of_measurement_results.copy()
+        measured_qubits = self._measured_qubits.copy()
         copies = {a: a.copy() for a in set(self.args.values())}
         for copy in copies.values():
             copy._log_of_measurement_results = logs
+            copy._measured_qubits = measured_qubits
         args = {q: copies[a] for q, a in self.args.items()}
         return ActOnArgsContainer(args, self.qubits, self.split_untangled_states, logs)
 
@@ -146,6 +153,10 @@ def qubits(self) -> Tuple['cirq.Qid', ...]:
     def log_of_measurement_results(self) -> Dict[str, Any]:
         return self._log_of_measurement_results
 
+    @property
+    def measured_qubits(self) -> Mapping[str, Tuple['cirq.Qid', ...]]:
+        return self._measured_qubits
+
     def sample(
         self,
         qubits: List['cirq.Qid'],

cirq-core/cirq/sim/act_on_density_matrix_args.py

@@ -41,6 +41,7 @@ def __init__(
         prng: np.random.RandomState = None,
         log_of_measurement_results: Dict[str, Any] = None,
         qubits: Sequence['cirq.Qid'] = None,
+        measured_qubits: Dict[str, Tuple['cirq.Qid', ...]] = None,
         ignore_measurement_results: bool = False,
     ):
         """Inits ActOnDensityMatrixArgs.
@@ -61,12 +62,20 @@ def __init__(
                 effects.
             log_of_measurement_results: A mutable object that measurements are
                 being recorded into.
+            measured_qubits: A dictionary that contains the qubits that were
+                measured in each measurement.
             ignore_measurement_results: If True, then the simulation
                 will treat measurement as dephasing instead of collapsing
                 process. This is only applicable to simulators that can
                 model dephasing.
         """
-        super().__init__(prng, qubits, log_of_measurement_results, ignore_measurement_results)
+        super().__init__(
+            prng=prng,
+            qubits=qubits,
+            log_of_measurement_results=log_of_measurement_results,
+            measured_qubits=measured_qubits,
+            ignore_measurement_results=ignore_measurement_results,
+        )
         self.target_tensor = target_tensor
         self.available_buffer = available_buffer
         self.qid_shape = qid_shape

cirq-core/cirq/sim/act_on_state_vector_args.py

@@ -44,6 +44,7 @@ def __init__(
         prng: np.random.RandomState = None,
         log_of_measurement_results: Dict[str, Any] = None,
         qubits: Sequence['cirq.Qid'] = None,
+        measured_qubits: Dict[str, Tuple['cirq.Qid', ...]] = None,
     ):
         """Inits ActOnStateVectorArgs.
 
@@ -63,8 +64,10 @@ def __init__(
                 effects.
             log_of_measurement_results: A mutable object that measurements are
                 being recorded into.
+            measured_qubits: A dictionary that contains the qubits that were
+                measured in each measurement.
         """
-        super().__init__(prng, qubits, log_of_measurement_results)
+        super().__init__(prng, qubits, log_of_measurement_results, measured_qubits)
         self.target_tensor = target_tensor
         self.available_buffer = available_buffer
 

cirq-core/cirq/sim/clifford/act_on_clifford_tableau_args.py

@@ -14,7 +14,7 @@
 """A protocol for implementing high performance clifford tableau evolutions
  for Clifford Simulator."""
 
-from typing import Any, Dict, TYPE_CHECKING, List, Sequence, Union
+from typing import Any, Dict, TYPE_CHECKING, List, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -42,6 +42,7 @@ def __init__(
         prng: np.random.RandomState,
         log_of_measurement_results: Dict[str, Any],
         qubits: Sequence['cirq.Qid'] = None,
+        measured_qubits: Dict[str, Tuple['cirq.Qid', ...]] = None,
     ):
         """Inits ActOnCliffordTableauArgs.
 
@@ -55,8 +56,10 @@ def __init__(
                 effects.
             log_of_measurement_results: A mutable object that measurements are
                 being recorded into.
+            measured_qubits: A dictionary that contains the qubits that were
+                measured in each measurement.
         """
-        super().__init__(prng, qubits, log_of_measurement_results)
+        super().__init__(prng, qubits, log_of_measurement_results, measured_qubits)
         self.tableau = tableau
 
     def _act_on_fallback_(

cirq-core/cirq/sim/clifford/act_on_stabilizer_ch_form_args.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, TYPE_CHECKING, List, Sequence, Union
+from typing import Any, Dict, TYPE_CHECKING, List, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -41,6 +41,7 @@ def __init__(
         prng: np.random.RandomState,
         log_of_measurement_results: Dict[str, Any],
         qubits: Sequence['cirq.Qid'] = None,
+        measured_qubits: Dict[str, Tuple['cirq.Qid', ...]] = None,
     ):
         """Initializes with the given state and the axes for the operation.
         Args:
@@ -53,8 +54,10 @@ def __init__(
                 effects.
             log_of_measurement_results: A mutable object that measurements are
                 being recorded into.
+            measured_qubits: A dictionary that contains the qubits that were
+                measured in each measurement.
         """
-        super().__init__(prng, qubits, log_of_measurement_results)
+        super().__init__(prng, qubits, log_of_measurement_results, measured_qubits)
         self.state = state
 
     def _act_on_fallback_(