import inspect
import json
import time
from copy import deepcopy
from dataclasses import asdict, dataclass, fields
from functools import wraps
from pathlib import Path
from typing import Callable, Generic, NewType, Optional, TypeVar, Union
import numpy as np
import numpy.typing as npt
from qibolab.platform import Platform
from qibolab.qubits import Qubit, QubitId, QubitPair, QubitPairId
from qibocal.config import log
from .serialize import deserialize, load, serialize
OperationId = NewType("OperationId", str)
"""Identifier for a calibration routine."""
ParameterValue = Union[float, int]
"""Valid value for a routine and runcard parameter."""
Qubits = dict[QubitId, Qubit]
"""Convenient way of passing qubit pairs in the routines."""
QubitsPairs = dict[tuple[QubitId, QubitId], QubitPair]
DATAFILE = "data"
"""Name of the file where data are dumped."""
RESULTSFILE = "results"
"""Name of the file where results are dumped."""
[docs]def show_logs(func):
"""Decorator to add logs."""
@wraps(func)
# necessary to maintain the function signature
def wrapper(*args, **kwds):
start = time.perf_counter()
out = func(*args, **kwds)
end = time.perf_counter()
if end - start < 1:
message = " in less than 1 second."
else:
message = f" in {end-start:.2f} seconds"
log.info(f"Finished {func.__name__[1:]}" + message)
return out, end - start
return wrapper
DEFAULT_PARENT_PARAMETERS = {
"nshots": None,
"relaxation_time": None,
}
"""Default values of the parameters of `Parameters`"""
[docs]class Parameters:
"""Generic action parameters.
Implement parameters as Algebraic Data Types (similar to), by subclassing
this marker in actual parameters specification for each calibration routine.
The actual parameters structure is only used inside the routines themselves.
"""
nshots: int
"""Number of executions on hardware"""
relaxation_time: float
"""Wait time for the qubit to decohere back to the `gnd` state"""
[docs] @classmethod
def load(cls, input_parameters):
"""Load parameters from runcard.
Possibly looking into previous steps outputs.
Parameters defined in Parameters class are removed from `parameters`
before `cls` is created.
Then `nshots` and `relaxation_time` are assigned to cls.
.. todo::
move the implementation to History, since it is required to resolve
the linked outputs
"""
default_parent_parameters = deepcopy(DEFAULT_PARENT_PARAMETERS)
parameters = deepcopy(input_parameters)
for parameter, value in default_parent_parameters.items():
default_parent_parameters[parameter] = parameters.pop(parameter, value)
instantiated_class = cls(**parameters)
for parameter, value in default_parent_parameters.items():
setattr(instantiated_class, parameter, value)
return instantiated_class
[docs]class AbstractData:
"""Abstract data class."""
def __init__(
self, data: dict[Union[tuple[QubitId, int], QubitId], npt.NDArray] = None
):
self.data = data if data is not None else {}
def __getitem__(self, qubit: Union[QubitId, tuple[QubitId, int]]):
"""Access data attribute member."""
return self.data[qubit]
@property
def params(self) -> dict:
"""Convert non-arrays attributes into dict."""
global_dict = asdict(self)
if hasattr(self, "data"):
global_dict.pop("data")
return global_dict
[docs] def save(self, path: Path, filename: str):
"""Dump class to file."""
self._to_json(path, filename)
self._to_npz(path, filename)
def _to_npz(self, path: Path, filename: str):
"""Helper function to use np.savez while converting keys into strings."""
if hasattr(self, "data"):
np.savez(
path / f"{filename}.npz",
**{json.dumps(i): self.data[i] for i in self.data},
)
def _to_json(self, path: Path, filename: str):
"""Helper function to dump to json."""
if self.params:
(path / f"{filename}.json").write_text(
json.dumps(serialize(self.params), indent=4)
)
[docs] @classmethod
def load(cls, path: Path, filename: str):
"""Generic load method."""
data_dict = cls.load_data(path, filename)
params = cls.load_params(path, filename)
if data_dict is not None:
if params is not None:
return cls(data=data_dict, **params)
else:
return cls(data=data_dict)
elif params is not None:
return cls(**params)
[docs] @staticmethod
def load_data(path: Path, filename: str):
"""Load data stored in a npz file."""
file = path / f"{filename}.npz"
if file.is_file():
raw_data_dict = dict(np.load(file))
data_dict = {}
for data_key, array in raw_data_dict.items():
data_dict[load(data_key)] = np.rec.array(array)
return data_dict
[docs] @staticmethod
def load_params(path: Path, filename: str):
"""Load parameters stored in a json file."""
file = path / f"{filename}.json"
if file.is_file():
params = json.loads(file.read_text())
params = deserialize(params)
return params
[docs]class Data(AbstractData):
"""Data resulting from acquisition routine."""
@property
def qubits(self):
"""Access qubits from data structure."""
if set(map(type, self.data)) == {tuple}:
return list({q[0] for q in self.data})
return [q for q in self.data]
@property
def pairs(self):
"""Access qubit pairs ordered alphanumerically from data structure."""
return list({tuple(sorted(q[:2])) for q in self.data})
[docs] def register_qubit(self, dtype, data_keys, data_dict):
"""Store output for single qubit.
Args:
data_keys (tuple): Keys of Data.data.
data_dict (dict): The keys are the fields of `dtype` and
the values are the related arrays.
"""
# to be able to handle the non-sweeper case
ar = np.empty(np.shape(data_dict[list(data_dict.keys())[0]]), dtype=dtype)
for key, value in data_dict.items():
ar[key] = value
if data_keys in self.data:
self.data[data_keys] = np.rec.array(
np.concatenate((self.data[data_keys], ar))
)
else:
self.data[data_keys] = np.rec.array(ar)
[docs] def save(self, path: Path):
"""Store data to file."""
super()._to_json(path, DATAFILE)
super()._to_npz(path, DATAFILE)
[docs] @classmethod
def load(cls, path: Path):
"""Load data and parameters."""
return super().load(path, filename=DATAFILE)
[docs]@dataclass
class Results(AbstractData):
"""Generic runcard update."""
def __contains__(self, key: Union[QubitId, QubitPairId, tuple[QubitId, ...]]):
"""Checking if qubit is in Results.
If key is not present means that fitting failed or
was not performed.
"""
return all(
key in getattr(self, field.name)
for field in fields(self)
if isinstance(getattr(self, field.name), dict)
)
[docs] @classmethod
def load(cls, path: Path):
"""Load results."""
return super().load(path, filename=RESULTSFILE)
[docs] def save(self, path: Path):
"""Store results to file."""
super()._to_json(path, RESULTSFILE)
super()._to_npz(path, RESULTSFILE)
# Internal types, in particular `_ParametersT` is used to address function
# contravariance on parameter type
_ParametersT = TypeVar("_ParametersT", bound=Parameters, contravariant=True)
_DataT = TypeVar("_DataT", bound=Data)
_ResultsT = TypeVar("_ResultsT", bound=Results)
[docs]@dataclass
class Routine(Generic[_ParametersT, _DataT, _ResultsT]):
"""A wrapped calibration routine."""
acquisition: Callable[[_ParametersT], _DataT]
"""Data acquisition function."""
fit: Callable[[_DataT], _ResultsT] = None
"""Post-processing function."""
report: Callable[[_DataT, _ResultsT], None] = None
"""Plotting function."""
update: Callable[[_ResultsT, Platform], None] = None
"""Update function platform."""
two_qubit_gates: Optional[bool] = False
"""Flag to determine whether to allocate list of Qubits or Pairs."""
def __post_init__(self):
# add decorator to show logs
self.acquisition = show_logs(self.acquisition)
self.fit = show_logs(self.fit)
if self.update is None:
self.update = _dummy_update
@property
def parameters_type(self):
"""Input parameters type."""
sig = inspect.signature(self.acquisition)
param = next(iter(sig.parameters.values()))
return param.annotation
@property
def data_type(self):
""" "Data object type return by data acquisition."""
return inspect.signature(self.acquisition).return_annotation
@property
def results_type(self):
"""Results object type returned by data acquisition."""
return inspect.signature(self.fit).return_annotation
# TODO: I don't like these properties but it seems to work
@property
def platform_dependent(self):
"""Check if acquisition involves platform."""
return "platform" in inspect.signature(self.acquisition).parameters
@property
def targets_dependent(self):
"""Check if acquisition involves qubits."""
return "targets" in inspect.signature(self.acquisition).parameters
[docs]@dataclass
class DummyPars(Parameters):
"""Dummy parameters."""
[docs]@dataclass
class DummyData(Data):
"""Dummy data."""
[docs] def save(self, path):
"""Dummy method for saving data"""
[docs]@dataclass
class DummyRes(Results):
"""Dummy results."""
def _dummy_acquisition(pars: DummyPars, platform: Platform) -> DummyData:
"""Dummy data acquisition."""
return DummyData()
def _dummy_update(
results: DummyRes, platform: Platform, qubit: Union[QubitId, QubitPairId]
) -> None:
"""Dummy update function"""
dummy_operation = Routine(_dummy_acquisition)
"""Example of a dummy operation."""