Single spin: Bayesian optimization of a gate

This is similar to the 02B_Single_qubit_gate example, except that it uses Bayesian optimization (Shahriari et al., 2016) [13] instead of gradient descent.

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

from paraqeet.eom.schroedinger_equation import SchroedingerEquation
from paraqeet.hamiltonian.drive import Drive
from paraqeet.hamiltonian.qubit import QubitHamiltonian
from paraqeet.logger import Logger
from paraqeet.measurement.unitary_fidelity import UnitaryFidelity
from paraqeet.optimization_map import OptimizationMap
from paraqeet.optimizers.bayesian_optimizer import BayesianOptimizer
from paraqeet.propagation import Expm
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

qubit_hamiltonian = QubitHamiltonian(
    frequency=Quantity(omega_q, 0.8 * omega_q, 1.2 * omega_q, unit="Hz", two_pi=True), drives=[]
)
model = SchroedingerEquation(
    hamiltonian_func=qubit_hamiltonian.get_value,
    hamiltonian_gradient_func=qubit_hamiltonian.get_gradient,
)

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 of the drive. We add a drive on the qubit.

freq = 4.8e9 * 2 * np.pi
sigma_x = qubit_hamiltonian.sigma_x
drive = Drive(sigma_x, gen)
qubit_hamiltonian.drives = [drive]
model = SchroedingerEquation(
    hamiltonian_func=qubit_hamiltonian.get_value,
    hamiltonian_gradient_func=qubit_hamiltonian.get_gradient,
)

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 an 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\).

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

prop = Expm(eom_func=model.get_value, resolution=100e9, initial_state=np.identity(2))
gate_fid = UnitaryFidelity(propagation_func=prop.get_value, propagation_gradient_func=None, gate=sigma_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$"])

array([<Axes: ylabel='Amplitude n[MHz / $2\pi$]'>,
       <Axes: xlabel='Time [ns]', ylabel='Population'>], dtype=object)

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

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

Gate fidelity: 0.008842872531522971

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(measure_func=gate_fid.get_value, optimization_map=optmap, initial_samples=10, iterations=100)
opt.logger = CustomLogger()

opt.optimize(times)

|   iter    |  target   |     0     |     1     |     2     |
-------------------------------------------------------------
| 1         | 0.0036278 | -0.165955 | 0.4406489 | -0.999771 |
| 2         | 0.0001104 | -0.395334 | -0.706488 | -0.815322 |
| 3         | 0.0003939 | -0.627479 | -0.308878 | -0.206465 |
| 4         | 0.0094778 | 0.0776334 | -0.161610 | 0.3704390 |
| 5         | -.085e-06 | -0.591095 | 0.7562348 | -0.945224 |
| 6         | 0.2028943 | 0.3409350 | -0.165390 | 0.1173796 |
| 7         | 0.0007805 | -0.719226 | -0.603797 | 0.6014891 |
| 8         | 0.0434133 | 0.9365231 | -0.373151 | 0.3846452 |
| 9         | -.729e-06 | 0.7527783 | 0.7892133 | -0.829911 |
| 10        | -.574e-06 | -0.921890 | -0.660339 | 0.7562850 |

| 11        | 0.2425745 | 0.3864646 | -0.166284 | 0.0747446 |

| 12        | 0.3210977 | 0.3995182 | -0.075493 | -0.066396 |
| 13        | 0.1647017 | 0.5639544 | -0.194260 | -0.341354 |
| 14        | 0.3215744 | 0.5552035 | 0.1456268 | -0.017306 |

| 15        | -.246e-06 | 0.2755778 | 0.3197374 | -0.141076 |
| 16        | 0.2212795 | 0.5554470 | 0.1939444 | -0.022655 |
| 17        | 0.1641519 | 0.5808638 | -0.020411 | -0.013947 |

| 18        | 0.3386069 | 0.3983636 | -0.078108 | -0.062977 |

| 19        | 0.4735536 | 0.3644688 | -0.102168 | -0.005663 |
| 20        | 0.3445419 | 0.3365286 | -0.152913 | -0.047864 |

| 21        | 0.2996809 | 0.3799188 | -0.030457 | 0.0350806 |
| 22        | 0.3506612 | 0.2964954 | -0.068823 | -0.020524 |
| 23        | 0.0011528 | 0.9867827 | -0.644633 | 0.8585118 |
| 24        | -.894e-05 | -0.472166 | -0.869049 | -0.553717 |

| 25        | -.132e-08 | -0.877018 | -0.657881 | 0.0207137 |
| 26        | 0.0011949 | -0.795254 | -0.508318 | 0.9524448 |

| 27        | 0.4317754 | 0.4324396 | -0.137114 | -0.016619 |
| 28        | 0.0843465 | 0.5621222 | 0.1172306 | 0.1098338 |

| 29        | 0.4592509 | 0.6068948 | 0.1109552 | -0.122137 |

| 30        | 0.4500461 | 0.7011788 | 0.1225493 | -0.096994 |
| 31        | 0.5526225 | 0.6880545 | 0.1237988 | -0.212676 |

| 32        | 0.2963422 | 0.7218859 | 0.0311079 | -0.204660 |
| 33        | 0.0725034 | 0.6894703 | 0.2209747 | -0.202374 |

| 34        | 0.0137368 | 0.3362308 | -0.840353 | -0.837670 |
| 35        | 0.6050386 | 0.6234342 | 0.0913049 | -0.226812 |

| 36        | 0.1168021 | -0.049565 | -0.108954 | -0.202153 |
| 37        | 0.6613792 | 0.6515588 | 0.0986467 | -0.300028 |

| 38        | 0.6427795 | 0.5720024 | 0.0964943 | -0.330979 |
| 39        | 0.7138039 | 0.6342141 | 0.0810681 | -0.410679 |

| 40        | -.035e-07 | -0.831439 | 0.3584873 | -0.030378 |

| 41        | 0.1030228 | 0.6117614 | 0.1681040 | -0.428471 |
| 42        | 0.6177010 | 0.6223520 | 0.0172091 | -0.367927 |

| 43        | 0.8360415 | 0.7096159 | 0.0516334 | -0.406770 |
| 44        | 0.0008451 | -0.324396 | 0.8964272 | -0.695403 |
| 45        | 0.7703155 | 0.9743153 | 0.0468998 | 0.6726593 |

| 46        | 0.8642772 | 0.7028993 | 0.0132453 | -0.486736 |

| 47        | 0.9247866 | 0.8027911 | 0.0301315 | -0.485857 |
| 48        | 0.1162966 | -0.455283 | 0.0556825 | -0.294502 |
| 49        | 0.0005492 | -0.896720 | -0.199863 | -0.487141 |

| 50        | 0.2695574 | 0.8031122 | -0.066686 | -0.485512 |

| 51        | 0.4891204 | 0.7696168 | 0.0946999 | -0.517352 |

| 52        | 0.8797236 | 0.8181353 | 0.0512836 | -0.418533 |

| 53        | 0.8835306 | 0.8908719 | 0.0559286 | -0.473369 |

| 54        | 0.5029458 | 0.9895918 | 0.1379915 | 0.7384613 |

| 55        | 0.5246147 | 0.8660061 | 0.0128609 | 0.6739320 |
| 56        | -.541e-05 | -0.879484 | -0.954218 | 0.7154813 |

| 57        | 0.8529050 | 0.9997355 | 0.0692861 | 0.5637266 |

| 58        | 0.3471668 | 1.0       | -0.039209 | 0.5793447 |
| 59        | 0.2922228 | 0.9356930 | 0.1417109 | 0.5855986 |

| 60        | 0.7454913 | 0.9328742 | 0.0904889 | -0.391478 |

| 61        | 0.6214151 | 1.0       | 0.0882262 | -0.496075 |
| 62        | 0.5161665 | 0.9811929 | 0.1365475 | 0.7444498 |

| 63        | 0.6789449 | 0.9501645 | -0.002108 | -0.416996 |
| 64        | 0.0028832 | 0.6475185 | -0.866986 | 0.2989424 |
| 65        | 0.0003793 | -0.712240 | 0.6268556 | 0.2073338 |

| 66        | 0.0001554 | 0.5744876 | 0.9758150 | 0.6370951 |
| 67        | 0.0004002 | 0.3575027 | -0.296749 | -0.023575 |

| 68        | 0.6433174 | 1.0       | 0.0747461 | 0.4536843 |
| 69        | 0.0719186 | 0.9819331 | -0.005330 | 0.7978891 |

| 70        | 0.7855935 | 0.5960585 | -0.004125 | -0.512132 |

| 71        | 0.8401320 | 0.6270568 | -0.026322 | -0.622177 |

| 72        | 0.6730900 | 0.5174064 | -0.050482 | -0.632539 |

| 73        | 0.1059564 | 0.6089376 | -0.119924 | -0.588921 |
| 74        | 0.5234386 | 0.5791422 | 0.0552139 | -0.637371 |

| 75        | 0.8687634 | 0.7092606 | -0.000667 | -0.678066 |

| 76        | 0.8322264 | 0.6351141 | -0.030434 | -0.747049 |
| 77        | 0.1196956 | 0.4775626 | 0.1868229 | 0.1878511 |
| 78        | 0.8588845 | 0.8171016 | 0.0613024 | -0.428494 |

| 79        | 0.7013931 | 0.7520536 | -0.012540 | -0.801662 |

| 80        | 0.8251733 | 0.8749484 | 0.0230114 | -0.656802 |

| 81        | 0.0214752 | 0.8787445 | 0.0980545 | -0.762461 |

| 82        | 0.7054977 | 0.8168963 | -0.073347 | -0.696622 |
| 83        | 0.7031685 | 0.9319049 | 0.0973267 | -0.384528 |
| 84        | 0.0002128 | -0.873608 | -0.590417 | -0.322908 |

| 85        | 0.0019041 | 0.5024768 | -0.860733 | 0.3299947 |

| 86        | 0.9104421 | 0.9543916 | -0.028352 | -0.608714 |

| 87        | 0.5084796 | 1.0       | -0.097431 | -0.697569 |

| 88        | 0.5954764 | 0.6538100 | -0.111085 | -0.872570 |
| 89        | 0.0074612 | -0.127203 | -0.291009 | -0.601715 |

| 90        | 0.7212891 | 0.4793080 | -0.047544 | -0.801404 |

| 91        | 0.0744798 | 0.5848782 | 0.0427544 | -0.879393 |
| 92        | 0.0009172 | -0.723588 | 0.4703690 | -0.997180 |

| 93        | 0.9739757 | 0.8810222 | 0.0100648 | -0.571485 |

| 94        | 0.4856802 | 0.3563276 | -0.089477 | -0.729510 |

| 95        | 0.0776399 | 0.5006620 | -0.185180 | -0.817513 |
| 96        | 0.2675017 | 0.8040816 | -0.159214 | -0.860836 |

| 97        | 0.9507493 | 0.7930022 | -0.000173 | -0.605471 |
| 98        | 0.0001739 | -0.765238 | -0.898005 | 0.5169565 |
| 99        | 0.0165422 | -0.227062 | -0.174050 | 0.2779688 |

| 100       | 0.2512042 | 0.3942487 | 0.0544503 | -0.736694 |
| 101       | 0.6713142 | 0.4207872 | 0.0082869 | -0.461412 |

| 102       | -.840e-06 | -0.838311 | 0.7147163 | -0.902185 |
| 103       | 0.1634979 | 0.2705901 | -0.073235 | -0.490517 |

| 104       | 0.8174814 | 0.8789252 | -0.043038 | -0.618030 |
| 105       | 0.7748375 | 1.0       | 0.0417764 | -0.621004 |

| 106       | 0.0628136 | -0.231786 | -0.098097 | -0.206950 |
| 107       | 0.0144301 | 0.7975162 | 0.2655924 | 0.9053020 |

| 108       | -.685e-05 | -0.454330 | 0.6388374 | -0.363639 |
| 109       | 0.0194496 | 1.0       | 0.1289968 | 0.2124224 |

| 110       | -.449e-06 | 0.2988408 | -0.849121 | 0.4582089 |
=============================================================

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

The plot shows 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$"])

array([<Axes: ylabel='Amplitude n[MHz / $2\pi$]'>,
       <Axes: xlabel='Time [ns]', ylabel='Population'>], dtype=object)

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

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

Gate fidelity: 0.973975731118761

References

• (Shahriari et al., 2016) B. Shahriari et al., “Taking the human out of the loop: A review of Bayesian optimization,” Proceedings of the IEEE 104, 148–175 (2016).