Sample independent qubit sets without merging state space

cirq-core/cirq/__init__.py

@@ -398,6 +398,7 @@
     StateVectorStepResult,
     StateVectorTrialResult,
     StepResult,
+    StepResultBase,
 )
 
 from cirq.study import (

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

@@ -117,12 +117,9 @@ def _create_partial_act_on_args(
 
     def _create_step_result(
         self,
-        sim_state: 'MPSState',
-        qubit_map: Dict['cirq.Qid', int],
+        sim_state: 'cirq.OperationTarget[MPSState]',
     ):
-        return MPSSimulatorStepResult(
-            measurements=sim_state.log_of_measurement_results, state=sim_state
-        )
+        return MPSSimulatorStepResult(sim_state)
 
     def _create_simulator_trial_result(
         self,
@@ -169,22 +166,22 @@ def __str__(self) -> str:
         return f'measurements: {samples}\noutput state: {final}'
 
 
-class MPSSimulatorStepResult(simulator.StepResult['MPSState']):
+class MPSSimulatorStepResult(simulator_base.StepResultBase['MPSState', 'MPSState']):
     """A `StepResult` that can perform measurements."""
 
-    def __init__(self, state, measurements):
+    def __init__(
+        self,
+        sim_state: 'cirq.OperationTarget[MPSState]',
+    ):
         """Results of a step of the simulator.
         Attributes:
-            state: A MPSState
-            measurements: A dictionary from measurement gate key to measurement
-                results, ordered by the qubits that the measurement operates on.
-            qubit_map: A map from the Qubits in the Circuit to the the index
-                of this qubit for a canonical ordering. This canonical ordering
-                is used to define the state vector (see the state_vector()
-                method).
+            sim_state: The qubit:ActOnArgs lookup for this step.
         """
-        self.measurements = measurements
-        self.state = state.copy()
+        super().__init__(sim_state)
+
+    @property
+    def state(self):
+        return self._merged_sim_state
 
     def __str__(self) -> str:
         def bitstring(vals):
@@ -204,24 +201,6 @@ def bitstring(vals):
     def _simulator_state(self):
         return self.state
 
-    def sample(
-        self,
-        qubits: List[ops.Qid],
-        repetitions: int = 1,
-        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
-    ) -> np.ndarray:
-
-        measurements: List[int] = []
-
-        for _ in range(repetitions):
-            measurements.append(
-                self.state.perform_measurement(
-                    qubits, value.parse_random_state(seed), collapse_state_vector=False
-                )
-            )
-
-        return np.array(measurements, dtype=int)
-
 
 @value.value_equality
 class MPSState(ActOnArgs):
@@ -537,3 +516,21 @@ def perform_measurement(
     def _perform_measurement(self, qubits: Sequence['cirq.Qid']) -> List[int]:
         """Measures the axes specified by the simulator."""
         return self.perform_measurement(qubits, self.prng)
+
+    def sample(
+        self,
+        qubits: Sequence[ops.Qid],
+        repetitions: int = 1,
+        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    ) -> np.ndarray:
+
+        measurements: List[List[int]] = []
+
+        for _ in range(repetitions):
+            measurements.append(
+                self.perform_measurement(
+                    qubits, value.parse_random_state(seed), collapse_state_vector=False
+                )
+            )
+
+        return np.array(measurements, dtype=int)

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

@@ -274,11 +274,11 @@ def test_trial_result_str():
 
 def test_empty_step_result():
     q0 = cirq.LineQubit(0)
-    state = ccq.mps_simulator.MPSState(qubits=(q0,), prng=value.parse_random_state(0))
-    step_result = ccq.mps_simulator.MPSSimulatorStepResult(state, measurements={'0': [1]})
+    sim = ccq.mps_simulator.MPSSimulator()
+    step_result = next(sim.simulate_moment_steps(cirq.Circuit(cirq.measure(q0))))
     assert (
         str(step_result)
-        == """0=1
+        == """0=0
 TensorNetwork([
     Tensor(shape=(2,), inds=('i_0',), tags=set()),
 ])"""

cirq-core/cirq/experiments/grid_parallel_two_qubit_xeb_test.py

@@ -20,7 +20,7 @@ def test_estimate_parallel_two_qubit_xeb_fidelity_on_grid_no_noise(tmpdir):
     two_qubit_gate = cirq.ISWAP ** 0.5
     cycles = [5, 10, 15]
     data_collection_id = collect_grid_parallel_two_qubit_xeb_data(
-        sampler=cirq.Simulator(seed=34310),
+        sampler=cirq.Simulator(seed=34310, split_untangled_states=False),
         qubits=qubits,
         two_qubit_gate=two_qubit_gate,
         num_circuits=2,
@@ -53,7 +53,9 @@ def test_estimate_parallel_two_qubit_xeb_fidelity_on_grid_depolarizing(tmpdir):
     cycles = [5, 10, 15]
     e = 0.01
     data_collection_id = collect_grid_parallel_two_qubit_xeb_data(
-        sampler=cirq.DensityMatrixSimulator(noise=cirq.depolarize(e), seed=65008),
+        sampler=cirq.DensityMatrixSimulator(
+            noise=cirq.depolarize(e), seed=65008, split_untangled_states=False
+        ),
         qubits=qubits,
         two_qubit_gate=two_qubit_gate,
         num_circuits=2,

cirq-core/cirq/experiments/single_qubit_readout_calibration_test.py

@@ -31,7 +31,7 @@ def __init__(self, p0: float, p1: float, seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE'
         self.p0 = p0
         self.p1 = p1
         self.prng = cirq.value.parse_random_state(seed)
-        self.simulator = cirq.Simulator(seed=self.prng)
+        self.simulator = cirq.Simulator(seed=self.prng, split_untangled_states=False)
 
     def run_sweep(
         self,

cirq-core/cirq/protocols/act_on_protocol_test.py

@@ -37,6 +37,9 @@ def copy(self):
     def _act_on_fallback_(self, action, qubits, allow_decompose):
         return self.fallback_result
 
+    def sample(self, qubits, repetitions=1, seed=None):
+        pass
+
 
 op = cirq.X(cirq.LineQubit(0))
 

cirq-core/cirq/sim/__init__.py

@@ -64,6 +64,7 @@
 )
 
 from cirq.sim.simulator_base import (
+    StepResultBase,
     SimulatorBase,
 )
 

cirq-core/cirq/sim/act_on_args.py

@@ -72,7 +72,7 @@ def __init__(
         if log_of_measurement_results is None:
             log_of_measurement_results = {}
         self._qubits = tuple(qubits)
-        self.qubit_map = {q: i for i, q in enumerate(self.qubits)}
+        self.qubit_map = {q: i for i, q in enumerate(qubits)}
         self._axes = tuple(axes)
         self.prng = prng
         self._log_of_measurement_results = log_of_measurement_results
@@ -89,9 +89,9 @@ def measure(self, qubits: Sequence['cirq.Qid'], key: str, invert_mask: Sequence[
         """
         bits = self._perform_measurement(qubits)
         corrected = [bit ^ (bit < 2 and mask) for bit, mask in zip(bits, invert_mask)]
-        if key in self.log_of_measurement_results:
+        if key in self._log_of_measurement_results:
             raise ValueError(f"Measurement already logged to key {key!r}")
-        self.log_of_measurement_results[key] = corrected
+        self._log_of_measurement_results[key] = corrected
 
     def get_axes(self, qubits: Sequence['cirq.Qid']) -> List[int]:
         return [self.qubit_map[q] for q in qubits]
@@ -139,6 +139,15 @@ def log_of_measurement_results(self) -> Dict[str, Any]:
     def qubits(self) -> Tuple['cirq.Qid', ...]:
         return self._qubits
 
+    @abc.abstractmethod
+    def sample(
+        self,
+        qubits: Sequence['cirq.Qid'],
+        repetitions: int = 1,
+        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    ) -> np.ndarray:
+        """Samples the qubits."""
+
     def __getitem__(self: TSelf, item: Optional['cirq.Qid']) -> TSelf:
         if item not in self.qubit_map:
             raise IndexError(f'{item} not in {self.qubits}')

cirq-core/cirq/sim/act_on_args_container.py

@@ -20,10 +20,14 @@
     Sequence,
     Optional,
     Iterator,
-    Tuple,
     Any,
+    Tuple,
+    Set,
+    List,
 )
 
+import numpy as np
+
 from cirq import ops
 from cirq.sim.operation_target import OperationTarget
 from cirq.sim.simulator import (
@@ -58,7 +62,8 @@ def __init__(
                 unentangled qubit sets independently and merging those states
                 at the end.
             log_of_measurement_results: A mutable object that measurements are
-                being recorded into.
+                being recorded into. Edit it easily by calling
+                `ActOnStateVectorArgs.record_measurement_result`.
         """
         self.args = args
         self._qubits = tuple(qubits)
@@ -122,6 +127,26 @@ def qubits(self) -> Tuple['cirq.Qid', ...]:
     def log_of_measurement_results(self) -> Dict[str, Any]:
         return self._log_of_measurement_results
 
+    def sample(
+        self,
+        qubits: List[ops.Qid],
+        repetitions: int = 1,
+        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    ) -> np.ndarray:
+        columns = []
+        selected_order: List[ops.Qid] = []
+        q_set = set(qubits)
+        for v in dict.fromkeys(self.args.values()):
+            qs = [q for q in v.qubits if q in q_set]
+            if any(qs):
+                column = v.sample(qs, repetitions, seed)
+                columns.append(column)
+                selected_order += qs
+        stacked = np.column_stack(columns)
+        qubit_map = {q: i for i, q in enumerate(selected_order)}
+        index_order = [qubit_map[q] for q in qubits]
+        return stacked[:, index_order]
+
     def __getitem__(self, item: Optional['cirq.Qid']) -> TActOnArgs:
         return self.args[item]
 

cirq-core/cirq/sim/act_on_args_container_test.py

@@ -64,6 +64,9 @@ def transpose_to_qubit_order(self, qubits: Sequence['cirq.Qid']) -> 'EmptyActOnA
             logs=self.log_of_measurement_results,
         )
 
+    def sample(self, qubits, repetitions=1, seed=None):
+        pass
+
 
 q0, q1 = qs2 = cirq.LineQubit.range(2)
 

cirq-core/cirq/sim/act_on_args_test.py

@@ -25,6 +25,9 @@ def __init__(self):
     def copy(self):
         pass
 
+    def sample(self, qubits, repetitions=1, seed=None):
+        pass
+
     def _perform_measurement(self, qubits):
         return [5, 3]
 

cirq-core/cirq/sim/act_on_density_matrix_args.py

@@ -197,6 +197,21 @@ def transpose_to_qubit_order(
             log_of_measurement_results=self.log_of_measurement_results,
         )
 
+    def sample(
+        self,
+        qubits: Sequence['cirq.Qid'],
+        repetitions: int = 1,
+        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    ) -> np.ndarray:
+        indices = [self.qubit_map[q] for q in qubits]
+        return sim.sample_density_matrix(
+            self.target_tensor,
+            indices,
+            qid_shape=tuple(q.dimension for q in self.qubits),
+            repetitions=repetitions,
+            seed=seed,
+        )
+
 
 def _strat_apply_channel_to_state(
     action: Any, args: ActOnDensityMatrixArgs, qubits: Sequence['cirq.Qid']

cirq-core/cirq/sim/act_on_state_vector_args.py

@@ -255,6 +255,21 @@ def transpose_to_qubit_order(self, qubits: Sequence['cirq.Qid']) -> 'cirq.ActOnS
         )
         return new_args
 
+    def sample(
+        self,
+        qubits: Sequence['cirq.Qid'],
+        repetitions: int = 1,
+        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    ) -> np.ndarray:
+        indices = [self.qubit_map[q] for q in qubits]
+        return sim.sample_state_vector(
+            self.target_tensor,
+            indices,
+            qid_shape=tuple(q.dimension for q in self.qubits),
+            repetitions=repetitions,
+            seed=seed,
+        )
+
 
 def _strat_act_on_state_vector_from_apply_unitary(
     unitary_value: Any,

cirq-core/cirq/sim/clifford/act_on_clifford_tableau_args.py

@@ -111,6 +111,15 @@ def copy(self) -> 'cirq.ActOnCliffordTableauArgs':
             log_of_measurement_results=self.log_of_measurement_results.copy(),
         )
 
+    def sample(
+        self,
+        qubits: Sequence['cirq.Qid'],
+        repetitions: int = 1,
+        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    ) -> np.ndarray:
+        # Unnecessary for now but can be added later if there is a use case.
+        raise NotImplementedError()
+
 
 def _strat_act_on_clifford_tableau_from_single_qubit_decompose(
     val: Any, args: 'cirq.ActOnCliffordTableauArgs', qubits: Sequence['cirq.Qid']

cirq-core/cirq/sim/clifford/act_on_stabilizer_ch_form_args.py

@@ -16,6 +16,7 @@
 
 import numpy as np
 
+from cirq import value, ops, protocols
 from cirq._compat import deprecated_parameter
 from cirq.ops import common_gates, pauli_gates
 from cirq.ops.clifford_gate import SingleQubitCliffordGate
@@ -106,6 +107,21 @@ def copy(self) -> 'cirq.ActOnStabilizerCHFormArgs':
             log_of_measurement_results=self.log_of_measurement_results.copy(),
         )
 
+    def sample(
+        self,
+        qubits: Sequence['cirq.Qid'],
+        repetitions: int = 1,
+        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+    ) -> np.ndarray:
+        measurements: Dict[str, List[np.ndarray]] = {}
+        prng = value.parse_random_state(seed)
+        for i in range(repetitions):
+            op = ops.measure(*qubits, key=str(i))
+            state = self.state.copy()
+            ch_form_args = ActOnStabilizerCHFormArgs(state, prng, measurements, self.qubits)
+            protocols.act_on(op, ch_form_args)
+        return np.array(list(measurements.values()), dtype=bool)
+
 
 def _strat_act_on_stabilizer_ch_form_from_single_qubit_decompose(
     val: Any, args: 'cirq.ActOnStabilizerCHFormArgs', qubits: Sequence['cirq.Qid']