!pip install -U -r grading_tools/requirements.txt

from IPython.display import clear_output
clear_output()

# our backend is the Pulse Simulator
from resources import helper
from qiskit.providers.aer import PulseSimulator
backend_sim = PulseSimulator()

# sample duration for pulse instructions 
dt = 1e-9

# create the model
duffing_model = helper.get_transmon(dt)

# get qubit frequency from Duffing model
qubit_lo_freq = duffing_model.hamiltonian.get_qubit_lo_from_drift()

import numpy as np

# visualization tools
import matplotlib.pyplot as plt
plt.style.use('dark_background')

# unit conversion factors -> all backend properties returned in SI (Hz, sec, etc)
GHz = 1.0e9 # Gigahertz
MHz = 1.0e6 # Megahertz
kHz = 1.0e3 # kilohertz
us = 1.0e-6 # microseconds
ns = 1.0e-9 # nanoseconds

from qiskit import pulse
from qiskit.pulse import Play, Acquire
from qiskit.pulse.pulse_lib import GaussianSquare

# qubit to be used throughout the notebook
qubit = 0

### Collect the necessary channels
drive_chan = pulse.DriveChannel(qubit)
meas_chan = pulse.MeasureChannel(qubit)
acq_chan = pulse.AcquireChannel(qubit)

# Construct a measurement schedule and add it to an InstructionScheduleMap
meas_samples = 1200
meas_pulse = GaussianSquare(duration=meas_samples, amp=0.025, sigma=4, width=1150)
measure_sched = Play(meas_pulse, meas_chan) | Acquire(meas_samples, acq_chan, pulse.MemorySlot(qubit))

inst_map = pulse.InstructionScheduleMap()
inst_map.add('measure', [qubit], measure_sched)

# save the measurement/acquire pulse for later
measure = inst_map.get('measure', qubits=[qubit])

from qiskit.pulse import pulse_lib

def build_rabi_pulse_schedule(drive_duration, drive_amp, drive_sigma):
    ### create a Rabi schedule (already done)
    ### create a Gaussian Rabi pulse using pulse_lib
    ### play Rabi pulse on the Rabi schedule and return
    
    rabi_schedule = pulse.Schedule(name='rabi_experiment')
    
    ### WRITE YOUR CODE BETWEEN THESE LINES - START
    
    rabi_pulse = pulse_lib.gaussian(duration=drive_duration, amp=drive_amp,
                                        sigma=drive_sigma)
    rabi_schedule = pulse.Schedule()
    rabi_schedule += Play(rabi_pulse, drive_chan)
    
    ### WRITE YOUR CODE BETWEEN THESE LINES - END
    
    # add measurement to rabi_schedule
    # << indicates time shift the beginning to the start of the schedule
    rabi_schedule += measure << rabi_schedule.duration
    return rabi_schedule

# Gaussian pulse parameters, with varying amplitude
drive_duration = 128
num_rabi_points = 41
drive_amps = np.linspace(0, 0.9, num_rabi_points)
drive_sigma = 16

# now vary the amplitude for each drive amp
rabi_schedules = []
for drive_amp in drive_amps:
    rabi_schedules.append(build_rabi_pulse_schedule(drive_duration, drive_amp, drive_sigma))

rabi_schedules[-1].draw() 

# assemble the schedules into a Qobj
from qiskit import assemble

rabi_qobj = assemble(**helper.get_params('rabi', globals()))
answer1a = rabi_qobj

# run the simulation
rabi_result = backend_sim.run(rabi_qobj, duffing_model).result()

# retrieve the data from the experiment
rabi_values = helper.get_values_from_result(rabi_result, qubit)

fit_params, y_fit = helper.fit_sinusoid(drive_amps, rabi_values, [1, 0, 0.5, 0])

plt.scatter(drive_amps, rabi_values, color='white')
plt.plot(drive_amps, y_fit, color='red')

drive_period = fit_params[2] # get period of rabi oscillation

plt.axvline(0, color='red', linestyle='--')
plt.axvline(drive_period/2, color='red', linestyle='--')
plt.xlabel("Drive amp [a.u.]", fontsize=15)
plt.ylabel("Measured signal [a.u.]", fontsize=15)
plt.show()

print("Pi pulse amplitude is %f"%float(drive_period/2))

# x_90 is a concise way to say pi_over_2; i.e., an X rotation of 90 degrees
x90_pulse = pulse_lib.gaussian(duration=drive_duration,
                               amp=drive_period/4, 
                               sigma=drive_sigma,
                               name='x90_pulse')

# Ramsey experiment parameters
time_max_us = 0.4
time_step_us = 0.0035
times_us = np.arange(0.1, time_max_us, time_step_us)
# Convert to units of dt
delay_times_dt = times_us * us / dt

def build_ramsey_pulse_schedule(delay):
    ### create a Ramsey pulse schedule (already done)
    ### play an x90 pulse on the drive channel
    ### play another x90 pulse after delay
    ### add measurement pulse to schedule
    
    ramsey_schedule = pulse.Schedule(name='ramsey_experiment')
    
    ### HINT: include delay by adding it to the duration of the schedule
    ###       round delay to nearest integer with int(delay)
    
    ### WRITE YOUR CODE BETWEEN THESE LINES - START
    
    ramsey_schedule = pulse.Schedule()
    ramsey_schedule += Play(x90_pulse, drive_chan)
    ramsey_schedule += Play(x90_pulse, drive_chan) << ramsey_schedule.duration + int(delay)
    ramsey_schedule += measure << ramsey_schedule.duration
    
    ### WRITE YOUR CODE BETWEEN THESE LINES - END
    
    return ramsey_schedule

# create schedules for Ramsey experiment 
ramsey_schedules = []
for delay in delay_times_dt:
    ramsey_schedules.append(build_ramsey_pulse_schedule(delay))

ramsey_schedules[-1].draw()

# assemble the schedules into a Qobj
# the helper will drive the pulses off-resonantly by an unknown value
ramsey_qobj = assemble(**helper.get_params('ramsey', globals()))
answer1b = ramsey_qobj

# run the simulation
ramsey_result = backend_sim.run(ramsey_qobj, duffing_model).result()

# retrieve the data from the experiment
ramsey_values = helper.get_values_from_result(ramsey_result, qubit)

# off-resonance component
fit_params, y_fit = helper.fit_sinusoid(times_us, ramsey_values, [1, 0.7, 0.1, 0.25])
_, _, ramsey_period_us, _, = fit_params 
del_f_MHz = 1/ramsey_period_us # freq is MHz since times in us

plt.scatter(times_us, np.real(ramsey_values), color='white')
plt.plot(times_us, y_fit, color='red', label=f"df = {del_f_MHz:.6f} MHz")
plt.xlim(np.min(times_us), np.max(times_us))
plt.xlabel('Delay between X90 pulses [$\mu$s]', fontsize=15)
plt.ylabel('Measured Signal [a.u.]', fontsize=15)
plt.title('Ramsey Experiment', fontsize=15)
plt.legend(loc=3)
plt.show()

print("Drive is off-resonant by %f MHz"%float(del_f_MHz))

name = 'Saasha Joshi'
email = 'saashajoshi08@gmail.com'

from grading_tools import grade
grade(answer1a, name, email, 'lab6', 'ex1a')
grade(answer1b, name, email, 'lab6', 'ex1b')

from IPython.display import display, Javascript;display(Javascript('IPython.notebook.save_checkpoint();'));
from grading_tools import send_code;send_code('ex1.ipynb')

