Disentangle MS gate and XXPowGate. (#3015)

* Add separate MSGate and distangle from XXPowGate

* Revert accidental deletion of quil related changes.

* Fix formatting changes and add tests for coverage

* Add equality group test and resolved comments

Co-authored-by: Dave Bacon <dabacon@gmail.com>

cirq/ion/ion_gates.py

@@ -14,12 +14,17 @@
 
 """Operations native to iontrap systems."""
 
+from typing import Union, TYPE_CHECKING
 import numpy as np
 
-from cirq import ops
+from cirq import ops, value
+from cirq import protocols
 
+if TYPE_CHECKING:
+    import cirq
 
-def ms(rads: float) -> ops.XXPowGate:
+
+class MSGate(ops.XXPowGate):
     """The Mølmer–Sørensen gate, a native two-qubit operation in ion traps.
 
     A rotation around the XX axis in the two-qubit bloch sphere.
@@ -30,11 +35,42 @@ def ms(rads: float) -> ops.XXPowGate:
                        [ 0        cos(t)  -isin(t)  0      ]
                        [ 0       -isin(t)  cos(t)   0      ]
                        [-isin(t)  0        0        cos(t) ]
+    """
+
+    def __init__(
+            self,
+            *,  # Forces keyword args.
+            rads: float):
+        ops.XXPowGate.__init__(self,
+                               exponent=rads * 2 / np.pi,
+                               global_shift=-0.5)
+
+    def _with_exponent(self: 'MSGate', exponent: value.TParamVal) -> 'MSGate':
+        return type(self)(rads=exponent * np.pi / 2)
+
+    def _circuit_diagram_info_(self, args: 'cirq.CircuitDiagramInfoArgs'
+                              ) -> Union[str, 'protocols.CircuitDiagramInfo']:
+        angle_str = self._format_exponent_as_angle(args, order=4)
+        symbol = f'MS({angle_str})'
+        return protocols.CircuitDiagramInfo(wire_symbols=(symbol, symbol))
 
+    def __str__(self) -> str:
+        if self._exponent == 1:
+            return 'MS(π/2)'
+        return f'MS({self._exponent!r}π/2)'
+
+    def __repr__(self) -> str:
+        if self._exponent == 1:
+            return 'cirq.ms(np.pi/2)'
+        return f'cirq.ms({self._exponent!r}*np.pi/2)'
+
+
+def ms(rads: float) -> MSGate:
+    """
     Args:
         rads: The rotation angle in radians.
 
     Returns:
         Mølmer–Sørensen gate rotating by the desired amount.
     """
-    return ops.XXPowGate(exponent=rads * 2 / np.pi, global_shift=-0.5)
+    return MSGate(rads=rads)

cirq/ion/ion_gates_test.py

@@ -20,12 +20,17 @@
 def test_ms_arguments():
     eq_tester = cirq.testing.EqualsTester()
     eq_tester.add_equality_group(cirq.ms(np.pi / 2),
-                                 cirq.XXPowGate(global_shift=-0.5))
+                                 cirq.ion.ion_gates.MSGate(rads=np.pi / 2))
+    eq_tester.add_equality_group(cirq.XXPowGate(global_shift=-0.5))
 
 
 def test_ms_str():
-    assert str(cirq.ms(np.pi / 2)) == 'MS(π/2)'
+    ms = cirq.ms(np.pi / 2)
+    assert str(ms) == 'MS(π/2)'
     assert str(cirq.ms(np.pi)) == 'MS(2.0π/2)'
+    assert str(ms**0.5) == 'MS(0.5π/2)'
+    assert str(ms**2) == 'MS(2.0π/2)'
+    assert str(ms**-1) == 'MS(-1.0π/2)'
 
 
 def test_ms_matrix():
@@ -47,6 +52,9 @@ def test_ms_repr():
     assert repr(cirq.ms(np.pi / 2)) == 'cirq.ms(np.pi/2)'
     assert repr(cirq.ms(np.pi / 4)) == 'cirq.ms(0.5*np.pi/2)'
     cirq.testing.assert_equivalent_repr(cirq.ms(np.pi / 4))
+    ms = cirq.ms(np.pi / 2)
+    assert (repr(ms**2) == 'cirq.ms(2.0*np.pi/2)')
+    assert (repr(ms**-0.5) == 'cirq.ms(-0.5*np.pi/2)')
 
 
 def test_ms_diagrams():

cirq/ops/parity_gates.py

@@ -39,7 +39,7 @@ class XXPowGate(eigen_gate.EigenGate,
               [0 1 0 0]
               [1 0 0 0]
 
-    See also: `cirq.MS` (the Mølmer–Sørensen gate), which is implemented via
+    See also: `cirq.MSGate` (the Mølmer–Sørensen gate), which is implemented via
         this class.
     """
 
@@ -86,13 +86,6 @@ def _decompose_into_clifford_with_qubits_(self, qubits):
 
     def _circuit_diagram_info_(self, args: 'cirq.CircuitDiagramInfoArgs'
                               ) -> Union[str, 'protocols.CircuitDiagramInfo']:
-        if self._global_shift == -0.5:
-            # Mølmer–Sørensen gate.
-            angle_str = self._format_exponent_as_angle(args, order=4)
-            symbol = f'MS({angle_str})'
-            return protocols.CircuitDiagramInfo(
-                                wire_symbols=(symbol, symbol))
-
         return protocols.CircuitDiagramInfo(
             wire_symbols=('XX', 'XX'),
             exponent=self._diagram_exponent(args))
@@ -105,19 +98,11 @@ def _quil_(self, qubits: Tuple['cirq.Qid', ...],
                                 qubits[0], self._exponent, qubits[1])
 
     def __str__(self) -> str:
-        if self._global_shift == -0.5:
-            if self._exponent == 1:
-                return 'MS(π/2)'
-            return f'MS({self._exponent!r}π/2)'
         if self.exponent == 1:
             return 'XX'
         return f'XX**{self._exponent!r}'
 
     def __repr__(self) -> str:
-        if self._global_shift == -0.5 and not protocols.is_parameterized(self):
-            if self._exponent == 1:
-                return 'cirq.ms(np.pi/2)'
-            return f'cirq.ms({self._exponent!r}*np.pi/2)'
         if self._global_shift == 0:
             if self._exponent == 1:
                 return 'cirq.XX'

cirq/ops/parity_gates_test.py

@@ -63,23 +63,10 @@ def test_xx_str():
     assert str(cirq.XX**0.5) == 'XX**0.5'
     assert str(cirq.XXPowGate(global_shift=0.1)) == 'XX'
 
-    ms = cirq.XXPowGate(global_shift=-0.5)
-    assert str(ms) == 'MS(π/2)'
-    assert str(ms**0.5) == 'MS(0.5π/2)'
-    assert str(ms**2) == 'MS(2.0π/2)'
-    assert str(ms**-1) == 'MS(-1.0π/2)'
-
-
 def test_xx_repr():
     assert repr(cirq.XXPowGate()) == 'cirq.XX'
     assert repr(cirq.XXPowGate(exponent=0.5)) == '(cirq.XX**0.5)'
 
-    ms = cirq.XXPowGate(global_shift=-0.5)
-    assert (repr(ms) == 'cirq.ms(np.pi/2)')
-    assert (repr(ms**2) == 'cirq.ms(2.0*np.pi/2)')
-    assert (repr(ms**-0.5) == 'cirq.ms(-0.5*np.pi/2)')
-
-
 def test_xx_matrix():
     np.testing.assert_allclose(cirq.unitary(cirq.XX),
                                np.array([[0, 0, 0, 1],
@@ -109,13 +96,12 @@ def test_xx_diagrams():
         cirq.XX(a, b),
         cirq.XX(a, b)**3,
         cirq.XX(a, b)**0.5,
-        cirq.XXPowGate(global_shift=-0.5).on(a, b),
     )
     cirq.testing.assert_has_diagram(
         circuit, """
-a: ───XX───XX───XX───────MS(0.5π)───
-      │    │    │        │
-b: ───XX───XX───XX^0.5───MS(0.5π)───
+a: ───XX───XX───XX───────
+      │    │    │
+b: ───XX───XX───XX^0.5───
 """)