Allow clifford simulator to run as product state

cirq-core/cirq/sim/act_on_args.py

@@ -155,6 +155,11 @@ def factor(
         remainder._set_qubits([q for q in self.qubits if q not in qubits])
         return extracted, remainder
 
+    @property
+    def allows_factoring(self):
+        """Subclasses that allow factorization should override this."""
+        return False
+
     def _on_factor(
         self: TSelf,
         qubits: Sequence['cirq.Qid'],

cirq-core/cirq/sim/act_on_args_container.py

@@ -123,8 +123,9 @@ def _act_on_fallback_(
             or (isinstance(gate, ops.MeasurementGate) and not op_args.ignore_measurement_results)
         ):
             for q in qubits:
-                q_args, op_args = op_args.factor((q,), validate=False)
-                self.args[q] = q_args
+                if op_args.allows_factoring:
+                    q_args, op_args = op_args.factor((q,), validate=False)
+                    self.args[q] = q_args
 
             # (Backfill the args map with the new value)
             for q in op_args.qubits:

cirq-core/cirq/sim/act_on_args_container_test.py

@@ -52,6 +52,10 @@ def _on_transpose_to_qubit_order(self, qubits, target):
     def _on_factor(self, qubits, extracted, remainder, validate=True, atol=1e-07):
         pass
 
+    @property
+    def allows_factoring(self):
+        return True
+
     def sample(self, qubits, repetitions=1, seed=None):
         pass
 

cirq-core/cirq/sim/act_on_density_matrix_args.py

@@ -147,6 +147,10 @@ def _on_factor(
         ]
         remainder.qid_shape = remainder_tensor.shape[: int(remainder_tensor.ndim / 2)]
 
+    @property
+    def allows_factoring(self):
+        return True
+
     def _on_transpose_to_qubit_order(
         self, qubits: Sequence['cirq.Qid'], target: 'cirq.ActOnDensityMatrixArgs'
     ):

cirq-core/cirq/sim/act_on_state_vector_args.py

@@ -204,6 +204,10 @@ def _on_factor(
         remainder.target_tensor = remainder_tensor
         remainder.available_buffer = np.empty_like(remainder_tensor)
 
+    @property
+    def allows_factoring(self):
+        return True
+
     def _on_transpose_to_qubit_order(
         self, qubits: Sequence['cirq.Qid'], target: 'cirq.ActOnStateVectorArgs'
     ):

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

@@ -81,6 +81,17 @@ def _perform_measurement(self, qubits: Sequence['cirq.Qid']) -> List[int]:
     def _on_copy(self, target: 'ActOnStabilizerCHFormArgs'):
         target.state = self.state.copy()
 
+    def _on_kronecker_product(
+        self, other: 'cirq.ActOnStabilizerCHFormArgs', target: 'cirq.ActOnStabilizerCHFormArgs'
+    ):
+        target.state = self.state.kron(other.state)
+
+    def _on_transpose_to_qubit_order(
+        self, qubits: Sequence['cirq.Qid'], target: 'cirq.ActOnStabilizerCHFormArgs'
+    ):
+        axes = [self.qubit_map[q] for q in qubits]
+        target.state = self.state.reindex(axes)
+
     def sample(
         self,
         qubits: Sequence['cirq.Qid'],

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

@@ -49,14 +49,23 @@ class CliffordSimulator(
 ):
     """An efficient simulator for Clifford circuits."""
 
-    def __init__(self, seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None):
+    def __init__(
+        self,
+        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+        split_untangled_states: bool = False,
+    ):
         """Creates instance of `CliffordSimulator`.
 
         Args:
             seed: The random seed to use for this simulator.
+            split_untangled_states: Optimizes simulation by running separable
+                states independently and merging those states at the end.
         """
         self.init = True
-        super().__init__(seed=seed)
+        super().__init__(
+            seed=seed,
+            split_untangled_states=split_untangled_states,
+        )
 
     @staticmethod
     def is_supported_operation(op: 'cirq.Operation') -> bool:

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

@@ -186,10 +186,11 @@ def test_simulate_moment_steps_sample():
                 )
 
 
-def test_simulate_moment_steps_intermediate_measurement():
+@pytest.mark.parametrize('split', [True, False])
+def test_simulate_moment_steps_intermediate_measurement(split):
     q0 = cirq.LineQubit(0)
     circuit = cirq.Circuit(cirq.H(q0), cirq.measure(q0), cirq.H(q0))
-    simulator = cirq.CliffordSimulator()
+    simulator = cirq.CliffordSimulator(split_untangled_states=split)
     for i, step in enumerate(simulator.simulate_moment_steps(circuit)):
         if i == 1:
             result = int(step.measurements['0'][0])
@@ -330,7 +331,8 @@ def test_clifford_circuit_SHSYSHS():
     )
 
 
-def test_clifford_circuit():
+@pytest.mark.parametrize('split', [True, False])
+def test_clifford_circuit(split):
     (q0, q1) = (cirq.LineQubit(0), cirq.LineQubit(1))
     circuit = cirq.Circuit()
 
@@ -354,7 +356,7 @@ def test_clifford_circuit():
         elif x == 6:
             circuit.append(cirq.CZ(q0, q1))
 
-    clifford_simulator = cirq.CliffordSimulator()
+    clifford_simulator = cirq.CliffordSimulator(split_untangled_states=split)
     state_vector_simulator = cirq.Simulator()
 
     np.testing.assert_almost_equal(
@@ -364,7 +366,8 @@ def test_clifford_circuit():
 
 
 @pytest.mark.parametrize("qubits", [cirq.LineQubit.range(2), cirq.LineQubit.range(4)])
-def test_clifford_circuit_2(qubits):
+@pytest.mark.parametrize('split', [True, False])
+def test_clifford_circuit_2(qubits, split):
     circuit = cirq.Circuit()
 
     np.random.seed(2)
@@ -388,13 +391,14 @@ def test_clifford_circuit_2(qubits):
             circuit.append(cirq.CZ(qubits[0], qubits[1]))  # coverage: ignore
 
     circuit.append(cirq.measure(qubits[0]))
-    result = cirq.CliffordSimulator().run(circuit, repetitions=100)
+    result = cirq.CliffordSimulator(split_untangled_states=split).run(circuit, repetitions=100)
 
     assert sum(result.measurements['0'])[0] < 80
     assert sum(result.measurements['0'])[0] > 20
 
 
-def test_clifford_circuit_3():
+@pytest.mark.parametrize('split', [True, False])
+def test_clifford_circuit_3(split):
     # This test tests the simulator on arbitrary 1-qubit Clifford gates.
     (q0, q1) = (cirq.LineQubit(0), cirq.LineQubit(1))
     circuit = cirq.Circuit()
@@ -412,7 +416,7 @@ def random_clifford_gate():
         else:
             circuit.append(random_clifford_gate()(np.random.choice((q0, q1))))
 
-    clifford_simulator = cirq.CliffordSimulator()
+    clifford_simulator = cirq.CliffordSimulator(split_untangled_states=split)
     state_vector_simulator = cirq.Simulator()
 
     np.testing.assert_almost_equal(
@@ -549,14 +553,16 @@ def test_valid_apply_measurement():
     assert measurements == {'0': [1]}
 
 
-def test_reset():
+@pytest.mark.parametrize('split', [True, False])
+def test_reset(split):
     q = cirq.LineQubit(0)
     c = cirq.Circuit(cirq.X(q), cirq.reset(q), cirq.measure(q, key="out"))
-    assert cirq.CliffordSimulator().sample(c)["out"][0] == 0
+    sim = cirq.CliffordSimulator(split_untangled_states=split)
+    assert sim.sample(c)["out"][0] == 0
     c = cirq.Circuit(cirq.H(q), cirq.reset(q), cirq.measure(q, key="out"))
-    assert cirq.CliffordSimulator().sample(c)["out"][0] == 0
+    assert sim.sample(c)["out"][0] == 0
     c = cirq.Circuit(cirq.reset(q), cirq.measure(q, key="out"))
-    assert cirq.CliffordSimulator().sample(c)["out"][0] == 0
+    assert sim.sample(c)["out"][0] == 0
 
 
 def test_state_copy():

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Union
+from typing import Any, Dict, Sequence, Union
 import numpy as np
 
 import cirq
@@ -273,3 +273,32 @@ def project_Z(self, q, z):
             self.omega /= np.sqrt(2)
 
         self.update_sum(t, u, delta=delta)
+
+    def kron(self, other: 'cirq.StabilizerStateChForm') -> 'cirq.StabilizerStateChForm':
+        n = self.n + other.n
+        copy = StabilizerStateChForm(n)
+
+        def diag(a, b):
+            a = np.pad(a, [0, other.n])
+            b = np.pad(b, [self.n, 0])
+            return a + b
+
+        copy.G = diag(self.G, other.G)
+        copy.F = diag(self.F, other.F)
+        copy.M = diag(self.M, other.M)
+        copy.gamma = np.concatenate([self.gamma, other.gamma])
+        copy.v = np.concatenate([self.v, other.v])
+        copy.s = np.concatenate([self.s, other.s])
+        copy.omega = self.omega * other.omega
+        return copy
+
+    def reindex(self, axes: Sequence[int]) -> 'cirq.StabilizerStateChForm':
+        copy = StabilizerStateChForm(self.n)
+        copy.G = self.G[axes][:, axes]
+        copy.F = self.F[axes][:, axes]
+        copy.M = self.M[axes][:, axes]
+        copy.gamma = self.gamma[axes]
+        copy.v = self.v[axes]
+        copy.s = self.s[axes]
+        copy.omega = self.omega
+        return copy

cirq-core/cirq/sim/simulator_base_test.py

@@ -70,6 +70,10 @@ def _on_factor(self, qubits, extracted, remainder, validate=True, atol=1e-07):
         remainder.gate_count = 0
         remainder.measurement_count = 0
 
+    @property
+    def allows_factoring(self):
+        return True
+
     def _on_transpose_to_qubit_order(
         self, qubits: Sequence['cirq.Qid'], target: 'SplittableCountingActOnArgs'
     ):