Refactor

Author: babacry

cirq-core/cirq/__init__.py

@@ -380,6 +380,7 @@
     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,
     optimize_for_target_gateset as optimize_for_target_gateset,
     parameterized_2q_op_to_sqrt_iswap_operations as parameterized_2q_op_to_sqrt_iswap_operations,

cirq-core/cirq/transformers/merge_single_qubit_gates.py

@@ -16,25 +16,24 @@
 
 from __future__ import annotations
 
-from typing import Callable, cast, Hashable, List, Tuple, TYPE_CHECKING
-
-import sympy
+from typing import Callable, cast, Hashable, TYPE_CHECKING
 
 from cirq import circuits, ops, protocols
 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,
-    transformer_api,
-    transformer_primitives,
     symbolize,
     tag_transformers,
+    transformer_api,
+    transformer_primitives,
 )
 from cirq.transformers.analytical_decompositions import single_qubit_decompositions
 
 if TYPE_CHECKING:
     import cirq
+    import sympy
 
 
 @transformer_api.transformer
@@ -78,9 +77,9 @@ def merge_single_qubit_gates_to_phxz(
     circuit: cirq.AbstractCircuit,
     *,
     context: cirq.TransformerContext | None = None,
-    merge_tags_fn: Callable[[cirq.CircuitOperation], List[Hashable]] | None = None,
+    merge_tags_fn: Callable[[cirq.CircuitOperation], list[Hashable]] | None = None,
     atol: float = 1e-8,
-) -> 'cirq.Circuit':
+) -> cirq.Circuit:
     """Replaces runs of single qubit rotations with a single optional `cirq.PhasedXZGate`.
 
     Specifically, any run of non-parameterized single-qubit unitaries will be replaced by an
@@ -97,7 +96,7 @@ def merge_single_qubit_gates_to_phxz(
         Copy of the transformed input circuit.
     """
 
-    def rewriter(circuit_op: 'cirq.CircuitOperation') -> 'cirq.OP_TREE':
+    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()
@@ -170,58 +169,43 @@ def merge_func(m1: cirq.Moment, m2: cirq.Moment) -> cirq.Moment | None:
     ).unfreeze(copy=False)
 
 
-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 _sweep_on_symbols(sweep: Sweep, symbols: set[sympy.Symbol]) -> Sweep:
-    new_resolvers: List[cirq.ParamResolver] = []
+    new_resolvers: list[cirq.ParamResolver] = []
     for resolver in sweep:
-        param_dict: 'cirq.ParamMappingType' = {s: resolver.value_of(s) for s in symbols}
+        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,
+def _calc_phxz_sweeps(
+    symbolized_circuit: cirq.Circuit, resolved_circuits: list[cirq.Circuit]
 ) -> 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: str | None = 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}."
+    """Return the phxz sweep of the symbolized_circuit on resolved_circuits.
+
+    Raises:
+        ValueError: Structural mismatch: A `resolved_circuit` contains an unexpected gate type.
+            Expected a `PhasedXZGate` or `IdentityGate` at a position corresponding to a
+            symbolic `PhasedXZGate` in the `symbolized_circuit`.
+    """
+
+    def _extract_axz(op: ops.Operation) -> tuple[float, float, float]:
+        if not op.gate or not isinstance(op.gate, ops.IdentityGate | ops.PhasedXZGate):
+            raise ValueError(f"Expect a PhasedXZGate or IdentityGate on op {op}.")
+        if isinstance(op.gate, ops.IdentityGate):
+            return 0.0, 0.0, 0.0  # Identity gate's a, x, z in PhasedXZ
+        phxz = cast(ops.PhasedXZGate, op.gate)
+        return phxz.axis_phase_exponent, phxz.x_exponent, phxz.z_exponent
+
+    values_by_params: dict[str, list[float]] = {}
+    for mid, moment in enumerate(symbolized_circuit):
+        for op in moment.operations:
+            if op.gate and isinstance(op.gate, ops.PhasedXZGate) and protocols.is_parameterized(op):
+                sa, sx, sz = op.gate.axis_phase_exponent, op.gate.x_exponent, op.gate.z_exponent
+                values_by_params[sa], values_by_params[sx], values_by_params[sz] = zip(
+                    *[_extract_axz(c[mid].operation_at(op.qubits[0])) for c in resolved_circuits]
                 )
-            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),
-    )
+    return dict_to_zip_sweep(cast(ProductOrZipSweepLike, values_by_params))
 
 
 def merge_single_qubit_gates_to_phxz_symbolized(
@@ -230,7 +214,7 @@ def merge_single_qubit_gates_to_phxz_symbolized(
     context: cirq.TransformerContext | None = None,
     sweep: Sweep,
     atol: float = 1e-8,
-) -> Tuple[cirq.Circuit, Sweep]:
+) -> tuple[cirq.Circuit, Sweep]:
     """Merges consecutive single qubit gates as PhasedXZ Gates. Symbolizes if any of
       the consecutive gates is symbolized.
 
@@ -288,6 +272,10 @@ def merge_single_qubit_gates_to_phxz_symbolized(
             for op in circuit_tagged.all_operations()
         ]
     )
+    # 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
+    )
     # If all single qubit gates are not parameterized, call the nonparamerized version of
     # the transformer.
     if not single_qubit_gate_symbols:
@@ -299,61 +287,43 @@ def merge_single_qubit_gates_to_phxz_symbolized(
     ]
 
     # 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 = tag_transformers.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,
+    #  the symbolized_single_tag to indicate the operator is a merged one.
+    merged_circuits: list[cirq.Circuit] = [
+        merge_single_qubit_gates_to_phxz(
+            c,
+            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 []
             ),
-            context=transformer_api.TransformerContext(deep=deep),
-            target_tags={symbolized_single_tag},
+            atol=atol,
         )
-        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.")
+        for c in resolved_circuits
+    ]
 
-    # Step 2, get the new symbolized circuit by symbolization on indexed symbolized_single_tag.
+    # Step 2, get the new symbolized circuit by symbolizing on indexed symbolized_single_tag.
     new_circuit = align.align_right(
-        tag_transformers.remove_tags(
+        tag_transformers.remove_tags(  # remove the temp tags used to track merges
             symbolize.symbolize_single_qubit_gates_by_indexed_tags(
-                merged_circuits[0],
+                tag_transformers.index_tags(  # index all single qubit ops merged from ops with symbols
+                    merged_circuits[0],
+                    context=transformer_api.TransformerContext(deep=deep),
+                    target_tags={symbolized_single_tag},
+                ),
                 symbolize_tag=symbolize.SymbolizeTag(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
+    new_sweep = Zip(
+        _calc_phxz_sweeps(new_circuit, merged_circuits),  # phxz sweeps
+        _sweep_on_symbols(sweep, remaining_symbols),  # remaining sweeps
     )
 
     return new_circuit, new_sweep

cirq-core/cirq/transformers/merge_single_qubit_gates_test.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 from typing import List
-from unittest.mock import Mock, patch
 from unittest import TestCase
+from unittest.mock import Mock, patch
 
 import pytest
 import sympy
@@ -24,7 +24,7 @@
 
 def assert_optimizes(optimized: cirq.AbstractCircuit, expected: cirq.AbstractCircuit):
     # Ignore differences that would be caught by follow-up optimizations.
-    followup_transformers: List[cirq.TRANSFORMER] = [
+    followup_transformers: list[cirq.TRANSFORMER] = [
         cirq.drop_negligible_operations,
         cirq.drop_empty_moments,
     ]
@@ -241,7 +241,7 @@ def test_merge_single_qubit_gates_to_phased_x_and_z_global_phase():
 class TestMergeSingleQubitGatesSymbolized(TestCase):
     """Test suite for merge_single_qubit_gates_to_phxz_symbolized."""
 
-    def case1(self):
+    def test_case1(self):
         """Test case diagram.
         Input circuit:
         # pylint: disable=line-too-long
@@ -298,7 +298,7 @@ def case1(self):
                 {q: q for q in input_circuit.all_qubits()},
             )
 
-    def case_non_parameterized_singles(self):
+    def test_case_non_parameterized_singles(self):
         """Test merge_single_qubit_gates_to_phxz_symbolized when all single qubit gates are not
         parameterized."""
 
@@ -310,27 +310,24 @@ def case_non_parameterized_singles(self):
         )
         assert_optimizes(output_circuit, expected_circuit)
 
-    def fail_different_structures_error(self):
-        """Tests that the function raises a RuntimeError if merged structures of the circuit differ
+    def test_fail_different_structures_error(self):
+        """Tests that the function raises a ValueError if merged structures of the circuit differ
         for different parameterizations."""
-        a = cirq.NamedQubit("a")
-        circuit = cirq.Circuit(cirq.H(a) ** sympy.Symbol("exp"))
+        q0, q1 = cirq.LineQubit.range(2)
+        circuit = cirq.Circuit(cirq.H(q0) ** sympy.Symbol("exp"))
         sweep = cirq.Points(key="exp", points=[0.1, 0.2])
 
         with patch(
             "cirq.protocols.resolve_parameters",
-            side_effect=[
-                cirq.Circuit(cirq.H(a).with_tags("_temp_symbolize_tag")),
-                cirq.Circuit(cirq.H(a)),
+            side_effect=[  # Mock the return values of resolve_parameters
+                cirq.Circuit(cirq.I(q0).with_tags("_tmp_symbolize_tag")),
+                cirq.Circuit(cirq.CZ(q0, q1)),
             ],
         ):
-            with pytest.raises(
-                RuntimeError,
-                match="Different resolvers in sweep resulted in different merged structures.",
-            ):
+            with pytest.raises(ValueError, match="Expect a PhasedXZGate or IdentityGate.*"):
                 cirq.merge_single_qubit_gates_to_phxz_symbolized(circuit, sweep=sweep)
 
-    def fail_unexpected_gate_error(self):
+    def test_fail_unexpected_gate_error(self):
         """Tests that the function raises a RuntimeError of unexpected gate."""
         a, b = cirq.LineQubit.range(2)
         circuit = cirq.Circuit(
@@ -352,7 +349,5 @@ def fail_unexpected_gate_error(self):
             ".single_qubit_decompositions.single_qubit_matrix_to_phxz",
             return_value=cirq.H,
         ):
-            with pytest.raises(
-                RuntimeError, match="Expected the merged gate to be a PhasedXZGate or IdentityGate."
-            ):
+            with pytest.raises(ValueError, match="Expect a PhasedXZGate or IdentityGate.*"):
                 cirq.merge_single_qubit_gates_to_phxz_symbolized(circuit, sweep=sweep)