Override gate_op controlled_by and three_qubit gates controlled methods

Author: tanujkhattar

cirq-core/cirq/interop/quirk/cells/input_rotation_cells_test.py

@@ -64,9 +64,9 @@ def test_input_rotation_cells():
     assert_url_to_circuit_returns(
         '{"cols":[["•","Z^(A/2^n)","inputA2"]]}',
         diagram="""
-0: ───@───────────
+0: ───@^(A/2^2)───
       │
-1: ───Z^(A/2^2)───
+1: ───@───────────
       │
 2: ───A0──────────
       │
@@ -89,9 +89,9 @@ def test_input_rotation_cells():
     assert_url_to_circuit_returns(
         '{"cols":[["•","X^(-A/2^n)","inputA2"]]}',
         diagram="""
-0: ───@────────────
+0: ───@^(-A/2^2)───
       │
-1: ───X^(-A/2^2)───
+1: ───X────────────
       │
 2: ───A0───────────
       │

cirq-core/cirq/interop/quirk/url_to_circuit_test.py

@@ -193,6 +193,26 @@ def test_no_specialized_control_for_global_shift_non_zero(gate, specialized_type
     assert not isinstance(gate.controlled(), specialized_type)
 
 
+@pytest.mark.parametrize(
+    'input_gate, q, expected_gate',
+    [
+        (cirq.Z, cirq.LineQubit.range(2), cirq.CZ),
+        (cirq.Z, cirq.LineQubit.range(3), cirq.CCZ),
+        (cirq.X, cirq.LineQubit.range(2), cirq.CX),
+        (cirq.X, cirq.LineQubit.range(3), cirq.CCX),
+        (cirq.Z ** 0.5, cirq.LineQubit.range(2), cirq.CZ ** 0.5),
+        (cirq.Z ** 0.5, cirq.LineQubit.range(3), cirq.CCZ ** 0.5),
+        (cirq.X ** 0.5, cirq.LineQubit.range(2), cirq.CX ** 0.5),
+        (cirq.X ** 0.5, cirq.LineQubit.range(3), cirq.CCX ** 0.5),
+    ],
+)
+def test_specialized_control_gate_and_op_same(input_gate, q, expected_gate):
+    num_controls = len(q) - 1
+    op1 = input_gate.controlled(num_controls).on(*q[1:], q[0])  # Way-1: gate.controlled().on()
+    op2 = input_gate.on(q[0]).controlled_by(*q[1:])  # Way-2: gate.on().controlled_by()
+    assert op1 == op2 == expected_gate(*q[1:], q[0])
+
+
 @pytest.mark.parametrize(
     'gate, matrix',
     [

cirq-core/cirq/ops/common_gates_test.py

@@ -19,6 +19,7 @@
     AbstractSet,
     Any,
     cast,
+    Collection,
     Dict,
     FrozenSet,
     Iterable,
@@ -333,5 +334,17 @@ def _equal_up_to_global_phase_(
             return False
         return protocols.equal_up_to_global_phase(self.gate, other.gate, atol=atol)
 
+    def controlled_by(
+        self,
+        *control_qubits: 'cirq.Qid',
+        control_values: Optional[Sequence[Union[int, Collection[int]]]] = None,
+    ) -> 'cirq.Operation':
+        qubits = tuple(control_qubits)
+        return self._gate.controlled(
+            num_controls=len(qubits),
+            control_values=control_values,
+            control_qid_shape=tuple(q.dimension for q in qubits),
+        ).on(*(qubits + self._qubits))
+
 
 TV = TypeVar('TV', bound=raw_types.Gate)

cirq-core/cirq/ops/gate_operation.py

@@ -127,13 +127,13 @@ def test_gate():
 
 
 def test_op():
-    a, b, c = cirq.LineQubit.range(3)
+    a, b, c, d = cirq.LineQubit.range(4)
     g = ValiGate()
-    op = g(a)
-    assert op.controlled_by() is op
-    controlled_op = op.controlled_by(b, c)
+    op = g(a, b)
+    assert op.controlled_by() == op
+    controlled_op = op.controlled_by(c, d)
     assert controlled_op.sub_operation == op
-    assert controlled_op.controls == (b, c)
+    assert controlled_op.controls == (c, d)
 
 
 def test_op_validate():

cirq-core/cirq/ops/raw_types_test.py

@@ -14,7 +14,17 @@
 
 """Common quantum gates that target three qubits."""
 
-from typing import AbstractSet, Any, List, Optional, Tuple, TYPE_CHECKING
+from typing import (
+    AbstractSet,
+    Any,
+    Collection,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    TYPE_CHECKING,
+    Union,
+)
 
 import numpy as np
 import sympy
@@ -29,6 +39,7 @@
     gate_features,
     pauli_gates,
     swap_gates,
+    raw_types,
 )
 
 if TYPE_CHECKING:
@@ -167,6 +178,32 @@ def __str__(self) -> str:
             return 'CCZ'
         return f'CCZ**{self._exponent}'
 
+    def controlled(
+        self,
+        num_controls: int = None,
+        control_values: Optional[Sequence[Union[int, Collection[int]]]] = None,
+        control_qid_shape: Optional[Tuple[int, ...]] = None,
+    ) -> raw_types.Gate:
+        """
+        Returns a controlled `ZPowGate` with two additional controls.
+
+        The `controlled` method of the `Gate` class, of which this class is a
+        child, returns a `ControlledGate` with `sub_gate = self`. This method
+        overrides this behavior to return a `ControlledGate` with
+        `sub_gate = ZPowGate`.
+        """
+        if num_controls == 0:
+            return self
+        return controlled_gate.ControlledGate(
+            controlled_gate.ControlledGate(
+                common_gates.ZPowGate(exponent=self._exponent, global_shift=self._global_shift),
+                num_controls=2,
+            ),
+            num_controls=num_controls,
+            control_values=control_values,
+            control_qid_shape=control_qid_shape,
+        )
+
 
 @value.value_equality()
 class ThreeQubitDiagonalGate(gate_features.ThreeQubitGate):
@@ -424,6 +461,32 @@ def __str__(self) -> str:
             return 'TOFFOLI'
         return f'TOFFOLI**{self._exponent}'
 
+    def controlled(
+        self,
+        num_controls: int = None,
+        control_values: Optional[Sequence[Union[int, Collection[int]]]] = None,
+        control_qid_shape: Optional[Tuple[int, ...]] = None,
+    ) -> raw_types.Gate:
+        """
+        Returns a controlled `XPowGate` with two additional controls.
+
+        The `controlled` method of the `Gate` class, of which this class is a
+        child, returns a `ControlledGate` with `sub_gate = self`. This method
+        overrides this behavior to return a `ControlledGate` with
+        `sub_gate = XPowGate`.
+        """
+        if num_controls == 0:
+            return self
+        return controlled_gate.ControlledGate(
+            controlled_gate.ControlledGate(
+                common_gates.XPowGate(exponent=self._exponent, global_shift=self._global_shift),
+                num_controls=2,
+            ),
+            num_controls=num_controls,
+            control_values=control_values,
+            control_qid_shape=control_qid_shape,
+        )
+
 
 @value.value_equality()
 class CSwapGate(gate_features.ThreeQubitGate, gate_features.InterchangeableQubitsGate):
@@ -569,6 +632,29 @@ def __str__(self) -> str:
     def __repr__(self) -> str:
         return 'cirq.FREDKIN'
 
+    def controlled(
+        self,
+        num_controls: int = None,
+        control_values: Optional[Sequence[Union[int, Collection[int]]]] = None,
+        control_qid_shape: Optional[Tuple[int, ...]] = None,
+    ) -> raw_types.Gate:
+        """
+        Returns a controlled `SWAP` with one additional control.
+
+        The `controlled` method of the `Gate` class, of which this class is a
+        child, returns a `ControlledGate` with `sub_gate = self`. This method
+        overrides this behavior to return a `ControlledGate` with
+        `sub_gate = SWAP`.
+        """
+        if num_controls == 0:
+            return self
+        return controlled_gate.ControlledGate(
+            controlled_gate.ControlledGate(swap_gates.SWAP, num_controls=1),
+            num_controls=num_controls,
+            control_values=control_values,
+            control_qid_shape=control_qid_shape,
+        )
+
 
 CCZ = CCZPowGate()
 document(

cirq-core/cirq/ops/three_qubit_gates.py

@@ -319,3 +319,10 @@ def test_resolve(resolve_fn):
     diagonal_gate = resolve_fn(diagonal_gate, {'b': 19})
     assert diagonal_gate == cirq.ThreeQubitDiagonalGate(diagonal_angles)
     assert not cirq.is_parameterized(diagonal_gate)
+
+
+def test_controlled_ops_consistency():
+    a, b, c, d = cirq.LineQubit.range(4)
+    assert cirq.CCX(a, b, c).controlled_by(d) == cirq.CCX(d, b, c).controlled_by(a)
+    assert cirq.CCZ(a, b, c).controlled_by(d) == cirq.CCZ(d, b, c).controlled_by(a)
+    assert cirq.CSWAP(a, b, c).controlled_by(d) == cirq.CSWAP(d, b, c).controlled_by(a)