Fix dropped state collapse in FreeQuantumState.measure_multiple

sequence/kernel/quantum_state.py

@@ -361,11 +361,12 @@ def measure_multiple(basis, states, rng: Generator):
         possible_results = np.arange(0, basis_dimension, 1)
         # result gives index of the basis vector that will be projected to
         res = rng.choice(possible_results, p=probabilities)
-        # project to new state, then reassign quantum state and entangled photons
-        new_state = new_states[res]
-        for state in entangled_list:
-            state.quantum_state = new_state
-            state.entangled_photons = entangled_list
+        # project to new state, then reassign quantum state and entangled list
+        new_state = tuple(new_states[res])
+        for qs in entangled_list:
+            if qs is not None:
+                qs.state = new_state
+                qs.entangled_states = entangled_list
 
         return res
 

tests/components/test_photon.py

@@ -80,6 +80,27 @@ def test_measure_multiple():
     assert Photon.measure_multiple(basis, [photon1, photon2], rng) == 0
 
 
+def test_measure_multiple_collapses_state():
+    """After measure_multiple, the shared FreeQuantumState must collapse to the
+    measured basis vector (and stay shared between the entangled photons)."""
+    tl = Timeline()
+    photon1 = Photon("p1", tl, quantum_state=(complex(2 ** -0.5), complex(2 ** -0.5)))
+    photon2 = Photon("p2", tl, quantum_state=(complex(2 ** -0.5), complex(2 ** -0.5)))
+    photon1.combine_state(photon2)
+
+    basis = ((complex(1), complex(0), complex(0), complex(0)),
+             (complex(0), complex(1), complex(0), complex(0)),
+             (complex(0), complex(0), complex(1), complex(0)),
+             (complex(0), complex(0), complex(0), complex(1)))
+
+    res = Photon.measure_multiple(basis, [photon1, photon2], np.random.default_rng(0))
+
+    expected = tuple(complex(1) if i == res else complex(0) for i in range(4))
+    assert tuple(photon1.quantum_state.state) == expected, "state did not collapse to measured basis vector"
+    # both photons must keep sharing the same collapsed state object
+    assert photon1.quantum_state.state is photon2.quantum_state.state
+
+
 def test_add_loss():
     tl = Timeline()
     photon = Photon("", tl, encoding_type=single_atom)