Group singleshot emulator result extraction

src/qibolab/_core/instruments/emulator/results.py

@@ -130,6 +130,20 @@ def _marginalize_probability(
     return np.stack(res)
 
 
+def _sampled_measurements(
+    sampled: NDArray, dims: list[int], inverse_map: NDArray, indices: list[int]
+) -> NDArray:
+    """Extract measured subsystem states from sampled full-system states."""
+    indices = np.asarray(indices)
+    res = np.empty((len(indices), *sampled.shape[1:]), dtype=sampled.dtype)
+    for sample in np.unique(inverse_map):
+        states = np.stack(np.unravel_index(sampled[sample], dims))
+        measurements = np.flatnonzero(inverse_map == sample)
+        res[measurements] = states[indices[measurements]]
+
+    return res
+
+
 def select_acquisitions(
     states: list[Operator], acquisitions: Iterable[float], times: NDArray
 ) -> NDArray:
@@ -214,20 +228,15 @@ def _singleshot_results(
     # the shape now is: (M_unique, Nshots, *S, *H_dim)
     sampled = np.moveaxis(sampled, 1, 0)
 
-    acq_id = acquisitions(sequence).keys()
+    acq_id = list(acquisitions(sequence).keys())
     # from every acquisition pulse id we get the corresponding channel, and from the channel we get the
-    # corresponding qubit index, which is then used to correctly permute the rows of states_computational_idx.
+    # corresponding Hilbert-space index to extract from the sampled full-system state.
     qubit_indices = [
         index(sequence.pulse_channels(ro_id)[0], hamiltonian) for ro_id in acq_id
     ]
 
-    # we use inverse_map to expand back the sampled results
-    # res is a (M, M, Nshots, *S, ...) array
-    res = np.stack(np.unravel_index(sampled[inverse_map], hamiltonian.dims))[
-        qubit_indices
-    ]
-    # using np.einsum, so res is a (M, Nshots, *S, ...) array
-    res = np.einsum(res, np.array([0, 0] + [...]), np.array([0] + [...]))
+    # res is a (M, Nshots, *S, ...) array
+    res = _sampled_measurements(sampled, hamiltonian.dims, inverse_map, qubit_indices)
     res = np.clip(res, 0, 1)
 
     if options.acquisition_type is AcquisitionType.INTEGRATION:

tests/instruments/emulator/test_results.py

@@ -8,6 +8,7 @@
 from qibolab._core.instruments.emulator.results import (
     _cyclic_results,
     _marginalize_probability,
+    _sampled_measurements,
 )
 from qibolab._core.pulses import Acquisition
 
@@ -65,6 +66,27 @@ def test_cyclic_integration_results_marginalize_probabilities(monkeypatch):
     np.testing.assert_allclose(results[acq1.id], [0.85, 0.0])
 
 
+def test_sampled_measurements_groups_acquisitions_by_sample():
+    sampled = np.array(
+        [
+            [[0, 1], [4, 5]],
+            [[2, 3], [1, 0]],
+        ]
+    )
+    inverse_map = np.array([0, 0, 1])
+
+    measured = _sampled_measurements(
+        sampled=sampled,
+        dims=[2, 3],
+        inverse_map=inverse_map,
+        indices=[0, 1, 1],
+    )
+
+    np.testing.assert_allclose(measured[0], [[0, 0], [1, 1]])
+    np.testing.assert_allclose(measured[1], [[0, 1], [1, 2]])
+    np.testing.assert_allclose(measured[2], [[2, 0], [1, 0]])
+
+
 class _Sequence:
     def __init__(self, channels):
         self._channels = channels