diff --git a/src/qibocal/protocols/__init__.py b/src/qibocal/protocols/__init__.py index f6aa87e61..37c9cea57 100644 --- a/src/qibocal/protocols/__init__.py +++ b/src/qibocal/protocols/__init__.py @@ -11,6 +11,7 @@ from .dispersive_shift import dispersive_shift from .dispersive_shift_qutrit import dispersive_shift_qutrit from .drag import drag_tuning +from .drag_simple import drag_simple from .flipping import flipping from .flux_dependence.qubit_crosstalk import qubit_crosstalk from .flux_dependence.qubit_flux_dependence import qubit_flux @@ -104,4 +105,5 @@ "standard_rb_2q_inter", "optimize_two_qubit_gate", "ramsey_zz", + "drag_simple", ] diff --git a/src/qibocal/protocols/drag_simple.py b/src/qibocal/protocols/drag_simple.py new file mode 100644 index 000000000..c0f2a824b --- /dev/null +++ b/src/qibocal/protocols/drag_simple.py @@ -0,0 +1,259 @@ +from dataclasses import dataclass + +import numpy as np +import plotly.graph_objects as go +from qibolab import AcquisitionType, AveragingMode, Delay, Drag, Pulse, PulseSequence +from scipy.optimize import curve_fit + +from qibocal import update +from qibocal.auto.operation import QubitId, Routine +from qibocal.calibration import CalibrationPlatform +from qibocal.protocols import drag +from qibocal.result import probability +from qibocal.update import replace + +from .utils import COLORBAND, COLORBAND_LINE, table_dict, table_html + +SEQUENCES = ["YpX9", "XpY9"] +"""Sequences used to fit drag parameter.""" + + +@dataclass +class DragTuningSimpleParameters(drag.DragTuningParameters): + """DragTuningSimple runcard inputs.""" + + +@dataclass +class DragTuningSimpleResults(drag.DragTuningResults): + """DragTuningSimple outputs.""" + + # TODO: to be removed + def __contains__(self, key): + return True + + +@dataclass +class DragTuningSimpleData(drag.DragTuningData): + """DragTuningSimple acquisition outputs.""" + + def __getitem__(self, key: tuple[QubitId, str]): + qubit, setup = key + if setup == "YpX9": + return self.data[qubit][::2] + return self.data[qubit][1::2] + + +def add_drag(pulse: Pulse, beta: float) -> Pulse: + """Add DRAG component to Gaussian Pulse.""" + return replace(pulse, envelope=Drag(rel_sigma=pulse.envelope.rel_sigma, beta=beta)) + + +def _acquisition( + params: DragTuningSimpleParameters, + platform: CalibrationPlatform, + targets: list[QubitId], +) -> DragTuningSimpleData: + """Acquisition function for DRAG experiments. + + We execute two sequences YpX9 and XpY9 following + https://rsl.yale.edu/sites/default/files/2024-08/2011-RSL-Thesis-Matthew-Reed.pdf + for different value of the DRAG parameter. + """ + + data = DragTuningSimpleData() + beta_range = np.arange(params.beta_start, params.beta_end, params.beta_step) + + sequences, all_ro_pulses = [], [] + for beta in beta_range: + for setup in SEQUENCES: + sequence = PulseSequence() + ro_pulses = {} + for q in targets: + natives = platform.natives.single_qubit[q] + ro_channel, ro_pulse = natives.MZ()[0] + if setup == "YpX9": + qd_channel, ry = natives.R(phi=np.pi / 2)[0] + _, rx90 = natives.R(theta=np.pi / 2)[0] + sequence.append((qd_channel, add_drag(ry, beta=beta))) + sequence.append((qd_channel, add_drag(rx90, beta=beta))) + sequence.append( + (ro_channel, Delay(duration=rx90.duration + ry.duration)) + ) + else: + qd_channel, rx = natives.RX()[0] + _, ry90 = natives.R(theta=np.pi / 2, phi=np.pi / 2)[0] + sequence.append((qd_channel, add_drag(rx, beta=beta))) + sequence.append((qd_channel, add_drag(ry90, beta=beta))) + sequence.append( + (ro_channel, Delay(duration=rx.duration + ry90.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(np.repeat(beta_range, 2), all_ro_pulses): + for qubit in targets: + result = results[ro_pulses[qubit].id] + prob = probability(result, state=1) + # store the results + data.register_qubit( + drag.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) + setup = "YpX9" if i % 2 == 2 else "XpY9" + for qubit in targets: + ro_pulse = list(sequence.channel(platform.qubits[qubit].acquisition))[ + -1 + ] + prob = probability(result[ro_pulse.id], state=1) + # store the results + data.register_qubit( + drag.DragTuningType, + (qubit), + dict( + prob=np.array([prob]), + error=np.array([np.sqrt(prob * (1 - prob) / params.nshots)]), + beta=np.array([beta_range[i // 2]]), + ), + ) + + return data + + +def _fit(data: DragTuningSimpleData) -> DragTuningSimpleResults: + """Post-processing for DRAG protocol. + + A linear fit is applied for the probability of both sequences. + The optimal is determined as the point in which the two lines met. + """ + qubits = data.qubits + fitted_parameters = {} + betas_optimal = {} + for qubit in qubits: + for setup in SEQUENCES: + qubit_data = data[qubit, setup] + popt, _ = curve_fit( + lambda x, a, b: a * x + b, qubit_data.beta, qubit_data.prob + ) + fitted_parameters[qubit, setup] = popt.tolist() + betas_optimal[qubit] = -( + fitted_parameters[qubit, "YpX9"][1] - fitted_parameters[qubit, "XpY9"][1] + ) / (fitted_parameters[qubit, "YpX9"][0] - fitted_parameters[qubit, "XpY9"][0]) + + return DragTuningSimpleResults(betas_optimal, fitted_parameters) + + +def _plot(data: DragTuningSimpleData, target: QubitId, fit: DragTuningSimpleResults): + """Plotting function for DragTuning.""" + + figures = [] + fitting_report = "" + + fig = go.Figure() + for setup in SEQUENCES: + qubit_data = data[target, setup] + fig.add_trace( + go.Scatter( + x=qubit_data.beta, + y=qubit_data.prob, + opacity=1, + mode="lines", + name=setup, + showlegend=True, + legendgroup=setup, + ) + ) + + fig.add_trace( + go.Scatter( + x=np.concatenate((qubit_data.beta, qubit_data.beta[::-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), + name=setup, + showlegend=False, + legendgroup=setup, + ) + ) + + # # add fitting traces + if fit is not None: + for setup in SEQUENCES: + qubit_data = data[target, setup] + betas = qubit_data.beta + beta_range = np.linspace( + min(betas), + max(betas), + 20, + ) + + fig.add_trace( + go.Scatter( + x=beta_range, + y=fit.fitted_parameters[target, setup][0] * betas + + fit.fitted_parameters[target, setup][1], + name=f"Fit {setup}", + line=go.scatter.Line(dash="dot"), + ), + ) + fitting_report = table_html( + table_dict( + target, + ["Best DRAG parameter"], + [np.round(fit.betas[target], 4)], + ) + ) + + fig.update_layout( + showlegend=True, + xaxis_title="Beta parameter", + yaxis_title="Excited state probability", + ) + + figures.append(fig) + + return figures, fitting_report + + +def _update( + results: DragTuningSimpleResults, platform: CalibrationPlatform, target: QubitId +): + update.drag_pulse_beta( + results.betas[target], + platform, + target, + ) + + +drag_simple = Routine(_acquisition, _fit, _plot, _update) +"""DragTuning Routine object."""