from dataclasses import dataclass, field
from typing import Optional
import numpy as np
import numpy.typing as npt
import plotly.graph_objects as go
from qibolab import AcquisitionType, AveragingMode, Delay, Drag, PulseSequence
from scipy.optimize import curve_fit
from qibocal import update
from qibocal.auto.operation import Data, Parameters, QubitId, Results, Routine
from qibocal.calibration import CalibrationPlatform
from qibocal.config import log
from qibocal.result import probability
from qibocal.update import replace
from .utils import (
COLORBAND,
COLORBAND_LINE,
chi2_reduced,
fallback_period,
guess_period,
table_dict,
table_html,
)
# TODO: add errors in fitting
[docs]@dataclass
class DragTuningParameters(Parameters):
"""DragTuning runcard inputs."""
beta_start: float
"""DRAG pulse beta start sweep parameter."""
beta_end: float
"""DRAG pulse beta end sweep parameter."""
beta_step: float
"""DRAG pulse beta sweep step parameter."""
unrolling: bool = False
"""If ``True`` it uses sequence unrolling to deploy multiple sequences in a single instrument call.
Defaults to ``False``."""
nflips: int = 1
"""Repetitions of (Xpi - Xmpi)."""
[docs]@dataclass
class DragTuningResults(Results):
"""DragTuning outputs."""
betas: dict[QubitId, float]
"""Optimal beta paramter for each qubit."""
fitted_parameters: dict[QubitId, dict[str, float]]
"""Raw fitting output."""
chi2: dict[QubitId, tuple[float, Optional[float]]] = field(default_factory=dict)
"""Chi2 calculation."""
DragTuningType = np.dtype(
[("prob", np.float64), ("error", np.float64), ("beta", np.float64)]
)
[docs]@dataclass
class DragTuningData(Data):
"""DragTuning acquisition outputs."""
data: dict[QubitId, npt.NDArray[DragTuningType]] = field(default_factory=dict)
"""Raw data acquired."""
[docs]def _acquisition(
params: DragTuningParameters,
platform: CalibrationPlatform,
targets: list[QubitId],
) -> DragTuningData:
r"""
Data acquisition for drag pulse tuning experiment.
See https://arxiv.org/pdf/1504.06597.pdf Fig. 2 (c).
"""
data = DragTuningData()
beta_param_range = np.arange(params.beta_start, params.beta_end, params.beta_step)
sequences, all_ro_pulses = [], []
for beta_param in beta_param_range:
sequence = PulseSequence()
ro_pulses = {}
for q in targets:
natives = platform.natives.single_qubit[q]
qd_channel, qd_pulse = natives.RX()[0]
ro_channel, ro_pulse = natives.MZ()[0]
drag = replace(
qd_pulse,
envelope=Drag(
rel_sigma=qd_pulse.envelope.rel_sigma,
beta=beta_param,
),
)
drag_negative = replace(drag, relative_phase=np.pi)
for _ in range(params.nflips):
sequence.append((qd_channel, drag))
sequence.append((qd_channel, drag_negative))
sequence.append(
(
ro_channel,
Delay(
duration=params.nflips
* (drag.duration + drag_negative.duration)
),
)
)
sequence.append((ro_channel, ro_pulse))
sequences.append(sequence)
all_ro_pulses.append(
{
qubit: list(sequence.channel(platform.qubits[q].acquisition))[-1]
for qubit in targets
}
)
options = {
"nshots": params.nshots,
"relaxation_time": params.relaxation_time,
"acquisition_type": AcquisitionType.DISCRIMINATION,
"averaging_mode": AveragingMode.SINGLESHOT,
}
# execute the pulse sequence
if params.unrolling:
results = platform.execute(sequences, **options)
for beta, ro_pulses in zip(beta_param_range, all_ro_pulses):
for qubit in targets:
result = results[ro_pulses[qubit].id]
prob = probability(result, state=0)
# store the results
data.register_qubit(
DragTuningType,
(qubit),
dict(
prob=np.array([prob]),
error=np.array([np.sqrt(prob * (1 - prob) / params.nshots)]),
beta=np.array([beta]),
),
)
else:
for i, sequence in enumerate(sequences):
result = platform.execute([sequence], **options)
for qubit in targets:
ro_pulse = list(sequence.channel(platform.qubits[qubit].acquisition))[
-1
]
prob = probability(result[ro_pulse.id], state=0)
# store the results
data.register_qubit(
DragTuningType,
(qubit),
dict(
prob=np.array([prob]),
error=np.array([np.sqrt(prob * (1 - prob) / params.nshots)]),
beta=np.array([beta_param_range[i]]),
),
)
return data
[docs]def drag_fit(x, offset, amplitude, period, phase):
return offset + amplitude * np.cos(2 * np.pi * x / period + phase)
[docs]def _fit(data: DragTuningData) -> DragTuningResults:
qubits = data.qubits
betas_optimal = {}
fitted_parameters = {}
chi2 = {}
for qubit in qubits:
qubit_data = data[qubit]
# normalize prob
prob = qubit_data.prob
prob_min = np.min(prob)
prob_max = np.max(prob)
normalized_prob = (prob - prob_min) / (prob_max - prob_min)
# normalize beta
beta_params = qubit_data.beta
beta_min = np.min(beta_params)
beta_max = np.max(beta_params)
normalized_beta = (beta_params - beta_min) / (beta_max - beta_min)
# Guessing period using fourier transform
period = fallback_period(guess_period(normalized_beta, normalized_prob))
pguess = [0.5, 0.5, period, 0]
try:
popt, _ = curve_fit(
drag_fit,
normalized_beta,
normalized_prob,
p0=pguess,
maxfev=100000,
bounds=(
[0, 0, 0, -np.pi],
[1, 1, np.inf, np.pi],
),
sigma=qubit_data.error,
)
translated_popt = [
popt[0] * (prob_max - prob_min) + prob_min,
popt[1] * (prob_max - prob_min),
popt[2] * (beta_max - beta_min),
popt[3] - 2 * np.pi * beta_min / popt[2] / (beta_max - beta_min),
]
fitted_parameters[qubit] = translated_popt
predicted_prob = drag_fit(beta_params, *translated_popt)
betas_optimal[qubit] = beta_params[np.argmax(predicted_prob)]
chi2[qubit] = (
chi2_reduced(
prob,
predicted_prob,
qubit_data.error,
),
np.sqrt(2 / len(prob)),
)
except Exception as e:
log.warning(f"drag_tuning_fit failed for qubit {qubit} due to {e}.")
return DragTuningResults(betas_optimal, fitted_parameters, chi2=chi2)
[docs]def _plot(data: DragTuningData, target: QubitId, fit: DragTuningResults):
"""Plotting function for DragTuning."""
figures = []
fitting_report = ""
qubit_data = data[target]
betas = qubit_data.beta
fig = go.Figure(
[
go.Scatter(
x=qubit_data.beta,
y=qubit_data.prob,
opacity=1,
mode="lines",
name="Probability",
showlegend=True,
legendgroup="Probability",
),
go.Scatter(
x=np.concatenate((betas, betas[::-1])),
y=np.concatenate(
(
qubit_data.prob + qubit_data.error,
(qubit_data.prob - qubit_data.error)[::-1],
)
),
fill="toself",
fillcolor=COLORBAND,
line=dict(color=COLORBAND_LINE),
showlegend=True,
name="Errors",
),
]
)
# add fitting traces
if fit is not None:
beta_range = np.linspace(
min(betas),
max(betas),
20,
)
fig.add_trace(
go.Scatter(
x=beta_range,
y=drag_fit(beta_range, *fit.fitted_parameters[target]),
name="Fit",
line=go.scatter.Line(dash="dot"),
),
)
fitting_report = table_html(
table_dict(
target,
["Optimal Beta Param", "Chi2 reduced"],
[(np.round(fit.betas[target], 4), 0), fit.chi2[target]],
display_error=True,
)
)
fig.update_layout(
showlegend=True,
xaxis_title="Beta parameter",
yaxis_title="Ground State Probability",
)
figures.append(fig)
return figures, fitting_report
[docs]def _update(results: DragTuningResults, platform: CalibrationPlatform, target: QubitId):
update.drag_pulse_beta(
results.betas[target],
platform,
target,
)
drag_tuning = Routine(_acquisition, _fit, _plot, _update)
"""DragTuning Routine object."""