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

cirq-core/cirq/protocols/act_on_protocol.py

@@ -149,7 +149,7 @@ def act_on(
 
     arg_fallback = getattr(sim_state, '_act_on_fallback_', None)
     if arg_fallback is not None:
-        qubits = action.qubits if isinstance(action, ops.Operation) else qubits
+        qubits = action.qubits if is_op else qubits
         result = arg_fallback(action, qubits=qubits, allow_decompose=allow_decompose)
         if result is True:
             return

cirq-core/cirq/sim/density_matrix_simulation_state.py

@@ -285,6 +285,14 @@ 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 _act_on_fallback_(
         self, action: Any, qubits: Sequence['cirq.Qid'], allow_decompose: bool = True
     ) -> bool:

cirq-core/cirq/sim/simulation_state.py

@@ -166,6 +166,38 @@ def create_merged_state(self) -> Self:
         """Creates a final merged state."""
         return self
 
+    def add_qubits(self: Self, qubits: Sequence['cirq.Qid']) -> None:
+        """Add qubits to a new state space and take the kron product.
+
+        Note that only Density Matrix and State Vector simulators
+        override this function.
+
+        Args:
+            qubits: Sequence of qubits to be added.
+
+        Returns:
+            NotImplemented: If the subclass does not implement this method.
+
+        Raises:
+             ValueError: If a qubit being added is already tracked.
+        """
+        if any(q in self.qubits for q in qubits):
+            raise ValueError(f"Qubit to add {qubits} should not already be tracked.")
+        return NotImplemented
+
+    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 +326,20 @@ 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)
+    all_qubits = frozenset([q for op in operations for q in op.qubits])
+    qubit_map = dict(zip(all_qubits, all_qubits))
+    qubit_map.update(dict(zip(qubits1, qubits)))
+    new_ancilla = tuple(q for q in sorted(all_qubits.difference(qubits)) if q not in args.qubits)
+    args = args.add_qubits(new_ancilla)
+    if args is NotImplemented:
+        return NotImplemented
     for operation in operations:
         operation = operation.with_qubits(*[qubit_map[q] for q in operation.qubits])
         protocols.act_on(operation, args)
+    args.remove_qubits(new_ancilla)
     return True
 
 

cirq-core/cirq/sim/simulation_state_test.py

@@ -19,6 +19,7 @@
 
 import cirq
 from cirq.sim import simulation_state
+from cirq.testing import PhaseUsingCleanAncilla, PhaseUsingDirtyAncilla
 
 
 class DummyQuantumState(cirq.QuantumStateRepresentation):
@@ -31,17 +32,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], {})
@@ -57,8 +120,9 @@ def _decompose_(self, qubits):
             yield cirq.X(*qubits)
 
     args = DummySimulationState()
-    assert simulation_state.strat_act_on_from_apply_decompose(
-        Composite(), args, [cirq.LineQubit(0)]
+    assert (
+        simulation_state.strat_act_on_from_apply_decompose(Composite(), args, [cirq.LineQubit(0)])
+        is NotImplemented
     )
 
 
@@ -101,3 +165,85 @@ 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)
+
+
+@pytest.mark.parametrize('num_ancilla', [1, 2, 3])
+def test_phase_using_dirty_ancilla(num_ancilla: int):
+    q = cirq.LineQubit(0)
+    anc = cirq.NamedQubit.range(num_ancilla, prefix='anc')
+
+    u = cirq.MatrixGate(cirq.testing.random_unitary(2 ** (num_ancilla + 1)))
+    test_circuit = cirq.Circuit(
+        u.on(q, *anc), PhaseUsingDirtyAncilla(ancilla_bitsize=num_ancilla).on(q)
+    )
+    control_circuit = cirq.Circuit(u.on(q, *anc), cirq.Z(q))
+    assert_test_circuit_for_sv_dm_simulators(test_circuit, control_circuit)
+
+
+@pytest.mark.parametrize('num_ancilla', [1, 2, 3])
+@pytest.mark.parametrize('theta', np.linspace(0, 2 * np.pi, 10))
+def test_phase_using_clean_ancilla(num_ancilla: int, theta: float):
+    q = cirq.LineQubit(0)
+    u = cirq.MatrixGate(cirq.testing.random_unitary(2))
+    test_circuit = cirq.Circuit(
+        u.on(q), PhaseUsingCleanAncilla(theta=theta, ancilla_bitsize=num_ancilla).on(q)
+    )
+    control_circuit = cirq.Circuit(u.on(q), cirq.ZPowGate(exponent=theta).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)
+        cirq.testing.assert_allclose_up_to_global_phase(test_sim, control_sim, atol=1e-6)

cirq-core/cirq/sim/state_vector_simulation_state.py

@@ -355,6 +355,14 @@ 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 _act_on_fallback_(
         self, action: Any, qubits: Sequence['cirq.Qid'], allow_decompose: bool = True
     ) -> bool: