Support merge 1-qubit gates for parameterized circuits

Note: if users expect update sweeps together with merging, it's not
supported yet. It's a todo item to be supported in the followup PRs.

Author: babacry

cirq-core/cirq/__init__.py

@@ -378,6 +378,7 @@
     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_moments_to_phxz as merge_single_qubit_moments_to_phxz,
+    merge_into_symbolized_phxz as merge_into_symbolized_phxz,
     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,

cirq-core/cirq/transformers/__init__.py

@@ -101,6 +101,7 @@
     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_moments_to_phxz as merge_single_qubit_moments_to_phxz,
+    merge_into_symbolized_phxz as merge_into_symbolized_phxz,
 )
 
 from cirq.transformers.qubit_management_transformers import (

cirq-core/cirq/transformers/merge_single_qubit_gates.py

@@ -14,9 +14,13 @@
 
 """Transformer passes to combine adjacent single-qubit rotations."""
 
-from typing import Optional, TYPE_CHECKING
+import enum
+import warnings
+from typing import  Optional, TYPE_CHECKING
+
 
 from cirq import circuits, ops, protocols
+from cirq.study import sweepable
 from cirq.transformers.analytical_decompositions import single_qubit_decompositions
 from cirq.transformers import transformer_api, transformer_primitives, merge_k_qubit_gates
 
@@ -152,3 +156,87 @@ 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)
+
+
+@transformer_api.transformer
+def merge_into_symbolized_phxz(
+    circuit: 'cirq.AbstractCircuit',
+    *,
+    context: Optional['cirq.TransformerContext'] = None,
+    sweeps: Optional['sweepable.Sweepable'] = None,
+    atol: float = 1e-8,
+) -> 'cirq.Circuit':
+    """Merge consecutive single qubit gates into connected symbolized PhasedXZ gates.
+
+    Specifically, if at least one of the consecutive gates is symbolized, then the merged gate
+    will be a symbolized gate.
+
+    e.g., X-Y-H-phxz(sa, sx, sz) ---transform---> phxz(sa, sx, sz)
+
+    Note, we only consider merging non-parameterized gates to symbolized phxz with
+     3 degrees of freedom, meaning that gates like Z^exp_symbol will be considered non-mergable.
+
+    Args:
+        circuit: Input circuit to transform. It will not be modified.
+        sweeps: Sweeps of the symbols in the input circuit, updated Sweeps will be returned
+            based on the transformation.
+        context: `cirq.TransformerContext` storing common configurable options for transformers.
+        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.
+    """
+
+    # TODO(#6994): support returning update sweeps when sweeps are provided.
+    if sweeps is not None:
+        raise NotImplementedError("To be supported in #6994.")
+
+    if not protocols.is_parameterized(circuit):
+        warnings.warn(
+            "Expect parameterized circuits. "
+            "Please use cirq.merge_single_qubit_gates_to_phxz instead.",
+            UserWarning,
+        )
+        return merge_single_qubit_gates_to_phxz(circuit, context=context, atol=atol)
+
+    # Merge all non parameterized single qubit gates first.
+    circuit = merge_single_qubit_gates_to_phxz(circuit, context=context, atol=atol)
+
+    def _merge_func(op1: 'cirq.Operation', op2: 'cirq.Operation'):
+
+        class _MergeGateType(enum.Enum):
+            MERAGABLE_NON_PARAMETERIZED = 0
+            MERAGABLE_PARAMETERIZED_PHXZ = 1
+            NON_MERGEABLE = 2
+
+        def _categorize(op: 'cirq.Operation') -> _MergeGateType:
+            if protocols.has_unitary(op) and protocols.num_qubits(op) == 1:
+                return _MergeGateType.MERAGABLE_NON_PARAMETERIZED
+            if isinstance(op.gate, ops.PhasedXZGate) and protocols.is_parameterized(op):
+                return _MergeGateType.MERAGABLE_PARAMETERIZED_PHXZ
+            return _MergeGateType.NON_MERGEABLE
+
+        merge_type1 = _categorize(op1)
+        merge_type2 = _categorize(op2)
+
+        if (
+            merge_type1 == _MergeGateType.NON_MERGEABLE
+            or merge_type2 == _MergeGateType.NON_MERGEABLE
+        ):
+            return None
+
+        # absorb the non-parameterized gate into the parameterized gate.
+        if merge_type1 == _MergeGateType.MERAGABLE_PARAMETERIZED_PHXZ:
+            return op1
+        if merge_type2 == _MergeGateType.MERAGABLE_PARAMETERIZED_PHXZ:
+            return op2
+
+        return None  # pragma: no cover
+
+    return transformer_primitives.merge_operations(
+        circuit,
+        _merge_func,
+        deep=context.deep if context else False,
+        tags_to_ignore=context.tags_to_ignore if context else (),
+    ).unfreeze()

cirq-core/cirq/transformers/merge_single_qubit_gates_test.py

@@ -14,7 +14,10 @@
 
 from typing import List
 
+import pytest
+import sympy
 import cirq
+from cirq.study.sweeps import Points
 
 
 def assert_optimizes(optimized: cirq.AbstractCircuit, expected: cirq.AbstractCircuit):
@@ -231,3 +234,71 @@ def test_merge_single_qubit_moments_to_phased_x_and_z_global_phase():
     c = cirq.Circuit(cirq.GlobalPhaseGate(1j).on())
     c2 = cirq.merge_single_qubit_gates_to_phased_x_and_z(c)
     assert c == c2
+
+
+def test_merge_into_symbolized_phxz():
+    """Test case diagram.
+    Input circuit:
+    0: ───X───────@───H[ignore]───H───X───PhXZ(a=a1,x=x1,z=z1)───X───PhXZ(a=a2,x=x2,z=z2)───H───
+                  │                                                                         ║
+    1: ───Y^0.5───@───M─────────────────────────────────────────────────────────────────────╫───
+                      ║                                                                     ║
+    m: ═══════════════@═════════════════════════════════════════════════════════════════════^═══
+    Expected output:
+    0: ───PhXZ(a=-1,x=1,z=0)──────@───H[ignore]───PhXZ(a=a1,x=x1,z=z1)───H───
+                                  │                                      ║
+    1: ───PhXZ(a=0.5,x=0.5,z=0)───@───M──────────────────────────────────╫───
+                                      ║                                  ║
+    m: ═══════════════════════════════@══════════════════════════════════^═══
+    """
+    a, b = cirq.LineQubit.range(2)
+    sa1, sa2 = [sympy.Symbol(a) for a in ["a1", "a2"]]
+    sx1, sx2 = [sympy.Symbol(x) for x in ["x1", "x2"]]
+    sz1, sz2 = [sympy.Symbol(z) for z in ["z1", "z2"]]
+    input_circuit = cirq.Circuit(
+        cirq.X(a),
+        cirq.Y(b) ** 0.5,
+        cirq.CZ(a, b),
+        cirq.H(a).with_tags("ignore"),
+        cirq.H(a),
+        cirq.X(a),
+        _phxz(sa1, sx1, sz1).on(a),
+        cirq.X(a),
+        _phxz(sa2, sx2, sz2).on(a),
+        cirq.measure(b, key="m"),
+        cirq.H(a).with_classical_controls("m"),
+    )
+    context = cirq.TransformerContext(tags_to_ignore=["ignore"])
+    assert_optimizes(
+        optimized=cirq.merge_into_symbolized_phxz(input_circuit, context=context),
+        expected=cirq.Circuit(
+            _phxz(-1, 1, 0).on(a),
+            _phxz(0.5, 0.5, 0).on(b),
+            cirq.CZ(a, b),
+            cirq.H(a).with_tags("ignore"),
+            _phxz(sa1, sx1, sz1).on(a),
+            cirq.measure(b, key="m"),
+            cirq.H(a).with_classical_controls("m"),
+        ),
+    )
+
+
+def test_merge_into_symbolized_phxz_other_symbolized_gates():
+    a = cirq.NamedQubit('a')
+    input_circuit = cirq.Circuit(_phxz(1, 1, 1).on(a), cirq.H(a) ** sympy.Symbol("exp"))
+    assert_optimizes(
+        optimized=cirq.merge_into_symbolized_phxz(input_circuit), expected=input_circuit
+    )
+
+
+def test_merge_into_symbolized_phxz_non_symbolized_input():
+    a = cirq.NamedQubit('a')
+    with pytest.warns(UserWarning):
+        cirq.merge_into_symbolized_phxz(cirq.Circuit(cirq.H(a), cirq.H(a)))
+
+
+def test_merge_into_symbolized_phxz_with_sweeps():
+    with pytest.raises(NotImplementedError):
+        cirq.merge_into_symbolized_phxz(
+            cirq.Circuit(), sweeps=[Points(key="x", points=[0.1, 0.2, 0.5])]
+        )