Repository: quantumlib/Cirq

Issue:
Density matrix simulator should be optimized for terminal measurements
similar to other simulators. 


Discussion:
Instead of running the density simulation once per repetition, run it once and then sample from the final density matrix multiple times. Works almost exactly the same as in the sparse wavefunction simulator; there is existing code to sample from the density matrix.
I'm happy to look into this, it seems doable in 2 weeks with iterations. 

Resolving patch:
diff --git a/cirq/sim/density_matrix_simulator.py b/cirq/sim/density_matrix_simulator.py
--- a/cirq/sim/density_matrix_simulator.py
+++ b/cirq/sim/density_matrix_simulator.py
@@ -123,12 +123,35 @@ def _run(self,
             repetitions: int) -> Dict[str, np.ndarray]:
         """See definition in `cirq.SimulatesSamples`."""
         param_resolver = param_resolver or study.ParamResolver({})
-        resolved_circuit = protocols.resolve_parameters(circuit, param_resolver)
+        resolved_circuit = protocols.resolve_parameters(circuit,
+                                                        param_resolver)
+
+        if circuit.are_all_measurements_terminal():
+            return self._run_sweep_sample(resolved_circuit, repetitions)
+        else:
+            return self._run_sweep_repeat(resolved_circuit, repetitions)
+
+    def _run_sweep_sample(self,
+                          circuit: circuits.Circuit,
+                          repetitions: int) -> Dict[str, np.ndarray]:
+        for step_result in self._base_iterator(
+                circuit=circuit,
+                qubit_order=ops.QubitOrder.DEFAULT,
+                initial_state=0,
+                perform_measurements=False):
+            pass
+        measurement_ops = [op for _, op, _ in
+                           circuit.findall_operations_with_gate_type(
+                               ops.MeasurementGate)]
+        return step_result.sample_measurement_ops(measurement_ops, repetitions)
+
+    def _run_sweep_repeat(self,
+                          circuit: circuits.Circuit,
+                          repetitions: int) -> Dict[str, np.ndarray]:
         measurements = {}  # type: Dict[str, List[np.ndarray]]
-        # TODO: optimize for all terminal measurements.
         for _ in range(repetitions):
             all_step_results = self._base_iterator(
-                resolved_circuit,
+                circuit,
                 qubit_order=ops.QubitOrder.DEFAULT,
                 initial_state=0,
                 perform_measurements=True)


Test patch:
diff --git a/cirq/sim/density_matrix_simulator_test.py b/cirq/sim/density_matrix_simulator_test.py
--- a/cirq/sim/density_matrix_simulator_test.py
+++ b/cirq/sim/density_matrix_simulator_test.py
@@ -11,13 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Tuple
 
 import numpy as np
 import pytest
 import sympy
 
 import cirq
+from cirq.testing.mock import mock
 
 
 def test_invalid_dtype():
@@ -119,34 +119,40 @@ def test_run_channel(dtype):
 def test_run_repetitions_measure_at_end(dtype):
     q0, q1 = cirq.LineQubit.range(2)
     simulator = cirq.DensityMatrixSimulator(dtype=dtype)
-    for b0 in [0, 1]:
-        for b1 in [0, 1]:
-            circuit = cirq.Circuit.from_ops((cirq.X**b0)(q0),
-                                            (cirq.X**b1)(q1),
-                                            cirq.measure(q0),
-                                            cirq.measure(q1))
-            result = simulator.run(circuit, repetitions=3)
-            np.testing.assert_equal(result.measurements,
-                                    {'0': [[b0]] * 3, '1': [[b1]] * 3})
+    with mock.patch.object(simulator, '_base_iterator',
+                           wraps=simulator._base_iterator) as mock_sim:
+        for b0 in [0, 1]:
+            for b1 in [0, 1]:
+                circuit = cirq.Circuit.from_ops((cirq.X**b0)(q0),
+                                                (cirq.X**b1)(q1),
+                                                cirq.measure(q0),
+                                                cirq.measure(q1))
+                result = simulator.run(circuit, repetitions=3)
+                np.testing.assert_equal(result.measurements,
+                                        {'0': [[b0]] * 3, '1': [[b1]] * 3})
+                assert result.repetitions == 3
+        assert mock_sim.call_count == 4
 
 
 @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
 def test_run_repetitions_measurement_not_terminal(dtype):
     q0, q1 = cirq.LineQubit.range(2)
     simulator = cirq.DensityMatrixSimulator(dtype=dtype)
-    for b0 in [0, 1]:
-        for b1 in [0, 1]:
-            circuit = cirq.Circuit.from_ops((cirq.X**b0)(q0),
-                                            (cirq.X**b1)(q1),
-                                            cirq.measure(q0),
-                                            cirq.measure(q1),
-                                            cirq.H(q0),
-                                            cirq.H(q1))
-            result = simulator.run(circuit, repetitions=3)
-            np.testing.assert_equal(result.measurements,
-                                    {'0': [[b0]] * 3, '1': [[b1]] * 3})
-            assert result.repetitions == 3
-
+    with mock.patch.object(simulator, '_base_iterator',
+                           wraps=simulator._base_iterator) as mock_sim:
+        for b0 in [0, 1]:
+            for b1 in [0, 1]:
+                circuit = cirq.Circuit.from_ops((cirq.X**b0)(q0),
+                                                (cirq.X**b1)(q1),
+                                                cirq.measure(q0),
+                                                cirq.measure(q1),
+                                                cirq.H(q0),
+                                                cirq.H(q1))
+                result = simulator.run(circuit, repetitions=3)
+                np.testing.assert_equal(result.measurements,
+                                        {'0': [[b0]] * 3, '1': [[b1]] * 3})
+                assert result.repetitions == 3
+        assert mock_sim.call_count == 12
 
 @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
 def test_run_param_resolver(dtype):
diff --git a/cirq/sim/sparse_simulator_test.py b/cirq/sim/sparse_simulator_test.py
--- a/cirq/sim/sparse_simulator_test.py
+++ b/cirq/sim/sparse_simulator_test.py
@@ -12,11 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
 import numpy as np
+import pytest
 import sympy
 
 import cirq
+from cirq.testing.mock import mock
 
 
 def test_invalid_dtype():
@@ -70,33 +71,40 @@ def test_run_bit_flips(dtype):
 def test_run_repetitions_measure_at_end(dtype):
     q0, q1 = cirq.LineQubit.range(2)
     simulator = cirq.Simulator(dtype=dtype)
-    for b0 in [0, 1]:
-        for b1 in [0, 1]:
-            circuit = cirq.Circuit.from_ops((cirq.X**b0)(q0),
-                                            (cirq.X**b1)(q1),
-                                            cirq.measure(q0),
-                                            cirq.measure(q1))
-            result = simulator.run(circuit, repetitions=3)
-            np.testing.assert_equal(result.measurements,
-                                    {'0': [[b0]] * 3, '1': [[b1]] * 3})
+    with mock.patch.object(simulator, '_base_iterator',
+                           wraps=simulator._base_iterator) as mock_sim:
+        for b0 in [0, 1]:
+            for b1 in [0, 1]:
+                circuit = cirq.Circuit.from_ops((cirq.X**b0)(q0),
+                                                (cirq.X**b1)(q1),
+                                                cirq.measure(q0),
+                                                cirq.measure(q1))
+                result = simulator.run(circuit, repetitions=3)
+                np.testing.assert_equal(result.measurements,
+                                        {'0': [[b0]] * 3, '1': [[b1]] * 3})
+                assert result.repetitions == 3
+        assert mock_sim.call_count == 4
 
 
 @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])
 def test_run_repetitions_measurement_not_terminal(dtype):
     q0, q1 = cirq.LineQubit.range(2)
     simulator = cirq.Simulator(dtype=dtype)
-    for b0 in [0, 1]:
-        for b1 in [0, 1]:
-            circuit = cirq.Circuit.from_ops((cirq.X**b0)(q0),
-                                            (cirq.X**b1)(q1),
-                                            cirq.measure(q0),
-                                            cirq.measure(q1),
-                                            cirq.H(q0),
-                                            cirq.H(q1))
-            result = simulator.run(circuit, repetitions=3)
-            np.testing.assert_equal(result.measurements,
-                                    {'0': [[b0]] * 3, '1': [[b1]] * 3})
-            assert result.repetitions == 3
+    with mock.patch.object(simulator, '_base_iterator',
+                           wraps=simulator._base_iterator) as mock_sim:
+        for b0 in [0, 1]:
+            for b1 in [0, 1]:
+                circuit = cirq.Circuit.from_ops((cirq.X**b0)(q0),
+                                                (cirq.X**b1)(q1),
+                                                cirq.measure(q0),
+                                                cirq.measure(q1),
+                                                cirq.H(q0),
+                                                cirq.H(q1))
+                result = simulator.run(circuit, repetitions=3)
+                np.testing.assert_equal(result.measurements,
+                                        {'0': [[b0]] * 3, '1': [[b1]] * 3})
+                assert result.repetitions == 3
+        assert  mock_sim.call_count == 12
 
 
 @pytest.mark.parametrize('dtype', [np.complex64, np.complex128])