GRAPE for preparing a cat state using Adam Tutorial#

cat state
# ruff: noqa
import feedback_grape.grape as fg
import jax.numpy as jnp
from feedback_grape.utils.operators import *
from feedback_grape.utils.states import basis, coherent
from feedback_grape.utils.tensor import tensor

Step 1: initialize parameters#

T = 1  # microsecond
num_of_intervals = 100
N = 30  # dimension of hilbert space
alpha = 1.5
phi = jnp.pi
hconj = lambda a: jnp.swapaxes(a.conj(), -1, -2)
chi = 0.2385 * (2 * jnp.pi)
mu_qub = 4.0
mu_cav = 8.0

Step 2: define target and start states#

psi0 = tensor(basis(2), basis(N))
cat_target_state = coherent(N, alpha) + jnp.exp(-1j * phi) * coherent(
    N, -alpha
)
psi_target = tensor(basis(2), cat_target_state)

Step 3: Build the Hamiltonian#

# Using Jaynes-Cummings model for qubit + cavity
def build_grape_format_ham():
    """
    Build Hamiltonian for given (complex) e_qub and e_cav
    """

    a = tensor(identity(2), destroy(N))
    adag = tensor(identity(2), create(N))
    n_phot = adag @ a
    sigz = tensor(sigmaz(), identity(N))
    sigp = tensor(sigmap(), identity(N))
    one = tensor(identity(2), identity(N))

    H0 = +(chi / 2) * n_phot @ (sigz + one)
    H_ctrl_qub = mu_qub * sigp
    H_ctrl_qub_dag = hconj(H_ctrl_qub)
    H_ctrl_cav = mu_cav * adag
    H_ctrl_cav_dag = hconj(H_ctrl_cav)

    H_ctrl = [H_ctrl_qub, H_ctrl_qub_dag, H_ctrl_cav, H_ctrl_cav_dag]

    return H0, H_ctrl

Step 4: run GRAPE#

# playing with the learning rate a bit may lead to better results (try it out!)
H0, H_ctrl = build_grape_format_ham()
res = fg.optimize_pulse(
    H0,
    H_ctrl,
    psi0,
    psi_target,
    num_t_slots=num_of_intervals,
    total_evo_time=T,
    evo_type="state",
)
res.final_fidelity
Array(0.6765205, dtype=float64)