Add qubits to PauliStringPhasor

cirq-core/cirq/ops/gate_operation_test.py

@@ -488,6 +488,7 @@ def all_subclasses(cls):
 
     skip_classes = {
         # Abstract or private parent classes.
+        cirq.ArithmeticGate,
         cirq.BaseDensePauliString,
         cirq.EigenGate,
         cirq.Pauli,
@@ -503,17 +504,10 @@ def all_subclasses(cls):
         # Interop gates
         cirq.interop.quirk.QuirkQubitPermutationGate,
         cirq.interop.quirk.QuirkArithmeticGate,
-        # No reason given for missing json.
-        # TODO(#5353): Serialize these gates.
-        cirq.ArithmeticGate,
     }
 
-    # Gates that do not satisfy the contract.
-    # TODO(#5167): Fix this case.
-    exceptions = {cirq.PauliStringPhasorGate}
-
     skipped = set()
-    for gate_cls in gate_subclasses - exceptions:
+    for gate_cls in gate_subclasses:
         filename = test_module_spec.test_data_path.joinpath(f"{gate_cls.__name__}.json")
         if pathlib.Path(filename).is_file():
             gates = cirq.read_json(filename)

cirq-core/cirq/ops/pauli_string.py

@@ -391,7 +391,7 @@ def _repr_pretty_(self, p: Any, cycle: bool) -> None:
             p.text(str(self))
 
     def __repr__(self) -> str:
-        ordered_qubits = sorted(self.qubits)
+        ordered_qubits = self.qubits
         prefix = ''
 
         factors = []

cirq-core/cirq/ops/pauli_string_phasor.py

@@ -12,7 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import AbstractSet, cast, Dict, Iterable, Union, TYPE_CHECKING, Sequence, Iterator
+from typing import (
+    AbstractSet,
+    cast,
+    Dict,
+    Iterable,
+    Iterator,
+    Optional,
+    Sequence,
+    TYPE_CHECKING,
+    Union,
+)
+
 import numbers
 
 import sympy
@@ -35,16 +46,25 @@
 
 @value.value_equality(approximate=True)
 class PauliStringPhasor(gate_operation.GateOperation):
-    """An operation that phases the eigenstates of a Pauli string.
+    r"""An operation that phases the eigenstates of a Pauli string.
+
+    This class takes `PauliString`, which is a sequence of non-identity
+    Pauli operators, potentially with a $\pm 1$ valued coefficient,
+    acting on qubits.
 
     The -1 eigenstates of the Pauli string will have their amplitude multiplied
     by e^(i pi exponent_neg) while +1 eigenstates of the Pauli string will have
     their amplitude multiplied by e^(i pi exponent_pos).
+
+    The class also takes a list of qubits, which can be bigger than those
+    provided acted on by `PauliString`.  Those extra qubits are assumed to be
+    acted upon via identity.
     """
 
     def __init__(
         self,
         pauli_string: ps.PauliString,
+        qubits: Optional[Sequence['cirq.Qid']] = None,
         *,
         exponent_neg: Union[int, float, sympy.Expr] = 1,
         exponent_pos: Union[int, float, sympy.Expr] = 0,
@@ -54,20 +74,36 @@ def __init__(
         Args:
             pauli_string: The PauliString defining the positive and negative
                 eigenspaces that will be independently phased.
+            qubits: The qubits upon which the PauliStringPhasor acts. This
+                can be different from the qubits of `pauli_string` but if
+                it is not supplied it will use the qubits from `pauli_string`.
+                The `pauli_string` contains only the non-identity component
+                of the phasor, while the qubits supplied here and not in
+                `pauli_string` are acted upon by identity. The order of
+                these qubits must match the order in `pauli_string`.
             exponent_neg: How much to phase vectors in the negative eigenspace,
                 in the form of the t in (-1)**t = exp(i pi t).
             exponent_pos: How much to phase vectors in the positive eigenspace,
                 in the form of the t in (-1)**t = exp(i pi t).
 
         Raises:
-            ValueError: If coefficient is not 1 or -1.
+            ValueError: If coefficient is not 1 or -1 or the qubits of
+                `pauli_string` are not a subset of `qubits`.
         """
+        if qubits is not None:
+            it = iter(qubits)
+            if any(not any(q0 == q1 for q1 in it) for q0 in pauli_string.qubits):
+                raise ValueError(
+                    f"PauliStringPhasor's pauli string qubits ({pauli_string.qubits}) "
+                    f"are not a subset of the explicit qubits ({qubits})."
+                )
         gate = PauliStringPhasorGate(
-            pauli_string.dense(pauli_string.qubits),
+            pauli_string.dense(qubits or pauli_string.qubits),
             exponent_neg=exponent_neg,
             exponent_pos=exponent_pos,
         )
-        super().__init__(gate, pauli_string.qubits)
+        super().__init__(gate, qubits or pauli_string.qubits)
+        it = iter(self.qubits)
         self._pauli_string = gate.dense_pauli_string.on(*self.qubits)
 
     @property
@@ -96,41 +132,54 @@ def exponent_relative(self) -> Union[int, float, sympy.Expr]:
         return self.gate.exponent_relative
 
     def _value_equality_values_(self):
-        return (self.pauli_string, self.exponent_neg, self.exponent_pos)
+        return (self.pauli_string, self.qubits, self.exponent_neg, self.exponent_pos)
 
     def equal_up_to_global_phase(self, other):
         """Checks equality of two PauliStringPhasors, up to global phase."""
         if isinstance(other, PauliStringPhasor):
-            rel1 = self.exponent_relative
-            rel2 = other.exponent_relative
-            return rel1 == rel2 and self.pauli_string == other.pauli_string
+            return (
+                self.exponent_relative == other.exponent_relative
+                and self.pauli_string == other.pauli_string
+                and self.qubits == other.qubits
+            )
         return False
 
     def map_qubits(self, qubit_map: Dict[raw_types.Qid, raw_types.Qid]):
         """Maps the qubits inside the PauliString."""
         return PauliStringPhasor(
-            self.pauli_string.map_qubits(qubit_map),
+            pauli_string=self.pauli_string.map_qubits(qubit_map),
+            qubits=[qubit_map[q] for q in self.qubits],
             exponent_neg=self.exponent_neg,
             exponent_pos=self.exponent_pos,
         )
 
     def can_merge_with(self, op: 'PauliStringPhasor') -> bool:
         """Checks whether the underlying PauliStrings can be merged."""
-        return self.pauli_string.equal_up_to_coefficient(op.pauli_string)
+        return (
+            self.pauli_string.equal_up_to_coefficient(op.pauli_string) and self.qubits == op.qubits
+        )
 
     def merged_with(self, op: 'PauliStringPhasor') -> 'PauliStringPhasor':
         """Merges two PauliStringPhasors."""
         if not self.can_merge_with(op):
             raise ValueError(f'Cannot merge operations: {self}, {op}')
         pp = self.exponent_pos + op.exponent_pos
         pn = self.exponent_neg + op.exponent_neg
-        return PauliStringPhasor(self.pauli_string, exponent_pos=pp, exponent_neg=pn)
+        return PauliStringPhasor(
+            self.pauli_string, qubits=self.qubits, exponent_pos=pp, exponent_neg=pn
+        )
 
     def _circuit_diagram_info_(
         self, args: 'cirq.CircuitDiagramInfoArgs'
     ) -> 'cirq.CircuitDiagramInfo':
         qubits = self.qubits if args.known_qubits is None else args.known_qubits
-        syms = tuple(f'[{self.pauli_string[qubit]}]' for qubit in qubits)
+
+        def sym(qubit):
+            if qubit in self.pauli_string:
+                return f'[{self.pauli_string[qubit]}]'
+            return '[I]'
+
+        syms = tuple(sym(qubit) for qubit in qubits)
         return protocols.CircuitDiagramInfo(wire_symbols=syms, exponent=self.exponent_relative)
 
     def pass_operations_over(
@@ -170,6 +219,7 @@ def pass_operations_over(
     def __repr__(self) -> str:
         return (
             f'cirq.PauliStringPhasor({self.pauli_string!r}, '
+            f'qubits={self.qubits!r}, '
             f'exponent_neg={proper_repr(self.exponent_neg)}, '
             f'exponent_pos={proper_repr(self.exponent_pos)})'
         )
@@ -182,7 +232,9 @@ def __str__(self) -> str:
         return f'({self.pauli_string})**{self.exponent_relative}'
 
     def _json_dict_(self):
-        return protocols.obj_to_dict_helper(self, ['pauli_string', 'exponent_neg', 'exponent_pos'])
+        return protocols.obj_to_dict_helper(
+            self, ['pauli_string', 'qubits', 'exponent_neg', 'exponent_pos']
+        )
 
 
 @value.value_equality(approximate=True)
@@ -352,6 +404,7 @@ def on(self, *qubits: 'cirq.Qid') -> 'cirq.PauliStringPhasor':
         """Creates a PauliStringPhasor on the qubits."""
         return PauliStringPhasor(
             self.dense_pauli_string.on(*qubits),
+            qubits=qubits,
             exponent_pos=self.exponent_pos,
             exponent_neg=self.exponent_neg,
         )

cirq-core/cirq/ops/pauli_string_phasor_test.py

@@ -36,7 +36,6 @@ def test_init():
     a = cirq.LineQubit(0)
     with pytest.raises(ValueError, match='eigenvalues'):
         _ = cirq.PauliStringPhasor(1j * cirq.X(a))
-
     v1 = cirq.PauliStringPhasor(-cirq.X(a), exponent_neg=0.25, exponent_pos=-0.5)
     assert v1.pauli_string == cirq.X(a)
     assert v1.exponent_neg == -0.5
@@ -48,8 +47,20 @@ def test_init():
     assert v2.exponent_pos == -0.125
 
 
+def test_qubit_order_mismatch():
+    q0, q1 = cirq.LineQubit.range(2)
+    with pytest.raises(ValueError, match='are not a subset'):
+        _ = cirq.PauliStringPhasor(1j * cirq.X(q0), qubits=[q1])
+    with pytest.raises(ValueError, match='are not a subset'):
+        _ = cirq.PauliStringPhasor(1j * cirq.X(q0) * cirq.X(q1), qubits=[q1])
+    with pytest.raises(ValueError, match='are not a subset'):
+        _ = cirq.PauliStringPhasor(1j * cirq.X(q0), qubits=[])
+    with pytest.raises(ValueError, match='are not a subset'):
+        _ = cirq.PauliStringPhasor(1j * cirq.X(q0) * cirq.X(q1), qubits=[q1, q0])
+
+
 def test_eq_ne_hash():
-    q0, q1, q2 = _make_qubits(3)
+    q0, q1, q2, q3 = _make_qubits(4)
     eq = cirq.testing.EqualsTester()
     ps1 = cirq.X(q0) * cirq.Y(q1) * cirq.Z(q2)
     ps2 = cirq.X(q0) * cirq.Y(q1) * cirq.X(q2)
@@ -66,10 +77,11 @@ def test_eq_ne_hash():
     eq.add_equality_group(cirq.PauliStringPhasor(ps2, exponent_neg=0.5))
     eq.add_equality_group(cirq.PauliStringPhasor(-ps2, exponent_neg=-0.5))
     eq.add_equality_group(cirq.PauliStringPhasor(ps1, exponent_neg=sympy.Symbol('a')))
+    eq.add_equality_group(cirq.PauliStringPhasor(ps1, qubits=[q0, q1, q2, q3]))
 
 
 def test_equal_up_to_global_phase():
-    a, b = cirq.LineQubit.range(2)
+    a, b, c = cirq.LineQubit.range(3)
     groups = [
         [
             cirq.PauliStringPhasor(cirq.PauliString({a: cirq.X}), exponent_neg=0.25),
@@ -83,6 +95,11 @@ def test_equal_up_to_global_phase():
         [cirq.PauliStringPhasor(cirq.PauliString({a: cirq.X}))],
         [cirq.PauliStringPhasor(cirq.PauliString({a: cirq.Y}), exponent_neg=0.25)],
         [cirq.PauliStringPhasor(cirq.PauliString({a: cirq.X, b: cirq.Y}), exponent_neg=0.25)],
+        [
+            cirq.PauliStringPhasor(
+                cirq.PauliString({a: cirq.X, b: cirq.Y}), qubits=[a, b, c], exponent_neg=0.25
+            )
+        ],
     ]
     for g1 in groups:
         for e1 in g1:
@@ -93,12 +110,21 @@ def test_equal_up_to_global_phase():
 
 
 def test_map_qubits():
-    q0, q1, q2, q3 = _make_qubits(4)
+    q0, q1, q2, q3, q4, q5 = _make_qubits(6)
     qubit_map = {q1: q2, q0: q3}
     before = cirq.PauliStringPhasor(cirq.PauliString({q0: cirq.Z, q1: cirq.Y}), exponent_neg=0.1)
     after = cirq.PauliStringPhasor(cirq.PauliString({q3: cirq.Z, q2: cirq.Y}), exponent_neg=0.1)
     assert before.map_qubits(qubit_map) == after
 
+    qubit_map = {q1: q3, q0: q4, q2: q5}
+    before = cirq.PauliStringPhasor(
+        cirq.PauliString({q0: cirq.Z, q1: cirq.Y}), qubits=[q0, q1, q2], exponent_neg=0.1
+    )
+    after = cirq.PauliStringPhasor(
+        cirq.PauliString({q4: cirq.Z, q3: cirq.Y}), qubits=[q4, q3, q5], exponent_neg=0.1
+    )
+    assert before.map_qubits(qubit_map) == after
+
 
 def test_pow():
     a = cirq.LineQubit(0)
@@ -108,12 +134,16 @@ def test_pow():
     with pytest.raises(TypeError, match='unsupported operand'):
         _ = p ** object()
     assert p**1 == p
+    p = cirq.PauliStringPhasor(s, qubits=[a], exponent_neg=0.25, exponent_pos=0.5)
+    assert p**0.5 == cirq.PauliStringPhasor(s, exponent_neg=0.125, exponent_pos=0.25)
 
 
 def test_consistent():
     a, b = cirq.LineQubit.range(2)
     op = np.exp(1j * np.pi / 2 * cirq.X(a) * cirq.X(b))
     cirq.testing.assert_implements_consistent_protocols(op)
+    p = cirq.PauliStringPhasor(cirq.X(a), qubits=[a], exponent_neg=0.25, exponent_pos=0.5)
+    cirq.testing.assert_implements_consistent_protocols(p)
 
 
 def test_pass_operations_over():
@@ -167,7 +197,7 @@ def test_inverse():
 
 
 def test_can_merge_with():
-    (q0,) = _make_qubits(1)
+    q0, q1 = _make_qubits(2)
 
     op1 = cirq.PauliStringPhasor(cirq.PauliString({}), exponent_neg=0.25)
     op2 = cirq.PauliStringPhasor(cirq.PauliString({}), exponent_neg=0.75)
@@ -181,6 +211,12 @@ def test_can_merge_with():
     op2 = cirq.PauliStringPhasor(cirq.PauliString({q0: cirq.Y}, -1), exponent_neg=0.75)
     assert not op1.can_merge_with(op2)
 
+    op1 = cirq.PauliStringPhasor(
+        cirq.PauliString({q0: cirq.X}, +1), qubits=[q0, q1], exponent_neg=0.25
+    )
+    op2 = cirq.PauliStringPhasor(cirq.PauliString({q0: cirq.X}, -1), exponent_neg=0.75)
+    assert not op1.can_merge_with(op2)
+
 
 def test_merge_with():
     (q0,) = _make_qubits(1)
@@ -363,16 +399,17 @@ def test_text_diagram():
             cirq.PauliString({q0: cirq.Z, q1: cirq.Y, q2: cirq.X}, -1),
             exponent_neg=sympy.Symbol('b'),
         ),
+        cirq.PauliStringPhasor(cirq.PauliString({q0: cirq.Z}), qubits=[q0, q1], exponent_neg=0.5),
     )
 
     cirq.testing.assert_has_diagram(
         circuit,
         """
-q0: ───[Z]───[Y]^0.25───[Z]───[Z]────────[Z]─────[Z]────────
+q0: ───[Z]───[Y]^0.25───[Z]───[Z]────────[Z]─────[Z]────────[Z]───────
+                        │     │          │       │          │
+q1: ────────────────────[Z]───[Y]────────[Y]─────[Y]────────[I]^0.5───
                         │     │          │       │
-q1: ────────────────────[Z]───[Y]────────[Y]─────[Y]────────
-                        │     │          │       │
-q2: ────────────────────[Z]───[X]^-0.5───[X]^a───[X]^(-b)───
+q2: ────────────────────[Z]───[X]^-0.5───[X]^a───[X]^(-b)─────────────
 """,
     )
 
@@ -408,6 +445,9 @@ def test_str():
     assert str(np.exp(-0.25j * np.pi * cirq.X(q0) * cirq.Y(q1))) == 'exp(-iπ0.25*X(q0)*Y(q1))'
     assert str(np.exp(0.5j * np.pi * cirq.PauliString())) == 'exp(iπ0.5*I)'
 
+    ps = cirq.PauliStringPhasor(cirq.PauliString({q0: cirq.X}, +1), qubits=[q0, q1]) ** 0.5
+    assert str(ps) == '(X(q0))**0.5'
+
 
 def test_gate_init():
     a = cirq.LineQubit(0)

cirq-core/cirq/ops/pauli_string_test.py

@@ -310,6 +310,18 @@ def test_repr():
     cirq.testing.assert_equivalent_repr(cirq.PauliString())
 
 
+def test_repr_preserves_qubit_order():
+    q0, q1, q2 = _make_qubits(3)
+    pauli_string = cirq.PauliString({q2: cirq.X, q1: cirq.Y, q0: cirq.Z})
+    assert eval(repr(pauli_string)).qubits == pauli_string.qubits
+
+    pauli_string = cirq.PauliString(cirq.X(q2), cirq.Y(q1), cirq.Z(q0))
+    assert eval(repr(pauli_string)).qubits == pauli_string.qubits
+
+    pauli_string = cirq.PauliString(cirq.Z(q0), cirq.Y(q1), cirq.X(q2))
+    assert eval(repr(pauli_string)).qubits == pauli_string.qubits
+
+
 def test_repr_coefficient_of_one():
     pauli_string = cirq.Z(cirq.LineQubit(0)) * 1
     assert type(pauli_string) == type(eval(repr(pauli_string)))

cirq-core/cirq/protocols/json_test_data/PauliStringPhasor.json

@@ -34,5 +34,15 @@
       "real": 1.0,
       "imag": 0.0
     }
-  }
+  },
+  "qubits": [
+    {
+      "cirq_type": "LineQubit",
+      "x": 0
+    },
+    {
+      "cirq_type": "LineQubit",
+      "x": 1
+    }
+  ]
 }