Lazily create TrialResult.final_simulator_state (quantumlib#4317)

* Lazily create TrialResult.final_simulator_state

* Change TrialResult only to accept StepResult in constructor

* Backwards compatibility for TrialResult initializer

* lint

* Platform independent test

* Update density_matrix_simulator_test.py

Co-authored-by: Orion Martin <[email protected]>

Author: daxfohl

cirq/contrib/quimb/mps_simulator.py

@@ -125,23 +125,21 @@ def _create_simulator_trial_result(
         self,
         params: study.ParamResolver,
         measurements: Dict[str, np.ndarray],
-        final_simulator_state: 'MPSState',
+        final_step_result: 'MPSSimulatorStepResult',
     ) -> 'MPSTrialResult':
         """Creates a single trial results with the measurements.
 
         Args:
-            circuit: The circuit to simulate.
-            param_resolver: A ParamResolver for determining values of
-                Symbols.
+            params: A ParamResolver for determining values of Symbols.
             measurements: A dictionary from measurement key (e.g. qubit) to the
                 actual measurement array.
-            final_simulator_state: The final state of the simulator.
+            final_step_result: The final step result of the simulation.
 
         Returns:
             A single result.
         """
         return MPSTrialResult(
-            params=params, measurements=measurements, final_simulator_state=final_simulator_state
+            params=params, measurements=measurements, final_step_result=final_step_result
         )
 
 
@@ -152,13 +150,15 @@ def __init__(
         self,
         params: study.ParamResolver,
         measurements: Dict[str, np.ndarray],
-        final_simulator_state: 'MPSState',
+        final_step_result: 'MPSSimulatorStepResult',
     ) -> None:
         super().__init__(
-            params=params, measurements=measurements, final_simulator_state=final_simulator_state
+            params=params, measurements=measurements, final_step_result=final_step_result
         )
 
-        self.final_state = final_simulator_state
+    @property
+    def final_state(self):
+        return self._final_simulator_state
 
     def __str__(self) -> str:
         samples = super().__str__()

cirq/contrib/quimb/mps_simulator_test.py

@@ -1,14 +1,15 @@
 import itertools
 import math
+from unittest import mock
 
 import numpy as np
 import pytest
 import sympy
 
 import cirq
-from cirq import value
 import cirq.contrib.quimb as ccq
 import cirq.experiments.google_v2_supremacy_circuit as supremacy_v2
+from cirq import value
 
 
 def assert_same_output_as_dense(circuit, qubit_order, initial_state=0, grouping=None):
@@ -252,7 +253,8 @@ def test_measurement_str():
 
 def test_trial_result_str():
     q0 = cirq.LineQubit(0)
-    final_simulator_state = ccq.mps_simulator.MPSState(
+    final_step_result = mock.Mock(cirq.StepResult)
+    final_step_result._simulator_state.return_value = ccq.mps_simulator.MPSState(
         qubits=(q0,),
         prng=value.parse_random_state(0),
         simulation_options=ccq.mps_simulator.MPSOptions(),
@@ -262,7 +264,7 @@ def test_trial_result_str():
             ccq.mps_simulator.MPSTrialResult(
                 params=cirq.ParamResolver({}),
                 measurements={'m': np.array([[1]])},
-                final_simulator_state=final_simulator_state,
+                final_step_result=final_step_result,
             )
         )
         == """measurements: m=1

cirq/sim/clifford/clifford_simulator.py

@@ -105,11 +105,11 @@ def _create_simulator_trial_result(
         self,
         params: study.ParamResolver,
         measurements: Dict[str, np.ndarray],
-        final_simulator_state,
+        final_step_result: 'CliffordSimulatorStepResult',
     ):
 
         return CliffordTrialResult(
-            params=params, measurements=measurements, final_simulator_state=final_simulator_state
+            params=params, measurements=measurements, final_step_result=final_step_result
         )
 
 
@@ -118,13 +118,15 @@ def __init__(
         self,
         params: study.ParamResolver,
         measurements: Dict[str, np.ndarray],
-        final_simulator_state: 'CliffordState',
+        final_step_result: 'CliffordSimulatorStepResult',
     ) -> None:
         super().__init__(
-            params=params, measurements=measurements, final_simulator_state=final_simulator_state
+            params=params, measurements=measurements, final_step_result=final_step_result
         )
 
-        self.final_state = final_simulator_state
+    @property
+    def final_state(self):
+        return self._final_simulator_state
 
     def __str__(self) -> str:
         samples = super().__str__()

cirq/sim/clifford/clifford_simulator_test.py

@@ -1,4 +1,6 @@
 import itertools
+from unittest import mock
+
 import numpy as np
 import pytest
 import sympy
@@ -208,13 +210,14 @@ def test_clifford_state_initial_state():
 
 def test_clifford_trial_result_repr():
     q0 = cirq.LineQubit(0)
-    final_simulator_state = cirq.CliffordState(qubit_map={q0: 0})
+    final_step_result = mock.Mock(cirq.CliffordSimulatorStepResult)
+    final_step_result._simulator_state.return_value = cirq.CliffordState(qubit_map={q0: 0})
     assert (
         repr(
             cirq.CliffordTrialResult(
                 params=cirq.ParamResolver({}),
                 measurements={'m': np.array([[1]])},
-                final_simulator_state=final_simulator_state,
+                final_step_result=final_step_result,
             )
         )
         == "cirq.SimulationTrialResult(params=cirq.ParamResolver({}), "
@@ -225,13 +228,14 @@ def test_clifford_trial_result_repr():
 
 def test_clifford_trial_result_str():
     q0 = cirq.LineQubit(0)
-    final_simulator_state = cirq.CliffordState(qubit_map={q0: 0})
+    final_step_result = mock.Mock(cirq.CliffordSimulatorStepResult)
+    final_step_result._simulator_state.return_value = cirq.CliffordState(qubit_map={q0: 0})
     assert (
         str(
             cirq.CliffordTrialResult(
                 params=cirq.ParamResolver({}),
                 measurements={'m': np.array([[1]])},
-                final_simulator_state=final_simulator_state,
+                final_step_result=final_step_result,
             )
         )
         == "measurements: m=1\n"

cirq/sim/density_matrix_simulator.py

@@ -222,10 +222,10 @@ def _create_simulator_trial_result(
         self,
         params: study.ParamResolver,
         measurements: Dict[str, np.ndarray],
-        final_simulator_state: 'DensityMatrixSimulatorState',
+        final_step_result: 'DensityMatrixStepResult',
     ) -> 'DensityMatrixTrialResult':
         return DensityMatrixTrialResult(
-            params=params, measurements=measurements, final_simulator_state=final_simulator_state
+            params=params, measurements=measurements, final_step_result=final_step_result
         )
 
     # TODO(#4209): Deduplicate with identical code in sparse_simulator.
@@ -423,22 +423,27 @@ class DensityMatrixTrialResult(simulator.SimulationTrialResult):
             measurement gate.)
         final_simulator_state: The final simulator state of the system after the
             trial finishes.
-        final_density_matrix: The final density matrix of the system.
     """
 
     def __init__(
         self,
         params: study.ParamResolver,
         measurements: Dict[str, np.ndarray],
-        final_simulator_state: DensityMatrixSimulatorState,
+        final_step_result: DensityMatrixStepResult,
     ) -> None:
         super().__init__(
-            params=params, measurements=measurements, final_simulator_state=final_simulator_state
-        )
-        size = np.prod(protocols.qid_shape(self), dtype=np.int64)
-        self.final_density_matrix = np.reshape(
-            final_simulator_state.density_matrix.copy(), (size, size)
+            params=params, measurements=measurements, final_step_result=final_step_result
         )
+        self._final_density_matrix: Optional[np.ndarray] = None
+
+    @property
+    def final_density_matrix(self):
+        if self._final_density_matrix is None:
+            size = np.prod(protocols.qid_shape(self), dtype=np.int64)
+            self._final_density_matrix = np.reshape(
+                self._final_simulator_state.density_matrix.copy(), (size, size)
+            )
+        return self._final_density_matrix
 
     def _value_equality_values_(self) -> Any:
         measurements = {k: v.tolist() for k, v in sorted(self.measurements.items())}

cirq/sim/density_matrix_simulator_test.py

@@ -11,10 +11,11 @@
 # 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 Type
-from unittest import mock
 import itertools
 import random
+from typing import Type
+from unittest import mock
+
 import numpy as np
 import pytest
 import sympy
@@ -998,61 +999,65 @@ def test_density_matrix_simulator_state_repr():
 
 def test_density_matrix_trial_result_eq():
     q0 = cirq.LineQubit(0)
-    final_simulator_state = cirq.DensityMatrixSimulatorState(
+    final_step_result = mock.Mock(cirq.StepResult)
+    final_step_result._simulator_state.return_value = cirq.DensityMatrixSimulatorState(
         density_matrix=np.ones((2, 2)) * 0.5, qubit_map={q0: 0}
     )
     eq = cirq.testing.EqualsTester()
     eq.add_equality_group(
         cirq.DensityMatrixTrialResult(
             params=cirq.ParamResolver({}),
             measurements={},
-            final_simulator_state=final_simulator_state,
+            final_step_result=final_step_result,
         ),
         cirq.DensityMatrixTrialResult(
             params=cirq.ParamResolver({}),
             measurements={},
-            final_simulator_state=final_simulator_state,
+            final_step_result=final_step_result,
         ),
     )
     eq.add_equality_group(
         cirq.DensityMatrixTrialResult(
             params=cirq.ParamResolver({'s': 1}),
             measurements={},
-            final_simulator_state=final_simulator_state,
+            final_step_result=final_step_result,
         )
     )
     eq.add_equality_group(
         cirq.DensityMatrixTrialResult(
             params=cirq.ParamResolver({'s': 1}),
             measurements={'m': np.array([[1]])},
-            final_simulator_state=final_simulator_state,
+            final_step_result=final_step_result,
         )
     )
 
 
 def test_density_matrix_trial_result_qid_shape():
     q0, q1 = cirq.LineQubit.range(2)
+    final_step_result = mock.Mock(cirq.StepResult)
+    final_step_result._simulator_state.return_value = cirq.DensityMatrixSimulatorState(
+        density_matrix=np.ones((4, 4)) / 4, qubit_map={q0: 0, q1: 1}
+    )
     assert (
         cirq.qid_shape(
             cirq.DensityMatrixTrialResult(
                 params=cirq.ParamResolver({}),
                 measurements={},
-                final_simulator_state=cirq.DensityMatrixSimulatorState(
-                    density_matrix=np.ones((4, 4)) / 4, qubit_map={q0: 0, q1: 1}
-                ),
+                final_step_result=final_step_result,
             ),
         )
         == (2, 2)
     )
     q0, q1 = cirq.LineQid.for_qid_shape((3, 4))
+    final_step_result._simulator_state.return_value = cirq.DensityMatrixSimulatorState(
+        density_matrix=np.ones((12, 12)) / 12, qubit_map={q0: 0, q1: 1}
+    )
     assert (
         cirq.qid_shape(
             cirq.DensityMatrixTrialResult(
                 params=cirq.ParamResolver({}),
                 measurements={},
-                final_simulator_state=cirq.DensityMatrixSimulatorState(
-                    density_matrix=np.ones((12, 12)) / 12, qubit_map={q0: 0, q1: 1}
-                ),
+                final_step_result=final_step_result,
             ),
         )
         == (3, 4)
@@ -1061,15 +1066,16 @@ def test_density_matrix_trial_result_qid_shape():
 
 def test_density_matrix_trial_result_repr():
     q0 = cirq.LineQubit(0)
-    final_simulator_state = cirq.DensityMatrixSimulatorState(
+    final_step_result = mock.Mock(cirq.StepResult)
+    final_step_result._simulator_state.return_value = cirq.DensityMatrixSimulatorState(
         density_matrix=np.ones((2, 2)) * 0.5, qubit_map={q0: 0}
     )
     assert (
         repr(
             cirq.DensityMatrixTrialResult(
                 params=cirq.ParamResolver({'s': 1}),
                 measurements={'m': np.array([[1]])},
-                final_simulator_state=final_simulator_state,
+                final_step_result=final_step_result,
             )
         )
         == "cirq.DensityMatrixTrialResult("
@@ -1166,11 +1172,12 @@ def test_works_on_pauli_string():
 
 def test_density_matrix_trial_result_str():
     q0 = cirq.LineQubit(0)
-    final_simulator_state = cirq.DensityMatrixSimulatorState(
+    final_step_result = mock.Mock(cirq.StepResult)
+    final_step_result._simulator_state.return_value = cirq.DensityMatrixSimulatorState(
         density_matrix=np.ones((2, 2)) * 0.5, qubit_map={q0: 0}
     )
     result = cirq.DensityMatrixTrialResult(
-        params=cirq.ParamResolver({}), measurements={}, final_simulator_state=final_simulator_state
+        params=cirq.ParamResolver({}), measurements={}, final_step_result=final_step_result
     )
 
     # numpy varies whitespace in its representation for different versions
@@ -1527,3 +1534,26 @@ def test_density_matrices_same_with_or_without_split_untangled_states():
     sim = cirq.DensityMatrixSimulator(split_untangled_states=True)
     result2 = sim.simulate(circuit).final_density_matrix
     assert np.allclose(result1, result2)
+
+
+def test_large_untangled_okay():
+    circuit = cirq.Circuit()
+    for i in range(59):
+        for _ in range(9):
+            circuit.append(cirq.X(cirq.LineQubit(i)))
+        circuit.append(cirq.measure(cirq.LineQubit(i)))
+
+    # Validate this can't be allocated with entangled state
+    with pytest.raises(MemoryError, match='Unable to allocate'):
+        _ = cirq.DensityMatrixSimulator(split_untangled_states=False).simulate(circuit)
+
+    # Validate a simulation run
+    result = cirq.DensityMatrixSimulator(split_untangled_states=True).simulate(circuit)
+    assert set(result._final_step_result._qubits) == set(cirq.LineQubit.range(59))
+    # _ = result.final_density_matrix hangs (as expected)
+
+    # Validate a trial run and sampling
+    result = cirq.DensityMatrixSimulator(split_untangled_states=True).run(circuit, repetitions=1000)
+    assert len(result.measurements) == 59
+    assert len(result.measurements['0']) == 1000
+    assert (result.measurements['0'] == np.full(1000, 1)).all()