Source code for qblox_scheduler.segments.loops

# Repository: https://gitlab.com/qblox/packages/software/qblox-scheduler
# Licensed according to the LICENSE file on the main branch
#
# Copyright 2025, Qblox B.V.
"""Module containing the segment to a set a parameter."""

from __future__ import annotations

from typing import TYPE_CHECKING

from tqdm.autonotebook import trange

from qblox_scheduler.analysis.helpers import concat_acq_data
from qblox_scheduler.operations.loop_domains import unpack_complex_variables
from qblox_scheduler.segments.segment import Segment

if TYPE_CHECKING:
    from xarray import Dataset

    from qblox_scheduler.device_under_test import QuantumDevice
    from qblox_scheduler.operations.expressions import Expression
    from qblox_scheduler.operations.loop_domains import LinearDomain
    from qblox_scheduler.operations.variables import Variable
    from qblox_scheduler.schedules.schedule import TimeableScheduleBase




class Loop(Segment):
    """Experiment segment that loops other segments over some values."""

    def __init__(self, domains: dict[Variable, LinearDomain], segments: list[Segment]) -> None:
        super().__init__(f"loop segments over {domains}")
        self.data["loop_info"] = {
            "domains": domains,
            "segments": segments,
        }

    @property


    def domains(self) -> dict[Variable, LinearDomain]:
        """Domains to loop over."""
        return self.data["loop_info"]["domains"]


    @property


    def segments(self) -> list[Segment]:
        """Segments to execute."""
        return self.data["loop_info"]["segments"]




    def run(self, device: QuantumDevice, timeout: int = 10) -> Dataset | None:
        """Execute segment on quantum device."""
        if not self.domains:
            return

        domains = unpack_complex_variables(self.domains)

        data_set = None
        nsteps = len(next(iter(domains.values())))
        values = {var: iter(dom.values()) for var, dom in domains.items()}
        for _ in trange(nsteps, unit="step"):
            substitutions: dict[Expression, Expression | int | float | complex] = {
                var: next(values) for var, values in values.items()
            }
            for segment in self.segments:
                sub_segment = segment.substitute(substitutions)
                sub_data_set = sub_segment.run(device, timeout=timeout)
                if sub_data_set is not None:
                    if data_set is None:
                        data_set = sub_data_set
                    else:
                        data_set = concat_acq_data([data_set, sub_data_set])

        return data_set




    def _unrolled_segments(self) -> list[TimeableScheduleBase | Segment]:
        if not self.domains:
            return []

        timeable_schedule_segments = []

        domains = unpack_complex_variables(self.domains)
        nsteps = len(next(iter(domains.values())))
        values = {var: iter(dom.values()) for var, dom in domains.items()}
        for _ in range(nsteps):
            substitutions: dict[Expression, Expression | int | float | complex] = {
                var: next(values) for var, values in values.items()
            }
            for segment in self.segments:
                substituted_segment = segment.substitute(substitutions)
                timeable_schedule_segments += substituted_segment._unrolled_segments()

        return timeable_schedule_segments