Repository: quantumlib/Cirq

Issue:
Flatten circuits containing CircuitOperation
I want to use `CircuitOperation` to carefully construct a circuit while respecting moment structure. At some point I need to flatten/unnest it to a normal two-dimensional circuit.

@Strilanc suggests

```python
def flatten(obj: Union[cirq.Moment, cirq.Circuit, cirq.FrozenCircuit]) -> cirq.Circuit:
    if isinstance(obj, cirq.Moment):
        flattenable = [op for op in obj if isinstance(op, cirq.CircuitOperation)]
        if not flattenable:
            return cirq.Circuit(obj)
        return cirq.Circuit.zip(
            cirq.Circuit(obj - flattenable),
            *[flatten(cast(cirq.CircuitOperation, e).circuit) for e in flattenable]
        )
    return cirq.Circuit(
        [moment for e in obj for moment in flatten(e)]
    )
```

but I'm running into problems with `circuit_operation.circuit` respecting a qubit mapping.


Discussion:
On the original `CircuitOperation` PR I [suggested](https://github.com/quantumlib/Cirq/pull/3580#discussion_r544578017) that we should add a method to "apply" the various maps (qubit, measurement, parameter) to produce a new circuit operation with maps applied. That would then work with the code from @Strilanc. Currently, this map application only happens implicitly during decomposition. I'd be happy to take a crack at this.

Resolving patch:
diff --git a/cirq-core/cirq/circuits/circuit.py b/cirq-core/cirq/circuits/circuit.py
--- a/cirq-core/cirq/circuits/circuit.py
+++ b/cirq-core/cirq/circuits/circuit.py
@@ -224,7 +224,7 @@ def __getitem__(self, key):
     # pylint: enable=function-redefined
 
     @abc.abstractmethod
-    def _with_sliced_moments(self, moments: Sequence['cirq.Moment']):
+    def _with_sliced_moments(self: CIRCUIT_TYPE, moments: Iterable['cirq.Moment']) -> CIRCUIT_TYPE:
         """Helper method for constructing circuits from __getitem__."""
 
     def __str__(self) -> str:
@@ -883,6 +883,25 @@ def all_operations(self) -> Iterator[ops.Operation]:
         """
         return (op for moment in self for op in moment.operations)
 
+    def map_operations(
+        self: CIRCUIT_TYPE, func: Callable[['cirq.Operation'], 'cirq.OP_TREE']
+    ) -> CIRCUIT_TYPE:
+        """Applies the given function to all operations in this circuit.
+
+        Args:
+            func: a mapping function from operations to OP_TREEs.
+
+        Returns:
+            A circuit with the same basic structure as the original, but with
+            each operation `op` replaced with `func(op)`.
+        """
+
+        def map_moment(moment: 'cirq.Moment') -> 'cirq.Circuit':
+            """Apply func to expand each op into a circuit, then zip up the circuits."""
+            return Circuit.zip(*[Circuit(func(op)) for op in moment])
+
+        return self._with_sliced_moments(m for moment in self for m in map_moment(moment))
+
     def qid_shape(
         self, qubit_order: 'cirq.QubitOrderOrList' = ops.QubitOrder.DEFAULT
     ) -> Tuple[int, ...]:
@@ -1652,7 +1671,7 @@ def copy(self) -> 'Circuit':
         copied_circuit._moments = self._moments[:]
         return copied_circuit
 
-    def _with_sliced_moments(self, moments: Sequence['cirq.Moment']) -> 'Circuit':
+    def _with_sliced_moments(self, moments: Iterable['cirq.Moment']) -> 'Circuit':
         new_circuit = Circuit(device=self.device)
         new_circuit._moments = list(moments)
         return new_circuit
diff --git a/cirq-core/cirq/circuits/circuit_operation.py b/cirq-core/cirq/circuits/circuit_operation.py
--- a/cirq-core/cirq/circuits/circuit_operation.py
+++ b/cirq-core/cirq/circuits/circuit_operation.py
@@ -188,24 +188,65 @@ def _parameter_names_(self) -> AbstractSet[str]:
             )
         }
 
-    def _decompose_(self) -> 'cirq.OP_TREE':
-        result = self.circuit.unfreeze()
-        result = result.transform_qubits(lambda q: self.qubit_map.get(q, q))
+    def mapped_circuit(self, deep: bool = False) -> 'cirq.Circuit':
+        """Applies all maps to the contained circuit and returns the result.
+
+        Args:
+            deep: If true, this will also call mapped_circuit on any
+                CircuitOperations this object contains.
+
+        Returns:
+            The contained circuit with all other member variables (repetitions,
+            qubit mapping, parameterization, etc.) applied to it. This behaves
+            like `cirq.decompose(self)`, but preserving moment structure.
+        """
+        circuit = self.circuit.unfreeze()
+        circuit = circuit.transform_qubits(lambda q: self.qubit_map.get(q, q))
         if self.repetitions < 0:
-            result = result ** -1
-        result = protocols.with_measurement_key_mapping(result, self.measurement_key_map)
-        result = protocols.resolve_parameters(result, self.param_resolver, recursive=False)
-        # repetition_ids don't need to be taken into account if the circuit has no measurements
-        # or if repetition_ids are unset.
-        if self.repetition_ids is None or not protocols.is_measurement(result):
-            return list(result.all_operations()) * abs(self.repetitions)
-        # If it's a measurement circuit with repetitions/repetition_ids, prefix the repetition_ids
-        # to measurements. Details at https://tinyurl.com/measurement-repeated-circuitop.
-        ops = []  # type: List[cirq.Operation]
-        for repetition_id in self.repetition_ids:
-            path = self.parent_path + (repetition_id,)
-            ops += protocols.with_key_path(result, path).all_operations()
-        return ops
+            circuit = circuit ** -1
+        has_measurements = protocols.is_measurement(circuit)
+        if has_measurements:
+            circuit = protocols.with_measurement_key_mapping(circuit, self.measurement_key_map)
+        circuit = protocols.resolve_parameters(circuit, self.param_resolver, recursive=False)
+        if deep:
+
+            def map_deep(op: 'cirq.Operation') -> 'cirq.OP_TREE':
+                return op.mapped_circuit(deep=True) if isinstance(op, CircuitOperation) else op
+
+            if self.repetition_ids is None:
+                return circuit.map_operations(map_deep)
+            if not has_measurements:
+                return circuit.map_operations(map_deep) * abs(self.repetitions)
+
+            # Path must be constructed from the top down.
+            rekeyed_circuit = circuits.Circuit(
+                protocols.with_key_path(circuit, self.parent_path + (rep,))
+                for rep in self.repetition_ids
+            )
+            return rekeyed_circuit.map_operations(map_deep)
+
+        if self.repetition_ids is None:
+            return circuit
+        if not has_measurements:
+            return circuit * abs(self.repetitions)
+
+        def rekey_op(op: 'cirq.Operation', rep: str):
+            """Update measurement keys in `op` to include repetition ID `rep`."""
+            rekeyed_op = protocols.with_key_path(op, self.parent_path + (rep,))
+            if rekeyed_op is NotImplemented:
+                return op
+            return rekeyed_op
+
+        return circuits.Circuit(
+            circuit.map_operations(lambda op: rekey_op(op, rep)) for rep in self.repetition_ids
+        )
+
+    def mapped_op(self, deep: bool = False) -> 'cirq.CircuitOperation':
+        """As `mapped_circuit`, but wraps the result in a CircuitOperation."""
+        return CircuitOperation(circuit=self.mapped_circuit(deep=deep).freeze())
+
+    def _decompose_(self) -> 'cirq.OP_TREE':
+        return self.mapped_circuit(deep=False).all_operations()
 
     # Methods for string representation of the operation.
 
diff --git a/cirq-core/cirq/circuits/frozen_circuit.py b/cirq-core/cirq/circuits/frozen_circuit.py
--- a/cirq-core/cirq/circuits/frozen_circuit.py
+++ b/cirq-core/cirq/circuits/frozen_circuit.py
@@ -17,6 +17,7 @@
     AbstractSet,
     Callable,
     FrozenSet,
+    Iterable,
     Iterator,
     Optional,
     Sequence,
@@ -163,7 +164,7 @@ def __pow__(self, other) -> 'FrozenCircuit':
         except:
             return NotImplemented
 
-    def _with_sliced_moments(self, moments: Sequence['cirq.Moment']) -> 'FrozenCircuit':
+    def _with_sliced_moments(self, moments: Iterable['cirq.Moment']) -> 'FrozenCircuit':
         new_circuit = FrozenCircuit(device=self.device)
         new_circuit._moments = tuple(moments)
         return new_circuit


Test patch:
diff --git a/cirq-core/cirq/circuits/circuit_operation_test.py b/cirq-core/cirq/circuits/circuit_operation_test.py
--- a/cirq-core/cirq/circuits/circuit_operation_test.py
+++ b/cirq-core/cirq/circuits/circuit_operation_test.py
@@ -702,6 +702,8 @@ def test_decompose_nested():
         cirq.measure(d, key='md'),
     )
     assert cirq.Circuit(cirq.decompose(final_op)) == expected_circuit
+    # Verify that mapped_circuit gives the same operations.
+    assert final_op.mapped_circuit(deep=True) == expected_circuit
 
 
 def test_decompose_repeated_nested_measurements():
@@ -752,6 +754,33 @@ def test_decompose_repeated_nested_measurements():
     assert cirq.Circuit(cirq.decompose(op3)) == expected_circuit
     assert cirq.measurement_keys(expected_circuit) == set(expected_measurement_keys_in_order)
 
+    # Verify that mapped_circuit gives the same operations.
+    assert op3.mapped_circuit(deep=True) == expected_circuit
+
+
+def test_mapped_circuit_preserves_moments():
+    q0, q1 = cirq.LineQubit.range(2)
+    fc = cirq.FrozenCircuit(cirq.Moment(cirq.X(q0)), cirq.Moment(cirq.X(q1)))
+    op = cirq.CircuitOperation(fc)
+    assert op.mapped_circuit() == fc
+    assert op.repeat(3).mapped_circuit(deep=True) == fc * 3
+
+
+def test_mapped_op():
+    q0, q1 = cirq.LineQubit.range(2)
+    a, b = (sympy.Symbol(x) for x in 'ab')
+    fc1 = cirq.FrozenCircuit(cirq.X(q0) ** a, cirq.measure(q0, q1, key='m'))
+    op1 = (
+        cirq.CircuitOperation(fc1)
+        .with_params({'a': 'b'})
+        .with_qubits(q1, q0)
+        .with_measurement_key_mapping({'m': 'k'})
+    )
+    fc2 = cirq.FrozenCircuit(cirq.X(q1) ** b, cirq.measure(q1, q0, key='k'))
+    op2 = cirq.CircuitOperation(fc2)
+
+    assert op1.mapped_op() == op2
+
 
 def test_tag_propagation():
     # Tags are not propagated from the CircuitOperation to its components.