This is equivalent to Example C but contains only the code blocks as in the accompanying paper

# ruff: noqa
import os

os.sys.path.append("../../../..")

from feedback_grape.utils.states import basis, coherent
N_cav = 20
basis_state = basis(n=2)
coherent_state = coherent(n=N_cav, alpha=3)

from feedback_grape.utils.operators import create, destroy
number_operator = create(N_cav) @ destroy(N_cav)

from feedback_grape.utils.tensor import tensor
ground_qubit = basis(n=2)
ground_cavity = basis(n=N_cav)
ground_state = tensor(ground_cavity, ground_qubit)

import jax
import jax.numpy as jnp
from jax.scipy.linalg import expm
from feedback_grape.utils.operators import (
sigmap, sigmam, create, destroy, identity, cosm, sinm
)

def meas_operator(measurement_outcome, params):
    gamma, delta = params
    cav_operator = (
    gamma * create(N_cav) @ destroy(N_cav)
    + delta / 2 * identity(N_cav)
    )
    angle = tensor(cav_operator, identity(2))
    meas_op = jnp.where(
    measurement_outcome == 1,
    cosm(angle),
    sinm(angle),
    )
    return meas_op

def qubit_unitary(alphas):
    alpha_re, alpha_im = alphas
    alpha = alpha_re + 1j * alpha_im
    h_dt = (alpha*sigmap() + alpha.conjugate()*sigmam()) / 2
    return tensor(identity(N_cav), expm(-1j * h_dt))
    
def qubit_cavity_unitary(betas):
    beta_re, beta_im = betas
    beta = beta_re + 1j * beta_im
    a_sigp = tensor(destroy(N_cav), sigmap())
    adag_sigm = tensor(create(N_cav), sigmam())
    h_dt = (beta * a_sigp + beta.conjugate() * adag_sigm) / 2
    return expm(-1j * h_dt)

from feedback_grape.fgrape import Gate
key1, key2, key3 = jax.random.split(jax.random.PRNGKey(42), 3)
measure = Gate(
    gate=meas_operator,
    initial_params=jax.random.uniform(
        key1,
        shape=(2,),
        minval=-2 * jnp.pi,
        maxval=2 * jnp.pi,
    ),
    measurement_flag=True,
)
qub_unitary = Gate(
    gate=qubit_unitary,
    initial_params=jax.random.uniform(
        key2,
        shape=(2,),
        minval=-2 * jnp.pi,
        maxval=2 * jnp.pi,
    ),
    measurement_flag=False,
)
qub_cav = Gate(
    gate=qubit_cavity_unitary,
    initial_params=jax.random.uniform(
        key3,
        shape=(2,),
        minval=-2 * jnp.pi,
        maxval=2 * jnp.pi,
    ),
    measurement_flag=False,
)

n_average = 1
beta = jnp.log((1 / n_average) + 1)
diags = jnp.exp(-beta * jnp.arange(N_cav))
normalized_diags = diags / jnp.sum(diags, axis=0)
rho_cav0 = jnp.diag(normalized_diags)
rho_qub0 = basis(2, 0) @ basis(2, 0).conj().T
rho0 = tensor(rho_cav0, rho_qub0)

from feedback_grape.utils.states import fock
psi_cav_target = (
    fock(N_cav, 1) + fock(N_cav, 2) + fock(N_cav, 3)
) / jnp.sqrt(3)
psi_qub_target = basis(2, 0)
psi_target = tensor(psi_cav_target, psi_qub_target)
rho_target = psi_target @ psi_target.conj().T

system_params = [measure, qub_unitary, qub_cav]
num_time_steps = 5

from feedback_grape.fgrape import optimize_pulse
result = optimize_pulse(
    init_state=rho0,
    target_state=rho_target,
    system_params=system_params,
    num_time_steps=num_time_steps,
    mode="lookup",
    goal="fidelity",
    max_iter=1000,
    convergence_threshold=1e-6,
    learning_rate=0.02,
    evo_type="density",
    progress=True,
    batch_size=10,
    eval_time_steps=2*num_time_steps
)

Iteration 10, Loss: 0.619694, T=3s, eta=286s
Iteration 20, Loss: 0.559191, T=6s, eta=296s
Iteration 30, Loss: 0.654462, T=9s, eta=298s
Iteration 40, Loss: 0.690517, T=12s, eta=297s
Iteration 50, Loss: 0.606377, T=15s, eta=298s
Iteration 60, Loss: 0.688028, T=19s, eta=295s
Iteration 70, Loss: 0.736810, T=22s, eta=292s
Iteration 80, Loss: 0.865780, T=25s, eta=290s
Iteration 90, Loss: 0.895460, T=28s, eta=287s
Iteration 100, Loss: 1.038856, T=31s, eta=284s
Iteration 110, Loss: 1.158938, T=35s, eta=281s
Iteration 120, Loss: 1.002144, T=38s, eta=278s
Iteration 130, Loss: 0.755226, T=41s, eta=275s
Iteration 140, Loss: 0.877550, T=44s, eta=272s
Iteration 150, Loss: 1.289478, T=47s, eta=268s
Iteration 160, Loss: 0.713522, T=50s, eta=265s
Iteration 170, Loss: 1.031733, T=54s, eta=262s
Iteration 180, Loss: 1.096958, T=57s, eta=259s
Iteration 190, Loss: 1.128355, T=60s, eta=256s
Iteration 200, Loss: 0.994936, T=63s, eta=253s
Iteration 210, Loss: 1.328315, T=66s, eta=250s
Iteration 220, Loss: 1.221207, T=69s, eta=247s
Iteration 230, Loss: 1.432580, T=73s, eta=244s
Iteration 240, Loss: 1.133913, T=76s, eta=240s
Iteration 250, Loss: 1.102306, T=79s, eta=237s
Iteration 260, Loss: 1.132770, T=82s, eta=234s
Iteration 270, Loss: 1.355784, T=85s, eta=231s
Iteration 280, Loss: 1.338068, T=88s, eta=228s
Iteration 290, Loss: 0.966543, T=92s, eta=225s
Iteration 300, Loss: 1.274789, T=95s, eta=222s
Iteration 310, Loss: 1.058684, T=98s, eta=218s
Iteration 320, Loss: 1.207416, T=101s, eta=215s
Iteration 330, Loss: 1.172187, T=104s, eta=212s
Iteration 340, Loss: 0.795722, T=108s, eta=209s
Iteration 350, Loss: 0.934281, T=111s, eta=206s
Iteration 360, Loss: 0.876484, T=114s, eta=203s
Iteration 370, Loss: 0.976146, T=118s, eta=200s
Iteration 380, Loss: 0.924410, T=121s, eta=197s
Iteration 390, Loss: 1.057753, T=124s, eta=194s
Iteration 400, Loss: 1.177116, T=127s, eta=191s
Iteration 410, Loss: 1.000939, T=131s, eta=189s
Iteration 420, Loss: 1.092718, T=134s, eta=185s
Iteration 430, Loss: 1.275725, T=137s, eta=182s
Iteration 440, Loss: 0.881084, T=141s, eta=179s
Iteration 450, Loss: 1.116676, T=144s, eta=176s
Iteration 460, Loss: 0.993254, T=147s, eta=172s
Iteration 470, Loss: 1.205236, T=150s, eta=169s
Iteration 480, Loss: 0.983355, T=153s, eta=166s
Iteration 490, Loss: 1.001310, T=156s, eta=163s
Iteration 500, Loss: 0.760032, T=159s, eta=159s
Iteration 510, Loss: 0.938214, T=163s, eta=156s
Iteration 520, Loss: 1.165753, T=166s, eta=153s
Iteration 530, Loss: 0.697868, T=169s, eta=150s
Iteration 540, Loss: 0.815368, T=172s, eta=147s
Iteration 550, Loss: 1.022441, T=175s, eta=143s
Iteration 560, Loss: 1.295912, T=179s, eta=140s
Iteration 570, Loss: 1.107398, T=182s, eta=137s
Iteration 580, Loss: 1.335107, T=185s, eta=134s
Iteration 590, Loss: 1.032403, T=188s, eta=131s
Iteration 600, Loss: 1.406895, T=192s, eta=128s
Iteration 610, Loss: 1.069001, T=195s, eta=124s
Iteration 620, Loss: 0.852598, T=198s, eta=121s
Iteration 630, Loss: 0.854650, T=201s, eta=118s
Iteration 640, Loss: 0.750477, T=204s, eta=115s
Iteration 650, Loss: 1.008557, T=207s, eta=112s
Iteration 660, Loss: 1.044254, T=211s, eta=108s
Iteration 670, Loss: 0.986450, T=214s, eta=105s
Iteration 680, Loss: 0.951353, T=217s, eta=102s
Iteration 690, Loss: 0.928878, T=220s, eta=99s
Iteration 700, Loss: 0.824319, T=223s, eta=96s
Iteration 710, Loss: 0.972151, T=227s, eta=92s
Iteration 720, Loss: 0.963490, T=230s, eta=89s
Iteration 730, Loss: 0.908401, T=233s, eta=86s
Iteration 740, Loss: 0.970608, T=236s, eta=83s
Iteration 750, Loss: 0.639949, T=239s, eta=80s
Iteration 760, Loss: 0.391333, T=242s, eta=76s
Iteration 770, Loss: 0.332053, T=245s, eta=73s
Iteration 780, Loss: 0.670661, T=249s, eta=70s
Iteration 790, Loss: 0.698586, T=252s, eta=67s
Iteration 800, Loss: 0.777814, T=255s, eta=64s
Iteration 810, Loss: 0.466045, T=258s, eta=60s
Iteration 820, Loss: 0.809259, T=261s, eta=57s
Iteration 830, Loss: 0.826517, T=264s, eta=54s
Iteration 840, Loss: 0.377347, T=268s, eta=51s
Iteration 850, Loss: 1.004180, T=271s, eta=48s
Iteration 860, Loss: 0.861270, T=274s, eta=44s
Iteration 870, Loss: 1.018132, T=277s, eta=41s
Iteration 880, Loss: 0.685581, T=280s, eta=38s
Iteration 890, Loss: 0.949864, T=284s, eta=35s
Iteration 900, Loss: 0.805575, T=287s, eta=32s
Iteration 910, Loss: 0.566240, T=290s, eta=29s
Iteration 920, Loss: 0.672469, T=293s, eta=25s
Iteration 930, Loss: 0.746736, T=296s, eta=22s
Iteration 940, Loss: 0.909266, T=300s, eta=19s
Iteration 950, Loss: 0.814764, T=303s, eta=16s
Iteration 960, Loss: 0.745589, T=306s, eta=13s
Iteration 970, Loss: 0.746876, T=309s, eta=9s
Iteration 980, Loss: 0.851863, T=313s, eta=6s
Iteration 990, Loss: 0.779292, T=316s, eta=3s

# Fidelity vs. evaluation time step.
# `result.fidelity_each_timestep` is a list of length eval_time_steps + 1
# (index 0 = initial state); each entry holds the fidelity of every
# evaluation trajectory in the batch at that time step.
from matplotlib.ticker import AutoMinorLocator, MultipleLocator
import matplotlib.pyplot as plt
n_steps = 5  # only plot the initial state plus the first `n_steps` time steps
fid = jnp.stack(result.fidelity_each_timestep)[: n_steps + 1]  # (n_steps + 1, batch)
t = jnp.arange(fid.shape[0])

fid_mean = fid.mean(axis=1)
fid_min = fid.min(axis=1)
fid_max = fid.max(axis=1)

peak_idx = int(jnp.argmax(fid_mean))
peak_val = float(fid_mean[peak_idx])

with plt.rc_context({
    "font.family": "serif",
    "mathtext.fontset": "dejavuserif",
    "axes.grid": True,
    "grid.alpha": 0.35,
    "grid.linewidth": 0.6,
    "xtick.direction": "in",
    "ytick.direction": "in",
    "xtick.top": True,
    "ytick.right": True,
    "xtick.minor.visible": True,
    "ytick.minor.visible": True,
}):
    fig, ax = plt.subplots(figsize=(10, 4.2))

    # min-max envelope across the evaluation batch
    ax.fill_between(
        t, fid_min, fid_max,
        color="#a9c6e8", alpha=0.45, lw=0, label="Min-max range",
    )
    # individual trajectories, faint, inside the band
    ax.plot(t, fid, color="#7ba3cc", alpha=0.35, lw=0.9)
    # batch-mean fidelity
    ax.plot(
        t, fid_mean,
        color="black", lw=2.2, marker="o", ms=6,
        mfc="white", mec="black", mew=1.5, zorder=5,
        label="Mean fidelity",
    )

    ax.set_xlabel("Evaluation time step", fontsize=13)
    ax.set_ylabel("Fidelity", fontsize=13)
    ax.set_xlim(t[0], t[-1])
    ax.set_ylim(-0.02, 1.02)
    # numbered major ticks plus unnumbered minor ticks on both axes
    ax.set_xticks(t)
    ax.xaxis.set_minor_locator(AutoMinorLocator(4))
    ax.yaxis.set_major_locator(MultipleLocator(0.2))
    ax.yaxis.set_minor_locator(AutoMinorLocator(4))
    ax.legend(loc="upper left", framealpha=0.9)
    # leave room on the right for the outside annotation
    fig.tight_layout(rect=(0, 0, 0.84, 1))
    plt.show()