# 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 2026, Qblox B.V.
"""Local copy of quantify_core utility functions to reduce dependency on quantify_core."""
import functools
import inspect
import re
import uuid
import warnings
from collections.abc import Callable, Sequence
from datetime import datetime
from pathlib import Path
import matplotlib.pyplot as plt
import matplotlib.ticker
import numpy as np
import xarray as xr
from matplotlib.axes import Axes
[docs]
SI_PREFIXES = dict(zip(range(-24, 25, 3), "yzafpnμm kMGTPEZY", strict=False))
SI_PREFIXES[0] = ""
# N.B. not all of these are SI units, however, all of these support SI prefixes
[docs]
SI_UNITS = [
"SI_PREFIX_ONLY",
"m",
"s",
"g",
"W",
"J",
"V",
"A",
"F",
"T",
"Hz",
"Ohm",
"S",
"N",
"C",
"px",
"b",
"B",
"K",
"Bar",
r"Vpeak",
r"Vpp",
r"Vp",
r"Vrms",
r"A/s",
r"$\Phi_0$",
]
[docs]
_SI_PREFIX_TO_FACTOR_MAPPING = {v: 10**key for key, v in SI_PREFIXES.items()}
_SI_PREFIX_TO_FACTOR_MAPPING["u"] = 10**-6
[docs]
_prefix_regexp = "(" + "|".join(list("yzafpnμmkMGTPEZY")) + ")"
[docs]
_si_regex = "(" + "|".join(map(re.escape, SI_UNITS)) + ")"
[docs]
_prefixed_si_regex = re.compile(f"{_prefix_regexp}{_si_regex}$")
[docs]
def _find_stack_level() -> int:
"""
Find the first place in the stack that is not inside qblox-scheduler.
Returns
-------
:
The stack level
"""
frame = inspect.currentframe()
if frame is None:
return 1
stack = inspect.getouterframes(frame)
for i, frame_info in enumerate(stack):
if "qblox_scheduler" not in frame_info.filename:
return max(1, i - 1)
return 1
[docs]
def deprecated(
drop_version: str,
message_or_alias: str | Callable,
) -> Callable:
"""
A decorator for deprecating classes and methods.
Parameters
----------
drop_version
A version of the package when the deprecated function or class will be dropped.
message_or_alias
Either an instruction about the usage of deprecated calls (string),
or the new drop-in replacement to the deprecated class or function (callable).
Returns
-------
:
The decorated function or class
"""
def decorator(obj: Callable) -> Callable:
if callable(message_or_alias):
new_obj = message_or_alias
if isinstance(obj, type):
new_obj.__name__ = obj.__name__
new_obj.__qualname__ = obj.__qualname__
else:
new_obj.__name__ = obj.__name__
new_obj.__qualname__ = obj.__qualname__
@functools.wraps(obj)
def wrapper_func(*args, **kwargs): # noqa: ANN202
warnings.warn(
f"{obj.__name__} is deprecated and will be removed in "
f"qblox-scheduler>={drop_version}. Use {new_obj.__name__} instead.",
category=FutureWarning,
stacklevel=_find_stack_level(),
)
return new_obj(*args, **kwargs)
return wrapper_func
else:
msg = message_or_alias
if isinstance(obj, type):
original_init = obj.__init__
@functools.wraps(original_init)
def new_init(self, *args, **kwargs) -> None: # noqa: ANN001
warnings.warn(
f"{obj.__name__} is deprecated and will be removed in "
f"qblox-scheduler>={drop_version}. {msg}",
category=FutureWarning,
stacklevel=_find_stack_level(),
)
original_init(self, *args, **kwargs)
obj.__init__ = new_init
return obj
else:
@functools.wraps(obj)
def wrapper_class(*args, **kwargs): # noqa: ANN202
warnings.warn(
f"{obj.__name__} is deprecated and will be removed in "
f"qblox-scheduler>={drop_version}. {msg}",
category=FutureWarning,
stacklevel=_find_stack_level(),
)
return obj(*args, **kwargs)
return wrapper_class
return decorator
[docs]
def without(d: dict, keys: list) -> dict:
"""
Return a copy of a dictionary with the specified keys removed.
Parameters
----------
d
The dictionary to copy
keys
List of keys to remove
Returns
-------
:
A new dictionary without the specified keys
"""
result = dict(d)
for key in keys:
result.pop(key, None)
return result
[docs]
def load_json_schema(file_location: str, filename: str) -> dict:
"""
Load a JSON schema from a file relative to a module location.
Parameters
----------
file_location
The location of the file (typically __file__)
filename
The name of the schema file
Returns
-------
:
The loaded JSON schema as a dictionary
"""
import json
from pathlib import Path
schema_path = Path(file_location).parent / filename
with open(schema_path) as f:
return json.load(f)
[docs]
DATASET_NAME = "dataset.hdf5"
[docs]
PROCESSED_DATASET_NAME = "dataset_processed.hdf5"
[docs]
QUANTITIES_OF_INTEREST_NAME = "quantities_of_interest.json"
[docs]
def gen_tuid() -> str:
"""
Generate a time-based unique identifier (TUID).
Returns
-------
:
A TUID string of the form "YYYYmmDD-HHMMSS-sss-******
"""
now = datetime.now()
tuid = now.strftime("%Y%m%d-%H%M%S")
ms = now.strftime("%f")[:3]
random_suffix = uuid.uuid4().hex[:6]
return f"{tuid}-{ms}-{random_suffix}"
[docs]
class TUID(str): # noqa: SLOT000
"""
Time-based unique identifier (TUID) class.
A TUID is a string of the form "YYYYmmDD-HHMMSS-sss-******
"""
def __new__(cls, tuid: str) -> "TUID": # noqa: D102
if not cls.is_valid(tuid):
raise ValueError(f"Invalid TUID: {tuid}")
return super().__new__(cls, tuid)
@classmethod
[docs]
def is_valid(cls, tuid: str) -> bool:
"""
Check if a string is a valid TUID.
Parameters
----------
tuid
The string to check
Returns
-------
:
True if the string is a valid TUID, False otherwise
"""
if len(tuid) < cls._TUID_LENGTH:
return False
try:
datetime.strptime(tuid[:15], "%Y%m%d-%H%M%S")
if tuid[15] != "-":
return False
int(tuid[16:19])
return tuid[19] == "-"
except (ValueError, IndexError):
return False
@classmethod
[docs]
def datetime_seconds(cls, tuid: str) -> datetime:
"""
Get the datetime for a TUID.
Parameters
----------
tuid
The TUID string
Returns
-------
:
Datetime object
"""
return datetime.strptime(tuid[:15], "%Y%m%d-%H%M%S")
[docs]
def _cast_numpy_bool_attrs(dataset: xr.Dataset) -> None:
"""Cast any np.bool_ attributes to Python bool (in-place) to avoid h5netcdf errors."""
def _cast(attrs: dict) -> None:
for key, val in attrs.items():
if isinstance(val, np.bool_):
attrs[key] = bool(val)
for var in dataset.variables.values():
_cast(var.attrs)
_cast(dataset.attrs)
[docs]
def write_dataset(path: str | Path, dataset: xr.Dataset) -> None:
"""
Write an xarray Dataset to disk.
Parameters
----------
path
The path to write to
dataset
The dataset to write
"""
_cast_numpy_bool_attrs(dataset)
dataset.to_netcdf(path, engine="h5netcdf", invalid_netcdf=True)
[docs]
def snapshot() -> dict:
"""
Create a snapshot of the current state.
Returns
-------
:
A dictionary containing the snapshot
"""
import platform
import sys
snapshot = {
"timestamp": datetime.now().isoformat(),
"platform": platform.platform(),
"python_version": sys.version,
}
return snapshot
[docs]
def save_json(
directory: str | Path | None,
filename: str,
data: dict,
compression: str | None = None, # pyright: ignore[reportUnusedParameter] # noqa: ARG001
) -> None:
"""
Save a dictionary to a JSON file.
Parameters
----------
directory
The directory to save to
filename
The filename
data
The data to save
compression
The compression type ("bz2", "gzip", "lzma", or None)
"""
import json
if directory is None:
directory = Path(".")
filepath = Path(directory) / filename
with open(filepath, "w") as f:
json.dump(data, f, indent=2)
[docs]
def to_gridded_dataset(
dataset: xr.Dataset,
dimension: str = "dim_0",
coords_names: list[str] | None = None,
) -> xr.Dataset:
"""
Convert a flattened dataset (as generated by initialize_dataset) to a gridded
dataset in which the measured values are mapped onto a grid.
Parameters
----------
dataset
The input dataset in flattened format.
dimension
The flattened xarray Dimension.
coords_names
Optionally specify which Variables correspond to orthogonal coordinates.
Defaults to all variables starting with "x".
Returns
-------
:
A gridded dataset.
"""
if dimension not in (dims := tuple(dataset.dims)):
raise ValueError(f"Dimension {dimension} not in dims {dims}.")
if coords_names is None:
coords_names = sorted(
v for v in dataset.variables if isinstance(v, str) and v.startswith("x")
)
else:
for coord in coords_names:
vars_ = tuple(dataset.variables.keys())
if coord not in vars_:
raise ValueError(f"Coordinate {coord} not in coordinates {vars_}.")
attrs_coords = tuple(dataset[name].attrs for name in coords_names)
gridded = dataset.set_coords(coords_names)
if len(coords_names) == 1:
for var in dataset.data_vars:
if dimension in gridded[var].dims:
gridded = gridded.update(
{var: gridded[var].swap_dims({dimension: coords_names[0]})}
)
else:
gridded = gridded.set_index({dimension: coords_names})
gridded = gridded.unstack(dim=dimension)
for name, attrs in zip(coords_names, attrs_coords, strict=False):
gridded[name].attrs = attrs
if "grid_2d" in gridded.attrs:
gridded.attrs["grid_2d"] = False
return gridded
[docs]
def set_xlabel(
label: str | Axes,
unit: str | None = None,
axis: Axes | None = None,
auto_scale: bool = True,
**kw,
) -> Axes:
"""
Add a unit aware x-label to an axis object.
Parameters
----------
label
the desired label
unit
the unit
auto_scale
If True, then automatically scale the units
axis
matplotlib axis object to set label on
**kw
keyword argument to be passed to matplotlib.set_xlabel
"""
_label: str
if isinstance(label, Axes):
warnings.warn(
"Passing axis as a first argument is deprecated and will be removed "
"in quantify-core >= 0.10.0. Please use the new syntax "
"set_xlabel(label, unit = None, axis = None)",
FutureWarning,
stacklevel=2,
)
_old_axis = axis
axis = label
_label = unit if isinstance(unit, str) else ""
unit = _old_axis if isinstance(_old_axis, str) else None # type: ignore[assignment]
else:
_label = label
if axis is None:
axis = plt.gca()
if unit and auto_scale:
xticks = axis.get_xticks()
precision = 4
scale_factor, offset, unit = _get_scale_factor_and_offset_and_prefix(
xticks, unit, precision
)
formatter = matplotlib.ticker.FuncFormatter(
lambda x, _pos: f"{x * scale_factor - offset:.{precision}g}"
)
axis.xaxis.set_major_formatter(formatter)
if offset != 0:
offset_scale, offset_unit = SI_prefix_and_scale_factor(offset, unit)
disp_offset = offset * offset_scale
axis.xaxis.get_offset_text().set_text(f"{disp_offset:+g} {offset_unit}")
axis.set_xlabel(_label + f" [{unit}]", **kw)
elif unit:
axis.set_xlabel(_label + f" [{unit}]", **kw)
else:
axis.set_xlabel(_label, **kw)
return axis
[docs]
def set_ylabel(
label: str | Axes,
unit: str | None = None,
axis: Axes | None = None,
auto_scale: bool = True,
**kw,
) -> Axes | None:
"""
Add a unit aware y-label to an axis object.
Parameters
----------
label
the desired label
unit
the unit
axis
matplotlib axis object to set label on
auto_scale
If True, then automatically scale the units
**kw
keyword argument to be passed to matplotlib.set_ylabel
"""
_label: str
if isinstance(label, Axes):
warnings.warn(
"Passing axis as a first argument is deprecated and will be removed"
" in quantify-core >= 0.10.0. Please use the new syntax"
" set_ylabel(label, unit = None, axis = None)",
FutureWarning,
stacklevel=2,
)
_old_axis = axis
axis = label
_label = unit if isinstance(unit, str) else ""
unit = _old_axis if isinstance(_old_axis, str) else None # type: ignore[assignment]
else:
_label = label
if axis is None:
axis = plt.gca()
if unit and auto_scale:
yticks = axis.get_yticks()
precision = 6
scale_factor, offset, unit = _get_scale_factor_and_offset_and_prefix(
yticks, unit, precision=precision
)
formatter = matplotlib.ticker.FuncFormatter(
lambda x, _pos: f"{x * scale_factor - offset:.{precision}g}"
)
axis.yaxis.set_major_formatter(formatter)
if offset != 0:
offset_scale, offset_unit = SI_prefix_and_scale_factor(offset, unit)
disp_offset = offset * offset_scale
axis.yaxis.get_offset_text().set_text(f"{disp_offset:+g} {offset_unit}")
axis.set_ylabel(_label + f" [{unit}]", **kw)
elif unit:
axis.set_ylabel(_label + f" [{unit}]", **kw)
else:
axis.set_ylabel(_label, **kw)
return axis
[docs]
def _get_scale_factor_and_offset_and_prefix(
ticks: Sequence[float] | np.ndarray, unit: str | None = None, precision: int = 4
) -> tuple[float, float, str]:
"""
Return a convenient scale factor, offset and SI prefix based on the tick values.
This function uses the :func:`~.SI_prefix_and_scale_factor` function to determine a
scale factor such that the distance between ticks is in the range [0.1, 100.0), plus
the corresponding scaled SI unit (e.g. 'mT', 'kV'), deduced from the input unit, to
represent the tick values in those scaled units. In addition, an offset is
calculated such that the maximum absolute tick value is less than 10^precision.
Parameters
----------
ticks
A list of axis tick values.
unit
The unit of the tick values.
precision
The maximum amount of digits to display as tick labels.
Returns
-------
scale_factor
The scale factor to multiply the tick values with.
offset
The offset to subtract from the tick values.
unit
The unit including the SI prefix.
Examples
--------
>>> _get_scale_factor_and_offset_and_prefix(
... ticks=[2100000, 2100100, 2100200],
... unit="Hz",
... precision=4,
... )
(1.0, 2100000, 'Hz')
"""
max_v, min_v = max(ticks), min(ticks)
resolution = (max_v - min_v) / len(ticks)
scale_factor, unit = SI_prefix_and_scale_factor(val=resolution * 10, unit=unit)
signed_max = max_v if abs(max_v) > abs(min_v) else min_v
factor = pow(10, precision - 1)
offset = int(signed_max * scale_factor / factor) * factor
return scale_factor, offset, unit
[docs]
def SI_prefix_and_scale_factor(val: float, unit: str | None = None) -> tuple[float, str]: # noqa: N802
"""
Takes in a value and unit, returns a scale factor and scaled unit.
It returns a scale factor to convert the input value to a value in the
range [1.0, 1000.0), plus the corresponding scaled SI unit (e.g. 'mT', 'kV'),
deduced from the input unit, to represent the input value in those scaled units.
The scaling is only applied if the unit is an unscaled or scaled unit present in
the variable :data::`SI_UNITS`.
If the unit is None, no scaling is done.
If the unit is "SI_PREFIX_ONLY", the value is scaled and an SI prefix is applied
without a base unit.
Parameters
----------
val
the value
unit
the unit of the value
Returns
-------
scale_factor
scale_factor needed to convert value
scaled_unit
unit including the prefix
"""
if unit and val is not None and (match := _prefixed_si_regex.match(unit)):
scale_part = match.group(1)
unit_part = match.group(2)
plus_scale = _SI_PREFIX_TO_FACTOR_MAPPING[scale_part]
scale_factor, scaled_unit = SI_prefix_and_scale_factor(plus_scale * val, unit_part)
return plus_scale * scale_factor, scaled_unit
if unit in SI_UNITS:
try:
with np.errstate(all="ignore"):
prefix_power = np.log10(abs(val)) // 3 * 3
prefix = SI_PREFIXES[prefix_power]
# Greek symbols not supported in tex
if plt.rcParams["text.usetex"] and prefix == "μ":
prefix = r"$\mu$"
if unit == "SI_PREFIX_ONLY":
scale_factor, scaled_unit = 10**-prefix_power, prefix
elif unit == "s" and val > 2 * 60:
if val > 2 * 3600: # Convert to hours if larger than 2 hours
scale_factor, scaled_unit = 1 / 3600, "hrs"
else: # Convert to minutes if between 2 minutes and 2 hours
scale_factor, scaled_unit = 1 / 60, "min"
else:
scale_factor, scaled_unit = 10**-prefix_power, prefix + unit
# this exception can be triggered in the pyqtgraph multi processing
except (KeyError, TypeError):
scale_factor, scaled_unit = 1, unit
elif unit is None:
scale_factor, scaled_unit = 1, ""
else:
scale_factor, scaled_unit = 1, unit
return scale_factor, scaled_unit
