Permit 2D _run output for backwards compatibility. (quantumlib#5014)

Fixes quantumlib#5000.

This PR reinstates support for 2D measurement data from `_run`, but logs a warning if that path is used. External simulators will need to modify their `_run` implementation before v0.15.

Author: 95-martin-orion

cirq/sim/simulator.py

@@ -44,6 +44,7 @@
     TypeVar,
     Union,
 )
+import warnings
 
 import numpy as np
 
@@ -110,6 +111,20 @@ def run_sweep_iter(
                 records = self._run(
                     circuit=program, param_resolver=param_resolver, repetitions=repetitions
                 )
+                flat_records = False
+                for k, v in records.items():
+                    if v.ndim == 2:
+                        flat_records = True
+                        records[k] = v.reshape((v.shape[0], 1, v.shape[1]))
+                if flat_records:
+                    warnings.warn(
+                        (
+                            'Starting in Cirq v0.15, values in the output of simulator._run must '
+                            'be 3D instead of 2D, with a new dimension between the existing two '
+                            'to capture "instances" of a key.'
+                        ),
+                        DeprecationWarning,
+                    )
             yield study.ResultDict(params=param_resolver, records=records)
 
     @abc.abstractmethod

cirq/sim/simulator_test.py

@@ -29,11 +29,40 @@
     SimulatesExpectationValues,
     SimulatesFinalState,
     SimulatesIntermediateState,
+    SimulatesSamples,
     SimulationTrialResult,
     TActOnArgs,
 )
 
 
+class FakeSimulatesSamples(SimulatesSamples):
+    """A SimulatesSamples that returns specified values from _run."""
+
+    def __init__(self, run_output: Dict[str, np.ndarray]):
+        self._run_output = run_output
+
+    def _run(self, *args, **kwargs) -> Dict[str, np.ndarray]:
+        return self._run_output
+
+
+class FakeStepResult(cirq.StepResult):
+    def __init__(self, *, ones_qubits=None, final_state=None):
+        self._ones_qubits = set(ones_qubits or [])
+        self._final_state = final_state
+
+    def _simulator_state(self):
+        return self._final_state
+
+    def state_vector(self):
+        pass
+
+    def __setstate__(self, state):
+        pass
+
+    def sample(self, qubits, repetitions=1, seed=None):
+        return np.array([[qubit in self._ones_qubits for qubit in qubits]] * repetitions)
+
+
 class SimulatesIntermediateStateImpl(
     Generic[TStepResult, TSimulatorState, TActOnArgs],
     SimulatesIntermediateState[TStepResult, 'SimulationTrialResult', TSimulatorState, TActOnArgs],
@@ -62,43 +91,48 @@ def _create_simulator_trial_result(
         )
 
 
-@mock.patch.multiple(cirq.SimulatesSamples, __abstractmethods__=set(), _run=mock.Mock())
 def test_run_simulator_run():
-    simulator = cirq.SimulatesSamples()
-    expected_measurements = {'a': np.array([[[1]]])}
-    simulator._run.return_value = expected_measurements
-    circuit = mock.Mock(cirq.Circuit)
-    circuit.__iter__ = mock.Mock(return_value=iter([]))
-    param_resolver = mock.Mock(cirq.ParamResolver)
-    param_resolver.param_dict = {}
-    expected_result = cirq.ResultDict(records=expected_measurements, params=param_resolver)
+    expected_records = {'a': np.array([[[1]]])}
+    simulator = FakeSimulatesSamples(expected_records)
+    circuit = cirq.Circuit(cirq.measure(cirq.LineQubit(0), key='k'))
+    param_resolver = cirq.ParamResolver({})
+    expected_result = cirq.ResultDict(records=expected_records, params=param_resolver)
     assert expected_result == simulator.run(
         program=circuit, repetitions=10, param_resolver=param_resolver
     )
-    simulator._run.assert_called_once_with(
-        circuit=circuit, repetitions=10, param_resolver=param_resolver
-    )
 
 
-@mock.patch.multiple(cirq.SimulatesSamples, __abstractmethods__=set(), _run=mock.Mock())
 def test_run_simulator_sweeps():
-    simulator = cirq.SimulatesSamples()
-    expected_measurements = {'a': np.array([[[1]]])}
-    simulator._run.return_value = expected_measurements
-    circuit = mock.Mock(cirq.Circuit)
-    circuit.__iter__ = mock.Mock(return_value=iter([]))
-    param_resolvers = [mock.Mock(cirq.ParamResolver), mock.Mock(cirq.ParamResolver)]
-    for resolver in param_resolvers:
-        resolver.param_dict = {}
+    expected_records = {'a': np.array([[[1]]])}
+    simulator = FakeSimulatesSamples(expected_records)
+    circuit = cirq.Circuit(cirq.measure(cirq.LineQubit(0), key='k'))
+    param_resolvers = [cirq.ParamResolver({}), cirq.ParamResolver({})]
     expected_results = [
-        cirq.ResultDict(records=expected_measurements, params=param_resolvers[0]),
-        cirq.ResultDict(records=expected_measurements, params=param_resolvers[1]),
+        cirq.ResultDict(records=expected_records, params=param_resolvers[0]),
+        cirq.ResultDict(records=expected_records, params=param_resolvers[1]),
     ]
     assert expected_results == simulator.run_sweep(
         program=circuit, repetitions=10, params=param_resolvers
     )
-    simulator._run.assert_called_with(circuit=circuit, repetitions=10, param_resolver=mock.ANY)
-    assert simulator._run.call_count == 2
+
+
+def test_run_simulator_sweeps_with_deprecated_run():
+    expected_measurements = {'a': np.array([[1]])}
+    simulator = FakeSimulatesSamples(expected_measurements)
+    circuit = cirq.Circuit(cirq.measure(cirq.LineQubit(0), key='k'))
+    param_resolvers = [cirq.ParamResolver({}), cirq.ParamResolver({})]
+    expected_records = {'a': np.array([[[1]]])}
+    expected_results = [
+        cirq.ResultDict(records=expected_records, params=param_resolvers[0]),
+        cirq.ResultDict(records=expected_records, params=param_resolvers[1]),
+    ]
+    with cirq.testing.assert_deprecated(
+        'values in the output of simulator._run must be 3D',
+        deadline='v0.15',
+    ):
+        assert expected_results == simulator.run_sweep(
+            program=circuit, repetitions=10, params=param_resolvers
+        )
 
 
 @mock.patch.multiple(
@@ -157,8 +191,7 @@ def steps(*args, **kwargs):
         program=circuit, params=param_resolvers, qubit_order=qubit_order, initial_state=2
     )
 
-    final_step_result = mock.Mock()
-    final_step_result._simulator_state.return_value = final_state
+    final_step_result = FakeStepResult(final_state=final_state)
     expected_results = [
         cirq.SimulationTrialResult(
             measurements={'a': np.array([True, True])},
@@ -174,27 +207,10 @@ def steps(*args, **kwargs):
     assert results == expected_results
 
 
-class FakeStepResult(cirq.StepResult):
-    def __init__(self, ones_qubits):
-        self._ones_qubits = set(ones_qubits)
-
-    def _simulator_state(self):
-        pass
-
-    def state_vector(self):
-        pass
-
-    def __setstate__(self, state):
-        pass
-
-    def sample(self, qubits, repetitions=1, seed=None):
-        return np.array([[qubit in self._ones_qubits for qubit in qubits]] * repetitions)
-
-
 def test_step_sample_measurement_ops():
     q0, q1, q2 = cirq.LineQubit.range(3)
     measurement_ops = [cirq.measure(q0, q1), cirq.measure(q2)]
-    step_result = FakeStepResult([q1])
+    step_result = FakeStepResult(ones_qubits=[q1])
 
     measurements = step_result.sample_measurement_ops(measurement_ops)
     np.testing.assert_equal(measurements, {'0,1': [[False, True]], '2': [[False]]})
@@ -203,7 +219,7 @@ def test_step_sample_measurement_ops():
 def test_step_sample_measurement_ops_repetitions():
     q0, q1, q2 = cirq.LineQubit.range(3)
     measurement_ops = [cirq.measure(q0, q1), cirq.measure(q2)]
-    step_result = FakeStepResult([q1])
+    step_result = FakeStepResult(ones_qubits=[q1])
 
     measurements = step_result.sample_measurement_ops(measurement_ops, repetitions=3)
     np.testing.assert_equal(measurements, {'0,1': [[False, True]] * 3, '2': [[False]] * 3})
@@ -215,29 +231,29 @@ def test_step_sample_measurement_ops_invert_mask():
         cirq.measure(q0, q1, invert_mask=(True,)),
         cirq.measure(q2, invert_mask=(False,)),
     ]
-    step_result = FakeStepResult([q1])
+    step_result = FakeStepResult(ones_qubits=[q1])
 
     measurements = step_result.sample_measurement_ops(measurement_ops)
     np.testing.assert_equal(measurements, {'0,1': [[True, True]], '2': [[False]]})
 
 
 def test_step_sample_measurement_ops_no_measurements():
-    step_result = FakeStepResult([])
+    step_result = FakeStepResult(ones_qubits=[])
 
     measurements = step_result.sample_measurement_ops([])
     assert measurements == {}
 
 
 def test_step_sample_measurement_ops_not_measurement():
     q0 = cirq.LineQubit(0)
-    step_result = FakeStepResult([q0])
+    step_result = FakeStepResult(ones_qubits=[q0])
     with pytest.raises(ValueError, match='MeasurementGate'):
         step_result.sample_measurement_ops([cirq.X(q0)])
 
 
 def test_step_sample_measurement_ops_repeated_qubit():
     q0, q1, q2 = cirq.LineQubit.range(3)
-    step_result = FakeStepResult([q0])
+    step_result = FakeStepResult(ones_qubits=[q0])
     with pytest.raises(ValueError, match='Measurement key 0 repeated'):
         step_result.sample_measurement_ops(
             [cirq.measure(q0), cirq.measure(q1, q2), cirq.measure(q0)]
@@ -246,8 +262,7 @@ def test_step_sample_measurement_ops_repeated_qubit():
 
 def test_simulation_trial_result_equality():
     eq = cirq.testing.EqualsTester()
-    final_step_result = mock.Mock(cirq.StepResult)
-    final_step_result._simulator_state.return_value = ()
+    final_step_result = FakeStepResult(final_state=())
     eq.add_equality_group(
         cirq.SimulationTrialResult(
             params=cirq.ParamResolver({}), measurements={}, final_step_result=final_step_result
@@ -270,7 +285,7 @@ def test_simulation_trial_result_equality():
             final_step_result=final_step_result,
         )
     )
-    final_step_result._simulator_state.return_value = (0, 1)
+    final_step_result._final_state = (0, 1)
     eq.add_equality_group(
         cirq.SimulationTrialResult(
             params=cirq.ParamResolver({'s': 1}),
@@ -281,8 +296,7 @@ def test_simulation_trial_result_equality():
 
 
 def test_simulation_trial_result_repr():
-    final_step_result = mock.Mock(cirq.StepResult)
-    final_step_result._simulator_state.return_value = (0, 1)
+    final_step_result = FakeStepResult(final_state=(0, 1))
     assert repr(
         cirq.SimulationTrialResult(
             params=cirq.ParamResolver({'s': 1}),
@@ -298,8 +312,7 @@ def test_simulation_trial_result_repr():
 
 
 def test_simulation_trial_result_str():
-    final_step_result = mock.Mock(cirq.StepResult)
-    final_step_result._simulator_state.return_value = (0, 1)
+    final_step_result = FakeStepResult(final_state=(0, 1))
     assert (
         str(
             cirq.SimulationTrialResult(
@@ -369,13 +382,10 @@ def text(self, to_print):
 @duet.sync
 async def test_async_sample():
     m = {'mock': np.array([[[0]], [[1]]])}
-
-    class MockSimulator(cirq.SimulatesSamples):
-        def _run(self, circuit, param_resolver, repetitions):
-            return m
+    simulator = FakeSimulatesSamples(m)
 
     q = cirq.LineQubit(0)
-    f = MockSimulator().run_async(cirq.Circuit(cirq.measure(q)), repetitions=10)
+    f = simulator.run_async(cirq.Circuit(cirq.measure(q)), repetitions=10)
     result = await f
     np.testing.assert_equal(result.records, m)
 
@@ -458,10 +468,8 @@ def _kraus_(self):
 
 
 def test_iter_definitions():
-    final_step_result = mock.Mock(cirq.StepResult)
-    final_step_result._simulator_state.return_value = []
     dummy_trial_result = SimulationTrialResult(
-        params={}, measurements={}, final_step_result=final_step_result
+        params={}, measurements={}, final_step_result=FakeStepResult(final_state=[])
     )
 
     class FakeNonIterSimulatorImpl(