Source code for qcodes.dataset.dond.do_nd

from __future__ import annotations

import itertools
import logging
import time
from import Mapping, Sequence
from contextlib import ExitStack
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Callable, Union, cast

import numpy as np
from opentelemetry import trace
from import tqdm
from typing_extensions import TypedDict

from qcodes import config
from qcodes.dataset.descriptions.detect_shapes import detect_shape_of_measurement
from qcodes.dataset.dond.do_nd_utils import (
from qcodes.dataset.measurements import Measurement
from qcodes.dataset.threading import (
from qcodes.parameters import ParameterBase

from .sweeps import AbstractSweep, TogetherSweep

LOG = logging.getLogger(__name__)

    from qcodes.dataset.descriptions.versioning.rundescribertypes import Shapes
    from qcodes.dataset.dond.do_nd_utils import (
    from qcodes.dataset.experiment_container import Experiment

SweepVarType = Any

TRACER = trace.get_tracer(__name__)

class ParameterGroup(TypedDict):
    params: tuple[ParamMeasT, ...]
    meas_name: str

class ParameterSetEvent:
    parameter: ParameterBase
    new_value: SweepVarType
    should_set: bool
    delay: float
    actions: ActionsT
    get_after_set: bool

class _Sweeper:
    def __init__(
        sweeps: Sequence[AbstractSweep | TogetherSweep],
        additional_setpoints: Sequence[ParameterBase],
        self._additional_setpoints = additional_setpoints
        self._sweeps = sweeps
        self._setpoints = self._make_setpoints_tuples()
        self._setpoints_dict = self._make_setpoints_dict()
        self._shape = self._make_shape(sweeps, additional_setpoints)
        self._iter_index = 0

    def setpoints_dict(self) -> dict[str, list[Any]]:
        return self._setpoints_dict

    def _make_setpoints_tuples(
    ) -> tuple[tuple[tuple[SweepVarType, ...] | SweepVarType, ...], ...]:
        sweeps = tuple(sweep.get_setpoints() for sweep in self._sweeps)
        return cast(
            tuple[tuple[Union[tuple[SweepVarType, ...], SweepVarType], ...], ...],

    def _make_single_point_setpoints_dict(self, index: int) -> dict[str, SweepVarType]:
        setpoint_dict = {}
        values = self._setpoints[index]
        for sweep, subvalues in zip(self._sweeps, values):
            if isinstance(sweep, TogetherSweep):
                for individual_sweep, value in zip(sweep.sweeps, subvalues):
                    setpoint_dict[individual_sweep.param.full_name] = value
                setpoint_dict[sweep.param.full_name] = subvalues
        return setpoint_dict

    def _make_setpoints_dict(self) -> dict[str, list[Any]]:
        setpoint_dict: dict[str, list[SweepVarType]] = {}

        for sweep in self._sweeps:
            if isinstance(sweep, TogetherSweep):
                for individual_sweep in sweep.sweeps:
                    setpoint_dict[individual_sweep.param.full_name] = []
                setpoint_dict[sweep.param.full_name] = []

        for setpoint_tuples in self._setpoints:
            for sweep, values in zip(self._sweeps, setpoint_tuples):
                if isinstance(sweep, TogetherSweep):
                    for individual_sweep, individual_value in zip(sweep.sweeps, values):
        return setpoint_dict

    def all_sweeps(self) -> tuple[AbstractSweep, ...]:
        sweeps: list[AbstractSweep] = []
        for sweep in self._sweeps:
            if isinstance(sweep, TogetherSweep):
        return tuple(sweeps)

    def all_setpoint_params(self) -> tuple[ParameterBase, ...]:
        return tuple(sweep.param for sweep in self.all_sweeps) + tuple(

    def sweep_groupes(self) -> tuple[tuple[ParameterBase, ...], ...]:
        These are all the combinations of setpoints we consider
        valid for setpoints in a dataset. A dataset must depend on
        at least one parameter from each dimension of the dond.
        For dimensions that uses TogetherSweep they may depend
        on one or more of the parameters of that sweep.
        param_tuple_list: list[tuple[ParameterBase, ...]] = []
        for sweep in self._sweeps:
            if isinstance(sweep, TogetherSweep):
                    tuple(sub_sweep.param for sub_sweep in sweep.sweeps)

            [(setpoint,) for setpoint in self._additional_setpoints]

        # in param_tuple_list there is a tuple of possible setpoints for each
        # dim in the dond. For regular sweeps this is a 1 tuple but for
        # a TogetherSweep it is of length of the number of parameters.

        # now we expand to a list of setpoints in a TogetherSweep
        # to a list of all possible combinations of these.

        expanded_param_tuples = tuple(
                    itertools.combinations(param_tuple, j + 1)
                    for j in range(len(param_tuple))
            for param_tuple in param_tuple_list

        # next we generate all valid combinations of picking one parameter from each
        # dimension in the setpoints.
        setpoint_combinations = itertools.product(*expanded_param_tuples)

        setpoint_combinations_expanded = tuple(
            for setpoint_combination in setpoint_combinations

        return setpoint_combinations_expanded

    def _make_shape(
        sweeps: Sequence[AbstractSweep | TogetherSweep],
        addtional_setpoints: Sequence[ParameterBase],
    ) -> tuple[int, ...]:
        loop_shape = tuple(sweep.num_points for sweep in sweeps) + tuple(
            1 for _ in addtional_setpoints
        return loop_shape

    def shape(self) -> tuple[int, ...]:
        return self._shape

    def __getitem__(self, index: int) -> tuple[ParameterSetEvent, ...]:
        setpoints = self._make_single_point_setpoints_dict(index)

        if index == 0:
            previous_setpoints: dict[str, SweepVarType | None] = {}
            for key in setpoints.keys():
                previous_setpoints[key] = None
            previous_setpoints = self._make_single_point_setpoints_dict(index - 1)

        parameter_set_events = []

        for sweep in self.all_sweeps:
            new_value = setpoints[sweep.param.full_name]
            old_value = previous_setpoints[sweep.param.full_name]
            if old_value is None or old_value != new_value:
                should_set = True
                should_set = False
            event = ParameterSetEvent(
        return tuple(parameter_set_events)

    def __len__(self) -> int:
        return int(

    def __iter__(self) -> _Sweeper:
        return self

    def __next__(self) -> tuple[ParameterSetEvent, ...]:
        if self._iter_index < len(self):
            return_val = self[self._iter_index]
            self._iter_index += 1
            return return_val
            raise StopIteration

class _Measurements:
    def __init__(
        sweeper: _Sweeper,
        measurement_name: str | Sequence[str],
        params_meas: Sequence[ParamMeasT | Sequence[ParamMeasT]],
        enter_actions: ActionsT,
        exit_actions: ActionsT,
        experiments: Experiment | Sequence[Experiment] | None,
        write_period: float | None,
        log_info: str | None,
        dataset_dependencies: Mapping[str, Sequence[ParamMeasT]] | None = None,
        self._sweeper = sweeper
        self._enter_actions = enter_actions
        self._exit_actions = exit_actions
        self._write_period = write_period
        self._log_info = log_info
        if log_info is not None:
            self._extra_log_info = log_info
            self._extra_log_info = "Using 'qcodes.dataset.dond'"

        ) = self._extract_parameters_by_type_and_group(params_meas)

        self._shapes = self._get_shapes()

        if dataset_dependencies and len(grouped_parameters) > 1:
            raise ValueError(
                "Measured parameters have been grouped both in input "
                "and given in dataset dependencies. This is not supported, "
                "group measurement parameters either in input or in "
                "dataset dependencies."

        if dataset_dependencies is None:
            self._groups = self._create_groups_from_grouped_parameters(
                grouped_parameters, experiments, measurement_name
        elif dataset_dependencies:
                dataset_dependencies, measurement_name
            dataset_dependencies_split = self._split_dateset_dependencies(
            self._groups = self._create_groups_from_dataset_dependencies(

    def measured_all(self) -> tuple[ParamMeasT, ...]:
        return self._measured_all

    def measured_parameters(self) -> tuple[ParameterBase, ...]:
        return self._measured_parameters

    def shapes(self) -> Shapes | None:
        return self._shapes

    def groups(self) -> tuple[_SweepMeasGroup, ...]:
        return self._groups

    def _create_measurement_names(
        measurement_name: str | Sequence[str], n_names_required: int
    ) -> tuple[str, ...]:
        if isinstance(measurement_name, str):
            return (measurement_name,) * n_names_required
            if len(measurement_name) != n_names_required:
                raise ValueError(
                    f"Got {len(measurement_name)} measurement names "
                    f"but should create {n_names_required} dataset(s)."
            return tuple(measurement_name)

    def _extract_parameters_by_type_and_group(
        params_meas: Sequence[ParamMeasT | Sequence[ParamMeasT]],
    ) -> tuple[
        tuple[ParamMeasT, ...],
        tuple[tuple[ParamMeasT, ...], ...],
        tuple[ParameterBase, ...],
        measured_parameters: list[ParameterBase] = []
        measured_all: list[ParamMeasT] = []
        single_group: list[ParamMeasT] = []
        multi_group: list[tuple[ParamMeasT, ...]] = []
        grouped_parameters: tuple[tuple[ParamMeasT, ...], ...]
        for param in params_meas:
            if not isinstance(param, Sequence):
                if isinstance(param, ParameterBase):
            elif not isinstance(param, str):
                for nested_param in param:
                    if isinstance(nested_param, ParameterBase):

        if single_group and multi_group:
            raise ValueError(
                f"Got both grouped and non grouped "
                f"parameters to measure in "
                f"{params_meas}. This is not supported."

        if single_group:
            grouped_parameters = (tuple(single_group),)
        elif multi_group:
            grouped_parameters = tuple(multi_group)
            raise ValueError("No parameters to measure supplied")

        return tuple(measured_all), grouped_parameters, tuple(measured_parameters)

    def _get_shapes(self) -> Shapes | None:
            shapes: Shapes | None = detect_shape_of_measurement(
                self.measured_parameters, self._sweeper.shape
            LOG.debug("Detected shapes to be %s", shapes)
        except TypeError:
                f"Could not detect shape of {self.measured_parameters} "
                f"falling back to unknown shape."
            shapes = None
        return shapes

    def _get_experiments(
        experiments: Experiment | Sequence[Experiment] | None,
        n_experiments_required: int,
    ) -> Sequence[Experiment | None]:
        if not isinstance(experiments, Sequence):
            experiments_internal: Sequence[Experiment | None] = [
            ] * n_experiments_required
            experiments_internal = experiments

        if len(experiments_internal) != n_experiments_required:
            raise ValueError(
                f"Inconsistent number of "
                f"datasets and experiments, "
                f"got {n_experiments_required} and {len(experiments_internal)}."
        return experiments_internal

    def _create_groups_from_grouped_parameters(
        grouped_parameters: tuple[tuple[ParamMeasT, ...], ...],
        experiments: Experiment | Sequence[Experiment] | None,
        meas_names: str | Sequence[str],
    ) -> tuple[_SweepMeasGroup, ...]:
        setpoints = self._sweeper.all_setpoint_params

        groups = []
        m_group: Sequence[ParamMeasT]

        experiments_internal = self._get_experiments(
            experiments, len(grouped_parameters)
        measurement_names = self._create_measurement_names(
            meas_names, len(grouped_parameters)

        for m_group, experiment, meas_name in zip(
            meas_ctx = self._create_measurement_ctx_manager(
                experiment, meas_name, setpoints, tuple(m_group)
            s_m_group = _SweepMeasGroup(setpoints, tuple(m_group), meas_ctx)
        return tuple(groups)

    def _create_groups_from_dataset_dependencies(
        dataset_dependencies: dict[
            str, tuple[Sequence[ParameterBase], Sequence[ParamMeasT]]
        all_measured_parameters: tuple[ParameterBase, ...],
        experiments: Experiment | Sequence[Experiment] | None,
        meas_names: str | Sequence[str],
    ) -> tuple[_SweepMeasGroup, ...]:
        potential_setpoint_groups = self._sweeper.sweep_groupes

        experiments_internal = self._get_experiments(
            experiments, len(dataset_dependencies)
        if meas_names == "":
            meas_names = tuple(dataset_dependencies.keys())

        measurement_names = self._create_measurement_names(
            meas_names, len(dataset_dependencies)

        all_dataset_dependencies_meas_parameters = tuple(
                output[1] for output in dataset_dependencies.values()

        for meas_param in all_measured_parameters:
            if meas_param not in all_dataset_dependencies_meas_parameters:
                raise ValueError(
                    f"Parameter {meas_param} is measured but not added "
                    f"to any dataset in dataset_dependencies."

        groups = []
        for experiment, meas_name in zip(
            (sp_group, m_group) = dataset_dependencies[meas_name]
            if tuple(sp_group) not in potential_setpoint_groups:
                raise ValueError(
                    f"dataset_dependencies contains {sp_group} "
                    f"which is not among the expected groups of setpoints "

                f"creating context manager for setpoints"
                f" {sp_group} and measurement parameters {m_group}"
            meas_ctx = self._create_measurement_ctx_manager(
                experiment, meas_name, tuple(sp_group), tuple(m_group)
            s_m_group = _SweepMeasGroup(tuple(sp_group), tuple(m_group), meas_ctx)
        return tuple(groups)

    def _create_measurement_ctx_manager(
        experiment: Experiment | None,
        measurement_name: str,
        sweep_parameters: Sequence[ParameterBase],
        measure_parameters: Sequence[ParamMeasT],
    ) -> Measurement:
        meas = Measurement(name=measurement_name, exp=experiment)
        _register_parameters(meas, sweep_parameters)
        meas._extra_log_info = self._extra_log_info
        _set_write_period(meas, self._write_period)
        _register_actions(meas, self._enter_actions, self._exit_actions)
        return meas

    def _split_dateset_dependencies(
        dataset_dependencies: Mapping[str, Sequence[ParamMeasT]],
    ) -> dict[str, tuple[Sequence[ParameterBase], Sequence[ParamMeasT]]]:
        # split measured parameters from setpoint parameters using param_meas
        dataset_dependencies_split: dict[
            str, tuple[Sequence[ParameterBase], Sequence[ParamMeasT]]
        ] = {}
        for name, dataset_parameters in dataset_dependencies.items():
            meas_parameters = tuple(
                param for param in dataset_parameters if param in self.measured_all
            setpoint_parameters = cast(
                    for param in dataset_parameters
                    if param not in self.measured_all
            dataset_dependencies_split[name] = (
        return dataset_dependencies_split

# idealy we would want this to be frozen but then postinit
# cannot calculate all the parameters
class _SweepMeasGroup:
    sweep_parameters: tuple[ParameterBase, ...]
    measure_parameters: tuple[ParamMeasT, ...]
    measurement_cxt: Measurement

    def __post_init__(self) -> None:
        meas_parameters = tuple(
            a for a in self.measure_parameters if isinstance(a, ParameterBase)
        self._parameters = self.sweep_parameters + meas_parameters

    def parameters(self) -> tuple[ParameterBase, ...]:
        return self._parameters

[docs] @TRACER.start_as_current_span("qcodes.dataset.dond") def dond( *params: AbstractSweep | TogetherSweep | ParamMeasT | Sequence[ParamMeasT], write_period: float | None = None, measurement_name: str | Sequence[str] = "", exp: Experiment | Sequence[Experiment] | None = None, enter_actions: ActionsT = (), exit_actions: ActionsT = (), do_plot: bool | None = None, show_progress: bool | None = None, use_threads: bool | None = None, additional_setpoints: Sequence[ParameterBase] = tuple(), log_info: str | None = None, break_condition: BreakConditionT | None = None, dataset_dependencies: Mapping[str, Sequence[ParamMeasT]] | None = None, in_memory_cache: bool | None = None, ) -> AxesTupleListWithDataSet | MultiAxesTupleListWithDataSet: """ Perform n-dimentional scan from slowest (first) to the fastest (last), to measure m measurement parameters. The dimensions should be specified as sweep objects, and after them the parameters to measure should be passed. Args: params: Instances of n sweep classes and m measurement parameters, e.g. if linear sweep is considered: .. code-block:: LinSweep(param_set_1, start_1, stop_1, num_points_1, delay_1), ..., LinSweep(param_set_n, start_n, stop_n, num_points_n, delay_n), param_meas_1, param_meas_2, ..., param_meas_m If multiple DataSets creation is needed, measurement parameters should be grouped, so one dataset will be created for each group. e.g.: .. code-block:: LinSweep(param_set_1, start_1, stop_1, num_points_1, delay_1), ..., LinSweep(param_set_n, start_n, stop_n, num_points_n, delay_n), [param_meas_1, param_meas_2], ..., [param_meas_m] If you want to sweep multiple parameters together. .. code-block:: TogetherSweep(LinSweep(param_set_1, start_1, stop_1, num_points, delay_1), LinSweep(param_set_2, start_2, stop_2, num_points, delay_2)) param_meas_1, param_meas_2, ..., param_meas_m write_period: The time after which the data is actually written to the database. measurement_name: Name(s) of the measurement. This will be passed down to the dataset produced by the measurement. If not given, a default value of 'results' is used for the dataset. If more than one is given, each dataset will have an individual name. exp: The experiment to use for this measurement. If you create multiple measurements using groups you may also supply multiple experiments. enter_actions: A list of functions taking no arguments that will be called before the measurements start. exit_actions: A list of functions taking no arguments that will be called after the measurements ends. do_plot: should png and pdf versions of the images be saved and plots are shown after the run. If None the setting will be read from ``qcodesrc.json`` show_progress: should a progress bar be displayed during the measurement. If None the setting will be read from ``qcodesrc.json`` use_threads: If True, measurements from each instrument will be done on separate threads. If you are measuring from several instruments this may give a significant speedup. additional_setpoints: A list of setpoint parameters to be registered in the measurement but not scanned/swept-over. log_info: Message that is logged during the measurement. If None a default message is used. break_condition: Callable that takes no arguments. If returned True, measurement is interrupted. dataset_dependencies: Optionally describe that measured datasets only depend on a subset of the setpoint parameters. Given as a mapping from measurement names to Sequence of Parameters. Note that a dataset must depend on at least one parameter from each dimension but can depend on one or more parameters from a dimension sweeped with a TogetherSweep. in_memory_cache: Should a cache of the data be kept available in memory for faster plotting and exporting. Useful to disable if the data is very large in order to save on memory consumption. If ``None``, the value for this will be read from ``qcodesrc.json`` config file. Returns: A tuple of QCoDeS DataSet, Matplotlib axis, Matplotlib colorbar. If more than one group of measurement parameters is supplied, the output will be a tuple of tuple(QCoDeS DataSet), tuple(Matplotlib axis), tuple(Matplotlib colorbar), in which each element of each sub-tuple belongs to one group, and the order of elements is the order of the supplied groups. """ if do_plot is None: do_plot = cast(bool, config.dataset.dond_plot) if show_progress is None: show_progress = config.dataset.dond_show_progress sweep_instances, params_meas = _parse_dond_arguments(*params) sweeper = _Sweeper(sweep_instances, additional_setpoints) measurements = _Measurements( sweeper, measurement_name, params_meas, enter_actions, exit_actions, exp, write_period, log_info, dataset_dependencies, ) "Starting a doNd with scan with\n setpoints: %s,\n measuring: %s", sweeper.all_setpoint_params, measurements.measured_all, ) LOG.debug( "dond has been grouped into the following datasets:\n%s", measurements.groups, ) datasets = [] plots_axes = [] plots_colorbar = [] if use_threads is None: use_threads = config.dataset.use_threads params_meas_caller = ( ThreadPoolParamsCaller(*measurements.measured_all) if use_threads else SequentialParamsCaller(*measurements.measured_all) ) datasavers = [] interrupted: Callable[ # noqa E731 [], KeyboardInterrupt | BreakConditionInterrupt | None ] = lambda: None try: with catch_interrupts() as interrupted, ExitStack() as stack, params_meas_caller as call_params_meas: datasavers = [ stack.enter_context( ) for group in measurements.groups ] additional_setpoints_data = process_params_meas(additional_setpoints) for set_events in tqdm(sweeper, disable=not show_progress): LOG.debug("Processing set events: %s", set_events) results: dict[ParameterBase, Any] = {} for set_event in set_events: if set_event.should_set: set_event.parameter(set_event.new_value) for act in set_event.actions: act() time.sleep(set_event.delay) if set_event.get_after_set: results[set_event.parameter] = set_event.parameter() else: results[set_event.parameter] = set_event.new_value meas_value_pair = call_params_meas() for meas_param, value in meas_value_pair: results[meas_param] = value for datasaver, group in zip(datasavers, measurements.groups): filtered_results_list = [ (param, value) for param, value in results.items() if param in group.parameters ] datasaver.add_result( *filtered_results_list, *additional_setpoints_data, ) if callable(break_condition): if break_condition(): raise BreakConditionInterrupt("Break condition was met.") finally: for datasaver in datasavers: ds, plot_axis, plot_color = _handle_plotting( datasaver.dataset, do_plot, interrupted() ) datasets.append(ds) plots_axes.append(plot_axis) plots_colorbar.append(plot_color) if len(measurements.groups) == 1: return datasets[0], plots_axes[0], plots_colorbar[0] else: return tuple(datasets), tuple(plots_axes), tuple(plots_colorbar)
def _validate_dataset_dependencies_and_names( dataset_dependencies: Mapping[str, Sequence[ParamMeasT]] | None, measurement_name: str | Sequence[str], ) -> None: if dataset_dependencies is not None and measurement_name != "": if isinstance(measurement_name, str): raise ValueError( "Creating multiple datasets but only one measurement name given." ) if set(dataset_dependencies.keys()) != set(measurement_name): raise ValueError( f"Inconsistent measurement names: measurement_name " f"contains {measurement_name} " f"but dataset_dependencies contains " f"{tuple(dataset_dependencies.keys())}." ) def _parse_dond_arguments( *params: AbstractSweep | TogetherSweep | ParamMeasT | Sequence[ParamMeasT], ) -> tuple[ list[AbstractSweep | TogetherSweep], list[ParamMeasT | Sequence[ParamMeasT]] ]: """ Parse supplied arguments into sweep objects and measurement parameters and their callables. """ sweep_instances: list[AbstractSweep | TogetherSweep] = [] params_meas: list[ParamMeasT | Sequence[ParamMeasT]] = [] for par in params: if isinstance(par, AbstractSweep): sweep_instances.append(par) elif isinstance(par, TogetherSweep): sweep_instances.append(par) else: params_meas.append(par) return sweep_instances, params_meas