fix checks

Author: babacry

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

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from typing import List
+from typing import cast, List
 
 import numpy as np
 
@@ -156,12 +156,12 @@ class _PhasedXYAndRz:
     The order is --(X|Y|I)--Rz(rad)--phase--.
     """
 
-    pauli: ops.X | ops.Y | ops.I
+    pauli: ops.Pauli | ops.IdentityGate  # X|Y|I
     rz_rads: float
     phase_exp: float  # phase of the qubit is e^{i*phase_exp*pi}
 
     def __init__(
-        self, pauli: ops.Pauli | ops.I = ops.I, rz_rads: float = 0, phase_exp: float = 0
+        self, pauli: ops.Pauli | ops.IdentityGate = ops.I, rz_rads: float = 0, phase_exp: float = 0
     ) -> None:
         if pauli == ops.Z:  # Merge Z gates to Rz where Z = Rz(π) * e^{iπ/2}
             self.pauli = ops.I
@@ -183,7 +183,7 @@ def _merge_right_rz(self, rads: float):
         """Merges Rz(rads) from right."""
         self.rz_rads += rads
 
-    def _merge_left_xy(self, other: ops.X | ops.Y):
+    def _merge_left_xy(self, other: ops.Pauli):
         """Merges --(X|Y)--self--."""
         if self.pauli == other:
             self.pauli = ops.I
@@ -202,7 +202,7 @@ def _merge_left_xy(self, other: ops.X | ops.Y):
             self.rz_rads -= np.pi
             return
 
-    def _merge_right_xy(self, other: ops.X | ops.Y):
+    def _merge_right_xy(self, other: ops.Pauli):
         """Merges --self--(X|Y)--."""
         self.rz_rads *= -1
         if self.pauli == other:
@@ -227,13 +227,13 @@ def merge_left(self, other: _PhasedXYAndRz) -> None:
         self._merge_left_rz(other.rz_rads)
         self.phase_exp += other.phase_exp
         if other.pauli != ops.I:
-            self._merge_left_xy(other.pauli)
+            self._merge_left_xy(cast(ops.Pauli, other.pauli))
 
     def merge_right(self, other: _PhasedXYAndRz) -> None:
         """Inplace merge other from right."""
         self.phase_exp += other.phase_exp
         if other.pauli != ops.I:
-            self._merge_right_xy(other.pauli)
+            self._merge_right_xy(cast(ops.Pauli, other.pauli))
         self._merge_right_rz(other.rz_rads)
 
     def after_cphase(
@@ -259,41 +259,47 @@ def after_cphase(
                 # Similarly for X|Y on qubit 0/1, the result is always flipping cphase and
                 # add an extra Rz rotation on the other qubit.
                 return (
-                    cphase**-1,
+                    cast(ops.CZPowGate, cphase**-1),
                     self,
                     _PhasedXYAndRz(rz_rads=cphase.exponent * np.pi, phase_exp=-cphase.exponent / 2),
                 )
 
     def __str__(self) -> str:
         return f"─{self.pauli}──Rz({self.rz_rads})──phase(e^{{i{self.phase_exp}π}})─"
 
-    def __eq__(self, other: _PhasedXYAndRz) -> bool:
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, _PhasedXYAndRz):
+            raise NotImplementedError
         return (
             self.pauli == other.pauli
-            and np.isclose(self.rz_rads, other.rz_rads, atol=1e-10)
-            and np.isclose(self.phase_exp, other.phase_exp, atol=1e-10)
+            and bool(np.isclose(self.rz_rads, other.rz_rads, atol=1e-10))
+            and bool(np.isclose(self.phase_exp, other.phase_exp, atol=1e-10))
         )
 
     def to_single_gate(self) -> ops.PhasedXZGate | ops.ZPowGate:
-        if self.pauli == ops.I:
-            rz_rads = self.rz_rads
-            if np.isclose(self.rz_rads, 0, atol=1e-2):
-                rz_rads = self.rz_rads + 4 * np.pi
-            return ops.ZPowGate(
-                exponent=rz_rads / np.pi, global_shift=np.pi * self.phase_exp / rz_rads - 0.5
-            )
-        if self.pauli == ops.X:
-            return ops.PhasedXZGate(
-                x_exponent=1,
-                z_exponent=2 * self.phase_exp,
-                axis_phase_exponent=self.rz_rads / 2 / np.pi - self.phase_exp,
-            )
-        if self.pauli == ops.Y:
-            return ops.PhasedXZGate(
-                x_exponent=1,
-                z_exponent=2 * self.phase_exp,
-                axis_phase_exponent=1 / 2 - self.phase_exp + self.rz_rads / 2 / np.pi,
-            )
+        """Converts the _PhasedXYAndRz to a single-qubit gate."""
+        match self.pauli:
+            case ops.I:
+                rz_rads = self.rz_rads
+                if np.isclose(self.rz_rads, 0, atol=1e-2):
+                    rz_rads = self.rz_rads + 4 * np.pi
+                return ops.ZPowGate(
+                    exponent=rz_rads / np.pi, global_shift=np.pi * self.phase_exp / rz_rads - 0.5
+                )
+            case ops.X:
+                return ops.PhasedXZGate(
+                    x_exponent=1,
+                    z_exponent=2 * self.phase_exp,
+                    axis_phase_exponent=self.rz_rads / 2 / np.pi - self.phase_exp,
+                )
+            case ops.Y:
+                return ops.PhasedXZGate(
+                    x_exponent=1,
+                    z_exponent=2 * self.phase_exp,
+                    axis_phase_exponent=1 / 2 - self.phase_exp + self.rz_rads / 2 / np.pi,
+                )
+            case _:
+                raise ValueError("Invalid self.pauli.")
 
 
 def _pull_through_single_cphase(
@@ -327,26 +333,31 @@ def _pull_through_single_cphase(
     return output_cphase, output0, output1
 
 
-def _multi_moment_pull_through(
+def _multi_moment_gauge_fn(
     moments: List[circuits.Moment], rng: np.random.Generator
 ) -> List[circuits.Moment]:
-    """TO FILL."""
+    """Generates a left layer with random generator, then pulling through all the moments."""
     all_qubits = [q for q in circuits.Circuit(moments).all_qubits()]
     if not all_qubits:
         return moments
     if not any(isinstance(op.gate, ops.CZPowGate) for moment in moments for op in moment):
         return moments
 
-    left_moment = circuits.Moment(
-        [rng.choice([ops.I, ops.X, ops.Y, ops.Z]).on(q) for q in all_qubits]
-    )
-    prev: map[ops.Qid, ops.Gate] = {
-        op.qubits[0]: _PhasedXYAndRz(pauli=op.gate) for op in left_moment
+    left_moment: List[ops.Operation] = [
+        rng.choice(
+            np.array([ops.I, ops.X, ops.Y, ops.Z], dtype=ops.Gate), p=[0.25, 0.25, 0.25, 0.25]
+        ).on(q)
+        for q in all_qubits
+    ]
+    prev: dict[ops.Qid, _PhasedXYAndRz] = {
+        op.qubits[0]: _PhasedXYAndRz(pauli=cast(ops.Pauli | ops.IdentityGate, op.gate))
+        for op in left_moment
+        if op.gate
     }
 
-    new_moments: List[circuits.Moment] = [left_moment]
+    new_moments: List[circuits.Moment] = [circuits.Moment(left_moment)]
 
-    pulled: map[ops.Qid, ops.Gate]
+    pulled: dict[ops.Qid, _PhasedXYAndRz]
     for moment in moments:
         pulled = {}
         new_moment: List[ops.Operation] = []
@@ -368,7 +379,7 @@ def _multi_moment_pull_through(
             if q not in pulled:
                 pulled[q] = prev[q]
         prev = pulled
-        new_moments.append(new_moment)
+        new_moments.append(circuits.Moment(new_moment))
 
     last_moment = circuits.Moment([pulled[q].to_single_gate().on(q) for q in all_qubits])
 
@@ -381,5 +392,5 @@ def _multi_moment_pull_through(
 CPhaseGaugeTransformerMM = GaugeTransformer(
     target=ops.Gateset(ops.CZPowGate, ops.ZPowGate),
     gauge_selector=CPhaseGaugeSelector,
-    multi_moment_pull_thourgh_fn=_multi_moment_pull_through,
+    multi_moment_gauge_fn=_multi_moment_gauge_fn,
 )

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

@@ -13,13 +13,14 @@
 # limitations under the License.
 
 from copy import deepcopy
+
 import numpy as np
 
 import cirq
 from cirq.transformers.gauge_compiling.cphase_gauge import (
+    _PhasedXYAndRz,
     CPhaseGaugeTransformer,
     CPhaseGaugeTransformerMM,
-    _PhasedXYAndRz,
 )
 from cirq.transformers.gauge_compiling.gauge_compiling_test_utils import GaugeTester
 
@@ -115,23 +116,23 @@ def test_multi_layer_pull_through():
     """Test case.
     Input:
                   ┌──┐
-    0: ───@────────@─────H───────────@───────@───────
-          │        │                 │       │
-    1: ───@^0.2────┼@────────────────@^0.1───@───────
+    0: ───@────────@─────H───Rz(-0.255π)───@───────@───────
+          │        │                       │       │
+    1: ───@^0.2────┼@──────────────────────@^0.1───@───────
                    ││
-    2: ───@────────@┼────────@───────@───────@───────
-          │         │        │       │       │
-    3: ───@─────────@────────@^0.2───@───────@^0.2───
+    2: ───@────────@┼────────@─────────────@───────@───────
+          │         │        │             │       │
+    3: ───@─────────@────────@^0.2─────────@───────@^0.2───
                   └──┘
     Example output:
                        ┌──┐
-    0: ───Y───@─────────@─────PhXZ(a=0.5,x=1,z=0)───H───Y────────────@────────@────────PhXZ(a=0.5,x=1,z=0)───
-              │         │                                            │        │
-    1: ───Z───@^-0.2────┼@────Z^-0.8────────────────────I────────────@^-0.1───@────────Z^-0.9────────────────
+    0: ───Z───@─────────@─────Z^0.2────────────────H───I───────────@────────@───────Z^0.845──────────────────  # pylint: disable=line-too-long
+              │         │                                          │        │
+    1: ───X───@^-0.2────┼@────PhXZ(a=0,x=1,z=1)────────X───────────@^-0.1───@───────PhXZ(a=0,x=1,z=0)────────  # pylint: disable=line-too-long
                         ││
-    2: ───Z───@─────────@┼────Z^0───────────────────────Y───@────────@────────@────────PhXZ(a=0.5,x=1,z=0)───
-              │          │                                  │        │        │
-    3: ───I───@──────────@────Z^0───────────────────────I───@^-0.2───@────────@^-0.2───Z^-0.6────────────────
+    2: ───X───@─────────@┼────PhXZ(a=0,x=1,z=1)────────Y───@───────@────────@───────PhXZ(a=0.5,x=1,z=1.4)────  # pylint: disable=line-too-long
+              │          │                                 │       │        │
+    3: ───X───@──────────@────PhXZ(a=2,x=1,z=-2)───────Y───@^0.2───@────────@^0.2───PhXZ(a=1.9,x=1,z=-1.4)───  # pylint: disable=line-too-long
                        └──┘
     """
     q0, q1, q2, q3 = cirq.LineQubit.range(4)
@@ -140,7 +141,7 @@ def test_multi_layer_pull_through():
             cirq.Moment(cirq.CZ(q0, q1) ** 0.2, cirq.CZ(q2, q3)),
             cirq.Moment(cirq.CZ(q0, q2), cirq.CZ(q1, q3)),
             cirq.Moment(cirq.H(q0)),
-            cirq.Moment(cirq.CZ(q2, q3) ** 0.2),
+            cirq.Moment(cirq.CZ(q2, q3) ** 0.2, cirq.Rz(rads=-0.8).on(q0)),
             cirq.Moment(cirq.CZ(q0, q1) ** 0.1, cirq.CZ(q2, q3)),
             cirq.Moment(cirq.CZ(q0, q1), cirq.CZ(q2, q3) ** 0.2),
         )

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

@@ -48,15 +48,17 @@
 )
 
 
-def _multi_moment_pull_through(
+def _multi_moment_gauge_fn(
     moments: List[circuits.Moment], rng: np.random.Generator
 ) -> List[circuits.Moment]:
     # Check all the ops are CZ first
     if not all(op.gate == CZ for moment in moments for op in moment):
         raise ValueError(f"Input moments must only contain CZ gates:\nmoments = {moments}.")
 
     left: List[ops.Operation] = [
-        rng.choice([ops.I, ops.X, ops.Y, ops.Z]).on(q)
+        rng.choice(
+            np.array([ops.I, ops.X, ops.Y, ops.Z], dtype=ops.Gate), p=[0.25, 0.25, 0.25, 0.25]
+        ).on(q)
         for q in circuits.Circuit(moments).all_qubits()
     ]
     if not left:
@@ -73,7 +75,5 @@ def _multi_moment_pull_through(
 
 # Multi-moments pull through version of CZGaugeTransformer
 CZGaugeTransformerMM = GaugeTransformer(
-    target=CZ,
-    gauge_selector=CZGaugeSelector,
-    multi_moment_pull_thourgh_fn=_multi_moment_pull_through,
+    target=CZ, gauge_selector=CZGaugeSelector, multi_moment_gauge_fn=_multi_moment_gauge_fn
 )

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

@@ -196,8 +196,8 @@ def __init__(
         target: Union[ops.Gate, ops.Gateset, ops.GateFamily],
         gauge_selector: Callable[[np.random.Generator], Gauge],
         two_qubit_gate_symbolizer: Optional[TwoQubitGateSymbolizer] = None,
-        multi_moment_pull_thourgh_fn: Optional[
-            Callable[[List[circuits.Moment], List[circuits.Moment]], List[circuits.Moment]]
+        multi_moment_gauge_fn: Optional[
+            Callable[[List[circuits.Moment], np.random.Generator], List[circuits.Moment]]
         ] = None,
     ) -> None:
         """Constructs a GaugeTransformer.
@@ -211,7 +211,7 @@ def __init__(
         self.target = ops.GateFamily(target) if isinstance(target, ops.Gate) else target
         self.gauge_selector = gauge_selector
         self.two_qubit_gate_symbolizer = two_qubit_gate_symbolizer
-        self.multi_moment_pull_thourgh_fn = multi_moment_pull_thourgh_fn
+        self.multi_moment_gauge_fn = multi_moment_gauge_fn
 
     def __call__(
         self,
@@ -225,22 +225,22 @@ def __call__(
             context = transformer_api.TransformerContext(deep=False)
         if context.deep:
             raise ValueError('GaugeTransformer cannot be used with deep=True')
-        new_moments = []
+        new_moments: List[circuits.Moment] = []
         left: List[List[ops.Operation]] = []
         right: List[List[ops.Operation]] = []
         all_target_moments: List[circuits.Moment] = []
 
         for moment in circuit:
-            if self.multi_moment_pull_thourgh_fn and all(
+            if self.multi_moment_gauge_fn and all(
                 [
                     op in self.target and not set(op.tags).intersection(context.tags_to_ignore)
                     for op in moment
                 ]
             ):  # all ops are target 2-qubit gates
                 all_target_moments.append(moment)
                 continue
-            if all_target_moments:
-                new_moments.extend(self.multi_moment_pull_thourgh_fn(all_target_moments, rng))
+            if all_target_moments and self.multi_moment_gauge_fn:
+                new_moments.extend(self.multi_moment_gauge_fn(all_target_moments, rng))
                 all_target_moments.clear()
 
             left.clear()
@@ -262,11 +262,11 @@ def __call__(
                     center.append(op)
             if left:
                 new_moments.extend(_build_moments(left))
-            new_moments.append(center)
+            new_moments.append(circuits.Moment(center))
             if right:
                 new_moments.extend(_build_moments(right))
-        if all_target_moments:
-            new_moments.extend(self.multi_moment_pull_thourgh_fn(all_target_moments, rng))
+        if all_target_moments and self.multi_moment_gauge_fn:
+            new_moments.extend(self.multi_moment_gauge_fn(all_target_moments, rng))
         return circuits.Circuit.from_moments(*new_moments)
 
     def as_sweep(
@@ -412,15 +412,15 @@ def two_qubit_gate_next_symbol_list(n: int) -> List[sympy.Symbol]:
         return circuits.Circuit.from_moments(*new_moments), Zip(*sweeps)
 
 
-def _build_moments(operation_by_qubits: List[List[ops.Operation]]) -> List[List[ops.Operation]]:
+def _build_moments(operation_by_qubits: List[List[ops.Operation]]) -> List[circuits.Moment]:
     """Builds moments from a list of operations grouped by qubits.
 
     Returns a list of moments from a list whose ith element is a list of operations applied
     to qubit i.
     """
     moments = []
     for moment in itertools.zip_longest(*operation_by_qubits):
-        moments.append([op for op in moment if op is not None])
+        moments.append(circuits.Moment([op for op in moment if op is not None]))
     return moments