Allow simulator quantum state representations to be used in qis

cirq-core/cirq/contrib/quimb/mps_simulator.py

@@ -231,7 +231,7 @@ def __init__(
             estimated_gate_error_list: The error estimations.
             simulation_options: Numerical options for the simulation.
         """
-        self._qid_shape = qid_shape
+        super().__init__(qid_shape)
         self._grouping = grouping
         self._M = M
         self._format_i = format_i

cirq-core/cirq/qis/clifford_tableau.py

@@ -14,9 +14,11 @@
 
 import abc
 from typing import Any, Dict, List, Sequence, Tuple, TYPE_CHECKING, TypeVar
+
 import numpy as np
 
 from cirq import protocols, value
+from cirq.qis import states
 from cirq.value import big_endian_int_to_digits, linear_dict
 
 if TYPE_CHECKING:
@@ -25,7 +27,7 @@
 TSelf = TypeVar('TSelf', bound='QuantumStateRepresentation')
 
 
-class QuantumStateRepresentation(metaclass=abc.ABCMeta):
+class QuantumStateRepresentation(states.HasQuantumState, metaclass=abc.ABCMeta):
     @abc.abstractmethod
     def copy(self: TSelf, deep_copy_buffers: bool = True) -> TSelf:
         """Creates a copy of the object.
@@ -97,11 +99,6 @@ def supports_factor(self) -> bool:
         """Subclasses that allow factorization should override this."""
         return False
 
-    @property
-    def can_represent_mixed_states(self) -> bool:
-        """Subclasses that can represent mixed states should override this."""
-        return False
-
 
 class StabilizerState(QuantumStateRepresentation, metaclass=abc.ABCMeta):
     """Interface for quantum stabilizer state representations.
@@ -214,13 +211,14 @@ class CliffordTableau(StabilizerState):
     an eigenoperator of the state vector with eigenvalue one: P|psi> = |psi>.
     """
 
-    def __init__(self, num_qubits, initial_state: int = 0):
+    def __init__(self, num_qubits: int, initial_state: int = 0):
         """Initializes CliffordTableau
         Args:
             num_qubits: The number of qubits in the system.
             initial_state: The computational basis representation of the
                 state as a big endian int.
         """
+        super().__init__((2,) * num_qubits)
         self.n = num_qubits
 
         # The last row (`2n+1`-th row) is the scratch row used in _measurement
@@ -663,3 +661,7 @@ def measure(
         self, axes: Sequence[int], seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None
     ) -> List[int]:
         return [self._measure(axis, seed) for axis in axes]
+
+    def state_vector(self):
+        # This function should exist.
+        raise NotImplementedError()

cirq-core/cirq/qis/measures.py

@@ -20,7 +20,7 @@
 
 from cirq import protocols, value, _import
 from cirq.qis.states import (
-    QuantumState,
+    HasQuantumState,
     infer_qid_shape,
     quantum_state,
     validate_density_matrix,
@@ -274,8 +274,8 @@ def von_neumann_entropy(
     Raises:
         ValueError: Invalid quantum state.
     """
-    if isinstance(state, QuantumState) and state._is_density_matrix():
-        state = state.data
+    if isinstance(state, HasQuantumState) and state.can_represent_mixed_states:
+        state = state.density_matrix()
     if isinstance(state, np.ndarray) and state.ndim == 2 and state.shape[0] == state.shape[1]:
         if validate:
             if qid_shape is None:

cirq-core/cirq/qis/measures_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for measures."""
+from typing import Sequence
+
 import numpy as np
 import pytest
 
@@ -164,6 +166,77 @@ def test_fidelity_product_states():
         )
 
 
+def test_fidelity_different_simulator_outputs_compatible():
+    from cirq.sim.act_on_state_vector_args import _BufferedStateVector
+    from cirq.sim.act_on_density_matrix_args import _BufferedDensityMatrix
+
+    gate_domain = {
+        cirq.CNOT: 2,
+        cirq.CZ: 2,
+        cirq.H: 1,
+        cirq.S: 1,
+        cirq.SWAP: 2,
+        cirq.X: 1,
+        cirq.Y: 1,
+        cirq.Z: 1,
+    }
+    circuits = [
+        cirq.testing.random_circuit(qubits=5, n_moments=30, op_density=1, gate_domain=gate_domain)
+        for _ in range(2)
+    ]
+
+    def final_state(sim: cirq.SimulatorBase):
+        return [
+            sim.simulate(c)._final_step_result_typed._sim_state.create_merged_state()._state
+            for c in circuits
+        ]
+
+    density_matrices: Sequence[_BufferedDensityMatrix] = final_state(cirq.DensityMatrixSimulator())
+    state_vectors: Sequence[_BufferedStateVector] = final_state(cirq.Simulator())
+    ch_forms: Sequence[cirq.StabilizerStateChForm] = final_state(cirq.CliffordSimulator())
+    density_matrix = density_matrices[0]
+    state_vector = state_vectors[0]
+    ch_form = ch_forms[0]
+
+    assert isinstance(density_matrix, _BufferedDensityMatrix)
+    assert isinstance(state_vector, _BufferedStateVector)
+    assert isinstance(ch_form, cirq.StabilizerStateChForm)
+
+    def assert_fidelity_one(x, y):
+        np.testing.assert_allclose(cirq.fidelity(x, y), 1, atol=1e-2)
+
+    assert_fidelity_one(density_matrix, density_matrix)
+    assert_fidelity_one(density_matrix, state_vector)
+    assert_fidelity_one(density_matrix, ch_form)
+    assert_fidelity_one(state_vector, density_matrix)
+    assert_fidelity_one(state_vector, state_vector)
+    assert_fidelity_one(state_vector, ch_form)
+    assert_fidelity_one(ch_form, density_matrix)
+    assert_fidelity_one(ch_form, state_vector)
+    assert_fidelity_one(ch_form, ch_form)
+
+    def assert_fidelity_equivalent(xs, ys):
+        atol = 1e-2
+        baseline = cirq.fidelity(xs[0], xs[1])
+        np.testing.assert_allclose(baseline, cirq.fidelity(ys[0], ys[1]), atol=atol)
+        np.testing.assert_allclose(baseline, cirq.fidelity(xs[0], ys[1]), atol=atol)
+        np.testing.assert_allclose(baseline, cirq.fidelity(ys[0], xs[1]), atol=atol)
+        np.testing.assert_allclose(baseline, cirq.fidelity(xs[1], xs[0]), atol=atol)
+        np.testing.assert_allclose(baseline, cirq.fidelity(ys[1], ys[0]), atol=atol)
+        np.testing.assert_allclose(baseline, cirq.fidelity(xs[1], ys[0]), atol=atol)
+        np.testing.assert_allclose(baseline, cirq.fidelity(ys[1], xs[0]), atol=atol)
+
+    assert_fidelity_equivalent(density_matrices, density_matrices)
+    assert_fidelity_equivalent(density_matrices, state_vectors)
+    assert_fidelity_equivalent(density_matrices, ch_forms)
+    assert_fidelity_equivalent(state_vectors, density_matrices)
+    assert_fidelity_equivalent(state_vectors, state_vectors)
+    assert_fidelity_equivalent(state_vectors, ch_forms)
+    assert_fidelity_equivalent(ch_forms, density_matrices)
+    assert_fidelity_equivalent(ch_forms, state_vectors)
+    assert_fidelity_equivalent(ch_forms, ch_forms)
+
+
 def test_fidelity_fail_inference():
     state_vector = cirq.one_hot(shape=(4,), dtype=np.complex128)
     state_tensor = np.reshape(state_vector, (2, 2))
@@ -224,6 +297,24 @@ def test_von_neumann_entropy():
     )
 
 
+@pytest.mark.parametrize(
+    'sim', (cirq.Simulator(), cirq.DensityMatrixSimulator(), cirq.CliffordSimulator())
+)
+def test_von_neumann_entropy_different_simulator_outputs(sim):
+    q0, q1 = cirq.LineQubit.range(2)
+    bell_circuit = cirq.Circuit(cirq.H(q0), cirq.CX(q0, q1))
+    bitflip_circuit = cirq.Circuit(cirq.H(q0), cirq.bit_flip(0.5)(q1))
+
+    def final_state(circuit):
+        return sim.simulate(circuit)._final_step_result._sim_state.create_merged_state()._state
+
+    assert cirq.von_neumann_entropy(final_state(bell_circuit), qid_shape=(2, 2)) == 0
+    np.testing.assert_allclose(
+        cirq.von_neumann_entropy(final_state(bitflip_circuit), qid_shape=(2, 2)),
+        1 if isinstance(sim, cirq.DensityMatrixSimulator) else 0,
+    )
+
+
 @pytest.mark.parametrize(
     'gate, expected_entanglement_fidelity',
     (

cirq-core/cirq/qis/states.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 """Classes and methods for quantum states."""
 
-from typing import Any, cast, Iterable, List, Optional, Sequence, Set, TYPE_CHECKING, Tuple, Union
+import abc
 import itertools
+from typing import Any, cast, Iterable, List, Optional, Sequence, Set, TYPE_CHECKING, Tuple, Union
 
 import numpy as np
 
@@ -46,14 +47,70 @@
     # density matrix
     np.ndarray,
     # quantum state object
-    'cirq.QuantumState',
+    'HasQuantumState',
 ]
 document(QUANTUM_STATE_LIKE, """An object representing a quantum state.""")  # type: ignore
 
 
-class QuantumState:
-    """A quantum state.
+class HasQuantumState(abc.ABC):
+    def __init__(self, qid_shape: Tuple[int, ...] = ()):
+        self._qid_shape = qid_shape
+        self._dim = np.prod(self.qid_shape, dtype=np.int64).item()
+
+    @property
+    def qid_shape(self) -> Tuple[int, ...]:
+        """The qid shape of the quantum state."""
+        return self._qid_shape
+
+    def state_vector(self) -> Optional[np.ndarray]:
+        """Return the state vector of this state.
+
+        A state vector stores the amplitudes of a pure state as a
+        one-dimensional array.
+        If the state is a density matrix, this method returns None.
+        """
+        return None
+
+    def state_tensor(self) -> Optional[np.ndarray]:
+        """Return the state tensor of this state.
+
+        A state tensor stores the amplitudes of a pure state as an array with
+        shape equal to the qid shape of the state.
+        If the state is a density matrix, this method returns None.
+        """
+        state_vector = self.state_vector()
+        if state_vector is None:
+            return None
+        return np.reshape(state_vector, self.qid_shape)
+
+    def density_matrix(self) -> np.ndarray:
+        """Return the density matrix of this state.
+
+        A density matrix stores the entries of a density matrix as a matrix
+        (a two-dimensional array).
+        """
+        state_vector = self.state_vector()
+        return np.outer(state_vector, np.conj(state_vector))
+
+    def state_vector_or_density_matrix(self) -> np.ndarray:
+        """Return the state vector or density matrix of this state.
+
+        If the state is a denity matrix, return the density matrix. Otherwise, return the state
+        vector.
+        """
+        state_vector = self.state_vector()
+        if state_vector is not None:
+            return state_vector
+        return self.density_matrix()
 
+    @property
+    def can_represent_mixed_states(self) -> bool:
+        """Whether this quantum state is a density matrix."""
+        return False
+
+
+class QuantumState(HasQuantumState):
+    """A quantum state.
     Can be a state vector, a state tensor, or a density matrix.
     """
 
@@ -83,9 +140,8 @@ def __init__(
         """
         if qid_shape is None:
             qid_shape = infer_qid_shape(data)
+        super().__init__(qid_shape)
         self._data = data
-        self._qid_shape = qid_shape
-        self._dim = np.prod(self.qid_shape, dtype=np.int64).item()
         if validate:
             self.validate(dtype=dtype, atol=atol)
 
@@ -94,11 +150,6 @@ def data(self) -> np.ndarray:
         """The data underlying the quantum state."""
         return self._data
 
-    @property
-    def qid_shape(self) -> Tuple[int, ...]:
-        """The qid shape of the quantum state."""
-        return self._qid_shape
-
     @property
     def dtype(self) -> np.dtype:
         """The data type of the quantum state."""
@@ -115,38 +166,14 @@ def state_vector(self) -> Optional[np.ndarray]:
             return None
         return np.reshape(self.data, (self._dim,))
 
-    def state_tensor(self) -> Optional[np.ndarray]:
-        """Return the state tensor of this state.
-
-        A state tensor stores the amplitudes of a pure state as an array with
-        shape equal to the qid shape of the state.
-        If the state is a density matrix, this method returns None.
-        """
-        if self._is_density_matrix():
-            return None
-        return np.reshape(self.data, self.qid_shape)
-
     def density_matrix(self) -> np.ndarray:
         """Return the density matrix of this state.
 
         A density matrix stores the entries of a density matrix as a matrix
         (a two-dimensional array).
         """
         if not self._is_density_matrix():
-            state_vector = self.state_vector()
-            assert state_vector is not None, 'only None if _is_density_matrix'
-            return np.outer(state_vector, np.conj(state_vector))
-        return self.data
-
-    def state_vector_or_density_matrix(self) -> np.ndarray:
-        """Return the state vector or density matrix of this state.
-
-        If the state is a denity matrix, return the density matrix. Otherwise, return the state
-        vector.
-        """
-        state_vector = self.state_vector()
-        if state_vector is not None:
-            return state_vector
+            return super().density_matrix()
         return self.data
 
     def _is_density_matrix(self) -> bool:
@@ -184,6 +211,11 @@ def validate(
                 f'compatible with qid shape of {self.qid_shape}.'
             )
 
+    @property
+    def can_represent_mixed_states(self) -> bool:
+        """Whether this quantum state is a density matrix."""
+        return self._is_density_matrix()
+
 
 def quantum_state(
     state: 'cirq.QUANTUM_STATE_LIKE',
@@ -263,6 +295,8 @@ def quantum_state(
             dtype = DEFAULT_COMPLEX_DTYPE
         data = one_hot(index=state, shape=(dim,), dtype=dtype)
     else:
+        if isinstance(state, HasQuantumState):
+            state = state.state_vector_or_density_matrix()
         data = np.array(state, copy=False)
         if qid_shape is None:
             qid_shape = infer_qid_shape(state)
@@ -344,7 +378,7 @@ def _infer_qid_shape_from_dimension(dim: int) -> Tuple[int, ...]:
     # Product state object
     'cirq.ProductState',
     # Quantum state object
-    'cirq.QuantumState',
+    'HasQuantumState',
 ]
 
 
@@ -410,10 +444,8 @@ def infer_qid_shape(*states: 'cirq.QUANTUM_STATE_LIKE') -> Tuple[int, ...]:
 
 def _potential_qid_shapes(state: _NON_INT_STATE_LIKE) -> '_QidShapeSet':
     """Return a set of qid shapes compatible with a given state."""
-    if isinstance(state, QuantumState):
+    if isinstance(state, HasQuantumState):
         return _QidShapeSet(explicit_qid_shapes={state.qid_shape})
-    if isinstance(state, value.ProductState):
-        return _QidShapeSet(explicit_qid_shapes={(2,) * len(state)})
 
     if isinstance(state, Sequence):
         state = np.array(state)