"""This module provides classes to define parameter sweeps.
Sweeps implement (and sublcass) :class:`SweepProtocol` s, which defines
binary operators to nest, parallelize, and concatenate sweeps.
TODO:
- labcore/quantify might be suitable to replace this
- compose sweeps with different parameters (broadcast)
- sweep arithmetic: +/-/*/**?
"""
from __future__ import annotations
import abc
import dataclasses
import numbers
import time
from collections.abc import Sequence, Callable, Iterable, Sized, Container
from typing import Any, TypeVar, Literal, Self, Protocol, runtime_checkable
import numpy as np
import numpy.typing as npt
from dataclassabc import dataclassabc
from qcodes.parameters import ParameterBase, DelegateGroupParameter
from qcodes_contrib_drivers.drivers.QDevil.QDAC2 import QDac2Channel
from qutil import functools, itertools
from ..helpers import find_param_source
from ..instruments.logical_instruments import Trap
from ..measurements.measures import MeasureSet, Id
from ..parameters import VirtualParameter, VirtualParameterContext
_S = TypeVar('_S')
_T = TypeVar('_T', bound=np.generic)
_InitActionsT = Sequence[Callable[[], None]]
_PostActionsT = Sequence[Callable[[Any], None]]
[docs]
@runtime_checkable
class SweepProtocol(Sized, Container[_T], Iterable[_T], Protocol[_T]):
"""Protocol defining a parameter sweep."""
param: tuple[ParameterBase, ...]
delay: tuple[float, ...]
num_points: int
init_actions: tuple[_InitActionsT, ...]
post_actions: tuple[_PostActionsT, ...]
get_after_set: tuple[bool, ...]
start: tuple[tuple[_T, ...], ...]
setpoints: tuple[npt.NDArray[_T], ...]
leak_measures: MeasureSet[_T]
transform: Callable[[_T], _T]
def __or__(self): ...
def __and__(self): ...
def __matmul__(self, other): ...
[docs]
def set(self, *vals: _T) -> tuple[_T, ...]: ...
[docs]
def initialize(self, i: int = 0): ...
[docs]
class AbstractSweep(Sequence[_T]):
"""Abstract base class or sweeps.
The intended implementation is a dataclass that accepts the
following attributes as initialization arguments. The attributes
will be tuples with length corresponding to the number of parameters
this sweep steps in parallel (see below). The only exception being
`num_points`, which must be the same for all parameters and is
therefore scalar.
Attributes
----------
param :
One or more parameters that provide a set method. If multiple,
all of them will be stepped in parallel (i.e., each parameter
will be set at each point of a sweep).
delay :
Time(s) to sleep after setting a value.
num_points :
The total number of points in the sweep (including all
directions).
init_actions :
A (tuple of) actions (callables that take no arguments) to run
at the start of the first sweep direction when calling
:meth:`initialize`.
post_actions :
A (tuple of) actions (callables that take the current setpoint
as sole argument) to run after setting each value. For mulitple
parameters, should also be a tuple of tuples of actions. If not,
some logic is applied to "broadcast" the supplied actions to the
parameters, but no guarantees.
get_after_set :
Should the parameter(s) be gotten after being set?
transform :
A callable with signature::
(val: _T) -> _T
that takes the set value and returns the transformed value.
Post-actions are applied to the transformed value.
The following should be properties, not init fields:
start :
A tuple of starting values at which the sweep starts. These
values are navigated to when calling :meth:`initialize`.
setpoints :
The parameter values that will be visited during the sweep.
Should be a tuple of arrays of shape (n_directions, n_pts,),
where the tuple is for each parameter and the first axis of the
array is for multiple directions of the sweep. That is, first
the first row of the array will be swept, then the second, etc.
Hence, it is the axis along which the concatenation operator
``@`` sequences sweeps.
leak_measures :
A :class:`~.measures.MeasureSet` of measures returning leakage
currents. Defaults to an empty set.
"""
param: ParameterBase | Sequence[ParameterBase]
delay: float | Sequence[float]
init_actions: _InitActionsT | Sequence[_InitActionsT]
post_actions: _PostActionsT | Sequence[_PostActionsT]
get_after_set: bool | Sequence[bool]
transform: Callable[[_T], _T] | Sequence[Callable[[_T], _T]] = Id
def __post_init__(self):
self.param = _to_tuple(self.param, ParameterBase)
self.delay = _to_tuple(self.delay, numbers.Real, (len(self.param),))
self.get_after_set = _to_tuple(self.get_after_set, bool, (len(self.param),))
self.transform = _to_tuple(self.transform, Callable, (len(self.param),))
self.init_actions = _to_tuple(self.init_actions, Callable)
self.post_actions = self._parse_post_actions(self.post_actions)
assert isinstance(self.num_points, numbers.Integral), 'Require integral num_points'
assert len(set(self.param)) == len(self.param), 'Require unique params'
def __len__(self) -> int:
return sum(len(direction) for direction in self)
def __getitem__(self, item) -> Sequence[Sequence[_T]]:
return self.setpoints[item]
def __matmul__(self, other) -> 'ArraySweep':
"""Sequence sweeps in time (concatenate)."""
if not isinstance(other, SweepProtocol):
return NotImplemented
if (
self.param != other.param
or self.delay != other.delay
or self.init_actions != other.init_actions
or self.post_actions != other.post_actions
or self.get_after_set != other.get_after_set
or self.transform != other.transform
):
raise TypeError(f'Incompatible sweeps for concatenation: {self}, {other}')
return ArraySweep(self.param, self.setpoints + other.setpoints, self.delay,
self.init_actions, self.post_actions, self.get_after_set, self.transform)
def __and__(self, other) -> Self:
"""Sweep in parallel (direct product)."""
if isinstance(other, self.__class__):
dct = dict.fromkeys((field.name for field in dataclasses.fields(self.__class__)))
for key in dct:
a = getattr(self, key)
b = getattr(other, key)
if np.iterable(a) and np.iterable(b):
dct[key] = tuple(a) + tuple(b)
elif a == b:
dct[key] = a
else:
raise ValueError(f'Incompatible sweep shapes: {key}.')
return self.__class__(**dct)
else:
return NotImplemented
def __or__(self, other) -> 'SweepList':
"""Nested loop (outer product)."""
if isinstance(other, SweepList):
return SweepList([self, *other])
return SweepList([self, other])
def __ror__(self, other) -> 'SweepList':
if isinstance(other, SweepList):
return SweepList([*other, self])
return SweepList([other, self])
def _parse_post_actions(self, actions: _PostActionsT) -> tuple[_PostActionsT, ...]:
actions = _to_tuple(actions, Callable)
if not actions:
actions = tuple(() for _ in range(len(self.param)))
elif all(callable(act) for act in actions):
if len(self.param) == 1:
actions = (actions,)
else:
actions = tuple((act,) for act in actions)
assert len(actions) == len(self.param)
assert all(isinstance(acts, Sequence) for acts in actions)
assert all(callable(act) for acts in actions for act in acts)
return actions
def _validate(self):
"""Validate every setpoint."""
for direction in self.setpoints:
for param, setpoints in zip(self.param, direction):
for setpoint in setpoints:
param.validate(setpoint)
[docs]
def set(self, *vals: _T) -> tuple[_T, ...]:
"""Set the parameters to vals.
Roughly equivalent to::
for param, trafo, val in zip(self.params, self.transforms, vals):
param.set(trafo(val))
.. todo::
Threaded?
"""
getvals = []
for i, (param, delay, trafo, post_actions, get_after_set, val) in enumerate(zip(
self.param, self.delay, self.transform, self.post_actions, self.get_after_set,
vals, strict=True
)):
if not any(val in direction[:, i] for direction in self):
raise ValueError(f'Can only set values from set of setpoints, not {val}')
val = trafo(val)
param.set(val)
tic = time.perf_counter()
for act in post_actions:
act(val)
if delay > (elapsed := time.perf_counter() - tic):
time.sleep(delay - elapsed)
if get_after_set:
getvals.append(param.get())
else:
getvals.append(val)
return tuple(getvals)
[docs]
def initialize(self, i: int = 0):
"""Go to starting point for direction *i* and perform custom
init tasks.
If *i* is 0, i.e, the initial direction, runs all init_actions.
Then sets params in reverse order so that going back to start
after sweeping will always be reproducible and not run into
parameter bounds.
"""
if i == 0:
for action in self.init_actions:
action()
for param, val in zip(self.param[::-1], self.start[i][::-1], strict=True):
param(val)
@property
@abc.abstractmethod
def num_points(self) -> int:
"""The total number of points that will be set."""
...
@functools.cached_property
@abc.abstractmethod
def setpoints(self) -> tuple[Sequence[Sequence[_T]], ...]:
"""A tuple of 2d-arrays of setpoints.
The last axis enumerates parameters (i.e., is of size
``len(self.param)``, the first axis enumerates setpoints (i.e.,
is of size :attr:`num_points`), and the tuple enumerates sweep
directions.
"""
...
@property
@abc.abstractmethod
def start(self) -> tuple[tuple[_T, ...], ...]:
"""The starting points of the sweep.
:meth:`initialize` will set parameters to these values.
"""
...
@property
def leak_measures(self) -> MeasureSet:
"""A :class:`~.measures.MeasureSet` of QDAC leakage parameters."""
return MeasureSet(_extract_leak_params(*self.param))
[docs]
@dataclassabc
class GridSweep(AbstractSweep[_T]):
"""Sweep values on a regularly spaced grid.
See base class for common parameters.
Parameters
----------
spacing :
The numpy function to use to grid the values.
- :func:`numpy:numpy.linspace`
- :func:`numpy:numpy.geomspace`
- :func:`numpy:numpy.logspace`
"""
param: ParameterBase | Sequence[ParameterBase]
rng: tuple[float, float] | Sequence[tuple[float, float]]
num_points: int
spacing: Literal['lin', 'geom', 'log'] | Sequence[Literal['lin', 'geom', 'log']] = 'lin'
delay: float | Sequence[float] = 0.0
init_actions: _InitActionsT | Sequence[_InitActionsT] = ()
post_actions: _PostActionsT | Sequence[_PostActionsT] = ()
get_after_set: bool | Sequence[bool] = False
transform: Callable[[_T], _T] | Sequence[Callable[[_T], _T]] = Id
def __post_init__(self):
super().__post_init__()
self.rng = tuple(np.broadcast_to(self.rng, (len(self.param), 2)))
self.spacing = _to_tuple(self.spacing, str, (len(self.param),))
self._validate()
@property
def start(self) -> tuple[tuple[float, ...]]:
return (tuple(rng[0] for rng in self.rng),)
@functools.cached_property
def setpoints(self) -> tuple[npt.NDArray[float]]:
return (np.array([getattr(np, f'{spacing}space')(*rng, self.num_points)
for spacing, rng in zip(self.spacing, self.rng)]).T,)
[docs]
@dataclassabc
class CenteredLinearSweep(AbstractSweep[_T]):
"""A linearly spaced sweep radiating out from a center position.
See base class for common parameters.
Parameters
----------
center :
The center (starting point) of the sweep.
"""
param: ParameterBase | Sequence[ParameterBase]
rng: float | Sequence[float]
num_points: int
center: float | Sequence[float] = np.nan
delay: float | Sequence[float] = 0.0
init_actions: _InitActionsT | Sequence[_InitActionsT] = ()
post_actions: _PostActionsT | Sequence[_PostActionsT] = ()
get_after_set: bool | Sequence[bool] = False
transform: Callable[[_T], _T] | Sequence[Callable[[_T], _T]] = Id
def __post_init__(self):
super().__post_init__()
self.rng = _to_tuple(self.rng, numbers.Real, (len(self.param),))
# Repeat center since both directions start at the same point, the center
self.center = tuple(np.repeat(np.array([[
center if not np.isnan(center) else param()
for center, param in itertools.zip_broadcast(self.center, self.param)
]]), 2, axis=0))
self._validate()
@property
def start(self) -> tuple[float, ...]:
return tuple(self.center)
@functools.cached_property
def setpoints(self) -> tuple[npt.NDArray[float], npt.NDArray[float]]:
# Both directions start at the same point
rng = np.array(self.rng)
points = self.start[0].T + np.linspace(-rng, rng, self.num_points)
a, b = np.split(points, [self.num_points // 2])
return b, a[::-1]
[docs]
@dataclassabc
class ArraySweep(AbstractSweep[_T]):
"""A sweep over arbitrary values.
See base class for common parameters.
Parameters
----------
array :
The setpoints to visit in the sweep.
"""
param: ParameterBase | Sequence[ParameterBase]
array: npt.ArrayLike | Sequence[npt.ArrayLike]
delay: float | Sequence[float] = 0.0
init_actions: _InitActionsT | Sequence[_InitActionsT] = ()
post_actions: _PostActionsT | Sequence[_PostActionsT] = ()
get_after_set: bool | Sequence[bool] = False
transform: Callable[[_T], _T] | Sequence[Callable[[_T], _T]] = Id
def __post_init__(self):
# Need to parse array before calling super's post_init since it requires num_points to be
# defined (which depends on array here).
try:
self.array = self._parse_array(self.array, max_ndim=3)
except ValueError:
# Directions are not homogeneous arrays. Parse each individually
self.array = tuple(self._parse_array(array, max_ndim=2) for array in self.array)
super().__post_init__()
# Finally broadcast arrays to expected shape
self.array = tuple(np.broadcast_to(array, (len(array), len(self.param)))
for array in self.array)
self._validate()
def _parse_array(self, array: npt.ArrayLike, max_ndim: Literal[2, 3]) -> npt.NDArray[float]:
# IMPORTANT! Array needs to be floating since otherwise its dtype is a numpy integer which
# is not JSON serializable
# TODO (08/05/24): But QCoDeS' NumpyJSONEncoder should be able to handle it?
# array = np.atleast_1d(array).astype(float)
array = np.atleast_1d(array)
if array.ndim == 1:
array = array[None, :, None]
elif array.ndim == 2:
try:
# Assume this dimension is meant as the one over parameters (i.e., the last)
ax = array.shape.index(len(self.param))
array = np.moveaxis(array, ax, -1)[None, :, :]
except ValueError:
# Assume it's the direction dimension
array = array[:, :, None]
elif array.ndim > max_ndim:
raise ValueError(f"Do not know how to interpret array shape {array.shape}")
if max_ndim == 3:
return array
return array[0]
@property
def start(self) -> tuple[npt.NDArray[float], ...]:
return tuple(array[0] for array in self.array) # noqa
@property
def num_points(self) -> int:
return sum(len(array) for array in self.array)
@functools.cached_property
def setpoints(self) -> tuple[npt.NDArray[float], ...]:
return tuple(self.array)
[docs]
class SweepList(list[SweepProtocol[_T]]):
"""A list of sweeps which will be executed nested.
The first item in the list will be the outermost loop. Hence,
measures will be gotten at each point of the sweep that is the last
item in this list.
:class:`SweepList` s can be composed via the usual list methods and
in-place using the |= operator.
"""
def __init__(self, sweeps: Iterable[SweepProtocol[_T]]):
super().__init__(sweeps)
self.__post_init__()
def __post_init__(self):
self._current_setpoints = [(None,)*len(sweep.param) for sweep in self]
# Get the setpoints of all sweeps to catch possible exceptions early on
for sweep in self:
_ = sweep.setpoints
def __ior__(self, other) -> Self:
if isinstance(other, SweepList):
self.extend(other)
elif isinstance(other, SweepProtocol):
self.append(other)
else:
return NotImplemented
# Update current_setpoints
self.__post_init__()
return self
def __add__(self, other) -> Self:
add = super().__add__(other)
return self.__class__(add)
def __iadd__(self, other) -> Self:
return self.__ior__(other)
def __mul__(self, other) -> Self:
return NotImplemented
def __imul__(self, other) -> Self:
return NotImplemented
def __getitem__(self, i) -> SweepProtocol[_T] | Self:
item = super().__getitem__(i)
if isinstance(item, list):
return self.__class__(item)
return item
@property
def params(self) -> tuple[ParameterBase, ...]:
"""All parameters of the complex sweep."""
return tuple(itertools.flatten(sweep.param for sweep in self))
@property
def leak_measures(self) -> MeasureSet:
"""All leakage measures of the complex sweep."""
return MeasureSet(itertools.flatten(sweep.leak_measures for sweep in self))
@property
def current_setpoints(self) -> list[tuple[_T | None, ...]]:
"""The current setpoints. Used internally."""
return self._current_setpoints
@property
def shape(self) -> tuple[int, ...]:
"""The shape (sweep_1.num_points, ..., sweep_n.num_points)."""
return tuple(len(sweep) for sweep in self)
def _extract_leak_params(*params: ParameterBase) -> list[ParameterBase]:
leak_params = set()
for param in params:
if isinstance(param.instrument, Trap):
leak_params.update(p for name, p in param.instrument.parameters.items() if
name.endswith('currents'))
elif isinstance(param := find_param_source(param), VirtualParameterContext):
# central / guard
# Should be covered by first branch
for current_param in param.current_parameters:
leak_params.add(current_param)
elif isinstance(param, VirtualParameter):
# difference_mode / common_mode.
# Should be covered by first branch
leak_params.update(_extract_leak_params(param.parameter_context))
elif isinstance(param, DelegateGroupParameter):
# top / bottom
# Should be covered by first branch
leak_params.add(getattr(param.instrument, f'{param.name}_current'))
elif isinstance(param.instrument, QDac2Channel):
leak_params.add(param.instrument.read_current_A)
return list(leak_params)
def _to_tuple(arg: _S | Iterable[_S], cls: type[_S],
shape: tuple[int, ...] | None = None) -> tuple[_S, ...]:
if isinstance(arg, cls):
res = (arg,)
else:
res = tuple(arg)
if shape is not None:
return tuple(np.broadcast_to(res, shape))
return res
