Fix global phase when removing qubits from simulators

cirq-core/cirq/sim/density_matrix_simulation_state.py

@@ -293,6 +293,14 @@ def add_qubits(self, qubits: Sequence['cirq.Qid']):
             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.py

@@ -196,7 +196,7 @@ def remove_qubits(self: Self, qubits: Sequence['cirq.Qid']) -> Self:
             `self` if there are no qubits to remove."""
         if qubits is None or not qubits:
             return self
-        return self.factor(qubits, inplace=True)[1]
+        return NotImplemented
 
     def kronecker_product(self, other: Self, *, inplace=False) -> Self:
         """Joins two state spaces together."""
@@ -339,7 +339,9 @@ def strat_act_on_from_apply_decompose(
     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)
+    args = args.remove_qubits(new_ancilla)
+    if args is NotImplemented:
+        raise TypeError(f"{type(args)} implements `add_qubits` but not `remove_qubits`.")
     return True
 
 

cirq-core/cirq/sim/simulation_state_test.py

@@ -246,4 +246,4 @@ def assert_test_circuit_for_sv_dm_simulators(test_circuit, control_circuit) -> N
     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)
+        assert np.allclose(test_sim, control_sim)

cirq-core/cirq/sim/state_vector_simulation_state.py

@@ -363,6 +363,14 @@ def add_qubits(self, qubits: Sequence['cirq.Qid']):
             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._state_vector *= extracted._state._state_vector.reshape((-1,))[0]
+        return remainder
+
     def _act_on_fallback_(
         self, action: Any, qubits: Sequence['cirq.Qid'], allow_decompose: bool = True
     ) -> bool: