!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

# the same spect pulse used in every schedule
drive_amp = 0.9
drive_sigma = 16
drive_duration = 128
spec_pulse = pulse_lib.gaussian(duration=drive_duration, amp=drive_amp, 
                                sigma=drive_sigma, name=f"Spec drive amplitude = {drive_amp}")

# Construct an np array of the frequencies for our experiment
spec_freqs_GHz = np.arange(5.0, 5.2, 0.005)

# Create the base schedule
# Start with drive pulse acting on the drive channel
spec_schedules = []
for freq in spec_freqs_GHz:
    sb_spec_pulse = helper.apply_sideband(spec_pulse, qubit_lo_freq[0]-freq*GHz, dt)
    
    spec_schedule = pulse.Schedule(name='SB Frequency = {}'.format(freq))
    spec_schedule += Play(sb_spec_pulse, drive_chan)
    # The left shift `<<` is special syntax meaning to shift the start time of the schedule by some duration
    spec_schedule += measure << spec_schedule.duration
    spec_schedules.append(spec_schedule)

spec_schedules[0].draw()

from qiskit import assemble

# assemble the schedules into a Qobj
spec01_qobj = assemble(**helper.get_params('spec01', globals()))

# run the simulation
spec01_result = backend_sim.run(spec01_qobj, duffing_model).result()

# retrieve the data from the experiment
spec01_values = helper.get_values_from_result(spec01_result, qubit)

fit_params, y_fit = helper.fit_lorentzian(spec_freqs_GHz, spec01_values, [5, 5, 1, 0])

f01 = fit_params[1]

plt.scatter(spec_freqs_GHz, np.real(spec01_values), color='white') # plot real part of sweep values
plt.plot(spec_freqs_GHz, y_fit, color='red')
plt.xlim([min(spec_freqs_GHz), max(spec_freqs_GHz)])

plt.xlabel("Frequency [GHz]")
plt.ylabel("Measured Signal [a.u.]")
plt.show()

print("01 Spectroscopy yields %f GHz"%f01)

x180_amp = 0.629070 #from lab 6 Rabi experiment

x_pulse = pulse_lib.gaussian(duration=drive_duration,
                             amp=x180_amp, 
                             sigma=drive_sigma,
                             name='x_pulse')

anharmonicity_guess_GHz = -0.3

def build_spec12_pulse_schedule(freq):
    sb12_spec_pulse = helper.apply_sideband(spec_pulse, (freq + anharmonicity_guess_GHz)*GHz, dt)
    
    ### create a 12 spectroscopy pulse schedule spec12_schedule (already done)
    ### play an x pulse on the drive channel
    ### play sidebanded spec pulse on the drive channel
    ### add measurement pulse to schedule
    
    spec12_schedule = pulse.Schedule()
    
    ### WRITE YOUR CODE BETWEEN THESE LINES - START

    spec12_schedule += Play(x_pulse, drive_chan)
    spec12_schedule += Play(sb12_spec_pulse, drive_chan)
    spec12_schedule += measure << spec12_schedule.duration

    ### WRITE YOUR CODE BETWEEN THESE LINES - STOP
    
    return spec12_schedule

sb_freqs_GHz = np.arange(-.1, .1, 0.005) # sweep +/- 100 MHz around guess

# now vary the sideband frequency for each spec pulse
spec_schedules = []
for freq in sb_freqs_GHz:
    spec_schedules.append(build_spec12_pulse_schedule(freq))

spec_schedules[0].draw()

# assemble the schedules into a Qobj
spec12_qobj = assemble(**helper.get_params('spec12', globals()))
answer1 = spec12_qobj

# run the simulation
spec12_result = backend_sim.run(spec12_qobj, duffing_model).result()

# retrieve the data from the experiment
spec12_values = helper.get_values_from_result(spec12_result, qubit)

anharm_offset = qubit_lo_freq[0]/GHz + anharmonicity_guess_GHz

fit_params, y_fit = helper.fit_lorentzian(anharm_offset + sb_freqs_GHz, spec12_values, [5, 4.5, .1, 3])

f12 = fit_params[1]

plt.scatter(anharm_offset + sb_freqs_GHz, np.real(spec12_values), color='white') # plot real part of sweep values
plt.plot(anharm_offset + sb_freqs_GHz, y_fit, color='red')
plt.xlim([anharm_offset + min(sb_freqs_GHz), anharm_offset + max(sb_freqs_GHz)])

plt.xlabel("Frequency [GHz]")
plt.ylabel("Measured Signal [a.u.]")
plt.show()

print("12 Spectroscopy yields %f GHz"%f12)
print("Measured transmon anharmonicity is %f MHz"%((f12-f01)*GHz/MHz))

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