Specialize ZPow.controlled() to return CZPow for single qubit control. (

#2402)

#2142

cirq/ops/common_gates.py

@@ -24,7 +24,7 @@
 Each of these are implemented as EigenGates, which means that they can be
 raised to a power (i.e. cirq.H**0.5). See the definition in EigenGate.
 """
-from typing import Any, cast, Optional, Tuple, Union
+from typing import Any, cast, Collection, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import sympy
@@ -380,6 +380,24 @@ def with_canonical_global_phase(self) -> 'ZPowGate':
         """Returns an equal-up-global-phase standardized form of the gate."""
         return ZPowGate(exponent=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:
+        """
+        Specialize controlled for ZPow to return corresponding CZPow when
+        controlled by a single qubit.
+        """
+        result = super().controlled(num_controls, control_values,
+                                    control_qid_shape)
+        if (result.control_values == ((1,),) and  # type: ignore
+                result.control_qid_shape == (2,)):  # type: ignore
+            return cirq.CZPowGate(exponent=self._exponent,
+                                  global_shift=self._global_shift)
+        return result
+
     def _eigen_components(self):
         return [
             (0, np.diag([1, 0])),

cirq/ops/common_gates_test.py

@@ -121,6 +121,15 @@ def test_z_init():
     assert cirq.Z**0.5 != cirq.Z**-0.5
     assert (cirq.Z**-1)**0.5 == cirq.Z**-0.5
     assert cirq.Z**-1 == cirq.Z
+    assert cirq.Z.controlled(num_controls=2) == cirq.ControlledGate(
+        cirq.Z, num_controls=2)
+    assert cirq.Z.controlled(num_controls=1) == cirq.CZ
+    assert cirq.Z.controlled(control_values=((1,),)) == cirq.CZ
+    assert cirq.Z.controlled(control_qid_shape=(2,)) == cirq.CZ
+    assert cirq.Z.controlled(num_controls=1,
+                             control_qid_shape=(3,)) == cirq.ControlledGate(
+                                 cirq.Z, num_controls=1, control_qid_shape=(3,))
+    assert z.controlled() == cirq.CZPowGate(exponent=5)
 
 
 def test_rot_gates_eq():