add back add/remove_qubit to density matrix sim state (#6259)

* wip: add back add/remove_qubit to density matrix sim state * factorize inplace while removing qubit in density matrix simulation state
quantumlib · Aug 28, 2023 · 56b5db2 · 56b5db2
1 parent ed26d2f
commit 56b5db2
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git a/cirq-core/cirq/sim/density_matrix_simulation_state.py b/cirq-core/cirq/sim/density_matrix_simulation_state.py
@@ -285,6 +285,22 @@ def __init__(
         )
         super().__init__(state=state, prng=prng, qubits=qubits, classical_data=classical_data)
 
+    def add_qubits(self, qubits: Sequence['cirq.Qid']):
+        ret = super().add_qubits(qubits)
+        return (
+            self.kronecker_product(type(self)(qubits=qubits), inplace=True)
+            if ret is NotImplemented
+            else ret
+        )
+
+    def remove_qubits(self, qubits: Sequence['cirq.Qid']):
+        ret = super().remove_qubits(qubits)
+        if ret is not NotImplemented:
+            return ret
+        extracted, remainder = self.factor(qubits, inplace=True)
+        remainder._state._density_matrix *= extracted._state._density_matrix.reshape(-1)[0]
+        return remainder
+
     def _act_on_fallback_(
         self, action: Any, qubits: Sequence['cirq.Qid'], allow_decompose: bool = True
     ) -> bool:

diff --git a/cirq-core/cirq/sim/simulation_state_test.py b/cirq-core/cirq/sim/simulation_state_test.py
@@ -164,9 +164,8 @@ def test_delegating_gate_channel(exp):
     control_circuit = cirq.Circuit(cirq.H(q))
     control_circuit.append(cirq.ZPowGate(exponent=exp).on(q))
 
-    with pytest.raises(TypeError, match="DensityMatrixSimulator doesn't support"):
-        # TODO: This test should pass once we extend support to DensityMatrixSimulator.
-        assert_test_circuit_for_dm_simulator(test_circuit, control_circuit)
+    assert_test_circuit_for_sv_simulator(test_circuit, control_circuit)
+    assert_test_circuit_for_dm_simulator(test_circuit, control_circuit)
 
 
 @pytest.mark.parametrize('num_ancilla', [1, 2, 3])