Add qubits as a param of PauliSum.matrix() (quantumlib#4203)

Author: madcpf

cirq-core/cirq/ops/linear_combinations.py

@@ -14,6 +14,7 @@
 from collections import defaultdict
 from typing import (
     AbstractSet,
+    Iterable,
     Mapping,
     Optional,
     Tuple,
@@ -416,18 +417,21 @@ def copy(self) -> 'PauliSum':
         factory = type(self)
         return factory(self._linear_dict.copy())
 
-    def matrix(self) -> np.ndarray:
-        """Reconstructs matrix of self from underlying Pauli operations.
+    def matrix(self, qubits: Optional[Iterable[raw_types.Qid]] = None) -> np.ndarray:
+        """Reconstructs matrix of self from underlying Pauli operations in
+        computational basis of qubits.
 
         Raises:
             TypeError: if any of the gates in self does not provide a unitary.
         """
-        num_qubits = len(self.qubits)
+
+        qubits = self.qubits if qubits is None else tuple(qubits)
+        num_qubits = len(qubits)
         num_dim = 2 ** num_qubits
         result = np.zeros((num_dim, num_dim), dtype=np.complex128)
         for vec, coeff in self._linear_dict.items():
             op = _pauli_string_from_unit(vec)
-            result += coeff * op.matrix(self.qubits)
+            result += coeff * op.matrix(qubits)
         return result
 
     def _has_unitary_(self) -> bool:

cirq-core/cirq/ops/linear_combinations_test.py

@@ -1035,6 +1035,7 @@ def test_non_pauli_sum_has_no_unitary(psum):
         (cirq.Z(q1), (q1,)),
         (cirq.X(q0) + cirq.Y(q0), (q0,)),
         (cirq.X(q0) + cirq.Y(q2), (q0, q2)),
+        (cirq.X(q2) + cirq.Y(q0), (q0, q2)),
         (cirq.X(q0) * cirq.Y(q1) + cirq.Y(q1) * cirq.Z(q3), (q0, q1, q3)),
     ),
 )
@@ -1189,6 +1190,35 @@ def test_pauli_sum_formatting():
     assert str(empty) == "0.000"
 
 
+def test_pauli_sum_matrix():
+    q = cirq.LineQubit.range(3)
+    paulisum = cirq.X(q[0]) * cirq.X(q[1]) + cirq.Z(q[0])
+    H1 = np.array(
+        [[1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0], [0.0, 1.0, -1.0, 0.0], [1.0, 0.0, 0.0, -1.0]]
+    )
+    assert np.allclose(H1, paulisum.matrix())
+    assert np.allclose(H1, paulisum.matrix([q[0], q[1]]))
+    # Expects a different matrix when change qubits order.
+    H2 = np.array(
+        [[1.0, 0.0, 0.0, 1.0], [0.0, -1.0, 1.0, 0.0], [0.0, 1.0, 1.0, 0.0], [1.0, 0.0, 0.0, -1.0]]
+    )
+    assert np.allclose(H2, paulisum.matrix([q[1], q[0]]))
+    # Expects matrix with a different size when add a new qubit.
+    H3 = np.array(
+        [
+            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
+            [0.0, -1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+            [0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 1.0, 0.0],
+            [0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, -1.0],
+        ]
+    )
+    assert np.allclose(H3, paulisum.matrix([q[1], q[2], q[0]]))
+
+
 def test_pauli_sum_repr():
     q = cirq.LineQubit.range(2)
     pstr1 = cirq.X(q[0]) * cirq.X(q[1])