Add handling for sympy conditions in deferred measurement transformer

cirq-core/cirq/transformers/measurement_transformers.py

@@ -13,7 +13,18 @@
 # limitations under the License.
 
 import itertools
-from typing import Any, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    TYPE_CHECKING,
+    Union,
+)
 
 import numpy as np
 
@@ -85,7 +96,6 @@ def defer_measurements(
         A circuit with equivalent logic, but all measurements at the end of the
         circuit.
     Raises:
-        ValueError: If sympy-based classical conditions are used.
         NotImplementedError: When attempting to defer a measurement with a
             confusion map. (https://github.com/quantumlib/Cirq/issues/5482)
     """
@@ -109,25 +119,34 @@ def defer(op: 'cirq.Operation', _) -> 'cirq.OP_TREE':
                 )
                 for indexes, m in gate.confusion_map.items()
             ]
-            cxs = [_mod_add(q, target) for q, target in zip(op.qubits, targets)]
             xs = [ops.X(targets[i]) for i, b in enumerate(gate.full_invert_mask()) if b]
             return cxs + confusions + xs
         elif protocols.is_measurement(op):
             return [defer(op, None) for op in protocols.decompose_once(op)]
         elif op.classical_controls:
-            new_op = op.without_classical_controls()
-            for c in op.classical_controls:
-                if isinstance(c, value.KeyCondition):
-                    if c.key not in measurement_qubits:
-                        raise ValueError(f'Deferred measurement for key={c.key} not found.')
-                    qs = measurement_qubits[c.key]
-                    all_values = itertools.product(*[range(q.dimension) for q in qs])
-                    anything_but_all_zeros = tuple(itertools.islice(all_values, 1, None))
-                    control_values = ops.SumOfProducts(anything_but_all_zeros)
-                    new_op = new_op.controlled_by(*qs, control_values=control_values)
-                else:
-                    raise ValueError('Only KeyConditions are allowed.')
-            return new_op
+            # Convert to a quantum control
+            keys = sorted(set(key for c in op.classical_controls for key in c.keys))
+            for key in keys:
+                if key not in measurement_qubits:
+                    raise ValueError(f'Deferred measurement for key={key} not found.')
+
+            # Try every possible datastore state against the condition, and the ones that work
+            # are the control values for the new op.
+            compatible_datastores = [
+                store
+                for store in _all_possible_datastore_states(keys, measurement_qubits)
+                if all(c.resolve(store) for c in op.classical_controls)
+            ]
+
+            # Rearrange these into the format expected by SumOfProducts
+            products = [
+                [i for key in keys for i in store.records[key][0]]
+                for store in compatible_datastores
+            ]
+
+            control_values = ops.SumOfProducts(products)
+            qs = [q for key in keys for q in measurement_qubits[key]]
+            return op.without_classical_controls().controlled_by(*qs, control_values=control_values)
         return op
 
     circuit = transformer_primitives.map_operations_and_unroll(
@@ -141,6 +160,39 @@ def defer(op: 'cirq.Operation', _) -> 'cirq.OP_TREE':
     return circuit
 
 
+def _all_possible_datastore_states(
+    keys: Iterable['cirq.MeasurementKey'],
+    measurement_qubits: Mapping['cirq.MeasurementKey', Iterable['cirq.Qid']],
+) -> Iterable['cirq.ClassicalDataStoreReader']:
+    """The full product of possible DatsStores states for the given keys."""
+    # First we get the list of all possible values. So if we have a key mapped to qubits of shape
+    # (2, 2) and a key mapped to a qutrit, the possible measurement values are:
+    # [((0, 0), (0,)),
+    #  ((0, 0), (1,)),
+    #  ((0, 0), (2,)),
+    #  ((0, 1), (0,)),
+    #  ((0, 1), (1,)),
+    #  ((0, 1), (2,)),
+    #  ((1, 0), (0,)),
+    #  ((1, 0), (1,)),
+    #  ((1, 0), (2,)),
+    #  ((1, 1), (0,)),
+    #  ((1, 1), (1,)),
+    #  ((1, 1), (2,))]
+    all_values = itertools.product(
+        *[
+            tuple(itertools.product(*[range(q.dimension) for q in measurement_qubits[k]]))
+            for k in keys
+        ]
+    )
+    # Then we create the ClassicalDataDictionaryStore for each of the above.
+    for sequences in all_values:
+        lookup = {k: [sequence] for k, sequence in zip(keys, sequences)}
+        yield value.ClassicalDataDictionaryStore(
+            _records=lookup, _measured_qubits={k: [tuple(measurement_qubits[k])] for k in keys}
+        )
+
+
 @transformer_api.transformer
 def dephase_measurements(
     circuit: 'cirq.AbstractCircuit',

cirq-core/cirq/transformers/measurement_transformers_test.py

@@ -15,6 +15,7 @@
 import numpy as np
 import pytest
 import sympy
+from sympy.parsing import sympy_parser
 
 import cirq
 from cirq.transformers.measurement_transformers import _ConfusionChannel, _MeasurementQid, _mod_add
@@ -79,6 +80,179 @@ def test_qudits():
     )
 
 
+def test_sympy_control():
+    q0, q1 = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.X(q1).with_classical_controls(sympy.Symbol('a')),
+        cirq.measure(q1, key='b'),
+    )
+    assert_equivalent_to_deferred(circuit)
+    deferred = cirq.defer_measurements(circuit)
+    q_ma = _MeasurementQid('a', q0)
+    cirq.testing.assert_same_circuits(
+        deferred,
+        cirq.Circuit(
+            cirq.CX(q0, q_ma),
+            cirq.CX(q_ma, q1),
+            cirq.measure(q_ma, key='a'),
+            cirq.measure(q1, key='b'),
+        ),
+    )
+
+
+def test_sympy_qudits():
+    q0, q1 = cirq.LineQid.range(2, dimension=3)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.XPowGate(dimension=3).on(q1).with_classical_controls(sympy.Symbol('a')),
+        cirq.measure(q1, key='b'),
+    )
+    assert_equivalent_to_deferred(circuit)
+    deferred = cirq.defer_measurements(circuit)
+    q_ma = _MeasurementQid('a', q0)
+    cirq.testing.assert_same_circuits(
+        deferred,
+        cirq.Circuit(
+            _mod_add(q0, q_ma),
+            cirq.XPowGate(dimension=3).on(q1).controlled_by(q_ma, control_values=[[1, 2]]),
+            cirq.measure(q_ma, key='a'),
+            cirq.measure(q1, key='b'),
+        ),
+    )
+
+
+def test_sympy_control_complex():
+    q0, q1, q2 = cirq.LineQubit.range(3)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.measure(q1, key='b'),
+        cirq.X(q2).with_classical_controls(sympy_parser.parse_expr('a >= b')),
+        cirq.measure(q2, key='c'),
+    )
+    assert_equivalent_to_deferred(circuit)
+    deferred = cirq.defer_measurements(circuit)
+    q_ma = _MeasurementQid('a', q0)
+    q_mb = _MeasurementQid('b', q1)
+    cirq.testing.assert_same_circuits(
+        deferred,
+        cirq.Circuit(
+            cirq.CX(q0, q_ma),
+            cirq.CX(q1, q_mb),
+            cirq.ControlledOperation(
+                [q_ma, q_mb], cirq.X(q2), cirq.SumOfProducts([[0, 0], [1, 0], [1, 1]])
+            ),
+            cirq.measure(q_ma, key='a'),
+            cirq.measure(q_mb, key='b'),
+            cirq.measure(q2, key='c'),
+        ),
+    )
+
+
+def test_sympy_control_complex_qudit():
+    q0, q1, q2 = cirq.LineQid.for_qid_shape((4, 2, 2))
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.measure(q1, key='b'),
+        cirq.X(q2).with_classical_controls(sympy_parser.parse_expr('a > b')),
+        cirq.measure(q2, key='c'),
+    )
+    assert_equivalent_to_deferred(circuit)
+    deferred = cirq.defer_measurements(circuit)
+    q_ma = _MeasurementQid('a', q0)
+    q_mb = _MeasurementQid('b', q1)
+    cirq.testing.assert_same_circuits(
+        deferred,
+        cirq.Circuit(
+            _mod_add(q0, q_ma),
+            cirq.CX(q1, q_mb),
+            cirq.ControlledOperation(
+                [q_ma, q_mb],
+                cirq.X(q2),
+                cirq.SumOfProducts([[1, 0], [2, 0], [3, 0], [2, 1], [3, 1]]),
+            ),
+            cirq.measure(q_ma, key='a'),
+            cirq.measure(q_mb, key='b'),
+            cirq.measure(q2, key='c'),
+        ),
+    )
+
+
+def test_multiple_sympy_control_complex():
+    q0, q1, q2 = cirq.LineQubit.range(3)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.measure(q1, key='b'),
+        cirq.X(q2)
+        .with_classical_controls(sympy_parser.parse_expr('a >= b'))
+        .with_classical_controls(sympy_parser.parse_expr('a <= b')),
+        cirq.measure(q2, key='c'),
+    )
+    assert_equivalent_to_deferred(circuit)
+    deferred = cirq.defer_measurements(circuit)
+    q_ma = _MeasurementQid('a', q0)
+    q_mb = _MeasurementQid('b', q1)
+    cirq.testing.assert_same_circuits(
+        deferred,
+        cirq.Circuit(
+            cirq.CX(q0, q_ma),
+            cirq.CX(q1, q_mb),
+            cirq.ControlledOperation(
+                [q_ma, q_mb], cirq.X(q2), cirq.SumOfProducts([[0, 0], [1, 1]])
+            ),
+            cirq.measure(q_ma, key='a'),
+            cirq.measure(q_mb, key='b'),
+            cirq.measure(q2, key='c'),
+        ),
+    )
+
+
+def test_sympy_and_key_control():
+    q0, q1 = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.X(q1).with_classical_controls(sympy.Symbol('a')).with_classical_controls('a'),
+        cirq.measure(q1, key='b'),
+    )
+    assert_equivalent_to_deferred(circuit)
+    deferred = cirq.defer_measurements(circuit)
+    q_ma = _MeasurementQid('a', q0)
+    cirq.testing.assert_same_circuits(
+        deferred,
+        cirq.Circuit(
+            cirq.CX(q0, q_ma),
+            cirq.CX(q_ma, q1),
+            cirq.measure(q_ma, key='a'),
+            cirq.measure(q1, key='b'),
+        ),
+    )
+
+
+def test_sympy_control_multiqubit():
+    q0, q1, q2 = cirq.LineQubit.range(3)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, q1, key='a'),
+        cirq.X(q2).with_classical_controls(sympy_parser.parse_expr('a >= 2')),
+        cirq.measure(q2, key='c'),
+    )
+    assert_equivalent_to_deferred(circuit)
+    deferred = cirq.defer_measurements(circuit)
+    q_ma0 = _MeasurementQid('a', q0)
+    q_ma1 = _MeasurementQid('a', q1)
+    cirq.testing.assert_same_circuits(
+        deferred,
+        cirq.Circuit(
+            cirq.CX(q0, q_ma0),
+            cirq.CX(q1, q_ma1),
+            cirq.ControlledOperation(
+                [q_ma0, q_ma1], cirq.X(q2), cirq.SumOfProducts([[1, 0], [1, 1]])
+            ),
+            cirq.measure(q_ma0, q_ma1, key='a'),
+            cirq.measure(q2, key='c'),
+        ),
+    )
+
+
 def test_nocompile_context():
     q0, q1 = cirq.LineQubit.range(2)
     circuit = cirq.Circuit(
@@ -316,15 +490,6 @@ def test_repr(qid: _MeasurementQid):
     test_repr(_MeasurementQid('0:1:a', cirq.LineQid(9, 4)))
 
 
-def test_sympy_control():
-    q0, q1 = cirq.LineQubit.range(2)
-    circuit = cirq.Circuit(
-        cirq.measure(q0, q1, key='a'), cirq.X(q1).with_classical_controls(sympy.Symbol('a'))
-    )
-    with pytest.raises(ValueError, match='Only KeyConditions are allowed'):
-        _ = cirq.defer_measurements(circuit)
-
-
 def test_confusion_map():
     q0, q1 = cirq.LineQubit.range(2)
     circuit = cirq.Circuit(