diff --git a/src/qibolab/_core/instruments/emulator/results.py b/src/qibolab/_core/instruments/emulator/results.py index 4ff7ce07c..fb1f38b8f 100644 --- a/src/qibolab/_core/instruments/emulator/results.py +++ b/src/qibolab/_core/instruments/emulator/results.py @@ -112,6 +112,24 @@ 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. This sum is done + # before stacking because measured subsystems may have different dimensions. + 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 +160,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) diff --git a/tests/instruments/emulator/test_results.py b/tests/instruments/emulator/test_results.py index 8e564f2ec..ef9580174 100644 --- a/tests/instruments/emulator/test_results.py +++ b/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