Support merging 1-qubit gates in transformers for parameterized circuits

cirq-core/cirq/__init__.py

@@ -363,6 +363,7 @@
     eject_z as eject_z,
     expand_composite as expand_composite,
     HardCodedInitialMapper as HardCodedInitialMapper,
+    index_tags as index_tags,
     is_negligible_turn as is_negligible_turn,
     LineInitialMapper as LineInitialMapper,
     MappingManager as MappingManager,
@@ -377,14 +378,17 @@
     merge_operations_to_circuit_op as merge_operations_to_circuit_op,
     merge_single_qubit_gates_to_phased_x_and_z as merge_single_qubit_gates_to_phased_x_and_z,
     merge_single_qubit_gates_to_phxz as merge_single_qubit_gates_to_phxz,
+    merge_single_qubit_gates_to_phxz_symbolized as merge_single_qubit_gates_to_phxz_symbolized,
     merge_single_qubit_moments_to_phxz as merge_single_qubit_moments_to_phxz,
+    symbolize_single_qubit_gates_by_indexed_tags as symbolize_single_qubit_gates_by_indexed_tags,
     optimize_for_target_gateset as optimize_for_target_gateset,
     parameterized_2q_op_to_sqrt_iswap_operations as parameterized_2q_op_to_sqrt_iswap_operations,
     prepare_two_qubit_state_using_cz as prepare_two_qubit_state_using_cz,
     prepare_two_qubit_state_using_iswap as prepare_two_qubit_state_using_iswap,
     prepare_two_qubit_state_using_sqrt_iswap as prepare_two_qubit_state_using_sqrt_iswap,
     quantum_shannon_decomposition as quantum_shannon_decomposition,
     RouteCQC as RouteCQC,
+    remove_tags as remove_tags,
     routed_circuit_with_mapping as routed_circuit_with_mapping,
     SqrtIswapTargetGateset as SqrtIswapTargetGateset,
     single_qubit_matrix_to_gates as single_qubit_matrix_to_gates,

cirq-core/cirq/transformers/__init__.py

@@ -100,9 +100,16 @@
 from cirq.transformers.merge_single_qubit_gates import (
     merge_single_qubit_gates_to_phased_x_and_z as merge_single_qubit_gates_to_phased_x_and_z,
     merge_single_qubit_gates_to_phxz as merge_single_qubit_gates_to_phxz,
+    merge_single_qubit_gates_to_phxz_symbolized as merge_single_qubit_gates_to_phxz_symbolized,
     merge_single_qubit_moments_to_phxz as merge_single_qubit_moments_to_phxz,
 )
 
+from cirq.transformers.tag_transformers import index_tags as index_tags, remove_tags as remove_tags
+from cirq.transformers.symbolize import (
+    symbolize_single_qubit_gates_by_indexed_tags as symbolize_single_qubit_gates_by_indexed_tags,
+)
+
+
 from cirq.transformers.qubit_management_transformers import (
     map_clean_and_borrowable_qubits as map_clean_and_borrowable_qubits,
 )

cirq-core/cirq/transformers/gauge_compiling/gauge_compiling.py

@@ -28,8 +28,7 @@
 from cirq import circuits, ops
 from cirq.protocols import unitary_protocol
 from cirq.protocols.has_unitary_protocol import has_unitary
-from cirq.study import sweepable
-from cirq.study.sweeps import Points, Zip
+from cirq.study.sweeps import Points, Sweep, Zip
 from cirq.transformers import transformer_api
 from cirq.transformers.analytical_decompositions import single_qubit_decompositions
 
@@ -256,7 +255,7 @@ def as_sweep(
         N: int,
         context: Optional[transformer_api.TransformerContext] = None,
         prng: Optional[np.random.Generator] = None,
-    ) -> Tuple[circuits.AbstractCircuit, sweepable.Sweepable]:
+    ) -> Tuple[circuits.AbstractCircuit, Sweep]:
         """Generates a parameterized circuit with *N* sets of sweepable parameters.
 
         Args:

cirq-core/cirq/transformers/merge_single_qubit_gates.py

@@ -14,11 +14,22 @@
 
 """Transformer passes to combine adjacent single-qubit rotations."""
 
-from typing import Optional, TYPE_CHECKING
+from typing import Callable, cast, Dict, Hashable, List, Optional, Tuple, TYPE_CHECKING
+
+import sympy
 
 from cirq import circuits, ops, protocols
-from cirq.transformers import merge_k_qubit_gates, transformer_api, transformer_primitives
+from cirq.study.resolver import ParamResolver
+from cirq.study.sweeps import dict_to_zip_sweep, ListSweep, ProductOrZipSweepLike, Sweep, Zip
+from cirq.transformers import (
+    align,
+    merge_k_qubit_gates,
+    symbolize,
+    transformer_api,
+    transformer_primitives,
+)
 from cirq.transformers.analytical_decompositions import single_qubit_decompositions
+from cirq.transformers.tag_transformers import index_tags, remove_tags
 
 if TYPE_CHECKING:
     import cirq
@@ -65,6 +76,7 @@ def merge_single_qubit_gates_to_phxz(
     circuit: 'cirq.AbstractCircuit',
     *,
     context: Optional['cirq.TransformerContext'] = None,
+    merge_tags_fn: Optional[Callable[['cirq.CircuitOperation'], List[Hashable]]] = None,
     atol: float = 1e-8,
 ) -> 'cirq.Circuit':
     """Replaces runs of single qubit rotations with a single optional `cirq.PhasedXZGate`.
@@ -75,19 +87,21 @@ def merge_single_qubit_gates_to_phxz(
     Args:
         circuit: Input circuit to transform. It will not be modified.
         context: `cirq.TransformerContext` storing common configurable options for transformers.
+        merge_tags_fn: A callable returns the tags to be added to the merged operation.
         atol: Absolute tolerance to angle error. Larger values allow more negligible gates to be
             dropped, smaller values increase accuracy.
 
     Returns:
         Copy of the transformed input circuit.
     """
 
-    def rewriter(op: 'cirq.CircuitOperation') -> 'cirq.OP_TREE':
-        u = protocols.unitary(op)
-        if protocols.num_qubits(op) == 0:
+    def rewriter(circuit_op: 'cirq.CircuitOperation') -> 'cirq.OP_TREE':
+        u = protocols.unitary(circuit_op)
+        if protocols.num_qubits(circuit_op) == 0:
             return ops.GlobalPhaseGate(u[0, 0]).on()
-        gate = single_qubit_decompositions.single_qubit_matrix_to_phxz(u, atol)
-        return gate(op.qubits[0]) if gate else []
+        gate = single_qubit_decompositions.single_qubit_matrix_to_phxz(u, atol) or ops.I
+        phxz_op = gate.on(circuit_op.qubits[0])
+        return phxz_op.with_tags(*merge_tags_fn(circuit_op)) if merge_tags_fn else phxz_op
 
     return merge_k_qubit_gates.merge_k_qubit_unitaries(
         circuit, k=1, context=context, rewriter=rewriter
@@ -152,3 +166,191 @@ def merge_func(m1: 'cirq.Moment', m2: 'cirq.Moment') -> Optional['cirq.Moment']:
         deep=context.deep if context else False,
         tags_to_ignore=tuple(tags_to_ignore),
     ).unfreeze(copy=False)
+
+
+def _sweep_on_symbols(sweep: Sweep, symbols: set[sympy.Symbol]) -> Sweep:
+    new_resolvers: List['cirq.ParamResolver'] = []
+    for resolver in sweep:
+        param_dict: 'cirq.ParamMappingType' = {s: resolver.value_of(s) for s in symbols}
+        new_resolvers.append(ParamResolver(param_dict))
+    return ListSweep(new_resolvers)
+
+
+def _parameterize_phxz_in_circuits(
+    circuit_list: List['cirq.Circuit'],
+    merge_tag_prefix: str,
+    phxz_symbols: set[sympy.Symbol],
+    remaining_symbols: set[sympy.Symbol],
+    sweep: Sweep,
+) -> Sweep:
+    """Parameterizes the circuits and returns a new sweep."""
+    values_by_params: Dict[str, List[float]] = {**{str(s): [] for s in phxz_symbols}}
+
+    for circuit in circuit_list:
+        for op in circuit.all_operations():
+            the_merge_tag: Optional[str] = None
+            for tag in op.tags:
+                if str(tag).startswith(merge_tag_prefix):
+                    the_merge_tag = str(tag)
+            if not the_merge_tag:
+                continue
+            sid = the_merge_tag.rsplit("_", maxsplit=-1)[-1]
+            x, z, a = 0.0, 0.0, 0.0  # Identity gate's parameters
+            if isinstance(op.gate, ops.PhasedXZGate):
+                x, z, a = op.gate.x_exponent, op.gate.z_exponent, op.gate.axis_phase_exponent
+            elif op.gate is not ops.I:
+                raise RuntimeError(
+                    f"Expected the merged gate to be a PhasedXZGate or IdentityGate,"
+                    f" but got {op.gate}."
+                )
+            values_by_params[f"x{sid}"].append(x)
+            values_by_params[f"z{sid}"].append(z)
+            values_by_params[f"a{sid}"].append(a)
+
+    return Zip(
+        dict_to_zip_sweep(cast(ProductOrZipSweepLike, values_by_params)),
+        _sweep_on_symbols(sweep, remaining_symbols),
+    )
+
+
+def _all_tags_startswith(circuit: 'cirq.AbstractCircuit', startswith: str):
+    tag_set: set[Hashable] = set()
+    for op in circuit.all_operations():
+        for tag in op.tags:
+            if str(tag).startswith(startswith):
+                tag_set.add(tag)
+    return tag_set
+
+
+def merge_single_qubit_gates_to_phxz_symbolized(
+    circuit: 'cirq.AbstractCircuit',
+    *,
+    context: Optional['cirq.TransformerContext'] = None,
+    sweep: Sweep,
+    atol: float = 1e-8,
+) -> Tuple['cirq.Circuit', Sweep]:
+    """Merges consecutive single qubit gates as PhasedXZ Gates. Symbolizes if any of
+      the consecutive gates is symbolized.
+
+    Example:
+        >>> q0, q1 = cirq.LineQubit.range(2)
+        >>> c = cirq.Circuit(\
+                    cirq.X(q0),\
+                    cirq.CZ(q0,q1)**sympy.Symbol("cz_exp"),\
+                    cirq.Y(q0)**sympy.Symbol("y_exp"),\
+                    cirq.X(q0))
+        >>> print(c)
+        0: ───X───@──────────Y^y_exp───X───
+                  │
+        1: ───────@^cz_exp─────────────────
+        >>> new_circuit, new_sweep = cirq.merge_single_qubit_gates_to_phxz_symbolized(\
+                c, sweep=cirq.Zip(cirq.Points(key="cz_exp", points=[0, 1]),\
+                                  cirq.Points(key="y_exp",  points=[0, 1])))
+        >>> print(new_circuit)
+        0: ───PhXZ(a=-1,x=1,z=0)───@──────────PhXZ(a=a0,x=x0,z=z0)───
+                                   │
+        1: ────────────────────────@^cz_exp──────────────────────────
+        >>> assert new_sweep[0] == cirq.ParamResolver({'a0': -1, 'x0': 1, 'z0': 0, 'cz_exp': 0})
+        >>> assert new_sweep[1] == cirq.ParamResolver({'a0': -0.5, 'x0': 0, 'z0': -1, 'cz_exp': 1})
+
+    Args:
+        circuit: Input circuit to transform. It will not be modified.
+        context: `cirq.TransformerContext` storing common configurable options for transformers.
+        sweep: Sweep of the symbols in the input circuit, updated Sweep will be returned
+            based on the transformation.
+        atol: Absolute tolerance to angle error. Larger values allow more negligible gates to be
+            dropped, smaller values increase accuracy.
+
+    Returns:
+        Copy of the transformed input circuit.
+    """
+    deep = context.deep if context else False
+
+    # Tag symbolized single-qubit op.
+    symbolized_single_tag = "TMP-TAG-symbolized-single"
+
+    circuit_tagged = transformer_primitives.map_operations(
+        circuit,
+        lambda op, _: (
+            op.with_tags(symbolized_single_tag)
+            if protocols.is_parameterized(op) and len(op.qubits) == 1
+            else op
+        ),
+        deep=deep,
+    )
+
+    # Step 0, isolate single qubit symbols and resolve the circuit on them.
+    single_qubit_gate_symbols: set[sympy.Symbol] = set().union(
+        *[
+            protocols.parameter_symbols(op) if symbolized_single_tag in op.tags else set()
+            for op in circuit_tagged.all_operations()
+        ]
+    )
+    # If all single qubit gates are not parameterized, call the nonparamerized version of
+    # the transformer.
+    if not single_qubit_gate_symbols:
+        return (merge_single_qubit_gates_to_phxz(circuit, context=context, atol=atol), sweep)
+    sweep_of_single: Sweep = _sweep_on_symbols(sweep, single_qubit_gate_symbols)
+    # Get all resolved circuits from all sets of resolvers in sweep_of_single.
+    resolved_circuits = [
+        protocols.resolve_parameters(circuit_tagged, resolver) for resolver in sweep_of_single
+    ]
+
+    # Step 1, merge single qubit gates per resolved circuit, preserving
+    #  the symbolized_single_tag with indexes.
+    merged_circuits: List['cirq.Circuit'] = []
+    for resolved_circuit in resolved_circuits:
+        merged_circuit = index_tags(
+            merge_single_qubit_gates_to_phxz(
+                resolved_circuit,
+                context=context,
+                merge_tags_fn=lambda circuit_op: (
+                    [symbolized_single_tag]
+                    if any(
+                        symbolized_single_tag in set(op.tags)
+                        for op in circuit_op.circuit.all_operations()
+                    )
+                    else []
+                ),
+                atol=atol,
+            ),
+            target_tags={symbolized_single_tag},
+            context=context,
+        )
+        merged_circuits.append(merged_circuit)
+
+    if not all(
+        _all_tags_startswith(merged_circuits[0], startswith=symbolized_single_tag)
+        == _all_tags_startswith(merged_circuit, startswith=symbolized_single_tag)
+        for merged_circuit in merged_circuits
+    ):
+        raise RuntimeError("Different resolvers in sweep resulted in different merged structures.")
+
+    # Step 2, get the new symbolized circuit by symbolization on indexed symbolized_single_tag.
+    new_circuit = align.align_right(
+        remove_tags(
+            symbolize.symbolize_single_qubit_gates_by_indexed_tags(
+                merged_circuits[0], tag_prefix=symbolized_single_tag
+            ),
+            remove_if=lambda tag: str(tag).startswith(symbolized_single_tag),
+        )
+    )
+
+    # Step 3, get N sets of parameterizations as new_sweep.
+    phxz_symbols: set[sympy.Symbol] = set().union(
+        *[
+            set(
+                [sympy.Symbol(tag.replace(f"{symbolized_single_tag}_", s)) for s in ["x", "z", "a"]]
+            )
+            for tag in _all_tags_startswith(merged_circuits[0], startswith=symbolized_single_tag)
+        ]
+    )
+    # Remaining symbols, e.g., 2 qubit gates' symbols. Sweep of those symbols keeps unchanged.
+    remaining_symbols: set[sympy.Symbol] = set(
+        protocols.parameter_symbols(circuit) - single_qubit_gate_symbols
+    )
+    new_sweep = _parameterize_phxz_in_circuits(
+        merged_circuits, symbolized_single_tag, phxz_symbols, remaining_symbols, sweep
+    )
+
+    return new_circuit, new_sweep