Update Density Matrix and State Vector Simulators to work when an operation allocates new qubits as part of its decomposition

cirq-core/cirq/sim/simulation_state.py

@@ -166,6 +166,37 @@ def create_merged_state(self) -> Self:
         """Creates a final merged state."""
         return self
 
+    def add_qubits(self: Self, qubits: Sequence['cirq.Qid']) -> Self:
+        """Add qubits to a new state space and take the kron product.
+        Note that only subclasses that support `kronecker_product`
+        will support this function. E.g Density Matrix and
+        State Vector simulators.
+
+        Args:
+            qubits: Sequence of qubits to be added.
+
+        Returns:
+            A new Simulation State with the new qubits added. Or
+            `self` if there are no qubits to add.
+        """
+        if qubits is None or not qubits:
+            return self
+        new_space = type(self)(qubits=qubits)
+        return self.kronecker_product(new_space, inplace=True)
+
+    def remove_qubits(self: Self, qubits: Sequence['cirq.Qid']) -> Self:
+        """Remove qubits from the state space.
+
+        Args:
+            qubits: Sequence of qubits to be added.
+
+        Returns:
+            A new Simulation State with qubits removed. Or
+            `self` if there are no qubits to remove."""
+        if qubits is None or not qubits:
+            return self
+        return self.factor(qubits, inplace=True)[1]
+
     def kronecker_product(self, other: Self, *, inplace=False) -> Self:
         """Joins two state spaces together."""
         args = self if inplace else copy.copy(self)
@@ -294,13 +325,18 @@ def strat_act_on_from_apply_decompose(
     val: Any, args: 'cirq.SimulationState', qubits: Sequence['cirq.Qid']
 ) -> bool:
     operations, qubits1, _ = _try_decompose_into_operations_and_qubits(val)
-    assert len(qubits1) == len(qubits)
-    qubit_map = {q: qubits[i] for i, q in enumerate(qubits1)}
     if operations is None:
         return NotImplemented
+    assert len(qubits1) == len(qubits)
+    qubit_map = {q: qubits[i] for i, q in enumerate(qubits1)}
+    ancillas = list(set(q for op in operations for q in op.qubits if q not in qubits1))
+    for q in ancillas:
+        qubit_map[q] = q
+    args.add_qubits(ancillas)
     for operation in operations:
         operation = operation.with_qubits(*[qubit_map[q] for q in operation.qubits])
         protocols.act_on(operation, args)
+    args.remove_qubits(ancillas)
     return True
 
 

cirq-core/cirq/sim/simulation_state_test.py

@@ -31,17 +31,79 @@ def measure(self, axes, seed=None):
     def reindex(self, axes):
         return self
 
+    def kron(self, other):
+        return self
+
+    def factor(self, axes, validate=True, atol=1e-07):
+        return (self, self)
+
 
 class DummySimulationState(cirq.SimulationState):
-    def __init__(self):
-        super().__init__(state=DummyQuantumState(), qubits=cirq.LineQubit.range(2))
+    def __init__(self, qubits=cirq.LineQubit.range(2)):
+        super().__init__(state=DummyQuantumState(), qubits=qubits)
 
     def _act_on_fallback_(
         self, action: Any, qubits: Sequence['cirq.Qid'], allow_decompose: bool = True
     ) -> bool:
         return True
 
 
+class AncillaZ(cirq.Gate):
+    def __init__(self, exponent=1):
+        self._exponent = exponent
+
+    def num_qubits(self) -> int:
+        return 1
+
+    def _decompose_(self, qubits):
+        ancilla = cirq.NamedQubit('Ancilla')
+        yield cirq.CX(qubits[0], ancilla)
+        yield cirq.Z(ancilla) ** self._exponent
+        yield cirq.CX(qubits[0], ancilla)
+
+
+class AncillaH(cirq.Gate):
+    def __init__(self, exponent=1):
+        self._exponent = exponent
+
+    def num_qubits(self) -> int:
+        return 1
+
+    def _decompose_(self, qubits):
+        ancilla = cirq.NamedQubit('Ancilla')
+        yield cirq.H(ancilla) ** self._exponent
+        yield cirq.CX(ancilla, qubits[0])
+        yield cirq.H(ancilla) ** self._exponent
+
+
+class AncillaY(cirq.Gate):
+    def __init__(self, exponent=1):
+        self._exponent = exponent
+
+    def num_qubits(self) -> int:
+        return 1
+
+    def _decompose_(self, qubits):
+        ancilla = cirq.NamedQubit('Ancilla')
+        yield cirq.Y(ancilla) ** self._exponent
+        yield cirq.CX(ancilla, qubits[0])
+        yield cirq.Y(ancilla) ** self._exponent
+
+
+class DelegatingAncillaZ(cirq.Gate):
+    def __init__(self, exponent=1):
+        self._exponent = exponent
+
+    def num_qubits(self) -> int:
+        return 1
+
+    def _decompose_(self, qubits):
+        a = cirq.NamedQubit('a')
+        yield cirq.CX(qubits[0], a)
+        yield AncillaZ(self._exponent).on(a)
+        yield cirq.CX(qubits[0], a)
+
+
 def test_measurements():
     args = DummySimulationState()
     args.measure([cirq.LineQubit(0)], "test", [False], {})
@@ -101,3 +163,59 @@ def test_field_getters():
     args = DummySimulationState()
     assert args.prng is np.random
     assert args.qubit_map == {q: i for i, q in enumerate(cirq.LineQubit.range(2))}
+
+
+@pytest.mark.parametrize('exp', [-3, -2, -1, 0, 1, 2, 3])
+def test_ancilla_z(exp):
+    q = cirq.LineQubit(0)
+    test_circuit = cirq.Circuit(AncillaZ(exp).on(q))
+
+    control_circuit = cirq.Circuit(cirq.ZPowGate(exponent=exp).on(q))
+
+    assert_test_circuit_for_sv_dm_simulators(test_circuit, control_circuit)
+
+
+@pytest.mark.parametrize('exp', [-3, -2, -1, 0, 1, 2, 3])
+def test_ancilla_y(exp):
+    q = cirq.LineQubit(0)
+    test_circuit = cirq.Circuit(AncillaY(exp).on(q))
+
+    control_circuit = cirq.Circuit(cirq.Y(q))
+    control_circuit.append(cirq.Y(q))
+    control_circuit.append(cirq.XPowGate(exponent=exp).on(q))
+
+    assert_test_circuit_for_sv_dm_simulators(test_circuit, control_circuit)
+
+
+@pytest.mark.parametrize('exp', [-3, -2, -1, 0, 1, 2, 3])
+def test_borrowable_qubit(exp):
+    q = cirq.LineQubit(0)
+    test_circuit = cirq.Circuit()
+    test_circuit.append(cirq.H(q))
+    test_circuit.append(cirq.X(q))
+    test_circuit.append(AncillaH(exp).on(q))
+
+    control_circuit = cirq.Circuit(cirq.H(q))
+
+    assert_test_circuit_for_sv_dm_simulators(test_circuit, control_circuit)
+
+
+@pytest.mark.parametrize('exp', [-3, -2, -1, 0, 1, 2, 3])
+def test_delegating_gate_qubit(exp):
+    q = cirq.LineQubit(0)
+
+    test_circuit = cirq.Circuit()
+    test_circuit.append(cirq.H(q))
+    test_circuit.append(DelegatingAncillaZ(exp).on(q))
+
+    control_circuit = cirq.Circuit(cirq.H(q))
+    control_circuit.append(cirq.ZPowGate(exponent=exp).on(q))
+
+    assert_test_circuit_for_sv_dm_simulators(test_circuit, control_circuit)
+
+
+def assert_test_circuit_for_sv_dm_simulators(test_circuit, control_circuit) -> None:
+    for test_simulator in ['cirq.final_state_vector', 'cirq.final_density_matrix']:
+        test_sim = eval(test_simulator)(test_circuit)
+        control_sim = eval(test_simulator)(control_circuit)
+        assert np.allclose(test_sim, control_sim)