# Source code for qiskit_experiments.library.calibration.rough_amplitude_cal

# This code is part of Qiskit.
#
# (C) Copyright IBM 2021.
#
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""Rough amplitude calibration using Rabi."""

from collections import namedtuple
from typing import Dict, Iterable, Optional
import numpy as np

from qiskit import QuantumCircuit
from qiskit.circuit import Parameter
from qiskit.providers.backend import Backend

from qiskit_experiments.framework import ExperimentData
from qiskit_experiments.calibration_management import BaseCalibrationExperiment, Calibrations
from qiskit_experiments.library.characterization import Rabi
from qiskit_experiments.calibration_management.update_library import BaseUpdater

AnglesSchedules = namedtuple(
"AnglesSchedules", ["target_angle", "parameter", "schedule", "previous_value"]
)

[docs]class RoughAmplitudeCal(BaseCalibrationExperiment, Rabi):
"""A calibration version of the Rabi experiment.

# section: see_also
qiskit_experiments.library.characterization.rabi.Rabi
"""

def __init__(
self,
qubit: int,
calibrations: Calibrations,
schedule_name: str = "x",
amplitudes: Iterable[float] = None,
cal_parameter_name: Optional[str] = "amp",
target_angle: float = np.pi,
auto_update: bool = True,
group: str = "default",
backend: Optional[Backend] = None,
):
r"""see class :class:Rabi for details.

Args:
qubit: The qubit for which to run the rough amplitude calibration.
calibrations: The calibrations instance with the schedules.
schedule_name: The name of the schedule to calibrate. Defaults to "x".
amplitudes: A list of amplitudes to scan. If None is given 51 amplitudes ranging
from -0.95 to 0.95 will be scanned.
cal_parameter_name: The name of the parameter in the schedule to update.
target_angle: The target angle of the gate to calibrate this will default to a
:math:\pi-pulse.
auto_update: Whether or not to automatically update the calibrations. By
default this variable is set to True.
group: The group of calibration parameters to use. The default value is "default".
backend: Optional, the backend to run the experiment on.
"""
schedule = calibrations.get_schedule(
schedule_name, qubit, assign_params={cal_parameter_name: Parameter("amp")}, group=group
)

self._validate_channels(schedule, [qubit])
self._validate_parameters(schedule, 1)

super().__init__(
calibrations,
qubit,
schedule=schedule,
amplitudes=amplitudes,
backend=backend,
schedule_name=schedule_name,
cal_parameter_name=cal_parameter_name,
auto_update=auto_update,
)

# Set the pulses to update.
prev_amp = calibrations.get_parameter_value(cal_parameter_name, qubit, schedule_name)
self.experiment_options.group = group
self.experiment_options.angles_schedules = [
AnglesSchedules(
target_angle=target_angle,
parameter=cal_parameter_name,
schedule=schedule_name,
previous_value=prev_amp,
)
]

@classmethod
def _default_experiment_options(cls):
"""Default values for the rough amplitude calibration experiment.

Experiment Options:
result_index (int): The index of the result from which to update the calibrations.
angles_schedules (list(float, str, str, float)): A list of parameter update information.
Each entry of the list is a tuple with four entries: the target angle of the
rotation, the name of the amplitude parameter to update, the name of the schedule
containing the amplitude parameter to update, and the previous value of the
amplitude parameter to update. This allows one experiment to update several
schedules, see for example :class:RoughXSXAmplitudeCal.
group (str): The calibration group to which the parameter belongs. This will default
to the value "default".
"""
options = super()._default_experiment_options()

options.angles_schedules = [
AnglesSchedules(target_angle=np.pi, parameter="amp", schedule="x", previous_value=None)
]

return options

def _metadata(self) -> Dict[str, any]:
"""Add metadata to the experiment data making it more self contained.

The following keys are added to each circuit's metadata:
angles_schedules: A list of parameter update information. Each entry of the list
is a tuple with four entries: the target angle of the rotation, the name of the
amplitude parameter to update, the name of the schedule containing the amplitude
parameter to update, and the previous value of the amplitude parameter to update.
cal_group: The calibration group to which the amplitude parameters belong.
"""
param_values = []
for angle, param_name, schedule_name, _ in self.experiment_options.angles_schedules:
param_val = self._cals.get_parameter_value(
param_name,
self._physical_qubits,
schedule_name,
group=self.experiment_options.group,
)

param_values.append(
AnglesSchedules(
target_angle=angle,
parameter=param_name,
schedule=schedule_name,
previous_value=param_val,
)
)

[docs]    def update_calibrations(self, experiment_data: ExperimentData):
r"""Update the amplitude of one or several schedules.

The update rule extracts the rate of the oscillation from the fit to the cosine function.
Recall that the amplitude is the x-axis in the analysis of the :class:Rabi experiment.
The value of the amplitude is thus the desired rotation-angle divided by the rate of
the oscillation:

.. math::

A_i \to \frac{\theta_i}{\omega}

where :math:\theta_i is the desired rotation angle (e.g. :math:\pi and :math:\pi/2
for "x" and "sx" gates, respectively) and :math:\omega is the rate of the oscillation.

Args:
experiment_data: The experiment data from which to extract the measured Rabi oscillation
used to set the pulse amplitude.
"""

result_index = self.experiment_options.result_index

rate = 2 * np.pi * BaseUpdater.get_value(experiment_data, self.__outcome__, result_index)

for angle, param, schedule, prev_amp in experiment_data.metadata["angles_schedules"]:

value = np.round(angle / rate, decimals=8) * np.exp(1.0j * np.angle(prev_amp))

self._cals, experiment_data, value, param, schedule, group
)

[docs]class RoughXSXAmplitudeCal(RoughAmplitudeCal):
"""A rough amplitude calibration of x and sx gates.

# section: see_also
qiskit_experiments.library.characterization.rabi.Rabi
"""

def __init__(
self,
qubit: int,
calibrations: Calibrations,
amplitudes: Iterable[float] = None,
backend: Optional[Backend] = None,
):
"""A rough amplitude calibration that updates both the sx and x pulses."""
super().__init__(
qubit,
calibrations,
schedule_name="x",
amplitudes=amplitudes,
backend=backend,
cal_parameter_name="amp",
target_angle=np.pi,
)

self.experiment_options.angles_schedules = [
AnglesSchedules(target_angle=np.pi, parameter="amp", schedule="x", previous_value=None),
AnglesSchedules(
target_angle=np.pi / 2, parameter="amp", schedule="sx", previous_value=None
),
]

[docs]class EFRoughXSXAmplitudeCal(RoughAmplitudeCal):
"""A rough amplitude calibration of x and sx gates on the 1<->2 transition.

# section: see_also
qiskit_experiments.library.characterization.rabi.Rabi
"""

__outcome__ = "rabi_rate_12"

def __init__(
self,
qubit: int,
calibrations: Calibrations,
amplitudes: Iterable[float] = None,
backend: Optional[Backend] = None,
ef_pulse_label: str = "12",
):
r"""A rough amplitude calibration that updates both the sx and x pulses on 1<->2.

Args:
qubit: The index of the qubit (technically a qutrit) to run on.
calibrations: The calibrations instance that stores the pulse schedules.
amplitudes: The amplitudes to scan.
backend: Optional, the backend to run the experiment on.
ef_pulse_label: A label that is post-pended to "x" and "sx" to obtain the name
of the pulses that drive a :math:\pi and :math:\pi/2 rotation on
the 1<->2 transition.
"""
super().__init__(
qubit,
calibrations,
schedule_name="x" + ef_pulse_label,
amplitudes=amplitudes,
backend=backend,
cal_parameter_name="amp",
target_angle=np.pi,
)

self.experiment_options.angles_schedules = [
AnglesSchedules(
target_angle=np.pi,
parameter="amp",
schedule="x" + ef_pulse_label,
previous_value=None,
),
AnglesSchedules(
target_angle=np.pi / 2,
parameter="amp",
schedule="sx" + ef_pulse_label,
previous_value=None,
),
]

def _pre_circuit(self) -> QuantumCircuit:
"""A circuit with operations to perform before the Rabi."""
circ = QuantumCircuit(1)
circ.x(0)
return circ