Single spin: Bayesian optimization of a gate

This is similar to the O2B_Single_qubit_gate example, except that it uses Bayesian optimization instead of gradient descent.

import matplotlib.pyplot as plt
import numpy as np
from jax import Array

from paraqeet.logger import Logger
from paraqeet.measurement.unitary_fidelity import UnitaryFidelity
from paraqeet.model.closed_system import ClosedSystem
from paraqeet.model.drive_operator import DriveOperator
from paraqeet.model.qubit import Qubit
from paraqeet.optimization_map import OptimizationMap
from paraqeet.optimizers.bayesian_optimizer import BayesianOptimizer
from paraqeet.propagation.scipy_expm_goat import ScipyExpmGOAT
from paraqeet.quantity import Quantity
from paraqeet.signal.envelopes import ConstantEnvelope
from paraqeet.signal.iq_mixer import IQMixer

Setup

We first set up the qubit system we want to control. We set the qubit frequency \(\omega_q / 2 \pi\) to be \(4.8\) GHz and define the Hamiltonian as

\[H(t)=H_\text{drift}+H_c(t)= \frac{\omega_q}{2} \sigma_z + \Omega(t)\sigma_x,\]

where \(\Omega(t)\) will be supplied by the generator.

freq_q = 4.8e9
omega_q = 2 * np.pi * freq_q

controlled_qubit = Qubit(frequency=Quantity(omega_q, 0.8 * omega_q, 1.2 * omega_q, unit="Hz", two_pi=True), drives=[])
model = ClosedSystem(controlled_qubit)

For signal generation, we define a simple cosine shaped tone generator \(A \cos(\omega t)\)

t_simu = 3e-9
tone = ConstantEnvelope()
tone.t_final.set_value(t_simu)
gen = IQMixer(envelopes=[tone])

We can inspect the pre-defined parameters with

params_gen = gen.get_parameters()

In this notebook, we would like to optimize the amplitude Amplitude and frequency lo_freq if the drive. We add a drive on the qubit.

freq = 4.8e9 * 2 * np.pi

drive = DriveOperator(gen, is_longitudinal=False)
controlled_qubit = Qubit(frequency=Quantity(freq, 0.8 * freq, 1.2 * freq), drives=[drive])
model = ClosedSystem(controlled_qubit)

Textbook values for implementing an \(X\) rotation on this system at a time \(T\) would be \(\omega=\omega_q\) and \(A=\pi/T\). We use some offset from these values as initial guess to demonstrate the optimization procedure.

params_gen[0].set_value(0.5 * np.pi / t_simu)
params_gen[2].set_value(1.01 * freq)

We select a propagation method, piecewise constant exponentiation, and configure an \(X\)-gate as a target gate. Also we initialize the identity at time \(0\).

prop = ScipyExpmGOAT(model, resolution=500e9)

times = np.array([0.0, t_simu])

pauli_x = np.array([[0.0, 1.0], [1.0, 0.0]])
pauli_y = np.array([[0.0, -1.0j], [1.0j, 0.0]])
pauli_z = np.array([[1.0, 0.0], [0.0, -1.0]])


prop.set_initial_state(np.identity(2))
gate_fid = UnitaryFidelity(
    propagation=prop,
    gate=pauli_x,
)

from plotting import plot_signal_and_dynamics

ts = np.linspace(0.0, t_simu, 301)
plot_signal_and_dynamics(gen, prop, ts, state_labels=[r"$|0\rangle$", r"$|1\rangle$"]);

As expected, we get a partial transfer and a low fidelity.

print(f"Gate fidelity: {gate_fid.measure(times)}")

Gate fidelity: 0.008892622545869415

Custom logger implementation

We define an optimizer and link our fidelity measure as a goal function and the parameters of the cosine tone. We also use a custom logger class to collect all samples that the optimizer takes

samples = []


class CustomLogger(Logger):
    """Custom logger class definition."""

    def log(self, params: list[Quantity], infid: Array):
        """Log the list of quantities and the fidelity.

        Parameters
        ----------
        params: list[Quantity]
            List of parameters of the system.
        infid: Array
            Inverse of fidelity.

        """
        samples.append((params[0].get_value(), params[1].get_value()))


optmap = OptimizationMap()
optmap.add(gen, [params_gen[0], params_gen[2], params_gen[3]])
opt = BayesianOptimizer(gate_fid, optimization_map=optmap, initial_samples=10, iterations=100)
opt.logger = CustomLogger()

opt.optimize(times)

|   iter    |  target   |     0     |     1     |     2     |
-------------------------------------------------------------
| 1         | 0.0036571 | -0.165955 | 0.4406489 | -0.999771 |
| 2         | 0.0001105 | -0.395334 | -0.706488 | -0.815322 |
| 3         | 0.0003969 | -0.627479 | -0.308878 | -0.206465 |
| 4         | 0.0095249 | 0.0776334 | -0.161610 | 0.3704390 |
| 5         | 3.107e-06 | -0.591095 | 0.7562348 | -0.945224 |
| 6         | 0.2036608 | 0.3409350 | -0.165390 | 0.1173796 |

| 7         | 0.0007858 | -0.719226 | -0.603797 | 0.6014891 |

| 8         | 0.0438829 | 0.9365231 | -0.373151 | 0.3846452 |
| 9         | 1.744e-06 | 0.7527783 | 0.7892133 | -0.829911 |
| 10        | 3.607e-06 | -0.921890 | -0.660339 | 0.7562850 |

| 11        | 0.2436705 | 0.3864240 | -0.166255 | 0.0746289 |

| 12        | 0.3152984 | 0.4014301 | -0.072420 | -0.063436 |
| 13        | 0.1466834 | 0.5649364 | -0.212259 | -0.332136 |
| 14        | 0.2810023 | 0.5635366 | 0.1276022 | 0.0006607 |

| 15        | 0.0152823 | 0.2674978 | 0.2561853 | -0.169559 |
| 16        | 0.3836315 | 0.4408709 | -0.098842 | -0.084145 |
| 17        | 0.4766822 | 0.6119481 | -0.095131 | -0.052578 |

| 18        | 0.2890354 | 0.4034507 | -0.068209 | -0.066917 |
| 19        | 0.4964185 | 0.5297423 | -0.121980 | -0.029561 |

| 20        | 0.2971542 | 0.5591593 | -0.177524 | -0.092768 |
| 21        | 0.4609845 | 0.5715910 | -0.053883 | 0.0224457 |

| 22        | 0.4405435 | 0.6169550 | -0.144028 | 0.0384503 |
| 23        | 0.0011515 | 0.9867827 | -0.644633 | 0.8585118 |
| 24        | 6.931e-05 | -0.472166 | -0.869049 | -0.553717 |
| 25        | 5.177e-08 | -0.877018 | -0.657881 | 0.0207137 |

| 26        | 0.0012021 | -0.795254 | -0.508318 | 0.9524448 |
| 27        | 0.2051465 | 0.3692185 | -0.191760 | -0.075047 |

| 28        | 0.4125911 | 0.6886774 | -0.046218 | 0.0114139 |
| 29        | 0.5436281 | 0.4663919 | -0.086370 | 0.0246949 |

| 30        | 0.6413501 | 0.4979543 | -0.041705 | 0.1428720 |
| 31        | 0.7819337 | 0.5710362 | -0.045724 | 0.2137782 |

| 32        | 0.4767379 | 0.5352728 | 0.0299197 | 0.2624907 |
| 33        | 0.3013880 | 0.5579259 | -0.130762 | 0.2397892 |

| 34        | 0.0138074 | 0.3362308 | -0.840353 | -0.837670 |
| 35        | 0.0013310 | -0.483655 | 0.3866039 | 0.6740750 |
| 36        | 0.1174572 | -0.049565 | -0.108954 | -0.202153 |

| 37        | 0.5255571 | 0.6097817 | -0.004734 | 0.1679875 |

| 38        | 0.7379100 | 0.5218613 | -0.034616 | 0.2090108 |
| 39        | 0.8490859 | 0.6323366 | -0.032667 | 0.2707394 |

| 40        | 3.044e-07 | -0.831439 | 0.3584873 | -0.030378 |
| 41        | 0.9025055 | 0.7107871 | -0.024897 | 0.3139254 |

| 42        | 0.8603514 | 0.6672859 | -0.024387 | 0.3857774 |

| 43        | 0.5009260 | 0.7225206 | 0.0629590 | 0.3653689 |
| 44        | 0.0008535 | -0.324396 | 0.8964272 | -0.695403 |
| 45        | 0.7708845 | 0.9743153 | 0.0468998 | 0.6726593 |

| 46        | 0.2796883 | 0.7344644 | -0.104430 | 0.3734581 |
| 47        | 0.5285330 | 0.4842338 | -0.073633 | 0.0276697 |
| 48        | 0.1170918 | -0.455283 | 0.0556825 | -0.294502 |

| 49        | 0.0005530 | -0.896720 | -0.199863 | -0.487141 |

| 50        | 0.6889618 | 1.0       | 0.0906087 | 0.7673260 |

| 51        | 0.6838936 | 0.8762485 | 0.0603446 | 0.7296905 |

| 52        | 0.0067735 | 0.9691647 | -0.042159 | 0.7514141 |

| 53        | 0.3685365 | 0.9467601 | 0.1473334 | 0.6715911 |

| 54        | 0.8365565 | 0.6495542 | 0.0045080 | 0.3293315 |

| 55        | 0.8216861 | 0.7354916 | -0.020858 | 0.2344069 |
| 56        | 5.577e-05 | -0.879484 | -0.954218 | 0.7154813 |

| 57        | 0.7832472 | 0.5852044 | -0.015138 | 0.4526613 |

| 58        | 0.7938804 | 0.6597925 | 0.0105868 | 0.5419519 |

| 59        | 0.2440360 | 0.5661377 | -0.041364 | 0.5962431 |

| 60        | 0.8320164 | 0.7867958 | 0.0324158 | 0.5959492 |

| 61        | 0.4959326 | 0.7079475 | 0.1158142 | 0.5953950 |
| 62        | 0.1669744 | 0.1229802 | 0.1805488 | 0.0361662 |

| 63        | 0.8216539 | 0.9007799 | 0.0125311 | 0.5672909 |
| 64        | 0.8316070 | 0.7870544 | 0.0327162 | 0.5968212 |

| 65        | 0.0003834 | -0.712240 | 0.6268556 | 0.2073338 |
| 66        | 0.0001571 | 0.5744876 | 0.9758150 | 0.6370951 |

| 67        | 0.0003994 | 0.3575027 | -0.296749 | -0.023575 |

| 68        | 0.0840021 | 0.7956127 | -0.065188 | 0.6093134 |

| 69        | 0.6736729 | 0.8465739 | 0.0872708 | 0.5378366 |

| 70        | 0.9500097 | 1.0       | 0.0343603 | 0.5404349 |

| 71        | 0.9937998 | 1.0       | 0.0120896 | 0.4337827 |

| 72        | 0.2720799 | 1.0       | 0.1068703 | 0.4237261 |

| 73        | 0.4196239 | 1.0       | -0.060757 | 0.4818069 |

| 74        | 0.9787179 | 0.9329835 | 0.0194547 | 0.4620656 |

| 75        | 0.9821017 | 0.9462511 | -0.011583 | 0.3572612 |

| 76        | 0.8554800 | 0.8962685 | -0.014639 | 0.2455088 |
| 77        | 0.1200786 | 0.4775626 | 0.1868229 | 0.1878511 |

| 78        | 0.9185996 | 1.0       | -0.047368 | 0.2737846 |

| 79        | 0.4297106 | 1.0       | 0.0093531 | 0.1550547 |

| 80        | 0.4254749 | 0.9266481 | -0.114392 | 0.2660849 |

| 81        | 0.7669252 | 1.0       | 0.0218863 | 0.3068366 |

| 82        | 0.8139725 | 0.8550330 | 0.0184427 | 0.3242011 |

| 83        | 0.3778610 | 0.9226054 | 0.1628879 | 0.9119624 |
| 84        | 0.0002142 | -0.873608 | -0.590417 | -0.322908 |

| 85        | 3.088e-06 | 0.1111606 | -1.0      | 1.0       |

| 86        | 0.0547661 | 0.8283657 | 0.0747173 | 0.1985936 |

| 87        | 0.0       | -1.0      | 1.0       | 1.0       |

| 88        | 0.6663733 | 0.6280386 | 0.0664456 | 0.4674776 |

| 89        | 0.0008108 | 1.0       | -1.0      | -0.261997 |

| 90        | 0.6253145 | 1.0       | -0.101278 | -1.0      |

| 91        | 0.0019486 | 1.0       | -0.273733 | -1.0      |
| 92        | 0.0009246 | -0.723588 | 0.4703690 | -0.997180 |

| 93        | 0.0230313 | 1.0       | 0.0522465 | -1.0      |

| 94        | 0.9625407 | 0.8660610 | -0.003961 | 0.4206869 |

| 95        | 0.9176635 | 0.7439016 | 0.0140624 | 0.4802423 |

| 96        | 0.0001725 | 1.0       | 1.0       | -0.071873 |

| 97        | 0.0001131 | -0.019492 | -1.0      | 0.1577813 |
| 98        | 0.0001749 | -0.765238 | -0.898005 | 0.5169565 |

| 99        | 0.0166398 | -0.227062 | -0.174050 | 0.2779688 |

| 100       | 0.5763802 | 0.7266887 | 0.1043723 | 0.8172724 |

| 101       | 0.3153232 | 0.5995976 | 0.1592729 | 0.9868910 |

| 102       | 0.6043643 | 0.7814017 | -0.106350 | -1.0      |

| 103       | 0.6357089 | 0.8468859 | -0.106238 | -0.833025 |

| 104       | 0.6316103 | 0.6694073 | -0.103468 | -0.847527 |

| 105       | 0.0003150 | 0.7309243 | -0.252831 | -0.851165 |

| 106       | 0.2091205 | 0.7374488 | 0.0266689 | -0.864756 |

| 107       | 0.5790649 | 0.5576477 | -0.091931 | -0.977306 |
| 108       | 8.780e-05 | -0.454330 | 0.6388374 | -0.363639 |

| 109       | 0.6700995 | 0.5047875 | -0.080379 | -0.799477 |

| 110       | 0.4629568 | 0.3737748 | -0.121167 | -0.896799 |
=============================================================

{'status': 0, 'value': 0.006200132733651831, 'iterations': 110}

The plot shows the all the samples that the optimization took in the two-dimensional parameter space. The red dot marks the best value.

plt.figure(figsize=(4, 4))
plt.scatter([s[0] for s in samples[:-1]], [s[1] for s in samples[:-1]], c="blue")
plt.scatter([params_gen[0].get_value()], [params_gen[2].get_value()], c="red", marker="o", s=100)
plt.xlim(params_gen[0].get_min_value()[0], params_gen[0].get_max_value()[0])
plt.ylim(params_gen[2].get_min_value()[0], params_gen[2].get_max_value()[0])
plt.xlabel(params_gen[0].get_name())
plt.ylabel(params_gen[2].get_name())
plt.show()

plot_signal_and_dynamics(gen, prop, ts, state_labels=[r"$|0\rangle$", r"$|1\rangle$"]);

We can see from the plot and optimizer output that we have found good controls.

print(f"Gate fidelity: {gate_fid.measure(times)}")

Gate fidelity: 0.9937998672663482