Add TensorProductState, a state-vector-like object (#3171)

This is factored out of #2781 with @Strilanc 's initial comments addressed, per comments on that PR.

`cirq.STATE_VECTOR_LIKE` accepts basis-state indices, sequences of per-qudit basis indices, or arrays of amplitudes. This PR extends this union type to accept `TensorProductState` -- an object representing a state vector which can be decomposed into a tensor product of single-qubit states. The structure shares a lot of similarity with `PauliString`. Whereas PauliString is a dictionary from qubits to (single-qubit) `Pauli` objects; here it is a mapping from qubits to single-qubit named states. Here we include `|0>, |1>, |+>, |->, |i>, |-i>`. As an example:

```python
# start from the |+++> state
q = cirq.LineQubit.range(3)
initial_state = cirq.KET_PLUS(q[0]) * cirq.KET_PLUS(q[1]) * cirq.KET_PLUS(q[2])
```

Since most things already use `STATE_VECTOR_LIKE` and its associated helper function `to_valid_state_vector`, only minor modifications to this function was necessary to enable specification of a tensor product state wherever initial states are sold.

Author: mpharrigan

cirq/__init__.py

@@ -413,6 +413,14 @@
     TParamVal,
     validate_probability,
     value_equality,
+    KET_PLUS,
+    KET_MINUS,
+    KET_IMAG,
+    KET_MINUS_IMAG,
+    KET_ZERO,
+    KET_ONE,
+    PAULI_STATES,
+    ProductState,
 )
 
 # pylint: disable=redefined-builtin

cirq/circuits/circuit_test.py

@@ -2411,6 +2411,28 @@ def test_apply_unitary_effect_to_state():
                                                     np.array([0, 0, 1, 0]),
                                                     atol=1e-8)
 
+    # Product state
+    cirq.testing.assert_allclose_up_to_global_phase(cirq.Circuit(cirq.CNOT(
+        a, b)).final_state_vector(initial_state=cirq.KET_ZERO(a) *
+                                  cirq.KET_ZERO(b)),
+                                                    np.array([1, 0, 0, 0]),
+                                                    atol=1e-8)
+    cirq.testing.assert_allclose_up_to_global_phase(cirq.Circuit(cirq.CNOT(
+        a, b)).final_state_vector(initial_state=cirq.KET_ZERO(a) *
+                                  cirq.KET_ONE(b)),
+                                                    np.array([0, 1, 0, 0]),
+                                                    atol=1e-8)
+    cirq.testing.assert_allclose_up_to_global_phase(cirq.Circuit(cirq.CNOT(
+        a, b)).final_state_vector(initial_state=cirq.KET_ONE(a) *
+                                  cirq.KET_ZERO(b)),
+                                                    np.array([0, 0, 0, 1]),
+                                                    atol=1e-8)
+    cirq.testing.assert_allclose_up_to_global_phase(cirq.Circuit(cirq.CNOT(
+        a,
+        b)).final_state_vector(initial_state=cirq.KET_ONE(a) * cirq.KET_ONE(b)),
+                                                    np.array([0, 0, 1, 0]),
+                                                    atol=1e-8)
+
     # Measurements.
     cirq.testing.assert_allclose_up_to_global_phase(cirq.Circuit(
         cirq.measure(a)).final_state_vector(),

cirq/ops/pauli_gates.py

@@ -12,9 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import abc
-from typing import Any, cast, Tuple, TYPE_CHECKING, Union
+from typing import Any, cast, Tuple, TYPE_CHECKING, Union, Dict
 
-from cirq import value
 from cirq._doc import document
 from cirq.ops import common_gates, raw_types, identity
 from cirq.type_workarounds import NotImplementedType
@@ -23,6 +22,8 @@
 if TYPE_CHECKING:
     import cirq
     from cirq.ops.pauli_string import SingleQubitPauliStringGateOperation
+    from cirq.value.product_state import (_XEigenState, _YEigenState,
+                                          _ZEigenState)  # coverage: ignore
 
 
 class Pauli(raw_types.Gate, metaclass=abc.ABCMeta):
@@ -98,19 +99,18 @@ def _canonical_exponent(self):
         """Overrides EigenGate._canonical_exponent in subclasses."""
         return 1
 
-
 class _PauliX(Pauli, common_gates.XPowGate):
 
     def __init__(self):
         Pauli.__init__(self, index=0, name='X')
         common_gates.XPowGate.__init__(self, exponent=1.0)
 
     def __pow__(self: '_PauliX',
-                exponent: value.TParamVal) -> common_gates.XPowGate:
+                exponent: 'cirq.TParamVal') -> common_gates.XPowGate:
         return common_gates.XPowGate(exponent=exponent)
 
     def _with_exponent(self: '_PauliX',
-                       exponent: value.TParamVal) -> common_gates.XPowGate:
+                       exponent: 'cirq.TParamVal') -> common_gates.XPowGate:
         return self.__pow__(exponent)
 
     @classmethod
@@ -119,6 +119,14 @@ def _from_json_dict_(cls, exponent, global_shift, **kwargs):
         assert exponent == 1
         return Pauli._XYZ[0]
 
+    @property
+    def basis(self: '_PauliX') -> Dict[int, '_XEigenState']:
+        from cirq.value.product_state import _XEigenState
+        return {
+            +1: _XEigenState(+1),
+            -1: _XEigenState(-1),
+        }
+
 
 class _PauliY(Pauli, common_gates.YPowGate):
 
@@ -127,11 +135,11 @@ def __init__(self):
         common_gates.YPowGate.__init__(self, exponent=1.0)
 
     def __pow__(self: '_PauliY',
-                exponent: value.TParamVal) -> common_gates.YPowGate:
+                exponent: 'cirq.TParamVal') -> common_gates.YPowGate:
         return common_gates.YPowGate(exponent=exponent)
 
     def _with_exponent(self: '_PauliY',
-                       exponent: value.TParamVal) -> common_gates.YPowGate:
+                       exponent: 'cirq.TParamVal') -> common_gates.YPowGate:
         return self.__pow__(exponent)
 
     @classmethod
@@ -140,6 +148,14 @@ def _from_json_dict_(cls, exponent, global_shift, **kwargs):
         assert exponent == 1
         return Pauli._XYZ[1]
 
+    @property
+    def basis(self: '_PauliY') -> Dict[int, '_YEigenState']:
+        from cirq.value.product_state import _YEigenState
+        return {
+            +1: _YEigenState(+1),
+            -1: _YEigenState(-1),
+        }
+
 
 class _PauliZ(Pauli, common_gates.ZPowGate):
 
@@ -148,11 +164,11 @@ def __init__(self):
         common_gates.ZPowGate.__init__(self, exponent=1.0)
 
     def __pow__(self: '_PauliZ',
-                exponent: value.TParamVal) -> common_gates.ZPowGate:
+                exponent: 'cirq.TParamVal') -> common_gates.ZPowGate:
         return common_gates.ZPowGate(exponent=exponent)
 
     def _with_exponent(self: '_PauliZ',
-                       exponent: value.TParamVal) -> common_gates.ZPowGate:
+                       exponent: 'cirq.TParamVal') -> common_gates.ZPowGate:
         return self.__pow__(exponent)
 
     @classmethod
@@ -161,6 +177,14 @@ def _from_json_dict_(cls, exponent, global_shift, **kwargs):
         assert exponent == 1
         return Pauli._XYZ[2]
 
+    @property
+    def basis(self: '_PauliZ') -> Dict[int, '_ZEigenState']:
+        from cirq.value.product_state import _ZEigenState
+        return {
+            +1: _ZEigenState(+1),
+            -1: _ZEigenState(-1),
+        }
+
 
 X = _PauliX()
 document(

cirq/ops/pauli_string_test.py

@@ -592,6 +592,31 @@ def test_to_z_basis_ops():
                     z_basis_state, expected_state, rtol=1e-7, atol=1e-7)
 
 
+def test_to_z_basis_ops_product_state():
+    q0, q1, q2, q3, q4, q5 = _make_qubits(6)
+    pauli_string = cirq.PauliString({
+        q0: cirq.X,
+        q1: cirq.X,
+        q2: cirq.Y,
+        q3: cirq.Y,
+        q4: cirq.Z,
+        q5: cirq.Z
+    })
+    circuit = cirq.Circuit(pauli_string.to_z_basis_ops())
+
+    initial_state = cirq.KET_PLUS(q0) * cirq.KET_MINUS(q1) * cirq.KET_IMAG(
+        q2) * cirq.KET_MINUS_IMAG(q3) * cirq.KET_ZERO(q4) * cirq.KET_ONE(q5)
+    z_basis_state = circuit.final_state_vector(initial_state)
+
+    expected_state = np.zeros(2**6)
+    expected_state[0b010101] = 1
+
+    cirq.testing.assert_allclose_up_to_global_phase(z_basis_state,
+                                                    expected_state,
+                                                    rtol=1e-7,
+                                                    atol=1e-7)
+
+
 def _assert_pass_over(ops: List[cirq.Operation],
                       before: cirq.PauliString,
                       after: cirq.PauliString):

cirq/protocols/json_serialization.py

@@ -80,6 +80,7 @@ def two_qubit_matrix_gate(matrix):
                 'AsymmetricDepolarizingChannel':
                 cirq.AsymmetricDepolarizingChannel,
                 'BitFlipChannel': cirq.BitFlipChannel,
+                'ProductState': cirq.ProductState,
                 'CCNotPowGate': cirq.CCNotPowGate,
                 'CCXPowGate': cirq.CCXPowGate,
                 'CCZPowGate': cirq.CCZPowGate,
@@ -120,6 +121,12 @@ def two_qubit_matrix_gate(matrix):
                 'MatrixGate': cirq.MatrixGate,
                 'MeasurementGate': cirq.MeasurementGate,
                 'Moment': cirq.Moment,
+                '_XEigenState':
+                cirq.value.product_state._XEigenState,  # type: ignore
+                '_YEigenState':
+                cirq.value.product_state._YEigenState,  # type: ignore
+                '_ZEigenState':
+                cirq.value.product_state._ZEigenState,  # type: ignore
                 '_NamedConstantXmonDevice': _NamedConstantXmonDevice,
                 '_NoNoiseModel': _NoNoiseModel,
                 'NamedQubit': cirq.NamedQubit,

cirq/protocols/json_serialization_test.py

@@ -138,6 +138,7 @@ def test_fail_to_resolve():
     # global objects
     'CONTROL_TAG',
     'PAULI_BASIS',
+    'PAULI_STATES',
 
     # abstract, but not inspect.isabstract():
     'Device',

cirq/protocols/json_test_data/KET_IMAG.json

@@ -0,0 +1,4 @@
+{
+  "cirq_type": "_YEigenState",
+  "eigenvalue": 1
+}

cirq/protocols/json_test_data/KET_IMAG.repr

@@ -0,0 +1 @@
+cirq.KET_IMAG

cirq/protocols/json_test_data/KET_MINUS.json

@@ -0,0 +1,4 @@
+{
+  "cirq_type": "_XEigenState",
+  "eigenvalue": -1
+}

cirq/protocols/json_test_data/KET_MINUS.repr

@@ -0,0 +1 @@
+cirq.KET_MINUS