Experiment Handler

ExperimentHandler

Bases: ABC

Source code in sqil_core/experiment/_experiment.py

class ExperimentHandler(ABC):
    setup: dict
    instruments: Instruments | None = None

    zi_setup: DeviceSetup
    zi_session: Session
    qpu: QPU

    db_schema: dict = None

    def __init__(
        self,
        params: dict = {},
        param_dict_path: str = "",
        setup_path: str = "",
        server=False,
    ):
        # Read setup file
        if not setup_path:
            config = read_yaml("config.yaml")
            setup_path = config.get("setup_path", "setup.py")
        self.setup = _extract_variables_from_module("setup", setup_path)

        # Get instruments through the server or connect locally
        if server:
            server, instrument_instances = link_instrument_server()
        else:
            instrument_dict = self.setup.get("instruments", None)
            if not instrument_dict:
                logger.warning(
                    f"Unable to find any instruments in {setup_path}"
                    + "Do you have an `instruments` entry in your setup file?"
                )
            # Reset event listeners
            clear_signal(before_experiment)
            clear_signal(before_sequence)
            clear_signal(after_sequence)
            clear_signal(after_experiment)
        instrument_instances = connect_instruments(instrument_dict)

        # Create Zurich Instruments session
        zi = cast(ZI_Instrument, instrument_instances.get("zi", None))
        if zi is not None:
            self.zi_setup = zi.generate_setup()
            # self.zi_setup = DeviceSetup.from_descriptor(zi.descriptor, zi.address)
            self.zi_session = Session(self.zi_setup)
            self.zi_session.connect()
            self._load_qpu(zi.generate_qpu)

        self.instruments = Instruments(instrument_instances)
        self._setup_instruments()

    def _load_qpu(self, generate_qpu: Callable):
        qpu_filename = self.setup["storage"].get("qpu_filename", "qpu.json")
        db_path_local = self.setup["storage"]["db_path_local"]
        try:
            self.qpu = serializers.load(os.path.join(db_path_local, qpu_filename))
        except FileNotFoundError:
            logger.warning(
                f"Cannot find QPU file name {qpu_filename} in {db_path_local}"
            )
            logger.warning(f" -> Creating a new QPU file")
            self.qpu = generate_qpu(self.zi_setup)
            os.makedirs(db_path_local, exist_ok=True)
            w_save(
                self.qpu,
                os.path.join(db_path_local, qpu_filename),
            )

    # Move to server
    def _setup_instruments(self):
        """Default setup for all instruments with support for custom setups"""
        logger.info("Setting up instruments")
        if not hasattr(self, "instruments"):
            logger.warning("No instruments to set up")
            return

        for instrument in self.instruments:
            if not hasattr(instrument, "setup"):
                continue
            instrument.setup()

    @abstractmethod
    def sequence(self, *args, **kwargs):
        """Experimental sequence defined by the user"""
        pass

    @abstractmethod
    def analyze(self, path, *args, **kwargs):
        pass

    def run(self, *args, **kwargs):
        try:
            db_type = self.setup.get("storage", {}).get("db_type", "")

            if db_type == "plottr":
                return self.run_with_plottr(*args, **kwargs)
            else:
                return self.run_raw(*args, **kwargs)

        finally:
            # Close and delete QCodes instances to avoid connection issues in following experiments
            QCodesInstrument.close_all()
            for instrument in self.instruments:
                del instrument

    def run_with_plottr(self, *args, **kwargs):
        logger.info("Before exp")
        before_experiment.send(sender=self)

        # Map input parameters index to their name
        params_map, _ = map_inputs(self.sequence)

        # Get information on sweeps
        sweeps: dict = kwargs.get("sweeps", None)
        sweep_keys = []
        sweep_grid = []
        sweep_schema = {}
        if sweeps is not None:
            # Name of the parameters to sweep
            sweep_keys = list(sweeps.keys())
            # Create a mesh grid of all the sweep parameters
            sweep_grid = list(itertools.product(*sweeps.values()))
            # Add sweeps to the database schema
            for i, key in enumerate(sweep_keys):
                # TODO: dynamically add unit
                sweep_schema[f"sweep{i}"] = {"role": "axis", "param_id": key}

        # Create the plotter datadict (database) using the inferred schema
        db_schema = {**self.db_schema, **sweep_schema}
        datadict = build_plottr_dict(db_schema)
        # Get local and server storage folders
        db_path = self.setup["storage"]["db_path"]
        db_path_local = self.setup["storage"]["db_path_local"]

        # TODO: dynamically assign self.exp_name to class name if not provided
        with DDH5Writer(datadict, db_path_local, name=self.exp_name) as writer:
            # Get the path to the folder where the data will be stored
            storage_path = get_plottr_path(writer, db_path)
            storage_path_local = get_plottr_path(writer, db_path_local)
            # Save helper files
            writer.save_text("paths.md", f"{storage_path_local}\n{storage_path}")
            # Save backup qpu
            old_qubits = self.qpu.copy_quantum_elements()
            serializers.save(self.qpu, os.path.join(storage_path_local, "qpu_old.json"))

            # TODO: for index sweep don't recompile laboneq
            for sweep_values in sweep_grid or [None]:
                data_to_save = {}

                # Run/create the experiment. Creates it for laboneq, runs it otherwise
                seq = self.sequence(*args, **kwargs)
                # Detect if the sequence created a laboneq experiment
                is_laboneq_exp = type(seq) == LaboneQExperiment

                if is_laboneq_exp:
                    qu_indices = kwargs.get("qu_idx", [0])
                    if type(qu_indices) == int:
                        qu_indices = [qu_indices]
                    used_qubits = [self.qpu.quantum_elements[i] for i in qu_indices]
                    qu_idx_by_uid = [qubit.uid for qubit in self.qpu.quantum_elements]
                    # TODO: save and re-apply old qubit params
                    # Reset to the first value of every sweep,
                    # then override current sweep value for all qubits
                    for qubit in used_qubits:
                        tmp = dict(zip(sweep_keys, sweep_values or []))
                        qubit.update(**tmp)
                    # Create the experiment (required to update params)
                    seq = self.sequence(*args, **kwargs)
                    compiled_exp = compile_experiment(self.zi_session, seq)
                    # pulse_sheet(self.zi_setup, compiled_exp, self.exp_name)
                    before_sequence.send(sender=self)
                    result = run_experiment(self.zi_session, compiled_exp)
                    after_sequence.send(sender=self)
                    # TODO: handle multiple qubits. Maybe different datadicts?
                    raw_data = result[qu_idx_by_uid[qu_indices[0]]].result.data
                    data_to_save["data"] = raw_data
                    result = raw_data
                else:
                    # TODO: handle results for different instrumets
                    data_to_save["data"] = seq

                # Add parameters to the data to save
                datadict_keys = datadict.keys()
                for key, value in params_map.items():
                    if key in datadict_keys:
                        data_to_save[key] = args[value]
                # Add sweeps to the data to save
                if sweeps is not None:
                    for i, key in enumerate(sweep_keys):
                        data_to_save[f"sweep{i}"] = sweep_values[i]

                # Save data using plottr
                writer.add_data(**data_to_save)

            after_experiment.send()

        # Reset the qpu to its previous state
        self.qpu.quantum_operations.detach_qpu()
        self.qpu = QPU(old_qubits, self.qpu.quantum_operations)

        # Run analysis script
        try:
            anal_res = self.analyze(storage_path_local, *args, **kwargs)
            if type(anal_res) == AnalysisResult:
                anal_res = cast(AnalysisResult, anal_res)
                anal_res.save_all(storage_path_local)
                # Update QPU
                if is_laboneq_exp and not kwargs.get("no_update", False):
                    for qu_id in anal_res.updated_params.keys():
                        qubit = self.qpu.quantum_element_by_uid(qu_id)
                        qubit.update(**anal_res.updated_params[qu_id])
                # writer.save_text("analysis.md", anal_res)
                plt.show()
        except Exception as e:
            logger.error(f"Error while analyzing the data {e}")

        w_save(self.qpu, os.path.join(storage_path_local, "qpu_new.json"))
        qpu_filename = self.setup["storage"].get("qpu_filename", "qpu.json")
        w_save(
            self.qpu,
            os.path.join(db_path_local, qpu_filename),
        )

        # Copy the local folder to the server
        copy_folder(storage_path_local, storage_path)

    def run_raw(self, *args, **kwargs):
        before_experiment.send(sender=self)

        seq = self.sequence(*args, **kwargs)
        is_laboneq_exp = type(seq) == LaboneQExperiment
        result = None

        if is_laboneq_exp:
            compiled_exp = compile_experiment(self.zi_session, seq)
            result = run_experiment(self.zi_session, compiled_exp)
        else:
            result = seq

        after_experiment.send(sender=self)

        return result

    def sweep_around(
        self,
        center: str | float,
        span: float | tuple[float, float],
        n_points: int = None,
        step: float = None,
        scale: str = "linear",
        qu_uid="q0",
    ):
        """
        Generates a sweep of values around a specified center, either numerically or by referencing
        a qubit parameter.

        Parameters
        ----------
        center : str or float
            Center of the sweep. If a string, it's interpreted as the name of a qubit parameter
            and resolved via `qubit_value`. If a float, used directly.
        span : float or tuple of float
            If a float, sweep will extend symmetrically by `span` on both sides of `center`.
            If a tuple `(left, right)`, creates an asymmetric sweep: `center - left` to `center + right`.
        n_points : int, optional
            Number of points in the sweep. Specify exactly one of `n_points` or `step`.
        step : float, optional
            Step size in the sweep. Specify exactly one of `n_points` or `step`.
        scale : {'linear', 'log'}, default 'linear'
            Whether to generate the sweep on a linear or logarithmic scale.
            For logarithmic sweeps, all generated values must be > 0.
        qu_uid : str, default "q0"
            Qubit identifier used to resolve `center` if it is a parameter name.

        Returns
        -------
        np.ndarray
            Array of sweep values.

        Raises
        ------
        AttributeError
            If `center` is a string and not found in the qubit's parameter set.
        ValueError
            If scale is not one of 'linear' or 'log'.
            If a log-scale sweep is requested with non-positive start/stop values.
            If both or neither of `n_points` and `step` are provided.

        Notes
        -----
        - For log scale and `step`-based sweeps, the step is interpreted in multiplicative terms,
          and an approximate number of points is derived.
        - Sweep boundaries are inclusive when using `step`, thanks to the `+ step / 2` adjustment.
        """

        if isinstance(center, str):
            value = self.qubit_value(param_id=center, qu_uid=qu_uid)
            if value is None:
                raise AttributeError(
                    f"No attribute {center} in qubit {qu_uid} parameters."
                )
            center = value

        # Handle symmetric or asymmetric span
        if isinstance(span, tuple):
            left, right = span
        else:
            left = right = span

        start = center - left
        stop = center + right

        if scale not in ("linear", "log"):
            raise ValueError("scale must be 'linear' or 'log'")

        if start <= 0 or stop <= 0:
            if scale == "log":
                raise ValueError("Logarithmic sweep requires all values > 0")

        if (n_points is None) == (step is None):
            raise ValueError("Specify exactly one of 'n_points' or 'step'")

        if scale == "linear":
            if step is not None:
                return np.arange(start, stop + step / 2, step)
            else:
                return np.linspace(start, stop, n_points)

        else:  # scale == "log"
            if step is not None:
                # Compute approximate number of points from step in log space
                log_start = np.log10(start)
                log_stop = np.log10(stop)
                num_steps = (
                    int(np.floor((log_stop - log_start) / np.log10(1 + step / start)))
                    + 1
                )
                return np.logspace(log_start, log_stop, num=num_steps)
            else:
                return np.logspace(np.log10(start), np.log10(stop), n_points)

    def qubit_value(self, param_id, qu_uid="q0"):
        """Get a qubit parameter value from the QPU."""
        params = self.qpu.quantum_element_by_uid(qu_uid).parameters
        return attrs.asdict(params).get(param_id)

qubit_value(param_id, qu_uid='q0')

Get a qubit parameter value from the QPU.

Source code in sqil_core/experiment/_experiment.py

def qubit_value(self, param_id, qu_uid="q0"):
    """Get a qubit parameter value from the QPU."""
    params = self.qpu.quantum_element_by_uid(qu_uid).parameters
    return attrs.asdict(params).get(param_id)

sequence(*args, **kwargs) abstractmethod

Experimental sequence defined by the user

Source code in sqil_core/experiment/_experiment.py

@abstractmethod
def sequence(self, *args, **kwargs):
    """Experimental sequence defined by the user"""
    pass

sweep_around(center, span, n_points=None, step=None, scale='linear', qu_uid='q0')

Generates a sweep of values around a specified center, either numerically or by referencing a qubit parameter.

Parameters:

Name	Type	Description	Default
center	str or float	Center of the sweep. If a string, it's interpreted as the name of a qubit parameter and resolved via qubit_value. If a float, used directly.	required
span	float or tuple of float	If a float, sweep will extend symmetrically by span on both sides of center. If a tuple (left, right), creates an asymmetric sweep: center - left to center + right.	required
n_points	int	Number of points in the sweep. Specify exactly one of n_points or step.	None
step	float	Step size in the sweep. Specify exactly one of n_points or step.	None
scale	('linear', 'log')	Whether to generate the sweep on a linear or logarithmic scale. For logarithmic sweeps, all generated values must be > 0.	'linear'
qu_uid	str	Qubit identifier used to resolve center if it is a parameter name.	"q0"

Returns:

Type	Description
ndarray	Array of sweep values.

Raises:

Type	Description
AttributeError	If center is a string and not found in the qubit's parameter set.
ValueError	If scale is not one of 'linear' or 'log'. If a log-scale sweep is requested with non-positive start/stop values. If both or neither of n_points and step are provided.

Notes

• For log scale and step-based sweeps, the step is interpreted in multiplicative terms, and an approximate number of points is derived.
• Sweep boundaries are inclusive when using step, thanks to the + step / 2 adjustment.

Source code in sqil_core/experiment/_experiment.py

def sweep_around(
    self,
    center: str | float,
    span: float | tuple[float, float],
    n_points: int = None,
    step: float = None,
    scale: str = "linear",
    qu_uid="q0",
):
    """
    Generates a sweep of values around a specified center, either numerically or by referencing
    a qubit parameter.

    Parameters
    ----------
    center : str or float
        Center of the sweep. If a string, it's interpreted as the name of a qubit parameter
        and resolved via `qubit_value`. If a float, used directly.
    span : float or tuple of float
        If a float, sweep will extend symmetrically by `span` on both sides of `center`.
        If a tuple `(left, right)`, creates an asymmetric sweep: `center - left` to `center + right`.
    n_points : int, optional
        Number of points in the sweep. Specify exactly one of `n_points` or `step`.
    step : float, optional
        Step size in the sweep. Specify exactly one of `n_points` or `step`.
    scale : {'linear', 'log'}, default 'linear'
        Whether to generate the sweep on a linear or logarithmic scale.
        For logarithmic sweeps, all generated values must be > 0.
    qu_uid : str, default "q0"
        Qubit identifier used to resolve `center` if it is a parameter name.

    Returns
    -------
    np.ndarray
        Array of sweep values.

    Raises
    ------
    AttributeError
        If `center` is a string and not found in the qubit's parameter set.
    ValueError
        If scale is not one of 'linear' or 'log'.
        If a log-scale sweep is requested with non-positive start/stop values.
        If both or neither of `n_points` and `step` are provided.

    Notes
    -----
    - For log scale and `step`-based sweeps, the step is interpreted in multiplicative terms,
      and an approximate number of points is derived.
    - Sweep boundaries are inclusive when using `step`, thanks to the `+ step / 2` adjustment.
    """

    if isinstance(center, str):
        value = self.qubit_value(param_id=center, qu_uid=qu_uid)
        if value is None:
            raise AttributeError(
                f"No attribute {center} in qubit {qu_uid} parameters."
            )
        center = value

    # Handle symmetric or asymmetric span
    if isinstance(span, tuple):
        left, right = span
    else:
        left = right = span

    start = center - left
    stop = center + right

    if scale not in ("linear", "log"):
        raise ValueError("scale must be 'linear' or 'log'")

    if start <= 0 or stop <= 0:
        if scale == "log":
            raise ValueError("Logarithmic sweep requires all values > 0")

    if (n_points is None) == (step is None):
        raise ValueError("Specify exactly one of 'n_points' or 'step'")

    if scale == "linear":
        if step is not None:
            return np.arange(start, stop + step / 2, step)
        else:
            return np.linspace(start, stop, n_points)

    else:  # scale == "log"
        if step is not None:
            # Compute approximate number of points from step in log space
            log_start = np.log10(start)
            log_stop = np.log10(stop)
            num_steps = (
                int(np.floor((log_stop - log_start) / np.log10(1 + step / start)))
                + 1
            )
            return np.logspace(log_start, log_stop, num=num_steps)
        else:
            return np.logspace(np.log10(start), np.log10(stop), n_points)

Instruments

Source code in sqil_core/experiment/_experiment.py

class Instruments:
    def __init__(self, data):
        self._instruments = data
        for key, value in data.items():
            setattr(self, key, value)

    def __iter__(self):
        """Allow iteration directly over instrument instances."""
        return iter(self._instruments.values())

__iter__()

Allow iteration directly over instrument instances.

Source code in sqil_core/experiment/_experiment.py

def __iter__(self):
    """Allow iteration directly over instrument instances."""
    return iter(self._instruments.values())

build_plottr_dict(db_schema)

Create a DataDict object from the given schema.

Source code in sqil_core/experiment/_experiment.py

def build_plottr_dict(db_schema):
    """Create a DataDict object from the given schema."""
    axes = []
    db = {}

    data_key = "data"
    data_unit = ""

    for key, value in db_schema.items():
        if value.get("role") in ("axis", "x-axis"):
            unit = value.get("unit", "")
            db[key] = dict(unit=unit)
            axes.append(key)
        elif value.get("role") == "data":
            data_key = key
            data_unit = value.get("unit", "")
    db[data_key] = dict(axes=axes, unit=data_unit)
    datadict = DataDict(**db)

    datadict.add_meta("schema", json.dumps(db_schema))

    return datadict

map_inputs(func)

Extracts parameter names and keyword arguments from a function signature.

Source code in sqil_core/experiment/_experiment.py

def map_inputs(func):
    """Extracts parameter names and keyword arguments from a function signature."""
    sig = inspect.signature(func)
    params = {}
    kwargs = []

    for index, (name, param) in enumerate(sig.parameters.items()):
        if param.default == inspect.Parameter.empty:
            # Positional or required argument
            params[name] = index
        else:
            # Keyword argument
            kwargs.append(name)

    return params, kwargs