Repository: quantumlib/Cirq

Issue:
Circuit.to_unitary_matrix() does not work on circuits with CompositeGates
If the CompositeGate subclass does not define a matrix, to_unitary_matrix() does not decompose it to Gates that have known matrices.


Discussion:

Resolving patch:
diff --git a/cirq/circuits/circuit.py b/cirq/circuits/circuit.py
--- a/cirq/circuits/circuit.py
+++ b/cirq/circuits/circuit.py
@@ -19,8 +19,8 @@
 Moment the Operations must all act on distinct Qubits.
 """
 
-from typing import Any, Dict, FrozenSet, Iterable, Optional, Sequence
-from typing import Union
+from typing import Any, Dict, FrozenSet, Generator, Iterable, Iterator
+from typing import Optional, Sequence, Union, TYPE_CHECKING
 
 import numpy as np
 
@@ -31,6 +31,10 @@
 from cirq.extension import Extensions
 from cirq.ops import QubitId
 
+if TYPE_CHECKING:
+    # pylint: disable=unused-import
+    from typing import Set
+
 
 class Circuit(object):
     """A mutable list of groups of operations to apply to some qubits.
@@ -142,6 +146,9 @@ def __len__(self):
     def __iter__(self):
         return iter(self.moments)
 
+    def iter_ops(self) -> Iterator[ops.Operation]:
+        return (op for moment in self for op in moment.operations)
+
     def __repr__(self):
         moment_lines = ('\n    ' + repr(moment) for moment in self.moments)
         return 'Circuit([{}])'.format(','.join(moment_lines))
@@ -452,30 +459,17 @@ def to_unitary_matrix(
                 by a Symbol, etc.
         """
 
-        def is_ignorable_measurement(moment_index, operation):
-            if not ignore_terminal_measurements:
-                return False
-            if not isinstance(operation.gate, ops.MeasurementGate):
-                return False
-            if self.next_moment_operating_on(operation.qubits,
-                                             moment_index + 1) is not None:
-                return False
-            return True
-
         if ext is None:
             ext = Extensions()
         qs = ops.QubitOrder.as_qubit_order(qubit_order).order_for(
             self.qubits().union(qubits_that_should_be_present))
         qubit_map = {i: q
                      for q, i in enumerate(qs)}  # type: Dict[QubitId, int]
-        total = np.eye(1 << len(qubit_map))
-        for k, moment in enumerate(self.moments):
-            for op in moment.operations:
-                if is_ignorable_measurement(k, op):
-                    continue
-                mat = _operation_to_unitary_matrix(op, qubit_map, ext)
-                total = np.matmul(mat, total)
-        return total
+        matrix_ops = _flatten_to_known_matrix_ops(self.iter_ops(), ext)
+        return _operations_to_unitary_matrix(matrix_ops,
+                                             qubit_map,
+                                             ignore_terminal_measurements,
+                                             ext)
 
     def to_text_diagram(
             self,
@@ -628,6 +622,75 @@ def _draw_moment_in_diagram(moment: Moment,
             out_diagram.write(x, y1, '^' + exponent)
 
 
+def _flatten_to_known_matrix_ops(iter_ops: Iterable[ops.Operation],
+                                 ext: Extensions
+                                 ) -> Generator[ops.Operation, None, None]:
+    for op in iter_ops:
+        # Check if the operation has a known matrix
+        known_matrix_gate = ext.try_cast(op.gate, ops.KnownMatrixGate)
+        if known_matrix_gate is not None:
+            yield op
+            continue
+
+        # If not, check if it has a decomposition
+        composite_gate = ext.try_cast(op.gate, ops.CompositeGate)
+        if composite_gate is not None:
+            # Recurse on composite gate decomposition to get known matrix gates.
+            op_tree = composite_gate.default_decompose(op.qubits)
+            op_list = ops.flatten_op_tree(op_tree)
+            for op in _flatten_to_known_matrix_ops(op_list, ext):
+                yield op
+            continue
+
+        # Pass measurement gates through
+        meas_gate = ext.try_cast(op.gate, ops.MeasurementGate)  # type: ops.Gate
+        if meas_gate is None:
+            meas_gate = op.gate
+        if isinstance(meas_gate, ops.MeasurementGate):
+            yield op
+            continue
+
+        # Otherwise, fail
+        raise TypeError(
+            'Operation without a known matrix or decomposition: {!r}'
+            .format(op))
+
+
+def _operations_to_unitary_matrix(iter_ops: Iterable[ops.Operation],
+                                  qubit_map: Dict[QubitId, int],
+                                  ignore_terminal_measurements: bool,
+                                  ext: Extensions) -> np.ndarray:
+    total = np.eye(1 << len(qubit_map))
+    measured_qubits = set()  # type: Set[QubitId]
+    for op in iter_ops:
+        meas_gate = ext.try_cast(op.gate, ops.MeasurementGate)  # type: ops.Gate
+        if meas_gate is None:
+            meas_gate = op.gate
+        if isinstance(meas_gate, ops.MeasurementGate):
+            if not ignore_terminal_measurements:
+                raise TypeError(
+                    'Measurement operation not supported: {!r}'.format(op))
+            measured_qubits.update(op.qubits)
+            continue
+        # Check if any of the op's qubits have been measured
+        measured_and_used = set(op.qubits) & measured_qubits
+        if measured_and_used:
+            if len(measured_and_used) == 1:
+                qubit = next(iter(measured_and_used))
+                raise TypeError(
+                    ('Non-terminal measurement on qubit {!r}. '
+                     + 'Was followed by {!r}')
+                    .format(qubit, op))
+            else:
+                raise TypeError(
+                    ('Non-terminal measurement on qubits {!r}. '
+                     + 'Was followed by {!r}')
+                    .format(tuple(measured_and_used), op))
+        mat = _operation_to_unitary_matrix(op, qubit_map, ext)
+        total = np.matmul(mat, total)
+    return total
+
+
 def _operation_to_unitary_matrix(op: ops.Operation,
                                  qubit_map: Dict[QubitId, int],
                                  ext: Extensions) -> np.ndarray:


Test patch:
diff --git a/cirq/circuits/circuit_test.py b/cirq/circuits/circuit_test.py
--- a/cirq/circuits/circuit_test.py
+++ b/cirq/circuits/circuit_test.py
@@ -992,18 +992,21 @@ def test_circuit_to_unitary_matrix():
     with pytest.raises(TypeError):
         _ = c.to_unitary_matrix()
 
-    # Ignore measurements by turning them into identities.
-    class IdentityGate(ops.KnownMatrixGate):
-        def matrix(self):
-            return np.eye(2)
-    ex = Extensions()
-    ex.add_cast(desired_type=ops.KnownMatrixGate,
-                actual_type=ops.MeasurementGate,
-                conversion=lambda _: IdentityGate())
-    c = Circuit.from_ops(ops.MeasurementGate()(a))
-    cirq.testing.assert_allclose_up_to_global_phase(
-        c.to_unitary_matrix(ext=ex, ignore_terminal_measurements=False),
-        np.eye(2))
+    # Non-terminal measurements are not ignored (multiple qubits).
+    c = Circuit.from_ops(
+            ops.MeasurementGate()(a),
+            ops.MeasurementGate()(b),
+            ops.CNOT(a, b))
+    with pytest.raises(TypeError):
+        _ = c.to_unitary_matrix()
+
+    # Gates without matrix or decomposition raise exception
+    class MysteryGate(ops.Gate):
+        pass
+    c = Circuit.from_ops(MysteryGate()(a, b))
+    with pytest.raises(TypeError):
+        _ = c.to_unitary_matrix()
+
 
 
 def test_simple_circuits_to_unitary_matrix():
@@ -1033,3 +1036,45 @@ def matrix(self):
         c = Circuit.from_ops(Passthrough()(a, b))
         m = c.to_unitary_matrix()
         cirq.testing.assert_allclose_up_to_global_phase(m, expected)
+
+
+def test_composite_gate_to_unitary_matrix():
+    class CNOT_composite(cirq.CompositeGate):
+        def default_decompose(self, qubits):
+            q0, q1 = qubits
+            return ops.Y(q1)**-0.5, ops.CZ(q0, q1), ops.Y(q1)**0.5
+
+    a = ops.NamedQubit('a')
+    b = ops.NamedQubit('b')
+    c = Circuit.from_ops(
+            ops.X(a),
+            CNOT_composite()(a, b),
+            ops.X(a),
+            ops.MeasurementGate()(a),
+            ops.X(b),
+            ops.MeasurementGate()(b))
+    mat = c.to_unitary_matrix()
+    mat_expected = ops.CNOT.matrix()
+
+    cirq.testing.assert_allclose_up_to_global_phase(mat, mat_expected)
+
+
+def test_iter_ops():
+    a = ops.NamedQubit('a')
+    b = ops.NamedQubit('b')
+    c = Circuit([
+            Moment([]),
+            Moment([ops.X(a), ops.Y(b)]),
+            Moment([]),
+            Moment([ops.CNOT(a, b)]),
+            Moment([ops.Z(b), ops.H(a)]),  # Different qubit order
+            Moment([])])
+
+    expected = [
+        ops.X(a),
+        ops.Y(b),
+        ops.CNOT(a, b),
+        ops.Z(b),
+        ops.H(a)]
+
+    assert list(c.iter_ops()) == expected