Defensively copy density matrices in simulator (#3109)

Author: dabacon

cirq/qis/states.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 """Utility methods for creating vectors and matrices."""
 
-from typing import (Any, Iterable, List, Optional, Sequence, Union,
-                    TYPE_CHECKING, Tuple, Type, cast)
+from typing import (Any, cast, Iterable, Optional, Sequence, TYPE_CHECKING,
+                    Tuple, Type, Union)
 
 import itertools
 
@@ -486,7 +486,8 @@ def to_valid_density_matrix(
 
     Returns:
         A numpy matrix corresponding to the density matrix on the given number
-        of qubits.
+        of qubits. Note that this matrix may share memory with the input
+        `density_matrix_rep`.
 
     Raises:
         ValueError if the density_matrix_rep is not valid.

cirq/sim/density_matrix_simulator.py

@@ -257,6 +257,8 @@ def _base_iterator(self,
                                                      len(qid_shape),
                                                      qid_shape=qid_shape,
                                                      dtype=self._dtype)
+        if np.may_share_memory(initial_matrix, initial_state):
+            initial_matrix = initial_matrix.copy()
         measured = collections.defaultdict(
             bool)  # type: Dict[Tuple[cirq.Qid, ...], bool]
         if len(circuit) == 0:
@@ -407,7 +409,7 @@ def set_density_matrix(self, density_matrix_repr: Union[int, np.ndarray]):
         density_matrix = np.reshape(density_matrix, sim_state_matrix.shape)
         np.copyto(dst=sim_state_matrix, src=density_matrix)
 
-    def density_matrix(self):
+    def density_matrix(self, copy=True):
         """Returns the density matrix at this step in the simulation.
 
         The density matrix that is stored in this result is returned in the
@@ -435,9 +437,16 @@ def density_matrix(self):
                 |  6  |   1    |   1    |   0    |
                 |  7  |   1    |   1    |   1    |
 
+        Args:
+            copy: If True, then the returned state is a copy of the density
+                matrix. If False, then the density matrix is not copied,
+                potentially saving memory. If one only needs to read derived
+                parameters from the density matrix and store then using False
+                can speed up simulation by eliminating a memory copy.
         """
         size = np.prod(self._qid_shape, dtype=int)
-        return np.reshape(self._density_matrix, (size, size))
+        matrix = self._density_matrix.copy() if copy else self._density_matrix
+        return np.reshape(matrix, (size, size))
 
     def sample(self,
                qubits: List[ops.Qid],
@@ -528,7 +537,7 @@ def __init__(self, params: study.ParamResolver,
                          final_simulator_state=final_simulator_state)
         size = np.prod(protocols.qid_shape(self), dtype=int)
         self.final_density_matrix = np.reshape(
-            final_simulator_state.density_matrix, (size, size))
+            final_simulator_state.density_matrix.copy(), (size, size))
 
     def _value_equality_values_(self) -> Any:
         measurements = {

cirq/sim/density_matrix_simulator_test.py

@@ -1112,3 +1112,41 @@ def test_simulate_noise_with_terminal_measurements():
     result2 = simulator.run(circuit2, repetitions=10)
 
     assert result1 == result2
+
+
+def test_density_matrix_copy():
+    sim = cirq.DensityMatrixSimulator()
+
+    q = cirq.LineQubit(0)
+    circuit = cirq.Circuit(cirq.H(q), cirq.H(q))
+
+    matrices = []
+    for step in sim.simulate_moment_steps(circuit):
+        matrices.append(step.density_matrix(copy=True))
+    assert all(np.isclose(np.trace(x), 1.0) for x in matrices)
+    for x, y in itertools.combinations(matrices, 2):
+        assert not np.shares_memory(x, y)
+
+    # If the density matrix is not copied, then applying second Hadamard
+    # causes old state to be modified.
+    matrices = []
+    traces = []
+    for step in sim.simulate_moment_steps(circuit):
+        matrices.append(step.density_matrix(copy=False))
+        traces.append(np.trace(step.density_matrix(copy=False)))
+    assert any(not np.isclose(np.trace(x), 1.0) for x in matrices)
+    assert all(np.isclose(x, 1.0) for x in traces)
+    assert all(not np.shares_memory(x, y)
+               for x, y in itertools.combinations(matrices, 2))
+
+
+def test_final_density_matrix_is_not_last_object():
+    sim = cirq.DensityMatrixSimulator()
+
+    q = cirq.LineQubit(0)
+    initial_state = np.array([[1, 0], [0, 0]], dtype=np.complex64)
+    circuit = cirq.Circuit(cirq.WaitGate(0)(q))
+    result = sim.simulate(circuit, initial_state=initial_state)
+    assert result.final_density_matrix is not initial_state
+    assert not np.shares_memory(result.final_density_matrix, initial_state)
+    np.testing.assert_equal(result.final_density_matrix, initial_state)