Bypass cyclic emulator state indexing

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

@@ -112,6 +112,23 @@ def index(ch: ChannelId, hconfig: HamiltonianConfig) -> int:
     return hconfig.hilbert_space_index(target)
 
 
+def _marginalize_probability(
+    probabilities: NDArray, dims: list[int], indices: Iterable[int]
+) -> NDArray:
+    """Marginalize full-system probabilities to measured Hilbert-space components."""
+    res = []
+    for measurement, index in enumerate(indices):
+        distribution = probabilities[..., measurement, :]
+        leading = distribution.shape[:-1]
+        distribution = distribution.reshape(*leading, *dims)
+        axes = tuple(len(leading) + i for i in range(len(dims)) if i != index)
+        marginalized = distribution.sum(axis=axes)
+        # States outside the ground state are classified as 1.
+        res.append(marginalized[..., 1:].sum(axis=-1))
+
+    return np.stack(res)
+
+
 def select_acquisitions(
     states: list[Operator], acquisitions: Iterable[float], times: NDArray
 ) -> NDArray:
@@ -142,25 +159,15 @@ def _cyclic_results(
     measurement subspaces and applying configured post-processing.
     """
 
-    # Through the entire function state_probs has dimensions:
-    # (*S, M *H_dim)
-    states_computational_idx = np.stack(
-        np.unravel_index(np.arange(state_probs.shape[-1]), hamiltonian.dims)
-    )
-
-    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 marginalize.
     qubit_indices = [
         index(sequence.pulse_channels(ro_id)[0], hamiltonian) for ro_id in acq_id
     ]
-    permuted_states_computational_idx = states_computational_idx[qubit_indices]
-
-    # applying a mask to select for each measurement the states that are outside the computational subspace, which are classified as 1
-    mask = permuted_states_computational_idx >= 1
 
     # res is a (M, *S, ...) array
-    res = np.moveaxis(np.sum(np.where(mask, state_probs, 0), axis=-1), -1, 0)
+    res = _marginalize_probability(state_probs, hamiltonian.dims, qubit_indices)
 
     if options.acquisition_type is AcquisitionType.INTEGRATION:
         res = np.random.normal(res, scale=0.001)

tests/instruments/emulator/test_results.py

@@ -1,5 +1,16 @@
 import numpy as np
 
+from qibolab._core.execution_parameters import (
+    AcquisitionType,
+    AveragingMode,
+    ExecutionParameters,
+)
+from qibolab._core.instruments.emulator.results import (
+    _cyclic_results,
+    _marginalize_probability,
+)
+from qibolab._core.pulses import Acquisition
+
 
 def random_states(space: tuple[int, ...], sweeps: tuple[int, ...] = (), nacq: int = 1):
     dimension = np.prod(space, dtype=int)
@@ -8,3 +19,72 @@ def random_states(space: tuple[int, ...], sweeps: tuple[int, ...] = (), nacq: in
     state = components / np.sqrt((components**2).sum(axis=-1))[..., np.newaxis]
 
     return np.einsum("...i,...j->...ij", state, state)
+
+
+def test_marginalize_probability_preserves_sweep_axes():
+    probabilities = np.arange(1, 49).reshape(2, 2, 12)
+
+    marginalized = _marginalize_probability(probabilities, [2, 3, 2], [1, 2])
+    expected = np.stack(
+        (
+            probabilities[:, 0].reshape(2, 2, 3, 2).sum(axis=(1, 3))[..., 1:].sum(-1),
+            probabilities[:, 1].reshape(2, 2, 3, 2).sum(axis=(1, 2))[..., 1:].sum(-1),
+        )
+    )
+
+    np.testing.assert_allclose(marginalized, expected)
+
+
+def test_cyclic_integration_results_marginalize_probabilities(monkeypatch):
+    monkeypatch.setattr(np.random, "normal", lambda loc, scale: loc)
+    acq0 = Acquisition(duration=1)
+    acq1 = Acquisition(duration=1)
+    sequence = _Sequence(
+        {
+            "0/acquisition": [acq0],
+            "1/acquisition": [acq1],
+        }
+    )
+    probabilities = np.array(
+        [
+            [[0.10, 0.20, 0.30], [0.05, 0.15, 0.20]],
+            [[0.10, 0.20, 0.30], [0.05, 0.15, 0.20]],
+        ]
+    )
+    results = _cyclic_results(
+        state_probs=probabilities.reshape(2, -1),
+        sequence=sequence,
+        hamiltonian=_HamiltonianConfig(dims=[2, 3]),
+        options=ExecutionParameters(
+            acquisition_type=AcquisitionType.INTEGRATION,
+            averaging_mode=AveragingMode.CYCLIC,
+        ),
+    )
+
+    np.testing.assert_allclose(results[acq0.id], [0.40, 0.0])
+    np.testing.assert_allclose(results[acq1.id], [0.85, 0.0])
+
+
+class _Sequence:
+    def __init__(self, channels):
+        self._channels = channels
+        self.channels = list(channels)
+        self._pulse_channels = {
+            event.id: channel
+            for channel, events in channels.items()
+            for event in events
+        }
+
+    def channel(self, channel):
+        return self._channels[channel]
+
+    def pulse_channels(self, pulse_id):
+        return [self._pulse_channels[pulse_id]]
+
+
+class _HamiltonianConfig:
+    def __init__(self, dims):
+        self.dims = dims
+
+    def hilbert_space_index(self, target):
+        return target