GRAPE on a single spin

This in an introductory example to the GRAPE method and its implementation in this software package.

1. Generate a PWC pulse shape

import matplotlib.pyplot as plt
import numpy as np
from paraqeet.quantity import Quantity

from paraqeet.signal.envelopes import GaussEnvelope
from paraqeet.signal.pwc_generator import PWCGenerator

First, let’s generate a piecewise constant (PWC) pulse envelope for the Gaussian pulse. For GRAPE, we need to specify an anstanz for piecewise constant controls at a given resolution. Here, we setup a Gaussian initial guess, sampling at 21 points during a gate time of 20ns.

t_final = 20e-9
tlist = np.linspace(0, t_final, 21)
tone = GaussEnvelope(amplitude=Quantity(np.pi / t_final / 3, -np.pi / t_final, np.pi / t_final))
tone.t_final.set_value(t_final)
gen = PWCGenerator(envelopes=[tone], tlist=tlist)
gen.multiply_flat_top = True
gen.max_amplitude = 2 * 1e8

We have added the option multiplyFlatTop, to ensure the pulse to start and end smoothly at 0 and t_final.

from plotting import plot_signal

ts = np.linspace(0, t_final, 501)
fig, ax = plt.subplots(1, figsize=(5, 3))
plot_signal(tone, ts, ax, linestyle="-", label="Smooth")
plot_signal(gen, ts, ax, linestyle="--", label="PWC")
ax.legend(loc=1, frameon=True)
plt.show()
../_images/04A_GRAPE_TLS_7_0.png

2. Define Hamiltonian in the rotating frame of drive

As a simple toy model, we use a single spin.

from paraqeet.model.closed_system import ClosedSystem
from paraqeet.model.rotating_frame_drive import RotatingFrameDrive
from paraqeet.model.hamiltonian import Hamiltonian


class SpinRWA(Hamiltonian):
    """A Single Spin."""

    def __init__(self, drives=None):
        super().__init__(drives)
        self.sigma_p = np.array([[0j, 1], [0, 0]])
        self.dim = 2

    def get_matrix_one_time(self, t):
        """Just sigma-X."""
        return self._drives[0].get_matrix_one_time(self.sigma_p, t)

    def gradient(self, t):
        """Gradient is just the drive matrix."""
        return self._drives[0].gradient(self.sigma_p, t)


drive = RotatingFrameDrive(gen)
spin = SpinRWA(drives=[drive])
model = ClosedSystem(spin)

from paraqeet.measurement.state_transfer_fidelity import StateTransferFidelityGRAPE
from paraqeet.propagation.scipy_expm_grape import ScipyExpmGRAPE


prop = ScipyExpmGRAPE(model, res=1e9)

init = np.array([[1.0], [0]])  # |0>
target = np.array([[0.0], [1]])  # |1>

prop.set_initial_state(init)
prop.target_state = target

zeroone = StateTransferFidelityGRAPE(
    propagation=prop,
    initial_state=init,
    target_state=target,
    times=tlist,
)

from plotting import plot_signal_and_dynamics

ts = np.linspace(0.0, t_final, 101)
plot_signal_and_dynamics(gen, prop, ts, state_labels=[r"$|0\rangle$", r"$|1\rangle$"]);
../_images/04A_GRAPE_TLS_11_0.png 
zeroone.measure()

0.10336679494545335

from paraqeet.optimisation_map import OptimisationMap
from paraqeet.optimisers.scipy_optimiser_gradient import ScipyOptimiserGradient


optmap = OptimisationMap()
optmap.add(gen)
optmap.register_params_with_optimisables()

opt_grad = ScipyOptimiserGradient(zeroone, optimisation_map=optmap)

opt_grad.optimise()

{'status': 1, 'value': 4.1100456371623295e-13, 'iterations': 7, 'message': 'CONVERGENCE: NORM OF PROJECTED GRADIENT <= PGTOL'}

For this simple example, we reach near perfect fidelity within a few iterations.

plot_signal_and_dynamics(gen, prop, ts, state_labels=[r"$|0\rangle$", r"$|1\rangle$"]);
../_images/04A_GRAPE_TLS_16_0.png

With open system

Lets first reset the pulse and create a open-system model

import matplotlib.pyplot as plt
import numpy as np
from paraqeet.quantity import Quantity

from paraqeet.signal.envelopes import GaussEnvelope
from paraqeet.signal.pwc_generator import PWCGenerator

t_final = 20e-9
tlist = np.linspace(0, t_final, 21)
tone = GaussEnvelope(amplitude=Quantity(np.pi / t_final / 3, -np.pi / t_final, np.pi / t_final))
tone.t_final.set_value(t_final)
gen = PWCGenerator(envelopes=[tone], tlist=tlist)
gen.multiply_flat_top = True
gen.max_amplitude = 5 * 1e8

from paraqeet.model.open_system import OpenSystem
from paraqeet.model.rotating_frame_drive import RotatingFrameDrive
from paraqeet.model.hamiltonian import Hamiltonian
import jax.numpy as jnp


class SpinRWA(Hamiltonian):
    """A Single Spin."""

    def __init__(self, drives=None):
        super().__init__(drives)
        self.sigma_p = np.array([[0j, 1], [0, 0]])
        self.sigma_m = np.array([[0j, 0], [1, 0]])
        self.sigma_z = np.array([[1, 0j], [0, -1]])
        self.dim = 2
        self.t1 = Quantity(50e-9, 1e-9, 100e-9)
        self.temp = Quantity(10e-3, 1e-3, 50e-3)
        self.t2star = Quantity(100e-9, 1e-9, 100e-9)

    def dimension(self):
        """Return TLS dimension."""
        return self.dim

    def get_matrix_one_time(self, t):
        """Just sigma-X."""
        return self._drives[0].get_matrix_one_time(self.sigma_p, t)

    def gradient(self, t):
        """Gradient is just the drive matrix."""
        return self._drives[0].gradient(self.sigma_p, t)

    def get_decay_rates(self) -> list[float]:
        """Return decay rate for T1, T2star and Temp respectively."""
        if (self.t1 is None) or (self.t2star is None) or (self.temp is None):
            raise Exception("Specify values of T1, T2star and Temp for Open system simulations.")

        gamma = 1 / self.t1.get_value()
        gamma_t2star = 0.5 / self.t2star.get_value()

        hbar_over_kb = 7.638232582257738e-12
        beta = hbar_over_kb / (self.temp.get_value())
        # inserting typical qubit freq here. TODO - CHECK
        nbar = jnp.exp(-beta * 5e9)
        gamma_temp = gamma * nbar
        gamma_t1 = gamma * (nbar + 1)
        return [gamma_t1, gamma_temp, gamma_t2star]

    def get_collapseops(self) -> list[jnp.ndarray]:
        """Return a list tuples of decay rates and collapse operators for each subsystem."""
        gamma_t1, gamma_temp, gamma_t2star = self.get_decay_rates()
        col_t1 = self.sigma_p
        col_temp = self.sigma_m
        col_t2star = 2 * self.sigma_z
        return [(gamma_t1, col_t1), (gamma_temp, col_temp), (gamma_t2star, col_t2star)]


drive = RotatingFrameDrive(gen)
spin = SpinRWA(drives=[drive])
model = OpenSystem(spin, ode_propagation=True)

Lets test GRAPE with ODE-propgation

from paraqeet.measurement.state_transfer_fidelity import StateTransferFidelityGRAPE
from paraqeet.propagation.vern7_grape import Vern7GRAPE

prop = Vern7GRAPE(model, res=1e9)

init = np.array([[1.0], [0.0j]])  # |0>
target = np.array([[0.0j], [1.0]])  # |1>

init = jnp.matmul(init, init.T.conj())
target = jnp.matmul(target, target.T.conj())

prop.set_initial_state(init)
prop.target_state = target

zeroone = StateTransferFidelityGRAPE(
    propagation=prop,
    initial_state=init,
    target_state=target,
    times=tlist,
)

ts = np.linspace(0, t_final, 101)
plot_signal_and_dynamics(gen, prop, times=ts, state_labels=[r"$\rho_0$", r"$\rho_1$"]);
../_images/04A_GRAPE_TLS_24_0.png 
zeroone.measure()

0.006921055509936116

from paraqeet.optimisation_map import OptimisationMap
from paraqeet.optimisers.scipy_optimiser_gradient import ScipyOptimiserGradient


optmap = OptimisationMap()
optmap.add(gen)
optmap.register_params_with_optimisables()

opt_grad = ScipyOptimiserGradient(zeroone, optimisation_map=optmap)
opt_grad.set_options({"disp": True})

opt_grad.optimise()

{'status': 2, 'value': 0.3986234856331238, 'iterations': 71, 'message': 'ABNORMAL: '}

plot_signal_and_dynamics(gen, prop, times=ts, state_labels=[r"$\rho_0$", r"$\rho_1$"]);
../_images/04A_GRAPE_TLS_29_0.png 
zeroone.measure()

0.6037315068515552