diff --git a/src/qibolab/_core/instruments/emulator/results.py b/src/qibolab/_core/instruments/emulator/results.py index fb1f38b8f..4220e3ff8 100644 --- a/src/qibolab/_core/instruments/emulator/results.py +++ b/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: diff --git a/tests/instruments/emulator/test_results.py b/tests/instruments/emulator/test_results.py index ef9580174..290ed9aa5 100644 --- a/tests/instruments/emulator/test_results.py +++ b/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