wip: add back add/remove_qubit to density matrix sim state

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)
+        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:

cirq-core/cirq/sim/simulation_state_test.py

@@ -153,7 +153,21 @@ def test_delegating_gate_unitary(exp):
     assert_test_circuit_for_sv_simulator(test_circuit, control_circuit)
 
 
-@pytest.mark.parametrize('exp', np.linspace(0, 2 * np.pi, 10))
+@pytest.mark.parametrize('exp', np.linspace(0, 2 * np.pi, 1))
+def test_delegating_gate_channel_sv(exp):
+    q = cirq.LineQubit(0)
+
+    test_circuit = cirq.Circuit()
+    test_circuit.append(cirq.H(q))
+    test_circuit.append(DelegatingAncillaZ(exp, True).on(q))
+
+    control_circuit = cirq.Circuit(cirq.H(q))
+    control_circuit.append(cirq.ZPowGate(exponent=exp).on(q))
+
+    assert_test_circuit_for_sv_simulator(test_circuit, control_circuit)
+
+
+@pytest.mark.parametrize('exp', np.linspace(0, 2 * np.pi, 1))
 def test_delegating_gate_channel(exp):
     q = cirq.LineQubit(0)
 
@@ -164,9 +178,7 @@ 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_dm_simulator(test_circuit, control_circuit)
 
 
 @pytest.mark.parametrize('num_ancilla', [1, 2, 3])