Add cirq.merge_k_qubit_unitaries transformer to replace cirq.MergeSingleQubitGates optimizer (quantumlib#4986)

* Replace cirq.MergeSingleQubitGates optimizer with cirq.merge_single_qubit_gates transformer

* Add merge_k_qubit_unitaries primitive

Author: tanujkhattar

cirq/__init__.py

@@ -372,10 +372,14 @@
     map_moments,
     map_operations,
     map_operations_and_unroll,
+    merge_k_qubit_unitaries,
     merge_k_qubit_unitaries_to_circuit_op,
     merge_moments,
     merge_operations,
     merge_operations_to_circuit_op,
+    merge_single_qubit_gates_to_phased_x_and_z,
+    merge_single_qubit_gates_to_phxz,
+    merge_single_qubit_moments_to_phxz,
     prepare_two_qubit_state_using_cz,
     prepare_two_qubit_state_using_sqrt_iswap,
     single_qubit_matrix_to_gates,

cirq/contrib/paulistring/convert_gate_set.py

@@ -28,7 +28,7 @@ def converted_gate_set(
     """
     conv_circuit = circuits.Circuit(circuit)
     optimizers.ConvertToCzAndSingleGates().optimize_circuit(conv_circuit)
-    optimizers.MergeSingleQubitGates().optimize_circuit(conv_circuit)
+    conv_circuit = transformers.merge_k_qubit_unitaries(conv_circuit, k=1)
     ConvertToPauliStringPhasors(
         ignore_failures=True,
         keep_clifford=not no_clifford_gates,

cirq/devices/noise_model_test.py

@@ -130,7 +130,6 @@ def test_noise_composition():
     a, b, c = cirq.LineQubit.range(3)
     noise_z = cirq.ConstantQubitNoiseModel(cirq.Z)
     noise_inv_s = cirq.ConstantQubitNoiseModel(cirq.S ** -1)
-    merge = cirq.optimizers.merge_single_qubit_gates_into_phased_x_z
     base_moments = [cirq.Moment([cirq.X(a)]), cirq.Moment([cirq.Y(b)]), cirq.Moment([cirq.H(c)])]
     circuit_z = cirq.Circuit(noise_z.noisy_moments(base_moments, [a, b, c]))
     circuit_s = cirq.Circuit(noise_inv_s.noisy_moments(base_moments, [a, b, c]))
@@ -147,9 +146,9 @@ def test_noise_composition():
     )
 
     # All of the gates will be the same, just out of order. Merging fixes this.
-    merge(actual_zs)
-    merge(actual_sz)
-    merge(expected_circuit)
+    actual_zs = cirq.merge_single_qubit_gates_to_phased_x_and_z(actual_zs)
+    actual_sz = cirq.merge_single_qubit_gates_to_phased_x_and_z(actual_sz)
+    expected_circuit = cirq.merge_single_qubit_gates_to_phased_x_and_z(expected_circuit)
     assert_equivalent_op_tree(actual_zs, actual_sz)
     assert_equivalent_op_tree(actual_zs, expected_circuit)
 

cirq/ion/convert_to_ion_gates.py

@@ -14,7 +14,7 @@
 
 import numpy as np
 
-from cirq import ops, protocols, optimizers, circuits, transformers
+from cirq import ops, protocols, circuits, transformers
 from cirq.ion import ms, two_qubit_matrix_to_ion_operations, ion_device
 
 
@@ -86,6 +86,4 @@ def convert_circuit(self, circuit: circuits.Circuit) -> circuits.Circuit:
         for moment in circuit:
             for op in moment.operations:
                 new_circuit.append(self.convert_one(op))
-        optimizers.merge_single_qubit_gates_into_phased_x_z(new_circuit)
-
-        return new_circuit
+        return transformers.merge_single_qubit_gates_to_phased_x_and_z(new_circuit)

cirq/ion/ion_decomposition.py

@@ -23,7 +23,7 @@
 
 import numpy as np
 
-from cirq import ops, linalg, protocols, optimizers, circuits, transformers
+from cirq import ops, linalg, protocols, circuits, transformers
 from cirq.ion import ms
 
 if TYPE_CHECKING:
@@ -52,7 +52,7 @@ def two_qubit_matrix_to_ion_operations(
 
 def _cleanup_operations(operations: List[ops.Operation]):
     circuit = circuits.Circuit(operations)
-    optimizers.merge_single_qubit_gates.merge_single_qubit_gates_into_phased_x_z(circuit)
+    circuit = transformers.merge_single_qubit_gates_to_phased_x_and_z(circuit)
     circuit = transformers.eject_phased_paulis(circuit)
     circuit = transformers.eject_z(circuit)
     circuit = circuits.Circuit(circuit.all_operations(), strategy=circuits.InsertStrategy.EARLIEST)

cirq/optimizers/merge_interactions_test.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, List
+from typing import List
 
 import pytest
 import sympy
@@ -26,14 +26,8 @@ def assert_optimizes(before: cirq.Circuit, expected: cirq.Circuit):
     opt.optimize_circuit(actual)
 
     # Ignore differences that would be caught by follow-up optimizations.
-    followup_optimizations: List[Callable[[cirq.Circuit], None]] = [
-        cirq.merge_single_qubit_gates_into_phased_x_z,
-    ]
-    for post in followup_optimizations:
-        post(actual)
-        post(expected)
-
     followup_transformers: List[cirq.TRANSFORMER] = [
+        cirq.merge_single_qubit_gates_to_phased_x_and_z,
         cirq.eject_phased_paulis,
         cirq.eject_z,
         cirq.drop_negligible_operations,

cirq/optimizers/merge_interactions_to_sqrt_iswap_test.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, List
+from typing import List
 
 import pytest
 
@@ -37,14 +37,8 @@ def assert_optimizes(before: cirq.Circuit, expected: cirq.Circuit, **kwargs):
     opt.optimize_circuit(actual)
 
     # Ignore differences that would be caught by follow-up optimizations.
-    followup_optimizations: List[Callable[[cirq.Circuit], None]] = [
-        cirq.merge_single_qubit_gates_into_phased_x_z,
-    ]
-    for post in followup_optimizations:
-        post(actual)
-        post(expected)
-
     followup_transformers: List[cirq.TRANSFORMER] = [
+        cirq.merge_single_qubit_gates_to_phased_x_and_z,
         cirq.eject_phased_paulis,
         cirq.eject_z,
         cirq.drop_negligible_operations,

cirq/optimizers/merge_single_qubit_gates.py

@@ -18,13 +18,13 @@
 
 import numpy as np
 
-from cirq import ops, linalg, protocols, circuits
-from cirq.transformers.analytical_decompositions import single_qubit_decompositions
+from cirq import ops, linalg, protocols, circuits, _compat, transformers
 
 if TYPE_CHECKING:
     import cirq
 
 
+@_compat.deprecated_class(deadline='v1.0', fix='Use cirq.merge_k_qubit_unitaries instead.')
 class MergeSingleQubitGates(circuits.PointOptimizer):
     """Optimizes runs of adjacent unitary 1-qubit operations."""
 
@@ -101,6 +101,9 @@ def optimization_at(
         )
 
 
+@_compat.deprecated(
+    deadline='v1.0', fix='Use cirq.merge_single_qubit_gates_to_phased_x_and_z instead.'
+)
 def merge_single_qubit_gates_into_phased_x_z(circuit: circuits.Circuit, atol: float = 1e-8) -> None:
     """Canonicalizes runs of single-qubit rotations in a circuit.
 
@@ -113,14 +116,12 @@ def merge_single_qubit_gates_into_phased_x_z(circuit: circuits.Circuit, atol: fl
         atol: Absolute tolerance to angle error. Larger values allow more
             negligible gates to be dropped, smaller values increase accuracy.
     """
-
-    def synth(qubit: 'cirq.Qid', matrix: np.ndarray) -> List[ops.Operation]:
-        out_gates = single_qubit_decompositions.single_qubit_matrix_to_phased_x_z(matrix, atol)
-        return [gate(qubit) for gate in out_gates]
-
-    MergeSingleQubitGates(synthesizer=synth).optimize_circuit(circuit)
+    circuit._moments = [
+        *transformers.merge_single_qubit_gates_to_phased_x_and_z(circuit, atol=atol)
+    ]
 
 
+@_compat.deprecated(deadline='v1.0', fix='Use cirq.merge_single_qubit_gates_to_phxz instead.')
 def merge_single_qubit_gates_into_phxz(
     circuit: circuits.Circuit,
     atol: float = 1e-8,
@@ -135,9 +136,4 @@ def merge_single_qubit_gates_into_phxz(
         atol: Absolute tolerance to angle error. Larger values allow more
             negligible gates to be dropped, smaller values increase accuracy.
     """
-
-    def synth(qubit: 'cirq.Qid', matrix: np.ndarray) -> List[ops.Operation]:
-        gate = single_qubit_decompositions.single_qubit_matrix_to_phxz(matrix, atol)
-        return [gate(qubit)] if gate else []
-
-    MergeSingleQubitGates(synthesizer=synth).optimize_circuit(circuit)
+    circuit._moments = [*transformers.merge_single_qubit_gates_to_phxz(circuit, atol=atol)]

cirq/optimizers/merge_single_qubit_gates_test.py

@@ -23,10 +23,12 @@ def assert_optimizes(
     before: cirq.Circuit,
     expected: cirq.Circuit,
     optimizer: Optional[Callable[[cirq.Circuit], None]] = None,
+    deprecated_msg: str = "Use cirq.merge_k_qubit_unitaries",
 ):
-    if optimizer is None:
-        optimizer = cirq.MergeSingleQubitGates().optimize_circuit
-    optimizer(before)
+    with cirq.testing.assert_deprecated(deprecated_msg, deadline='v1.0'):
+        if optimizer is None:
+            optimizer = cirq.MergeSingleQubitGates().optimize_circuit
+        optimizer(before)
 
     # Ignore differences that would be caught by follow-up optimizations.
     followup_transformers = [cirq.drop_negligible_operations, cirq.drop_empty_moments]
@@ -38,7 +40,8 @@ def assert_optimizes(
 
 
 def test_leaves_singleton():
-    m = cirq.MergeSingleQubitGates()
+    with cirq.testing.assert_deprecated("Use cirq.merge_k_qubit_unitaries", deadline='v1.0'):
+        m = cirq.MergeSingleQubitGates()
     q = cirq.NamedQubit('q')
     c = cirq.Circuit([cirq.Moment([cirq.X(q)])])
 
@@ -48,12 +51,16 @@ def test_leaves_singleton():
 
 
 def test_not_both():
-    with pytest.raises(ValueError):
-        _ = cirq.MergeSingleQubitGates(synthesizer=lambda *args: None, rewriter=lambda *args: None)
+    with cirq.testing.assert_deprecated("Use cirq.merge_k_qubit_unitaries", deadline='v1.0'):
+        with pytest.raises(ValueError):
+            _ = cirq.MergeSingleQubitGates(
+                synthesizer=lambda *args: None, rewriter=lambda *args: None
+            )
 
 
 def test_combines_sequence():
-    m = cirq.MergeSingleQubitGates()
+    with cirq.testing.assert_deprecated("Use cirq.merge_k_qubit_unitaries", deadline='v1.0'):
+        m = cirq.MergeSingleQubitGates()
     q = cirq.NamedQubit('q')
     c = cirq.Circuit(cirq.X(q) ** 0.5, cirq.Z(q) ** 0.5, cirq.X(q) ** -0.5)
 
@@ -83,7 +90,8 @@ def test_removes_identity_sequence():
 
 
 def test_stopped_at_2qubit():
-    m = cirq.MergeSingleQubitGates()
+    with cirq.testing.assert_deprecated("Use cirq.merge_k_qubit_unitaries", deadline='v1.0'):
+        m = cirq.MergeSingleQubitGates()
     q = cirq.NamedQubit('q')
     q2 = cirq.NamedQubit('q2')
     c = cirq.Circuit(
@@ -109,7 +117,8 @@ def test_stopped_at_2qubit():
 
 
 def test_ignores_2qubit_target():
-    m = cirq.MergeSingleQubitGates()
+    with cirq.testing.assert_deprecated("Use cirq.merge_k_qubit_unitaries", deadline='v1.0'):
+        m = cirq.MergeSingleQubitGates()
     q = cirq.NamedQubit('q')
     q2 = cirq.NamedQubit('q2')
     c = cirq.Circuit(
@@ -132,7 +141,8 @@ class UnsupportedDummy(cirq.SingleQubitGate):
         UnsupportedDummy()(q0),
     )
     c_orig = cirq.Circuit(circuit)
-    cirq.MergeSingleQubitGates().optimize_circuit(circuit)
+    with cirq.testing.assert_deprecated("Use cirq.merge_k_qubit_unitaries", deadline='v1.0'):
+        cirq.MergeSingleQubitGates().optimize_circuit(circuit)
 
     assert circuit == c_orig
 
@@ -147,9 +157,10 @@ def test_rewrite():
         cirq.CZ(q0, q1),
         cirq.Y(q1),
     )
-    cirq.MergeSingleQubitGates(rewriter=lambda ops: cirq.H(ops[0].qubits[0])).optimize_circuit(
-        circuit
-    )
+    with cirq.testing.assert_deprecated("Use cirq.merge_k_qubit_unitaries", deadline='v1.0'):
+        cirq.MergeSingleQubitGates(rewriter=lambda ops: cirq.H(ops[0].qubits[0])).optimize_circuit(
+            circuit
+        )
     circuit = cirq.drop_empty_moments(circuit)
 
     cirq.testing.assert_same_circuits(
@@ -180,6 +191,7 @@ def test_merge_single_qubit_gates_into_phased_x_z():
             (cirq.PhasedXPowGate(phase_exponent=-0.5)(a)) ** 0.5,
         ),
         optimizer=cirq.merge_single_qubit_gates_into_phased_x_z,
+        deprecated_msg="Use cirq.merge_single_qubit_gates_to_phased_x_and_z",
     )
 
 
@@ -207,4 +219,5 @@ def phxz(a, x, z):
             phxz(-0.5, 0.5, 0).on(a),
         ),
         optimizer=cirq.merge_single_qubit_gates_into_phxz,
+        deprecated_msg="Use cirq.merge_single_qubit_gates_to_phxz",
     )

cirq/transformers/__init__.py

@@ -60,6 +60,14 @@
     dephase_measurements,
 )
 
+from cirq.transformers.merge_k_qubit_gates import merge_k_qubit_unitaries
+
+from cirq.transformers.merge_single_qubit_gates import (
+    merge_single_qubit_gates_to_phased_x_and_z,
+    merge_single_qubit_gates_to_phxz,
+    merge_single_qubit_moments_to_phxz,
+)
+
 from cirq.transformers.synchronize_terminal_measurements import synchronize_terminal_measurements
 
 from cirq.transformers.transformer_api import (

cirq/transformers/analytical_decompositions/two_qubit_to_cz.py

@@ -23,9 +23,9 @@
 
 from cirq import ops, linalg, protocols, circuits
 from cirq.transformers.analytical_decompositions import single_qubit_decompositions
+from cirq.transformers.merge_single_qubit_gates import merge_single_qubit_gates_to_phased_x_and_z
 from cirq.transformers.eject_z import eject_z
 from cirq.transformers.eject_phased_paulis import eject_phased_paulis
-from cirq.optimizers import merge_single_qubit_gates
 
 if TYPE_CHECKING:
     import cirq
@@ -161,7 +161,7 @@ def _xx_yy_zz_interaction_via_full_czs(
 
 def _cleanup_operations(operations: Sequence[ops.Operation]):
     circuit = circuits.Circuit(operations)
-    merge_single_qubit_gates.merge_single_qubit_gates_into_phased_x_z(circuit)
+    circuit = merge_single_qubit_gates_to_phased_x_and_z(circuit)
     circuit = eject_phased_paulis(circuit)
     circuit = eject_z(circuit)
     circuit = circuits.Circuit(circuit.all_operations(), strategy=circuits.InsertStrategy.EARLIEST)