# Repository: https://gitlab.com/qblox/packages/software/qblox-scheduler
# Licensed according to the LICENSE file on the main branch
#
# Copyright 2020-2025, Quantify Consortium
# Copyright 2025, Qblox B.V.
"""Helper functions to generate acq_indices."""
from __future__ import annotations
import warnings
from collections import defaultdict, namedtuple
from copy import copy
from dataclasses import asdict, dataclass, field
from typing import TYPE_CHECKING, Union
from qblox_scheduler.enums import BinMode
from qblox_scheduler.helpers.importers import export_python_object_to_path_string
from qblox_scheduler.helpers.schedule import (
_is_acquisition_binned_average_append,
)
from qblox_scheduler.operations.control_flow_library import (
ConditionalOperation,
LoopOperation,
)
from qblox_scheduler.operations.expressions import Expression
from qblox_scheduler.schedules.schedule import (
AcquisitionChannelData,
AcquisitionChannelsData,
TimeableScheduleBase,
)
if TYPE_CHECKING:
from collections.abc import Hashable
from qblox_scheduler.operations.loop_domains import Domain
from qblox_scheduler.operations.operation import Operation
from qblox_scheduler.operations.variables import Variable
[docs]
SchedulableLabel = Union[str, None]
[docs]
FullSchedulableLabel = tuple[SchedulableLabel, ...]
@dataclass
[docs]
class AcquisitionIndices:
"""
The AcquisitionIndices stores data for the backend compiler
for each acquisition schedulable: the acquisition indices, and it's loop structure.
For example, if you have a multi-dimensional with 2, 3 and 4 repetitions,
and you average over the last nested level,
`loop_bin_modes=[APPEND, APPEND, AVERAGE]`, and number the indices if 2*3.
The `loop_bin_modes` element can be `None`, if the node is inside a conditional operation.
"""
[docs]
loop_bin_modes: list[BinMode]
def __getstate__(self) -> dict:
data = asdict(self)
return {
"deserialization_type": export_python_object_to_path_string(self.__class__),
"data": data,
}
def __setstate__(self, state: dict) -> None:
self.acq_index = state["data"]["acq_index"]
self.loop_bin_modes = [BinMode(x) for x in state["data"]["loop_bin_modes"]]
[docs]
SchedulableLabelToAcquisitionIndex = dict[FullSchedulableLabel, AcquisitionIndices]
"""
A mapping from schedulables to an acquisition index.
This mapping helps the backend to figure out which
binned acquisition corresponds to which acquisition index.
Note, it maps the full schedulable label to acquisition indices,
Only defined for binned acquisitions, and backend independent.
For control flows, the `None` in the schedulable label refers to the `body`
of the control flow. This is for future proofing, if control flows were extended
to include maybe multiple suboperations.
"""
@dataclass
[docs]
class LoopData:
"""
Data to contain relevant information from LoopOperation.
`repetitions` is `None` if and only if it's a conditional.
"""
[docs]
domain: dict[Variable, Domain] | None
[docs]
def _evaluate_coords_recursively(
loops: list[LoopData],
evaluated_coords: list[dict],
) -> list[dict]:
"""
Evaluate coords even if there are variables in it.
The accumulator is stored in evaluated_coords, which is returned.
"""
def evaluate_domain(
domain: dict[Variable, Domain],
) -> list[dict[Expression, Expression | int | float | complex]]:
steps = max(current_domain.num_steps for current_domain in domain.values())
# list of dictionaries from variable to a value.
evaluated_variables: list[dict[Expression, Expression | int | float | complex]] = [
{} for _ in range(steps)
]
for variable, current_domain in domain.items():
for i, value in enumerate(current_domain.values()):
assert isinstance(
value, Expression | int | float | complex
) # Currently only these are supported.
evaluated_variables[i][variable] = value
return evaluated_variables
def substitute_coords_variables(
coords: dict, evaluated_variables: dict[Expression, Expression | int | float | complex]
) -> dict:
coords = copy(coords)
for key, value in coords.items():
if isinstance(value, Expression):
coords[key] = value.substitute(evaluated_variables)
return coords
if len(loops) == 0:
return evaluated_coords
current_loop = loops[-1]
if current_loop.domain:
evaluated_coords = [
substitute_coords_variables(current_coords, evaluated_variables)
for evaluated_variables in evaluate_domain(current_loop.domain)
for current_coords in evaluated_coords
]
else:
evaluated_coords = [
copy(item) for _ in range(current_loop.repetitions) for item in evaluated_coords
]
return _evaluate_coords_recursively(loops[:-1], evaluated_coords)
[docs]
def _evaluate_coords(coords: dict, loops: list[LoopData]) -> list[dict]:
"""Evaluate coords even if there are variables in it."""
return _evaluate_coords_recursively(loops, [coords])
[docs]
def _get_loops_with_append_bin_mode_and_all_loop_bin_modes(
coords: dict,
loops: list[LoopData],
append_all_loops: bool,
) -> tuple[list[LoopData], list[BinMode]]:
def is_variable_in_expression(domain: dict[Variable, Domain] | None, coords: dict) -> bool:
if domain is None:
return False
for variable in domain:
for coords_value in coords.values():
if isinstance(coords_value, Expression) and (
(variable == coords_value) or (variable in coords_value)
):
return True
return False
append_loops = []
loop_bin_modes = []
for loop in loops:
if is_variable_in_expression(loop.domain, coords) or append_all_loops:
append_loops.append(loop)
loop_bin_modes.append(BinMode.APPEND)
else:
loop_bin_modes.append(BinMode.AVERAGE)
return append_loops, loop_bin_modes
@dataclass
[docs]
class _AcqIndexDimData:
"""
Data for each acquisition index dimension.
An acquisition index dimension can be shared between multiple acquisition channels.
This dataclass makes sure that two coords which have the same coords will
get the same acquisition index, otherwise they get unique acquisition indices.
For example a coords `{"amp": 0}` and coords `{"amp": 1}` will get two different
acquisition index generated, because they share at least one coords key `"amp"`,
but can be on different acquisition channels.
"""
[docs]
_next_acq_index: int = 0
[docs]
_coords_to_acq_index: dict[frozenset, int] = field(default_factory=dict)
[docs]
_coords_to_acq_channels: defaultdict[frozenset, set[Hashable]] = field(
default_factory=lambda: defaultdict(set)
)
[docs]
def update_and_get_acq_index(self, acq_channel: Hashable, coords: dict) -> int:
coords_frozenset = frozenset(coords.items())
if (
i := self._coords_to_acq_index.get(coords_frozenset)
) is not None and acq_channel not in self._coords_to_acq_channels[coords_frozenset]:
return i
else:
i = self._next_acq_index
self._next_acq_index += 1
if acq_channel not in self._coords_to_acq_channels[coords_frozenset]:
# Only record the first occurrence of a given `coords` value. If the same channel
# sees the same coords again, we still generate a new acquisition index
# for that repeated acquisition, but we do not overwrite `coords -> acq_index`.
# This preserves the first index as the canonical index for these coords, so that
# another acquisition channel with the same coords reuses that same index.
#
# Example:
# ch0: {"amp": 0} -> 0
# ch0: {"amp": 0} -> 1 (repeat on same channel; mapping is not overwritten)
# ch1: {"amp": 0} -> 0 (reuses the first canonical index across channels)
self._coords_to_acq_index[coords_frozenset] = i
self._coords_to_acq_channels[coords_frozenset].add(acq_channel)
return i
[docs]
def _generate_acq_channels_data_binned_append(
acq_channel_data: AcquisitionChannelData,
acq_channel: Hashable,
schedulable_label_to_acq_index: SchedulableLabelToAcquisitionIndex,
full_schedulable_label: FullSchedulableLabel,
loops: list[LoopData],
coords: dict,
append_all_loops: bool,
acq_index_dim_data: dict[str, _AcqIndexDimData],
) -> None:
"""
Generates the acquisition channel data, and updates acq_channel_data,
and updates schedulable_label_to_acq_index for average bin mode.
"""
evaluated_coords: list[dict] = []
if len(loops) == 0:
evaluated_coords = [coords]
loop_bin_modes = []
else:
append_loops, loop_bin_modes = _get_loops_with_append_bin_mode_and_all_loop_bin_modes(
coords, loops, append_all_loops
)
evaluated_coords = _evaluate_coords(coords, append_loops)
acq_index_dim_name = acq_channel_data.acq_index_dim_name
if acq_index_dim_name not in acq_index_dim_data:
acq_index_dim_data[acq_index_dim_name] = _AcqIndexDimData()
new_acq_indices = []
for evaluated_coords_i in evaluated_coords:
new_acq_index = acq_index_dim_data[acq_index_dim_name].update_and_get_acq_index(
acq_channel, evaluated_coords_i
)
new_acq_indices.append(new_acq_index)
schedulable_label_to_acq_index[full_schedulable_label] = AcquisitionIndices(
new_acq_indices, loop_bin_modes
)
acq_channel_data.coords |= dict(zip(new_acq_indices, evaluated_coords, strict=True))
[docs]
def _validate_trace_protocol(
acq_channel: Hashable,
acq_channels_data: AcquisitionChannelsData,
loops: list[LoopData], # noqa: ARG001
) -> None:
if acq_channel in acq_channels_data:
raise ValueError(
f"Multiple acquisitions found for acq_channel '{acq_channel}' "
f"which has a trace acquisition. "
f"Only one trace acquisition is allowed for each acq_channel.",
)
[docs]
def _generate_acq_channels_data_for_protocol(
acq_info: dict,
acq_channels_data: AcquisitionChannelsData,
schedulable_label_to_acq_index: SchedulableLabelToAcquisitionIndex,
full_schedulable_label: FullSchedulableLabel,
loops: list[LoopData],
acq_index_dim_names: dict[Hashable, str],
acq_index_dim_data: dict[str, _AcqIndexDimData],
) -> None:
"""
Generates the acquisition channel data, and updates acq_channel_data,
and updates schedulable_label_to_acq_index.
"""
acq_channel: Hashable = acq_info["acq_channel"]
protocol: str = acq_info["protocol"]
bin_mode: BinMode = acq_info["bin_mode"]
coords: dict = acq_info["coords"] or {}
if (acq_channel_data := acq_channels_data.get(acq_channel, None)) is not None:
if acq_channel_data.protocol != protocol:
raise ValueError(
f"Found different acquisition protocols "
f"('{acq_channel_data.protocol}' and '{protocol}') "
f"for acq_channel '{acq_channel}'. "
f"Make sure there is only one protocol for each acq_channel.",
)
if acq_channel_data.bin_mode != bin_mode:
raise ValueError(
f"Found different bin modes "
f"('{acq_channel_data.bin_mode}' and '{bin_mode}') "
f"for acq_channel '{acq_channel}'. "
f"Make sure there is only one bin mode for each acq_channel.",
)
if _is_acquisition_binned_average_append(protocol, bin_mode) or (
protocol == "TimetagTrace" and bin_mode == BinMode.APPEND
):
if acq_channel not in acq_channels_data:
acq_channels_data[acq_channel] = AcquisitionChannelData(
acq_index_dim_name=acq_index_dim_names[acq_channel],
protocol=protocol,
bin_mode=bin_mode,
coords={},
)
_generate_acq_channels_data_binned_append(
acq_channel_data=acq_channels_data[acq_channel],
acq_channel=acq_channel,
schedulable_label_to_acq_index=schedulable_label_to_acq_index,
full_schedulable_label=full_schedulable_label,
loops=loops,
coords=coords,
append_all_loops=(bin_mode == BinMode.APPEND),
acq_index_dim_data=acq_index_dim_data,
)
elif protocol == "Trace" and bin_mode in (BinMode.AVERAGE, BinMode.FIRST):
_validate_trace_protocol(
acq_channel=acq_channel,
acq_channels_data=acq_channels_data,
loops=loops,
)
acq_channels_data[acq_channel] = AcquisitionChannelData(
acq_index_dim_name=("acq_index_" + str(acq_channel)),
protocol=protocol,
bin_mode=bin_mode,
coords=coords,
)
elif protocol == "TriggerCount" and bin_mode == BinMode.DISTRIBUTION:
acq_channels_data[acq_channel] = AcquisitionChannelData(
acq_index_dim_name=("acq_index_" + str(acq_channel)),
protocol=protocol,
bin_mode=bin_mode,
coords=coords,
)
else:
raise ValueError(
f"Unsupported acquisition protocol '{protocol}' with bin mode '{bin_mode}' "
f"on acq_channel '{acq_channel}'.",
)
[docs]
def _generate_acq_channels_data(
operation: TimeableScheduleBase | Operation,
acq_channels_data: AcquisitionChannelsData,
schedulable_label_to_acq_index: SchedulableLabelToAcquisitionIndex,
full_schedulable_label: FullSchedulableLabel,
loops: list[LoopData],
acq_index_dim_names: dict[Hashable, str],
acq_index_dim_data: dict[str, _AcqIndexDimData],
) -> None:
"""
Adds mappings to acq_channels_data and schedulable_label_to_acq_index;
these are the output arguments; the others are input arguments.
It will also generate the acq_index.
"""
if isinstance(operation, TimeableScheduleBase):
sorted_schedulables = sorted(operation.schedulables.values(), key=lambda s: s["abs_time"])
for schedulable in sorted_schedulables:
schedulable_label = schedulable["name"]
new_full_schedulable_label = full_schedulable_label + (schedulable_label,)
inner_operation = operation.operations[schedulable["operation_id"]]
_generate_acq_channels_data(
operation=inner_operation,
acq_channels_data=acq_channels_data,
schedulable_label_to_acq_index=schedulable_label_to_acq_index,
full_schedulable_label=new_full_schedulable_label,
loops=loops,
acq_index_dim_names=acq_index_dim_names,
acq_index_dim_data=acq_index_dim_data,
)
elif isinstance(operation, LoopOperation):
# For control flows, `None` signifies we refer to the `body` of the control flow.
new_full_schedulable_label: FullSchedulableLabel = full_schedulable_label + (None,)
repetitions: int = operation.data["control_flow_info"]["repetitions"]
domain: dict[Variable, Domain] | None = operation.data["control_flow_info"]["domain"]
new_loops: list[LoopData] = loops + [LoopData(repetitions, domain)]
_generate_acq_channels_data(
operation=operation.body,
acq_channels_data=acq_channels_data,
schedulable_label_to_acq_index=schedulable_label_to_acq_index,
full_schedulable_label=new_full_schedulable_label,
loops=new_loops,
acq_index_dim_names=acq_index_dim_names,
acq_index_dim_data=acq_index_dim_data,
)
elif isinstance(operation, ConditionalOperation):
# For control flows, `None` signifies we refer to the `body` of the control flow.
new_full_schedulable_label = full_schedulable_label + (None,)
_generate_acq_channels_data(
operation=operation.body,
acq_channels_data=acq_channels_data,
schedulable_label_to_acq_index=schedulable_label_to_acq_index,
full_schedulable_label=new_full_schedulable_label,
loops=loops,
acq_index_dim_names=acq_index_dim_names,
acq_index_dim_data=acq_index_dim_data,
)
elif operation.valid_acquisition:
_generate_acq_channels_data_for_protocol(
acq_info=operation.data["acquisition_info"],
acq_channels_data=acq_channels_data,
schedulable_label_to_acq_index=schedulable_label_to_acq_index,
full_schedulable_label=full_schedulable_label,
loops=loops,
acq_index_dim_names=acq_index_dim_names,
acq_index_dim_data=acq_index_dim_data,
)
[docs]
[docs]
[docs]
_AcqChannelsCoordsKeys = namedtuple("_AcqChannelsCoordsKeys", ("acq_channels", "coords_keys"))
[docs]
def _get_acq_channels_coords_keys(
operation: TimeableScheduleBase | Operation,
acq_channels_coords_keys: list[_AcqChannelsCoordsKeys] | None = None,
) -> list[_AcqChannelsCoordsKeys]:
"""
Collects sets of acquisition channels in acq_channels_coords_keys.
Two acquisition channels will be at the same index in `acq_channels_coords_keys` if and only if
any acquisition within that acquisition channel share at least one coords key.
For example if two acquisitions on `acq_channel="ch0"` and `acq_channel="ch1"` have
`coords={"amp": 0}` and `coords={"amp": 1, "freq": 2}` respectively,
then they share the same `acq_channels_coords_key` index,
`acq_channels_coords_key=[{"ch0", "ch1"}, {"amp", "freq"}]`.
A different operation with different acquisition channel will be a different element:
`acq_channels_coords_key=[{"ch0", "ch1"}, {"amp", "freq"}, {"ch2", {"rep"}]`.
This function only applies to binned acquisitions.
"""
# Initializing the accumulator.
if acq_channels_coords_keys is None:
acq_channels_coords_keys = []
if isinstance(operation, TimeableScheduleBase):
for inner_operation in operation.operations.values():
_get_acq_channels_coords_keys(inner_operation, acq_channels_coords_keys)
elif isinstance(operation, (LoopOperation, ConditionalOperation)):
_get_acq_channels_coords_keys(operation.body, acq_channels_coords_keys)
elif operation.valid_acquisition:
acq_info = operation.data["acquisition_info"]
acq_channel: Hashable = acq_info["acq_channel"]
protocol: str = acq_info["protocol"]
bin_mode: BinMode = acq_info["bin_mode"]
coords: dict = acq_info["coords"] or {}
if _is_acquisition_binned_average_append(protocol, bin_mode) or (
protocol == "TimetagTrace" and bin_mode == BinMode.APPEND
):
for i in range(len(acq_channels_coords_keys)):
existing_acq_channels = acq_channels_coords_keys[i].acq_channels
existing_coords_keys = acq_channels_coords_keys[i].coords_keys
if acq_channel in existing_acq_channels or any(
k in existing_coords_keys for k in coords.keys()
):
existing_acq_channels.add(acq_channel)
existing_coords_keys |= set(coords.keys())
break
else:
acq_channels_coords_keys.append(
_AcqChannelsCoordsKeys({acq_channel}, set(coords.keys()))
)
else:
existing_coords_keys = {k for ac in acq_channels_coords_keys for k in ac.coords_keys}
if any(k in existing_coords_keys for k in coords.keys()):
warnings.warn(
f"A coords key is shared between "
f"a binned acquisition and acquisition channel '{acq_channel}'. "
f"Sharing a coords key between a binned and non-binned acquisition "
f"is not allowed."
)
return acq_channels_coords_keys
[docs]
def _generate_acq_index_dim_names(schedule: TimeableScheduleBase) -> dict[Hashable, str]:
"""
Generates an acquisition index dimension name for each acquisition channel.
Two acquisition channels share the same acquisition index dimension name if and only if
any acquisition within that acquisition channel share at least one coords key.
For example if three acquisitions on `acq_channel="ch0"`, `acq_channel="ch1"`
and `acq_channel="ch2"` have `coords={"amp": 0}, `coords={"amp": 1, "freq": 2}` and
`coords={"rep": 0}` respectively, then the generated acquisition dimension index names are
`"acq_index_ch0_ch1"` (for `"ch0"` and `"ch1"`), and `"acq_index_ch2"` (for `"ch2"`).
This function only applies to binned acquisitions.
"""
acq_channels_coords_keys = _get_acq_channels_coords_keys(schedule)
acq_index_dim_names: dict[Hashable, str] = {}
for acq_channels, _coords in acq_channels_coords_keys:
acq_dim_name = "acq_index_" + "_".join([str(acq_channel) for acq_channel in acq_channels])
for acq_channel in acq_channels:
acq_index_dim_names[acq_channel] = acq_dim_name
return acq_index_dim_names
[docs]
def generate_acq_channels_data(
schedule: TimeableScheduleBase,
) -> tuple[AcquisitionChannelsData, SchedulableLabelToAcquisitionIndex]:
"""
Generate acq_index for every schedulable,
and validate schedule regarding the acquisitions.
This function generates the ``AcquisitionChannelData`` for every ``acq_channel``,
and the ``SchedulableLabelToAcquisitionIndex``. It assumes the schedule is device-level.
"""
acq_channels_data: AcquisitionChannelsData = {}
schedulable_label_to_acq_index: SchedulableLabelToAcquisitionIndex = {}
acq_index_dim_names = _generate_acq_index_dim_names(schedule)
acq_index_dim_data = {}
_generate_acq_channels_data(
schedule,
acq_channels_data,
schedulable_label_to_acq_index,
full_schedulable_label=(),
loops=[],
acq_index_dim_names=acq_index_dim_names,
acq_index_dim_data=acq_index_dim_data,
)
return acq_channels_data, schedulable_label_to_acq_index
