Adds transmon calibration programs

sequencing/calibration.py

@@ -29,23 +29,22 @@ def sine(xs, amp=1, f0=1, phi=0, ofs=0.5):
     fft = np.fft.rfft(ys - ys.mean(), num_pts)
     ix = np.argmax(np.abs(fft))
     f0 = fs[ix]
-    phi0 = 2 * np.pi * f0 * xs[0]
-    phi = np.angle(fft[ix]) - phi0
-    phi = (phi + np.pi) % (2 * np.pi) - np.pi / 2
+    phi = np.angle(fft[ix])
+    phi = np.mod(phi, 2 * np.pi)
 
     model = lmfit.Model(sine)
-    model.set_param_hint("f0", value=f0, min=fs.min(), max=fs.max())
+    model.set_param_hint("f0", value=f0, min=fs.min(), max=2 * fs.max())
     model.set_param_hint("ofs", value=ys.mean(), min=ys.min(), max=ys.max())
     model.set_param_hint("amp", value=np.ptp(ys) / 2, min=0, max=np.ptp(ys))
-    model.set_param_hint("phi", value=phi, min=-2 * np.pi, max=2 * np.pi)
+    model.set_param_hint("phi", value=phi, min=0, max=2 * np.pi)
     return model.fit(ys, xs=xs)
 
 
 def fit_displacement(xs, ys):
     def displacement(xs, xscale=1.0, amp=1, ofs=0, n=0):
         alphas = xs * xscale
-        nbars = alphas ** 2
-        return ofs + amp * nbars ** n / factorial(int(n)) * np.exp(-nbars)
+        nbars = alphas**2
+        return ofs + amp * nbars**n / factorial(int(n)) * np.exp(-nbars)
 
     if xs[-1] > xs[0]:
         amp = ys[0] - ys[-1]
@@ -170,6 +169,230 @@ def tune_rabi(
     return (fig, ax), old_amp, new_amp
 
 
+def tune_repeated_pi_pulses(
+    system,
+    init_state,
+    e_ops=None,
+    mode_name="qubit",
+    pulse_name=None,
+    max_num_pulses=100,
+    progbar=True,
+    plot=True,
+    ax=None,
+    ylabel=None,
+    update=True,
+    verify=True,
+):
+    """Tune the amplitude of a Transmon pulse by playing train of pi pulses.
+
+    Args:
+        system (System): System containing the Transmon whose
+            pulse you want to tune.
+        init_state (qutip.Qobj): Initial state of the system.
+        e_ops (optional, list[qutip.Qobj]): List of expectation
+            operators. If none, defaults to init_state. Default: None.
+        mode_name (optional, str): Name of the Transmon mode. Default: 'qubit'.
+        pulse_name (optional, str): Name of the pulse to tune. If None,
+            will use transmon.default_pulse. Default: None.
+        max_num_pulses (optional, tuple[float, float, int]): Maximum number of
+            repeated pulses, Default: 100.
+        progbar (optional, bool): Whether to display a tqdm progress bar.
+            Default: True.
+        plot (optional, bool): Whether to plot the results: Default: True.
+        ax (optional, matplotlib axis): Axis on which to plot results. If None,
+            one is automatically created. Default: None.
+        ylabel (optional, str): ylabel for the plot. Default: None.
+        update (optional, bool): Whether to update the pulse amplitude based on
+            the fit result. Default: True.
+        verify (optional, bool): Whether to re-run the Rabi sequence with the
+            newly-determined amplitude to verify correctness. Default: True.
+
+    Returns:
+        tuple[tuple, float, float]: (fig, ax), old_amp, new_amp
+    """
+    init_state = ket2dm(init_state)
+    qubit = system.get_mode(mode_name)
+    pulse_name = pulse_name or qubit.default_pulse
+    pulse = getattr(qubit, pulse_name)
+
+    if e_ops is None:
+        e_ops = [init_state]
+    e_ops = ops2dms(e_ops)
+    e_pop = []
+    num_pulses = np.arange(max_num_pulses + 1)
+    current_state = init_state
+    prog = tqdm if progbar else lambda x, **kwargs: x
+    for _ in prog(num_pulses):
+        seq = get_sequence(system)
+        qubit.rotate_x(np.pi, pulse_name=pulse_name)
+        result = seq.run(current_state, e_ops=e_ops, only_final_state=False)
+        current_state = result.states[-1]
+        e_pop.append(result.expect[0][-1])
+    e_pop = np.array(e_pop)
+
+    fit_result = fit_sine(num_pulses, e_pop)
+    amp_scale = 0.5 / fit_result.params["f0"]
+    amp_scale = amp_scale**-1
+
+    old_amp = pulse.amp
+    new_amp = amp_scale * old_amp
+
+    if plot:
+        if ax is None:
+            fig, ax = plt.subplots(1, 1)
+        else:
+            fig = plt.gcf()
+        ax.plot(num_pulses, e_pop, "o")
+        ax.plot(num_pulses, fit_result.best_fit, "-")
+        ax.set_xlabel("Number of pulses")
+        if ylabel:
+            ax.set_ylabel(ylabel)
+        ax.grid(True)
+        ax.set_title(f"{pulse.name} Repeated pi pulses")
+        plt.pause(0.1)
+    else:
+        fig = None
+        ax = None
+    if update:
+        pulse.amp = new_amp
+        print(
+            f"Updating {qubit.name} unit amp from {old_amp:.2e} to {new_amp:.2e}.",
+            flush=True,
+        )
+    if verify:
+        _ = tune_repeated_pi_pulses(
+            system,
+            init_state,
+            e_ops=e_ops,
+            mode_name=mode_name,
+            pulse_name=pulse_name,
+            max_num_pulses=max_num_pulses,
+            progbar=progbar,
+            plot=True,
+            ax=ax,
+            update=False,
+            verify=False,
+        )
+    return (fig, ax), old_amp, new_amp
+
+
+def tune_repeated_pio2_pulses(
+    system,
+    init_state,
+    e_ops=None,
+    mode_name="qubit",
+    pulse_name=None,
+    max_num_pulses=100,
+    progbar=True,
+    plot=True,
+    ax=None,
+    ylabel=None,
+    update=True,
+    verify=True,
+):
+    """Tune the amplitude of a Transmon pulse by playing train of pi/2 pulses.
+
+    Args:
+        system (System): System containing the Transmon whose
+            pulse you want to tune.
+        init_state (qutip.Qobj): Initial state of the system.
+        e_ops (optional, list[qutip.Qobj]): List of expectation
+            operators. If none, defaults to init_state. Default: None.
+        mode_name (optional, str): Name of the Transmon mode. Default: 'qubit'.
+        pulse_name (optional, str): Name of the pulse to tune. If None,
+            will use transmon.default_pulse. Default: None.
+        max_num_pulses (optional, tuple[float, float, int]): Maximum number of
+            repeated pulses, Default: 100.
+        progbar (optional, bool): Whether to display a tqdm progress bar.
+            Default: True.
+        plot (optional, bool): Whether to plot the results: Default: True.
+        ax (optional, matplotlib axis): Axis on which to plot results. If None,
+            one is automatically created. Default: None.
+        ylabel (optional, str): ylabel for the plot. Default: None.
+        update (optional, bool): Whether to update the pulse amplitude based on
+            the fit result. Default: True.
+        verify (optional, bool): Whether to re-run the Rabi sequence with the
+            newly-determined amplitude to verify correctness. Default: True.
+
+    Returns:
+        tuple[tuple, float, float]: (fig, ax), old_amp, new_amp
+    """
+    init_state = ket2dm(init_state)
+    qubit = system.get_mode(mode_name)
+    pulse_name = pulse_name or qubit.default_pulse
+    pulse = getattr(qubit, pulse_name)
+
+    if e_ops is None:
+        e_ops = [init_state]
+    e_ops = ops2dms(e_ops)
+    e_pop = []
+    num_pulses = np.arange(max_num_pulses + 1)
+    current_state = init_state
+
+    def run_sim(current_state, N):
+        seq = get_sequence(system)
+        for _ in range(N):
+            qubit.rotate_x(np.pi / 2, pulse_name=pulse_name)
+        result = seq.run(current_state, e_ops=e_ops, only_final_state=False)
+        current_state = result.states[-1]
+        return result, current_state
+
+    result, current_state = run_sim(current_state, 1)
+    e_pop.append(result.expect[0][-1])
+
+    prog = tqdm if progbar else lambda x, **kwargs: x
+    for _ in prog(num_pulses[:-1]):
+        result, current_state = run_sim(current_state, 2)
+        e_pop.append(result.expect[0][-1])
+    e_pop = np.array(e_pop)
+
+    fit_result = fit_sine(num_pulses, e_pop)
+    print(fit_result)
+    amp_scale = 0.5 / fit_result.params["f0"]
+    amp_scale = amp_scale**-1
+
+    old_amp = pulse.amp
+    new_amp = amp_scale * old_amp
+
+    if plot:
+        if ax is None:
+            fig, ax = plt.subplots(1, 1)
+        else:
+            fig = plt.gcf()
+        ax.plot(num_pulses, e_pop, "o")
+        ax.plot(num_pulses, fit_result.best_fit, "-")
+        ax.set_xlabel("Number of pulses")
+        if ylabel:
+            ax.set_ylabel(ylabel)
+        ax.grid(True)
+        ax.set_title(f"{pulse.name} Repeated pi/2 pulses")
+        plt.pause(0.1)
+    else:
+        fig = None
+        ax = None
+    if 0:
+        pulse.amp = new_amp
+        print(
+            f"Updating {qubit.name} unit amp from {old_amp:.2e} to {new_amp:.2e}.",
+            flush=True,
+        )
+    if verify:
+        _ = tune_repeated_pio2_pulses(
+            system,
+            init_state,
+            e_ops=e_ops,
+            mode_name=mode_name,
+            pulse_name=pulse_name,
+            max_num_pulses=max_num_pulses,
+            progbar=progbar,
+            plot=True,
+            ax=ax,
+            update=False,
+            verify=False,
+        )
+    return (fig, ax), old_amp, new_amp
+
+
 def tune_drag(
     system,
     init_state,

sequencing/pulses.py

@@ -89,20 +89,20 @@ def array_pulse(
 
 def gaussian_wave(sigma, chop=4):
     ts = np.linspace(-chop // 2 * sigma, chop // 2 * sigma, int(chop * sigma // 4) * 4)
-    P = np.exp(-(ts ** 2) / (2.0 * sigma ** 2))
+    P = np.exp(-(ts**2) / (2.0 * sigma**2))
     ofs = P[0]
     return (P - ofs) / (1 - ofs)
 
 
 def gaussian_deriv_wave(sigma, chop=4):
     ts = np.linspace(-chop // 2 * sigma, chop // 2 * sigma, int(chop * sigma // 4) * 4)
-    ofs = np.exp(-ts[0] ** 2 / (2 * sigma ** 2))
-    return (0.25 / sigma ** 2) * ts * np.exp(-(ts ** 2) / (2 * sigma ** 2)) / (1 - ofs)
+    ofs = np.exp(-ts[0] ** 2 / (2 * sigma**2))
+    return (0.25 / sigma**2) * ts * np.exp(-(ts**2) / (2 * sigma**2)) / (1 - ofs)
 
 
 def gaussian_chop(t, sigma, t0):
-    P = np.exp(-(t ** 2) / (2.0 * sigma ** 2))
-    ofs = np.exp(-(t0 ** 2) / (2.0 * sigma ** 2))
+    P = np.exp(-(t**2) / (2.0 * sigma**2))
+    ofs = np.exp(-(t0**2) / (2.0 * sigma**2))
     return (P - ofs) / (1 - ofs)
 
 
@@ -135,9 +135,9 @@ def _ring_up(ts):
         if np.abs(ts) < ramp_offset:
             return 1.0
         elif ts > ramp_offset:
-            return np.exp(-((ts - ramp_offset) ** 2) / (2.0 * sigma ** 2))
+            return np.exp(-((ts - ramp_offset) ** 2) / (2.0 * sigma**2))
         else:  # ts < ramp_offset
-            return np.exp(-((ts + ramp_offset) ** 2) / (2.0 * sigma ** 2))
+            return np.exp(-((ts + ramp_offset) ** 2) / (2.0 * sigma**2))
 
     ts = np.linspace(-length + 1, 0, length)
     P = np.array([_ring_up(t) for t in ts])

sequencing/test/test_calibration.py

@@ -16,7 +16,7 @@ def tearDownClass(cls):
     def test_rabi_two_levels(self):
         qubit = Transmon("qubit", levels=2)
         system = System("system", modes=[qubit])
-        for _ in range(2):
+        for _ in range(5):
             _, old_amp, new_amp = tune_rabi(system, qubit.fock(0))
         self.assertLess(abs(old_amp - new_amp), 1e-7)
 
@@ -41,7 +41,7 @@ def test_drag(self):
         qubit = Transmon("qubit", levels=3, kerr=-200e-3)
         qubit.gaussian_pulse.sigma = 10
         system = System("system", modes=[qubit])
-        for _ in range(3):
+        for _ in range(5):
             _, old_amp, new_amp = tune_rabi(
                 system, qubit.fock(0), plot=False, verify=False
             )

sequencing/test/test_qasm.py

@@ -16,7 +16,7 @@ def setUpClass(cls):
 
         for qubit in [q0, q1]:
             init_state = system.fock()
-            for _ in range(1):
+            for _ in range(3):
                 _ = tune_rabi(
                     system, init_state=init_state, mode_name=qubit.name, verify=False
                 )