A2. State preparation with Jaynes-Cummings controls for 4-legged cat state

from feedback_grape.fgrape import optimize_pulse
from feedback_grape.utils.operators import (
    sigmap,
    sigmam,
    create,
    destroy,
    identity,
)
from feedback_grape.utils.states import basis, fock
from feedback_grape.utils.tensor import tensor
import jax.numpy as jnp
from jax.scipy.linalg import expm

As a preliminary step, we consider state preparation of a target state starting from a pure state. In addition, we assume that any coupling to an external environment is negligible and that the parametrized controls can be implemented perfectly.

Here no feedback is required, we are just testing the parameterized gates setup.

As a first example, we consider the state preparation of a cavity resonantly coupled to an externally driven qubit

Here, we consider a particular sequence of parametrized unitary gates originally introduced by Law and Eberly

N_cav = 40

def qubit_unitary(alphas):
    alpha_re = alphas[0]
    alpha_im = alphas[1]
    alpha = alpha_re + 1j * alpha_im
    return tensor(
        identity(N_cav),
        expm(-1j * (alpha * sigmap() + alpha.conjugate() * sigmam()) / 2),
    )

def qubit_cavity_unitary(beta_re):
    beta = beta_re
    return expm(
        -1j
        * (
            beta * (tensor(destroy(N_cav), sigmap()))
            + beta.conjugate() * (tensor(create(N_cav), sigmam()))
        )
        / 2
    )

In their groundbreaking work, Law and Eberly have shown that any arbitrary superposition of Fock states with maximal excitation number N can be prepared out of the ground state in a sequence of N such interleaved gates, also providing an algorithm to find the correct angles and interaction durations

First target is the 4 component cat state

time_steps = 20

from feedback_grape.utils.states import coherent


psi0 = tensor(basis(N_cav), basis(2))
psi0 = psi0 / jnp.linalg.norm(psi0)
average_photon_number = 9
alpha = jnp.sqrt(average_photon_number)
cat = (
    coherent(N_cav, alpha)
    + coherent(N_cav, -alpha)
    + coherent(N_cav, 1j * alpha)
    + coherent(N_cav, -1j * alpha)
)
psi_target = tensor(cat, basis(2))
psi_target = psi_target / jnp.linalg.norm(psi_target)

psi0.shape

(80, 1)

psi_target.shape

(80, 1)

print(fock(N_cav, 1))

[[0.+0.j]
 [1.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]
 [0.+0.j]]

Law and Eberly provided an algorithm to determine the correct parameters for state preparation. These include:

• The rotation angle \( |\alpha| \),

• The azimuthal angle \( \arg\left(\frac{\alpha}{|\alpha|}\right) \),

• The interaction duration \( |\beta| \).
  

So Goal is to find the best control vector (rather than control amplitudes, this time) that leads to finding the optimal state-preparation strategies. Performing as well as the Law-Eberly algorithm.

Optimizing

Currently l-bfgs with the same learning rate of 0.3 converges at a local minimum of 0.5, adam also converges at 0.5 but at smaller learning rates

from feedback_grape.fgrape import Gate
import jax

key = jax.random.PRNGKey(42)
# not provideing param_constraints just propagates the same initial_parameters for each time step
qub_unitary = Gate(
    gate=qubit_unitary,
    initial_params=jax.random.uniform(
        key,
        shape=(2,),  # 2 for gamma and delta
        minval=-5,
        maxval=5,
        dtype=jnp.float64,
    ),
    measurement_flag=False,
)

qub_cav = Gate(
    gate=qubit_cavity_unitary,
    initial_params=jax.random.uniform(
        key,
        shape=(1,),
        minval=-5,
        maxval=5,
        dtype=jnp.float64,
    ),
    measurement_flag=False,
)

system_params = [qub_unitary, qub_cav]


result = optimize_pulse(
    U_0=psi0,
    C_target=psi_target,
    system_params=system_params,
    num_time_steps=time_steps,
    max_iter=3000,
    convergence_threshold=None,  # None to reach max iterations
    evo_type="state",
    mode="no-measurement",
    goal="fidelity",
    learning_rate=0.01,
    batch_size=10,
    eval_batch_size=1,
    progress=True,
)

Iteration 10, Loss: -0.119766, T=0s, eta=113s
Iteration 20, Loss: -0.218807, T=0s, eta=118s
Iteration 30, Loss: -0.290923, T=1s, eta=119s
Iteration 40, Loss: -0.331432, T=1s, eta=120s
Iteration 50, Loss: -0.363328, T=2s, eta=120s
Iteration 60, Loss: -0.385187, T=2s, eta=120s
Iteration 70, Loss: -0.392812, T=2s, eta=119s
Iteration 80, Loss: -0.398971, T=3s, eta=119s
Iteration 90, Loss: -0.406213, T=3s, eta=118s
Iteration 100, Loss: -0.414206, T=4s, eta=118s
Iteration 110, Loss: -0.424084, T=4s, eta=117s
Iteration 120, Loss: -0.436752, T=4s, eta=117s
Iteration 130, Loss: -0.455166, T=5s, eta=116s
Iteration 140, Loss: -0.483714, T=5s, eta=116s
Iteration 150, Loss: -0.526234, T=6s, eta=115s
Iteration 160, Loss: -0.587400, T=6s, eta=115s
Iteration 170, Loss: -0.658891, T=6s, eta=115s
Iteration 180, Loss: -0.710771, T=7s, eta=114s
Iteration 190, Loss: -0.743793, T=7s, eta=114s
Iteration 200, Loss: -0.767970, T=8s, eta=113s
Iteration 210, Loss: -0.785470, T=8s, eta=113s
Iteration 220, Loss: -0.794869, T=8s, eta=113s
Iteration 230, Loss: -0.799156, T=9s, eta=112s
Iteration 240, Loss: -0.801985, T=9s, eta=112s
Iteration 250, Loss: -0.804416, T=10s, eta=111s
Iteration 260, Loss: -0.806428, T=10s, eta=111s
Iteration 270, Loss: -0.808274, T=11s, eta=111s
Iteration 280, Loss: -0.809989, T=11s, eta=110s
Iteration 290, Loss: -0.811600, T=11s, eta=110s
Iteration 300, Loss: -0.813111, T=12s, eta=109s
Iteration 310, Loss: -0.814524, T=12s, eta=109s
Iteration 320, Loss: -0.815841, T=13s, eta=108s
Iteration 330, Loss: -0.817068, T=13s, eta=108s
Iteration 340, Loss: -0.818218, T=13s, eta=108s
Iteration 350, Loss: -0.819303, T=14s, eta=108s
Iteration 360, Loss: -0.820341, T=14s, eta=107s
Iteration 370, Loss: -0.821346, T=15s, eta=107s
Iteration 380, Loss: -0.822334, T=15s, eta=106s
Iteration 390, Loss: -0.823330, T=15s, eta=106s
Iteration 400, Loss: -0.824376, T=16s, eta=105s
Iteration 410, Loss: -0.825539, T=16s, eta=105s
Iteration 420, Loss: -0.826936, T=17s, eta=105s
Iteration 430, Loss: -0.828753, T=17s, eta=104s
Iteration 440, Loss: -0.831279, T=17s, eta=104s
Iteration 450, Loss: -0.834815, T=18s, eta=103s
Iteration 460, Loss: -0.839446, T=18s, eta=103s
Iteration 470, Loss: -0.845438, T=19s, eta=103s
Iteration 480, Loss: -0.853808, T=19s, eta=102s
Iteration 490, Loss: -0.864693, T=19s, eta=102s
Iteration 500, Loss: -0.876147, T=20s, eta=101s
Iteration 510, Loss: -0.890039, T=20s, eta=101s
Iteration 520, Loss: -0.905138, T=21s, eta=101s
Iteration 530, Loss: -0.918734, T=21s, eta=100s
Iteration 540, Loss: -0.929061, T=22s, eta=100s
Iteration 550, Loss: -0.935838, T=22s, eta=100s
Iteration 560, Loss: -0.939936, T=22s, eta=99s
Iteration 570, Loss: -0.942359, T=23s, eta=100s
Iteration 580, Loss: -0.943913, T=23s, eta=99s
Iteration 590, Loss: -0.945100, T=24s, eta=99s
Iteration 600, Loss: -0.946135, T=24s, eta=99s
Iteration 610, Loss: -0.947034, T=25s, eta=99s
Iteration 620, Loss: -0.947801, T=25s, eta=98s
Iteration 630, Loss: -0.948456, T=26s, eta=98s
Iteration 640, Loss: -0.949025, T=26s, eta=98s
Iteration 650, Loss: -0.949528, T=27s, eta=97s
Iteration 660, Loss: -0.949979, T=27s, eta=97s
Iteration 670, Loss: -0.950387, T=27s, eta=97s
Iteration 680, Loss: -0.950758, T=28s, eta=96s
Iteration 690, Loss: -0.951097, T=28s, eta=96s
Iteration 700, Loss: -0.951405, T=29s, eta=95s
Iteration 710, Loss: -0.951687, T=29s, eta=95s
Iteration 720, Loss: -0.951943, T=29s, eta=94s
Iteration 730, Loss: -0.952177, T=30s, eta=94s
Iteration 740, Loss: -0.952389, T=30s, eta=93s
Iteration 750, Loss: -0.952581, T=31s, eta=93s
Iteration 760, Loss: -0.952754, T=31s, eta=93s
Iteration 770, Loss: -0.952908, T=32s, eta=92s
Iteration 780, Loss: -0.953045, T=32s, eta=92s
Iteration 790, Loss: -0.953168, T=32s, eta=92s
Iteration 800, Loss: -0.953279, T=33s, eta=91s
Iteration 810, Loss: -0.953380, T=33s, eta=91s
Iteration 820, Loss: -0.953474, T=34s, eta=90s
Iteration 830, Loss: -0.953563, T=34s, eta=90s
Iteration 840, Loss: -0.953649, T=34s, eta=89s
Iteration 850, Loss: -0.953732, T=35s, eta=89s
Iteration 860, Loss: -0.953814, T=35s, eta=88s
Iteration 870, Loss: -0.953894, T=36s, eta=88s
Iteration 880, Loss: -0.953972, T=36s, eta=88s
Iteration 890, Loss: -0.954047, T=37s, eta=88s
Iteration 900, Loss: -0.954119, T=37s, eta=87s
Iteration 910, Loss: -0.954189, T=38s, eta=87s
Iteration 920, Loss: -0.954255, T=38s, eta=87s
Iteration 930, Loss: -0.954318, T=38s, eta=86s
Iteration 940, Loss: -0.954379, T=39s, eta=86s
Iteration 950, Loss: -0.954437, T=39s, eta=85s
Iteration 960, Loss: -0.954492, T=40s, eta=85s
Iteration 970, Loss: -0.954545, T=40s, eta=84s
Iteration 980, Loss: -0.954595, T=40s, eta=84s
Iteration 990, Loss: -0.954643, T=41s, eta=84s
Iteration 1000, Loss: -0.954688, T=41s, eta=83s
Iteration 1010, Loss: -0.954731, T=42s, eta=83s
Iteration 1020, Loss: -0.954772, T=42s, eta=82s
Iteration 1030, Loss: -0.954811, T=43s, eta=82s
Iteration 1040, Loss: -0.954847, T=43s, eta=81s
Iteration 1050, Loss: -0.954882, T=43s, eta=81s
Iteration 1060, Loss: -0.954915, T=44s, eta=81s
Iteration 1070, Loss: -0.954946, T=44s, eta=80s
Iteration 1080, Loss: -0.954976, T=45s, eta=80s
Iteration 1090, Loss: -0.955004, T=45s, eta=79s
Iteration 1100, Loss: -0.955031, T=46s, eta=79s
Iteration 1110, Loss: -0.955057, T=46s, eta=79s
Iteration 1120, Loss: -0.955082, T=46s, eta=78s
Iteration 1130, Loss: -0.955105, T=47s, eta=78s
Iteration 1140, Loss: -0.955128, T=47s, eta=77s
Iteration 1150, Loss: -0.955150, T=48s, eta=77s
Iteration 1160, Loss: -0.955172, T=48s, eta=77s
Iteration 1170, Loss: -0.955193, T=48s, eta=76s
Iteration 1180, Loss: -0.955213, T=49s, eta=76s
Iteration 1190, Loss: -0.955233, T=49s, eta=75s
Iteration 1200, Loss: -0.955253, T=50s, eta=75s
Iteration 1210, Loss: -0.955272, T=50s, eta=75s
Iteration 1220, Loss: -0.955291, T=51s, eta=74s
Iteration 1230, Loss: -0.955310, T=51s, eta=74s
Iteration 1240, Loss: -0.955328, T=52s, eta=73s
Iteration 1250, Loss: -0.955347, T=52s, eta=73s
Iteration 1260, Loss: -0.955365, T=52s, eta=73s
Iteration 1270, Loss: -0.955384, T=53s, eta=72s
Iteration 1280, Loss: -0.955402, T=53s, eta=72s
Iteration 1290, Loss: -0.955421, T=54s, eta=71s
Iteration 1300, Loss: -0.955439, T=54s, eta=71s
Iteration 1310, Loss: -0.955457, T=54s, eta=70s
Iteration 1320, Loss: -0.955476, T=55s, eta=70s
Iteration 1330, Loss: -0.955396, T=55s, eta=70s
Iteration 1340, Loss: -0.955490, T=56s, eta=69s
Iteration 1350, Loss: -0.955524, T=56s, eta=69s
Iteration 1360, Loss: -0.955551, T=57s, eta=68s
Iteration 1370, Loss: -0.955571, T=57s, eta=68s
Iteration 1380, Loss: -0.955591, T=57s, eta=67s
Iteration 1390, Loss: -0.955611, T=58s, eta=67s
Iteration 1400, Loss: -0.955631, T=58s, eta=67s
Iteration 1410, Loss: -0.955652, T=59s, eta=66s
Iteration 1420, Loss: -0.955674, T=59s, eta=66s
Iteration 1430, Loss: -0.955695, T=60s, eta=65s
Iteration 1440, Loss: -0.955718, T=60s, eta=65s
Iteration 1450, Loss: -0.955740, T=60s, eta=65s
Iteration 1460, Loss: -0.955764, T=61s, eta=64s
Iteration 1470, Loss: -0.955788, T=61s, eta=64s
Iteration 1480, Loss: -0.955812, T=62s, eta=63s
Iteration 1490, Loss: -0.955837, T=62s, eta=63s
Iteration 1500, Loss: -0.955863, T=62s, eta=62s
Iteration 1510, Loss: -0.955890, T=63s, eta=62s
Iteration 1520, Loss: -0.955918, T=63s, eta=62s
Iteration 1530, Loss: -0.955946, T=64s, eta=61s
Iteration 1540, Loss: -0.955975, T=64s, eta=61s
Iteration 1550, Loss: -0.956005, T=65s, eta=60s
Iteration 1560, Loss: -0.956036, T=65s, eta=60s
Iteration 1570, Loss: -0.956068, T=65s, eta=60s
Iteration 1580, Loss: -0.956101, T=66s, eta=59s
Iteration 1590, Loss: -0.956136, T=66s, eta=59s
Iteration 1600, Loss: -0.956171, T=67s, eta=58s
Iteration 1610, Loss: -0.956208, T=67s, eta=58s
Iteration 1620, Loss: -0.956246, T=68s, eta=57s
Iteration 1630, Loss: -0.956285, T=68s, eta=57s
Iteration 1640, Loss: -0.956326, T=68s, eta=57s
Iteration 1650, Loss: -0.956315, T=69s, eta=56s
Iteration 1660, Loss: -0.956366, T=69s, eta=56s
Iteration 1670, Loss: -0.956448, T=70s, eta=55s
Iteration 1680, Loss: -0.956505, T=70s, eta=55s
Iteration 1690, Loss: -0.956556, T=71s, eta=55s
Iteration 1700, Loss: -0.956609, T=71s, eta=54s
Iteration 1710, Loss: -0.956664, T=71s, eta=54s
Iteration 1720, Loss: -0.956722, T=72s, eta=53s
Iteration 1730, Loss: -0.956783, T=72s, eta=53s
Iteration 1740, Loss: -0.956847, T=73s, eta=52s
Iteration 1750, Loss: -0.956915, T=73s, eta=52s
Iteration 1760, Loss: -0.956986, T=73s, eta=52s
Iteration 1770, Loss: -0.957062, T=74s, eta=51s
Iteration 1780, Loss: -0.957142, T=74s, eta=51s
Iteration 1790, Loss: -0.957227, T=75s, eta=50s
Iteration 1800, Loss: -0.957318, T=75s, eta=50s
Iteration 1810, Loss: -0.957414, T=76s, eta=49s
Iteration 1820, Loss: -0.957515, T=76s, eta=49s
Iteration 1830, Loss: -0.957544, T=76s, eta=49s
Iteration 1840, Loss: -0.957732, T=77s, eta=48s
Iteration 1850, Loss: -0.957855, T=77s, eta=48s
Iteration 1860, Loss: -0.957992, T=78s, eta=47s
Iteration 1870, Loss: -0.958128, T=78s, eta=47s
Iteration 1880, Loss: -0.958274, T=78s, eta=47s
Iteration 1890, Loss: -0.958427, T=79s, eta=46s
Iteration 1900, Loss: -0.958587, T=79s, eta=46s
Iteration 1910, Loss: -0.958753, T=80s, eta=45s
Iteration 1920, Loss: -0.958926, T=80s, eta=45s
Iteration 1930, Loss: -0.959104, T=80s, eta=44s
Iteration 1940, Loss: -0.959287, T=81s, eta=44s
Iteration 1950, Loss: -0.959475, T=81s, eta=44s
Iteration 1960, Loss: -0.959667, T=82s, eta=43s
Iteration 1970, Loss: -0.959863, T=82s, eta=43s
Iteration 1980, Loss: -0.960063, T=82s, eta=42s
Iteration 1990, Loss: -0.960265, T=83s, eta=42s
Iteration 2000, Loss: -0.960404, T=83s, eta=41s
Iteration 2010, Loss: -0.960693, T=84s, eta=41s
Iteration 2020, Loss: -0.960945, T=84s, eta=41s
Iteration 2030, Loss: -0.961196, T=85s, eta=40s
Iteration 2040, Loss: -0.961475, T=85s, eta=40s
Iteration 2050, Loss: -0.961783, T=86s, eta=39s
Iteration 2060, Loss: -0.962129, T=86s, eta=39s
Iteration 2070, Loss: -0.962529, T=86s, eta=39s
Iteration 2080, Loss: -0.962999, T=87s, eta=38s
Iteration 2090, Loss: -0.963567, T=87s, eta=38s
Iteration 2100, Loss: -0.964274, T=88s, eta=37s
Iteration 2110, Loss: -0.965175, T=88s, eta=37s
Iteration 2120, Loss: -0.966344, T=89s, eta=37s
Iteration 2130, Loss: -0.967844, T=89s, eta=36s
Iteration 2140, Loss: -0.969300, T=89s, eta=36s
Iteration 2150, Loss: -0.971686, T=90s, eta=35s
Iteration 2160, Loss: -0.973828, T=90s, eta=35s
Iteration 2170, Loss: -0.976316, T=91s, eta=34s
Iteration 2180, Loss: -0.979287, T=91s, eta=34s
Iteration 2190, Loss: -0.983004, T=92s, eta=34s
Iteration 2200, Loss: -0.987018, T=92s, eta=33s
Iteration 2210, Loss: -0.990289, T=92s, eta=33s
Iteration 2220, Loss: -0.992632, T=93s, eta=32s
Iteration 2230, Loss: -0.994462, T=93s, eta=32s
Iteration 2240, Loss: -0.995217, T=94s, eta=31s
Iteration 2250, Loss: -0.995547, T=94s, eta=31s
Iteration 2260, Loss: -0.995743, T=94s, eta=31s
Iteration 2270, Loss: -0.995900, T=95s, eta=30s
Iteration 2280, Loss: -0.996035, T=95s, eta=30s
Iteration 2290, Loss: -0.996151, T=96s, eta=29s
Iteration 2300, Loss: -0.996248, T=96s, eta=29s
Iteration 2310, Loss: -0.996330, T=96s, eta=28s
Iteration 2320, Loss: -0.996402, T=97s, eta=28s
Iteration 2330, Loss: -0.996465, T=97s, eta=28s
Iteration 2340, Loss: -0.996521, T=98s, eta=27s
Iteration 2350, Loss: -0.996571, T=98s, eta=27s
Iteration 2360, Loss: -0.996616, T=98s, eta=26s
Iteration 2370, Loss: -0.996656, T=99s, eta=26s
Iteration 2380, Loss: -0.996693, T=99s, eta=26s
Iteration 2390, Loss: -0.996726, T=100s, eta=25s
Iteration 2400, Loss: -0.996757, T=100s, eta=25s
Iteration 2410, Loss: -0.996785, T=101s, eta=24s
Iteration 2420, Loss: -0.996811, T=101s, eta=24s
Iteration 2430, Loss: -0.996835, T=101s, eta=23s
Iteration 2440, Loss: -0.996858, T=102s, eta=23s
Iteration 2450, Loss: -0.996879, T=102s, eta=23s
Iteration 2460, Loss: -0.996899, T=103s, eta=22s
Iteration 2470, Loss: -0.996918, T=103s, eta=22s
Iteration 2480, Loss: -0.996936, T=103s, eta=21s
Iteration 2490, Loss: -0.996953, T=104s, eta=21s
Iteration 2500, Loss: -0.996970, T=104s, eta=20s
Iteration 2510, Loss: -0.996986, T=105s, eta=20s
Iteration 2520, Loss: -0.997001, T=105s, eta=20s
Iteration 2530, Loss: -0.997016, T=105s, eta=19s
Iteration 2540, Loss: -0.997030, T=106s, eta=19s
Iteration 2550, Loss: -0.997034, T=106s, eta=18s
Iteration 2560, Loss: -0.997031, T=107s, eta=18s
Iteration 2570, Loss: -0.997055, T=107s, eta=18s
Iteration 2580, Loss: -0.997083, T=108s, eta=17s
Iteration 2590, Loss: -0.997093, T=108s, eta=17s
Iteration 2600, Loss: -0.997105, T=108s, eta=16s
Iteration 2610, Loss: -0.997118, T=109s, eta=16s
Iteration 2620, Loss: -0.997130, T=109s, eta=15s
Iteration 2630, Loss: -0.997141, T=110s, eta=15s
Iteration 2640, Loss: -0.997152, T=110s, eta=15s
Iteration 2650, Loss: -0.997163, T=110s, eta=14s
Iteration 2660, Loss: -0.997174, T=111s, eta=14s
Iteration 2670, Loss: -0.997184, T=111s, eta=13s
Iteration 2680, Loss: -0.997194, T=112s, eta=13s
Iteration 2690, Loss: -0.997203, T=112s, eta=13s
Iteration 2700, Loss: -0.997213, T=112s, eta=12s
Iteration 2710, Loss: -0.997222, T=113s, eta=12s
Iteration 2720, Loss: -0.997232, T=113s, eta=11s
Iteration 2730, Loss: -0.997240, T=114s, eta=11s
Iteration 2740, Loss: -0.997249, T=114s, eta=10s
Iteration 2750, Loss: -0.997258, T=114s, eta=10s
Iteration 2760, Loss: -0.997266, T=115s, eta=10s
Iteration 2770, Loss: -0.997274, T=115s, eta=9s
Iteration 2780, Loss: -0.997282, T=116s, eta=9s
Iteration 2790, Loss: -0.997290, T=116s, eta=8s
Iteration 2800, Loss: -0.997298, T=117s, eta=8s
Iteration 2810, Loss: -0.997305, T=117s, eta=7s
Iteration 2820, Loss: -0.997312, T=117s, eta=7s
Iteration 2830, Loss: -0.997319, T=118s, eta=7s
Iteration 2840, Loss: -0.997326, T=118s, eta=6s
Iteration 2850, Loss: -0.997332, T=119s, eta=6s
Iteration 2860, Loss: -0.997339, T=119s, eta=5s
Iteration 2870, Loss: -0.997345, T=119s, eta=5s
Iteration 2880, Loss: -0.997351, T=120s, eta=5s
Iteration 2890, Loss: -0.997356, T=120s, eta=4s
Iteration 2900, Loss: -0.997362, T=121s, eta=4s
Iteration 2910, Loss: -0.997321, T=121s, eta=3s
Iteration 2920, Loss: -0.997253, T=122s, eta=3s
Iteration 2930, Loss: -0.997341, T=122s, eta=2s
Iteration 2940, Loss: -0.997366, T=122s, eta=2s
Iteration 2950, Loss: -0.997377, T=123s, eta=2s
Iteration 2960, Loss: -0.997388, T=123s, eta=1s
Iteration 2970, Loss: -0.997393, T=124s, eta=1s
Iteration 2980, Loss: -0.997397, T=124s, eta=0s
Iteration 2990, Loss: -0.997401, T=125s, eta=0s

result.final_fidelity

Array(0.99740387, dtype=float64)

len(result.returned_params)

20

# here makes sense for each batch size we have a different set of parameters since there are no measurements and therefore no stochasticisty or randomness
result.returned_params

[[Array([[-4.63880455e-06, -5.89163813e+00]], dtype=float64),
  Array([[-3.27025652]], dtype=float64)],
 [Array([[ 5.30936785e-06, -2.78119760e+00]], dtype=float64),
  Array([[2.16693322]], dtype=float64)],
 [Array([[ 2.26132122e-06, -3.49741135e+00]], dtype=float64),
  Array([[1.57835183]], dtype=float64)],
 [Array([[-5.53203603e-07, -2.99806151e+00]], dtype=float64),
  Array([[-1.98438101]], dtype=float64)],
 [Array([[ 1.47626300e-05, -9.48264375e+00]], dtype=float64),
  Array([[-1.58191364]], dtype=float64)],
 [Array([[-2.53038158e-06, -2.97480749e+00]], dtype=float64),
  Array([[-1.1509593]], dtype=float64)],
 [Array([[ 3.40190087e-06, -2.96776457e+00]], dtype=float64),
  Array([[-1.06949831]], dtype=float64)],
 [Array([[ 2.91159759e-07, -3.72293771e+00]], dtype=float64),
  Array([[-1.68201837]], dtype=float64)],
 [Array([[ 1.07229655e-06, -2.75384585e+00]], dtype=float64),
  Array([[-0.98959449]], dtype=float64)],
 [Array([[ 4.21711601e-06, -4.55072344e+00]], dtype=float64),
  Array([[-0.21186256]], dtype=float64)],
 [Array([[-1.31300876e-06, -4.47787632e+00]], dtype=float64),
  Array([[-1.43772112]], dtype=float64)],
 [Array([[ 1.75143954e-06, -3.52033868e+00]], dtype=float64),
  Array([[-1.07387673]], dtype=float64)],
 [Array([[ 2.38186582e-06, -4.84808239e+00]], dtype=float64),
  Array([[-1.79592884]], dtype=float64)],
 [Array([[-2.29087240e-05, -5.37155881e+00]], dtype=float64),
  Array([[-0.95674839]], dtype=float64)],
 [Array([[-3.48485556e-06, -2.74685381e+00]], dtype=float64),
  Array([[0.54153511]], dtype=float64)],
 [Array([[-4.09400812e-06, -3.40237410e+00]], dtype=float64),
  Array([[0.4005666]], dtype=float64)],
 [Array([[-6.51121948e-06, -4.68318984e+00]], dtype=float64),
  Array([[-1.39839245]], dtype=float64)],
 [Array([[ 1.04878340e-05, -3.91073139e+00]], dtype=float64),
  Array([[-0.87290756]], dtype=float64)],
 [Array([[-8.10436237e-06, -9.49443035e+00]], dtype=float64),
  Array([[-0.79295764]], dtype=float64)],
 [Array([[ 1.12882672e-03, -3.13645801e+00]], dtype=float64),
  Array([[0.00098022]], dtype=float64)]]

print(result.final_state)

[[[ 2.21902316e-02+8.00569317e-06j]
  [-9.07085185e-05+1.48062715e-07j]
  [ 8.37049919e-08-1.60629316e-05j]
  [ 1.11667404e-07-1.20082580e-04j]
  [ 8.89964256e-05+1.31668532e-08j]
  [-2.55783178e-04+1.77021538e-07j]
  [-1.09140779e-07-1.25175086e-04j]
  [-1.88688772e-07-1.64900109e-04j]
  [ 3.67922871e-01+1.33257005e-04j]
  [-9.54586259e-05+1.27989948e-07j]
  [ 3.97951639e-08-4.83889940e-05j]
  [ 1.00821383e-07-2.20171882e-04j]
  [ 1.43144445e-04+1.48312850e-08j]
  [-1.68897309e-04+7.63629438e-10j]
  [-1.23867892e-07-5.65780041e-05j]
  [-1.08897864e-07-5.87678539e-04j]
  [ 7.27251528e-01+2.62989878e-04j]
  [ 3.78573488e-04-5.80239358e-07j]
  [-7.82849249e-08+4.85661120e-04j]
  [ 6.71442062e-07+2.49232286e-04j]
  [-4.02010005e-04+1.15718305e-08j]
  [ 1.54898507e-03-7.31482887e-07j]
  [-5.29097982e-08+5.29856337e-04j]
  [ 3.76315650e-07+1.50800408e-03j]
  [ 5.39566884e-01+1.95143210e-04j]
  [-1.55484269e-04-9.05215741e-08j]
  [ 1.07007240e-07-3.79505801e-04j]
  [ 6.24636449e-07+2.56979010e-04j]
  [-2.41706609e-04-1.70539021e-07j]
  [-3.86603553e-03+1.36891062e-06j]
  [ 6.80100377e-07-1.15293537e-03j]
  [-6.02621676e-07-1.67747281e-03j]
  [ 2.09879042e-01+7.60888057e-05j]
  [-8.04077224e-04+9.41947250e-08j]
  [-5.12140588e-08+9.70355331e-07j]
  [-8.88769414e-08-1.10448182e-03j]
  [ 6.03088939e-03+2.18926312e-06j]
  [ 1.22974986e-05-9.93866425e-09j]
  [-6.29870228e-08+1.75003094e-04j]
  [ 4.67700240e-09-1.33332647e-05j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 3.83577377e-07+1.38056970e-10j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]]]

from feedback_grape.utils.fidelity import fidelity, ket2dm


fidelity(
    U_final=result.final_state[0],
    C_target=ket2dm(psi_target),
    evo_type="density",
)

Array(0.99740423, dtype=float64)

import matplotlib.pyplot as plt
import jax.numpy as jnp
from feedback_grape.utils.operators import identity
from feedback_grape.utils.tensor import tensor
from feedback_grape.utils.states import (
    fock,
)  # Assuming this provides the Fock state function

# Parameters
target_levels = [0, 4, 8, 12, 16, 20]

# Projectors and probabilities
probs = []
for n in target_levels:
    # Projector: |n⟩⟨n| ⊗ I_qubit
    # probability that the final state is in the projection subspace of the
    # cavity Fock state |n⟩
    proj = tensor(ket2dm(fock(N_cav, n)), ket2dm(basis(2)))
    prob = jnp.real(
        jnp.vdot(result.final_state[0], proj @ result.final_state[0])
    )
    probs.append(prob)

# Plotting
plt.figure(figsize=(8, 5))
plt.bar([str(n) for n in target_levels], probs, color='orange')
plt.yscale('log')
plt.yticks(
    [1e-8, 1e-4, 1e-2, 1e0],
    labels=[
        '$10^{-8}$',
        '$10^{-4}$',
        '$10^{-2}$',
        '$10^0$',
    ],
)
plt.xlabel('Cavity Fock State |n⟩')
plt.ylabel('Excitation Probability $P_n$ (log scale)')
plt.title('Excitation Number Distribution $\\langle n | \\rho | n \\rangle$')
plt.grid(True, axis='y', linestyle='--', alpha=0.6)
plt.tight_layout()
plt.show()