Repository: quantumlib/Cirq

Issue:
PauliSum should support multiplication with PauliSum and raising to integer power



Discussion:
You beat me to this! An example of how we want this to behave:

```python

(cirq.X(bit_0) + cirq.Y(bit_0)) *  (cirq.X(bit_1) + cirq.Y(bit_1)) = (cirq.X(bit_0) * cirq.X(bit_1)) + 
(cirq.X(bit_0) * cirq.Y(bit_1)) + (cirq.Y(bit_0) * cirq.X(bit_1)) + (cirq.Y(bit_0) * cirq.y(bit)_1)


```

The multiplication symbol represents the tensor product here
In your example, `*` should denote plain (matrix) multiplication, not the tensor product. The tensor product is implicit in the fact that `cirq.X(bit_0)` represents `X ⊗ I`, `cirq.X(bit_1)` represents `I ⊗ X`, etc.

Resolving patch:
diff --git a/cirq/ops/linear_combinations.py b/cirq/ops/linear_combinations.py
--- a/cirq/ops/linear_combinations.py
+++ b/cirq/ops/linear_combinations.py
@@ -14,6 +14,7 @@
 from collections import defaultdict
 from typing import (Mapping, Optional, Tuple, Union, List, FrozenSet,
                     DefaultDict)
+import numbers
 
 import numpy as np
 
@@ -451,7 +452,7 @@ def __len__(self) -> int:
         return len(self._linear_dict)
 
     def __iadd__(self, other):
-        if isinstance(other, (float, int, complex)):
+        if isinstance(other, numbers.Complex):
             other = PauliSum.from_pauli_strings(
                 [PauliString(coefficient=other)])
         elif isinstance(other, PauliString):
@@ -464,7 +465,7 @@ def __iadd__(self, other):
         return self
 
     def __add__(self, other):
-        if not isinstance(other, (float, int, complex, PauliString, PauliSum)):
+        if not isinstance(other, (numbers.Complex, PauliString, PauliSum)):
             return NotImplemented
         result = self.copy()
         result += other
@@ -477,7 +478,7 @@ def __rsub__(self, other):
         return -self.__sub__(other)
 
     def __isub__(self, other):
-        if isinstance(other, (float, int, complex)):
+        if isinstance(other, numbers.Complex):
             other = PauliSum.from_pauli_strings(
                 [PauliString(coefficient=other)])
         if isinstance(other, PauliString):
@@ -490,7 +491,7 @@ def __isub__(self, other):
         return self
 
     def __sub__(self, other):
-        if not isinstance(other, (float, int, complex, PauliString, PauliSum)):
+        if not isinstance(other, (numbers.Complex, PauliString, PauliSum)):
             return NotImplemented
         result = self.copy()
         result -= other
@@ -500,17 +501,49 @@ def __neg__(self):
         factory = type(self)
         return factory(-self._linear_dict)
 
-    def __imul__(self, a: value.Scalar):
-        self._linear_dict *= a
+    def __imul__(self, other: PauliSumLike):
+        if not isinstance(other, (numbers.Complex, PauliString, PauliSum)):
+            return NotImplemented
+        if isinstance(other, numbers.Complex):
+            self._linear_dict *= other
+        elif isinstance(other, PauliString):
+            temp = PauliSum.from_pauli_strings([term * other for term in self])
+            self._linear_dict = temp._linear_dict
+        elif isinstance(other, PauliSum):
+            temp = PauliSum.from_pauli_strings(
+                [term * other_term for term in self for other_term in other])
+            self._linear_dict = temp._linear_dict
+
         return self
 
-    def __mul__(self, a: value.Scalar):
+    def __mul__(self, other: PauliSumLike):
+        if not isinstance(other, (numbers.Complex, PauliString, PauliSum)):
+            return NotImplemented
         result = self.copy()
-        result *= a
+        result *= other
         return result
 
-    def __rmul__(self, a: value.Scalar):
-        return self.__mul__(a)
+    def __rmul__(self, other: PauliSumLike):
+        if isinstance(other, numbers.Complex):
+            result = self.copy()
+            result *= other
+            return result
+        elif isinstance(other, PauliString):
+            result = self.copy()
+            return PauliSum.from_pauli_strings([other]) * result
+        return NotImplemented
+
+    def __pow__(self, exponent: int):
+        if not isinstance(exponent, numbers.Integral):
+            return NotImplemented
+        if exponent == 0:
+            return PauliSum(value.LinearDict({frozenset(): 1 + 0j}))
+        if exponent > 0:
+            base = self.copy()
+            for _ in range(exponent - 1):
+                base *= base
+            return base
+        return NotImplemented
 
     def __truediv__(self, a: value.Scalar):
         return self.__mul__(1 / a)


Test patch:
diff --git a/cirq/ops/linear_combinations_test.py b/cirq/ops/linear_combinations_test.py
--- a/cirq/ops/linear_combinations_test.py
+++ b/cirq/ops/linear_combinations_test.py
@@ -954,7 +954,6 @@ def test_bad_arithmetic():
 
     with pytest.raises(TypeError):
         _ = psum - 'hi mom'
-
     with pytest.raises(TypeError):
         psum *= [1, 2, 3]
 
@@ -967,6 +966,97 @@ def test_bad_arithmetic():
     with pytest.raises(TypeError):
         _ = psum / [1, 2, 3]
 
+    with pytest.raises(TypeError):
+        _ = psum**1.2
+
+    with pytest.raises(TypeError):
+        _ = psum**-2
+
+    with pytest.raises(TypeError):
+        _ = psum**"string"
+
+
+def test_paulisum_mul_paulistring():
+    q0, q1 = cirq.LineQubit.range(2)
+
+    psum1 = cirq.X(q0) + 2 * cirq.Y(q0) + 3 * cirq.Z(q0)
+    x0 = cirq.PauliString(cirq.X(q0))
+    y1 = cirq.PauliString(cirq.Y(q1))
+    assert x0 * psum1 == cirq.PauliString(cirq.I(q0)) \
+                         + 2j * cirq.PauliString(cirq.Z(q0)) \
+                         - 3j * cirq.PauliString(cirq.Y(q0))
+    assert y1 * psum1 == cirq.PauliString(cirq.X(q0) * cirq.Y(q1)) \
+                         + 2 * cirq.PauliString(cirq.Y(q0) * cirq.Y(q1)) \
+                         + 3 * cirq.PauliString(cirq.Z(q0) * cirq.Y(q1))
+    assert cirq.PauliString(cirq.I(q0)) * psum1 == psum1
+    assert psum1 * x0 == cirq.PauliString(cirq.I(q0)) \
+                         - 2j * cirq.PauliString(cirq.Z(q0)) \
+                         + 3j * cirq.PauliString(cirq.Y(q0))
+    assert psum1 * y1 == y1 * psum1
+
+    psum1 *= cirq.Z(q0)
+    assert psum1 == -1j * cirq.Y(q0) + 2j * cirq.X(q0) + 3
+
+
+def test_paulisum_mul_paulisum():
+    q0, q1, q2 = cirq.LineQubit.range(3)
+
+    psum1 = cirq.X(q0) + 2 * cirq.Y(q0) * cirq.Y(q1)
+    psum2 = cirq.X(q0) * cirq.Y(q1) + 3 * cirq.Z(q2)
+    assert psum1 * psum2 == cirq.Y(q1) + 3 * cirq.X(q0) * cirq.Z(q2) \
+                            - 2j * cirq.Z(q0) \
+                            + 6 * cirq.Y(q0) * cirq.Y(q1) * cirq.Z(q2)
+    assert psum2 * psum1 == cirq.Y(q1) + 3 * cirq.X(q0) * cirq.Z(q2) \
+                            + 2j * cirq.Z(q0) \
+                            + 6 * cirq.Y(q0) * cirq.Y(q1) * cirq.Z(q2)
+    psum3 = cirq.X(q1) + cirq.X(q2)
+    psum1 *= psum3
+    assert psum1 == cirq.X(q0) * cirq.X(q1) - 2j * cirq.Y(q0) * cirq.Z(q1) \
+        + cirq.X(q0) * cirq.X(q2) + 2 * cirq.Y(q0) * cirq.Y(q1) * cirq.X(q2)
+
+    psum4 = cirq.X(q0) + cirq.Y(q0) + cirq.Z(q1)
+    psum5 = cirq.Z(q0) + cirq.Y(q0) + cirq.PauliString(coefficient=1.2)
+    assert psum4 * psum5 == -1j * cirq.Y(q0) + 1j * (cirq.X(q0) + cirq.Z(q0)) \
+                            + (cirq.Z(q0) + cirq.Y(q0)) * cirq.Z(q1) + 1 \
+                            + 1.2 * psum4
+    assert psum5 * psum4 == 1j * cirq.Y(q0) + -1j * (cirq.X(q0) + cirq.Z(q0)) \
+                            + (cirq.Z(q0) + cirq.Y(q0)) * cirq.Z(q1) + 1 \
+                            + 1.2 * psum4
+
+
+def test_pauli_sum_pow():
+    identity = cirq.PauliSum.from_pauli_strings(
+        [cirq.PauliString(coefficient=1)])
+    psum1 = cirq.X(q0) + cirq.Y(q0)
+    assert psum1**2 == psum1 * psum1
+    assert psum1**2 == 2 * identity
+
+    psum2 = cirq.X(q0) + cirq.Y(q1)
+    assert psum2**2 == cirq.PauliString(
+        cirq.I(q0)) + 2 * cirq.X(q0) * cirq.Y(q1) + cirq.PauliString(cirq.I(q1))
+
+    psum3 = cirq.X(q0) * cirq.Z(q1) + 1.3 * cirq.Z(q0)
+    sqd = cirq.PauliSum.from_pauli_strings(
+        [2.69 * cirq.PauliString(cirq.I(q0))])
+    assert cirq.approx_eq(psum3**2, sqd, atol=1e-8)
+
+    psum4 = cirq.X(q0) * cirq.Z(q1) + 1.3 * cirq.Z(q1)
+    sqd2 = cirq.PauliSum.from_pauli_strings(
+        [2.69 * cirq.PauliString(cirq.I(q0)), 2.6 * cirq.X(q0)])
+    assert cirq.approx_eq(psum4**2, sqd2, atol=1e-8)
+
+    for psum in [psum1, psum2, psum3, psum4]:
+        assert cirq.approx_eq(psum**0, identity)
+
+
+def test_imul_aliasing():
+    q0, q1, q2 = cirq.LineQubit.range(3)
+    psum1 = cirq.X(q0) + cirq.Y(q1)
+    psum2 = psum1
+    psum2 *= (cirq.X(q0) * cirq.Y(q2))
+    assert psum1 is psum2
+    assert psum1 == psum2
+
 
 def test_expectation_from_wavefunction_invalid_input():
     q0, q1, q2, q3 = cirq.LineQubit.range(4)