# Copyright 2025 Qilimanjaro Quantum Tech
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from pathlib import Path
from typing import Callable, ClassVar, Literal
import matplotlib.pyplot as plt
import numpy as np
from loguru import logger
from matplotlib.figure import Figure
from qilisdk.functionals.functional_result import FunctionalResult
from qilisdk.yaml import yaml
@yaml.register_class
[docs]
class Dimension:
"""Represents a labeled dimension in an experiment sweep.
A `Dimension` defines one or more sweep parameters, such as drive
amplitude, frequency, or delay time, together with their associated
numerical values.
"""
def __init__(self, labels: list[str], values: list[np.ndarray]) -> None:
"""Initialize a Dimension object.
Args:
labels (list[str]): Labels describing each dimension (e.g. ``["Drive amplitude"]``).
values (list[np.ndarray]): Numerical arrays for the corresponding parameter values.
"""
def __repr__(self) -> str:
return f"Dimension(labels={self.labels}, values={self.values})"
_LABEL_AMPLITUDE = "Amplitude (V)"
_LABEL_PHASE = "Phase (rad)"
_LABEL_DB = "Amplitude (dB)"
[docs]
DimensionOverride = Callable[[Dimension], Dimension]
"""Callable that takes a Dimension and returns a transformed Dimension."""
def _secondary_axis_limits(
primary: np.ndarray, secondary: np.ndarray, primary_limits: tuple[float, float]
) -> tuple[float, float]:
"""Map the limits of a primary axis onto the values of its secondary twin axis.
A twin axis relabels the very same positions as the primary axis, so its limits are the
secondary values at the primary limits, taken from the linear map anchored on the first
and last swept point. Using the smallest and largest secondary values instead would flip
a secondary axis that runs opposite to the primary one, for example a current ramped down
while the bias it sets is swept up.
Args:
primary (np.ndarray): Values swept along the primary axis.
secondary (np.ndarray): Values swept along the secondary axis, paired point by point
with `primary`.
primary_limits (tuple[float, float]): Limits currently displayed by the primary axis.
Returns:
tuple[float, float]: The limits to display on the secondary axis.
"""
first, last = float(primary[0]), float(primary[-1])
secondary_first, secondary_last = float(secondary[0]), float(secondary[-1])
if first == last:
return secondary_first, secondary_last
slope = (secondary_last - secondary_first) / (last - first)
low, high = primary_limits
return secondary_first + (low - first) * slope, secondary_first + (high - first) * slope
@yaml.register_class
[docs]
class ExperimentResult(FunctionalResult):
"""Base class for storing and visualizing experiment results.
This class defines common utilities for handling experimental data,
including computation of S21 parameters and automatic 1D or 2D plotting.
Subclasses provide specific sweep parameters and plot titles.
"""
[docs]
plot_title: ClassVar[str]
"""Default plot title; subclasses provide the concrete label."""
[docs]
dims_override: ClassVar[list[DimensionOverride | None]] = []
"""Per-dimension overrides; each entry is a callable transforming that dimension, or None to use the default."""
[docs]
fit_by_default: ClassVar[bool] = False
"""Whether to perform fitting by default when plotting; can be overridden by subclasses if needed."""
def __init__(self, qubit: int, averages: int, data: np.ndarray, dims: list[Dimension]) -> None:
"""Initialize an experiment result.
Args:
qubit (int): The qubit index on which the experiment was performed.
averages (int): Number of averages acquired for the experiment.
data (np.ndarray): Raw experimental data array.
dims (list[Dimension]): Sweep dimensions of the experiment.
"""
[docs]
self.averages = averages
@property
[docs]
def s21(self) -> np.ndarray:
"""Complex S21 transmission parameter.
Returns:
np.ndarray: The complex-valued S21 response computed as ``Re + i * Im``.
"""
return self.data[..., 0] + 1j * self.data[..., 1]
@property
[docs]
def s21_modulus(self) -> np.ndarray:
"""Magnitude of the S21 parameter.
Returns:
np.ndarray: The absolute value of the S21 parameter.
"""
return np.abs(self.s21)
@property
[docs]
def s21_db(self) -> np.ndarray:
"""Magnitude of S21 in decibels (dB).
Returns:
np.ndarray: ``20 * log10(abs(S21))`` expressed in dB.
"""
return 20 * np.log10(self.s21_modulus)
@property
[docs]
def s21_phase(self) -> np.ndarray:
"""Phase of the S21 parameter in radians.
Returns:
np.ndarray: The angle of the complex S21 parameter.
"""
return np.unwrap(np.angle(self.s21))
@staticmethod
[docs]
def add_fit(x_values: np.ndarray, y_values: np.ndarray, initial_guess: list[float] | None = None) -> None:
"""
Fit a user-provided function to the experimental data.
This should be implemented by subclasses to provide specific fitting functionality relevant to the experiment type.
Args:
x_values (np.ndarray): The independent variable data (e.g., frequencies, drive durations).
y_values (np.ndarray): The dependent variable data (e.g., measured signal).
initial_guess (list[float] | None): Optional initial guess for the fit parameters. The specific parameters depend on the fit model used by the subclass.
"""
def _save_figure(self, figure: Figure, save_to: str | Path) -> None:
"""Save the figure to disk, handling both file and directory paths.
Args:
figure (Figure): The Matplotlib figure to save.
save_to (str | Path): The path or directory where the figure should be saved.
"""
save_to = Path(save_to)
if save_to.is_dir():
save_to /= f"{self.plot_title}_qubit{self.qubit}.png"
save_to.parent.mkdir(parents=True, exist_ok=True)
figure.savefig(save_to)
def _plot_1d(
self,
s21: np.ndarray,
dims: list[Dimension],
fit: bool = False,
save_to: str | None = None,
initial_guess: list[float] | None = None,
connect_points: bool = False,
default_y_label: str = _LABEL_AMPLITUDE,
apply_y_override: bool = True,
) -> None:
"""Plot 1D S21 data.
Args:
s21 (np.ndarray): The S21 data to plot.
dims (list[Dimension]): The dimensions of the experiment, used for labeling axes.
fit (bool): Whether to perform and plot the fit using the `add_fit` method.
save_to (str | None): Optional path or directory to save the figure.
initial_guess (list[float] | None): Optional initial guess passed to `add_fit`.
connect_points (bool): Whether to connect data points with a grey dashed line.
default_y_label (str): Default y-axis label used when no dims_override is set for the y dimension.
apply_y_override (bool): Whether to apply dims_override for the y dimension. Set to False for non-amplitude plot types.
"""
x_dim = self.dims_override[0](dims[0]) if len(self.dims_override) > 0 and self.dims_override[0] else dims[0]
x_labels, x_values = x_dim.labels, x_dim.values
y_override = self.dims_override[1] if apply_y_override and len(self.dims_override) > 1 else None
fig, ax1 = plt.subplots()
ax1.set_title(f"{self.plot_title} - Qubit {self.qubit}")
ax1.set_xlabel(x_labels[0])
y_dim_input = Dimension(labels=[default_y_label], values=[s21])
ax1.set_ylabel(y_override(y_dim_input).labels[0] if y_override else default_y_label)
if connect_points:
ax1.plot(x_values[0], s21, "--", color="grey", linewidth=0.8, zorder=1)
ax1.plot(x_values[0], s21, ".")
if fit:
self.add_fit(x_values[0], s21, initial_guess=initial_guess)
if len(x_labels) > 1:
ax2 = ax1.twiny()
ax2.set_xlabel(x_labels[1])
secondary_x_limits = _secondary_axis_limits(x_values[0], x_values[1], ax1.get_xlim())
ax2.set_xlim(secondary_x_limits)
ax2.set_xticks(np.linspace(*secondary_x_limits, num=6))
ax2.ticklabel_format(axis="x", style="sci", scilimits=(-3, 3))
fig.tight_layout()
if save_to:
self._save_figure(fig, save_to)
plt.show()
plt.close(fig)
def _plot_2d(
self,
s21: np.ndarray,
dims: list[Dimension],
fit: bool = False,
save_to: str | None = None,
initial_guess: list[float] | None = None,
default_z_label: str = _LABEL_AMPLITUDE,
apply_z_override: bool = True,
) -> None:
"""Plot 2D S21 data as a color mesh.
Args:
s21 (np.ndarray): The 2D S21 data to plot.
dims (list[Dimension]): The dimensions of the experiment, used for labeling axes.
fit (bool): Whether to perform and plot the fit using the `add_fit` method.
save_to (str | None): Optional path or directory to save the figure.
initial_guess (list[float] | None): Optional initial guess passed to `add_fit`.
default_z_label (str): Default colorbar label used when no dims_override is set for the z dimension.
apply_z_override (bool): Whether to apply dims_override for the z dimension. Set to False for non-amplitude plot types.
"""
x_dim = self.dims_override[0](dims[0]) if len(self.dims_override) > 0 and self.dims_override[0] else dims[0]
y_dim = self.dims_override[1](dims[1]) if len(self.dims_override) > 1 and self.dims_override[1] else dims[1]
z_override = (
self.dims_override[2]
if apply_z_override and len(self.dims_override) > 2 # ruff:ignore[magic-value-comparison]
else None
)
x_labels, x_values = x_dim.labels, x_dim.values
y_labels, y_values = y_dim.labels, y_dim.values
z_dim_input = Dimension(labels=[default_z_label], values=[s21])
z_dim = z_override(z_dim_input) if z_override else z_dim_input
z_values = z_dim.values[0]
fig, ax1 = plt.subplots()
ax1.set_title(f"{self.plot_title} - Qubit {self.qubit}")
ax1.set_xlabel(x_labels[0])
ax1.set_ylabel(y_labels[0])
ax1.ticklabel_format(axis="both", style="sci", scilimits=(-3, 3))
mesh = ax1.pcolormesh(x_values[0], y_values[0], z_values.T, cmap="viridis", shading="nearest")
colorbar_label = z_dim.labels[0]
fig.colorbar(mesh, ax=ax1, label=colorbar_label)
if len(x_labels) > 1:
ax2 = ax1.twiny()
ax2.set_xlabel(x_labels[1])
secondary_x_limits = _secondary_axis_limits(x_values[0], x_values[1], ax1.get_xlim())
ax2.set_xlim(secondary_x_limits)
ax2.set_xticks(np.linspace(*secondary_x_limits, num=6))
ax2.ticklabel_format(axis="x", style="sci", scilimits=(-3, 3))
if len(y_labels) > 1:
ax3 = ax1.twinx()
ax3.set_ylabel(y_labels[1])
secondary_y_limits = _secondary_axis_limits(y_values[0], y_values[1], ax1.get_ylim())
ax3.set_ylim(secondary_y_limits)
ax3.set_yticks(np.linspace(*secondary_y_limits, num=6))
ax3.ticklabel_format(axis="y", style="sci", scilimits=(-3, 3))
fig.tight_layout()
if save_to:
self._save_figure(fig, save_to)
plt.show()
plt.close(fig)
[docs]
def plot(
self,
save_to: str | None = None,
*,
initial_guess: list[float] | None = None,
fit: bool | None = None,
connect_points: bool = False,
plot_type: Literal["amplitude", "phase", "db"] = "amplitude",
) -> None:
"""Plot the S21 parameter from experiment results.
Automatically detects whether the dataset is 1D or 2D and creates
the appropriate figure. Optionally saves the figure to disk.
Args:
save_to (str | None): Optional path or directory to save the
generated plot. If a directory is provided, the filename is
automatically generated as ``{plot_title}_qubit{qubit}.png``.
initial_guess (list[float] | None): Optional initial guess for the fit parameters, passed to the `add_fit` method.
fit (bool | None): Whether to perform and plot the fit using the `add_fit` method. If None, the class-level `fit_by_default` is used.
connect_points (bool): Whether to connect data points with a grey dashed line (1D plots only).
plot_type (Literal["amplitude", "phase", "db"]): Whether to plot amplitude (default), phase, or magnitude in dB of the S21 parameter.
Raises:
NotImplementedError: If the experiment data has more than 2 dimensions.
"""
if plot_type == "phase":
to_plot = self.s21_phase
default_label = _LABEL_PHASE
elif plot_type == "db":
to_plot = self.s21_db
default_label = _LABEL_DB
else:
to_plot = self.s21_modulus
default_label = _LABEL_AMPLITUDE
is_amplitude = plot_type == "amplitude"
n_dimensions = len(to_plot.shape)
should_fit = fit if fit is not None else self.fit_by_default
if fit is not None and fit and not is_amplitude:
logger.warning(
"[ExperimentResult] Fitting is only implemented for amplitude plots. Ignoring fit request for non-amplitude plot."
)
should_fit = False
if n_dimensions == 1:
self._plot_1d(
to_plot,
self.dims,
fit=should_fit,
save_to=save_to,
initial_guess=initial_guess,
connect_points=connect_points,
default_y_label=default_label,
apply_y_override=is_amplitude,
)
elif n_dimensions == 2: # ruff: ignore[magic-value-comparison]
self._plot_2d(
to_plot,
self.dims,
fit=should_fit,
save_to=save_to,
initial_guess=initial_guess,
default_z_label=default_label,
apply_z_override=is_amplitude,
)
else:
raise NotImplementedError("3D and higher dimension plots are not supported yet.")
def __repr__(self) -> str:
return f"{self.__class__.__name__}(qubit={self.qubit}, averages={self.averages}, data={self.data}, dims={self.dims})"
