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(
    init_state=psi0,
    target_state=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=133s
Iteration 20, Loss: -0.218807, T=0s, eta=138s
Iteration 30, Loss: -0.290923, T=1s, eta=139s
Iteration 40, Loss: -0.331432, T=1s, eta=139s
Iteration 50, Loss: -0.363328, T=2s, eta=139s
Iteration 60, Loss: -0.385187, T=2s, eta=139s
Iteration 70, Loss: -0.392812, T=3s, eta=139s
Iteration 80, Loss: -0.398971, T=3s, eta=138s
Iteration 90, Loss: -0.406213, T=4s, eta=138s
Iteration 100, Loss: -0.414206, T=4s, eta=137s
Iteration 110, Loss: -0.424084, T=5s, eta=137s
Iteration 120, Loss: -0.436752, T=5s, eta=136s
Iteration 130, Loss: -0.455166, T=6s, eta=136s
Iteration 140, Loss: -0.483714, T=6s, eta=135s
Iteration 150, Loss: -0.526234, T=7s, eta=135s
Iteration 160, Loss: -0.587400, T=7s, eta=135s
Iteration 170, Loss: -0.658891, T=8s, eta=134s
Iteration 180, Loss: -0.710771, T=8s, eta=134s
Iteration 190, Loss: -0.743793, T=9s, eta=133s
Iteration 200, Loss: -0.767970, T=9s, eta=133s
Iteration 210, Loss: -0.785470, T=10s, eta=132s
Iteration 220, Loss: -0.794869, T=10s, eta=132s
Iteration 230, Loss: -0.799156, T=10s, eta=131s
Iteration 240, Loss: -0.801985, T=11s, eta=131s
Iteration 250, Loss: -0.804416, T=11s, eta=130s
Iteration 260, Loss: -0.806428, T=12s, eta=130s
Iteration 270, Loss: -0.808274, T=12s, eta=129s
Iteration 280, Loss: -0.809989, T=13s, eta=130s
Iteration 290, Loss: -0.811600, T=13s, eta=129s
Iteration 300, Loss: -0.813111, T=14s, eta=129s
Iteration 310, Loss: -0.814524, T=14s, eta=128s
Iteration 320, Loss: -0.815841, T=15s, eta=128s
Iteration 330, Loss: -0.817068, T=15s, eta=127s
Iteration 340, Loss: -0.818218, T=16s, eta=127s
Iteration 350, Loss: -0.819303, T=16s, eta=127s
Iteration 360, Loss: -0.820341, T=17s, eta=127s
Iteration 370, Loss: -0.821346, T=17s, eta=127s
Iteration 380, Loss: -0.822334, T=18s, eta=127s
Iteration 390, Loss: -0.823330, T=19s, eta=127s
Iteration 400, Loss: -0.824376, T=19s, eta=128s
Iteration 410, Loss: -0.825539, T=20s, eta=128s
Iteration 420, Loss: -0.826936, T=20s, eta=128s
Iteration 430, Loss: -0.828753, T=21s, eta=128s
Iteration 440, Loss: -0.831279, T=22s, eta=128s
Iteration 450, Loss: -0.834815, T=23s, eta=130s
Iteration 460, Loss: -0.839446, T=23s, eta=129s
Iteration 470, Loss: -0.845438, T=24s, eta=129s
Iteration 480, Loss: -0.853808, T=24s, eta=128s
Iteration 490, Loss: -0.864693, T=25s, eta=128s
Iteration 500, Loss: -0.876147, T=25s, eta=128s
Iteration 510, Loss: -0.890039, T=26s, eta=127s
Iteration 520, Loss: -0.905138, T=26s, eta=127s
Iteration 530, Loss: -0.918734, T=27s, eta=126s
Iteration 540, Loss: -0.929061, T=27s, eta=125s
Iteration 550, Loss: -0.935838, T=28s, eta=125s
Iteration 560, Loss: -0.939936, T=28s, eta=124s
Iteration 570, Loss: -0.942359, T=29s, eta=123s
Iteration 580, Loss: -0.943913, T=29s, eta=123s
Iteration 590, Loss: -0.945100, T=30s, eta=122s
Iteration 600, Loss: -0.946135, T=30s, eta=121s
Iteration 610, Loss: -0.947034, T=31s, eta=121s
Iteration 620, Loss: -0.947801, T=31s, eta=120s
Iteration 630, Loss: -0.948456, T=31s, eta=120s
Iteration 640, Loss: -0.949025, T=32s, eta=119s
Iteration 650, Loss: -0.949528, T=32s, eta=118s
Iteration 660, Loss: -0.949979, T=33s, eta=118s
Iteration 670, Loss: -0.950387, T=33s, eta=117s
Iteration 680, Loss: -0.950758, T=34s, eta=117s
Iteration 690, Loss: -0.951097, T=34s, eta=116s
Iteration 700, Loss: -0.951405, T=35s, eta=115s
Iteration 710, Loss: -0.951687, T=35s, eta=115s
Iteration 720, Loss: -0.951943, T=36s, eta=115s
Iteration 730, Loss: -0.952177, T=36s, eta=114s
Iteration 740, Loss: -0.952389, T=37s, eta=114s
Iteration 750, Loss: -0.952581, T=37s, eta=113s
Iteration 760, Loss: -0.952754, T=38s, eta=112s
Iteration 770, Loss: -0.952908, T=38s, eta=112s
Iteration 780, Loss: -0.953045, T=39s, eta=111s
Iteration 790, Loss: -0.953168, T=40s, eta=111s
Iteration 800, Loss: -0.953279, T=40s, eta=111s
Iteration 810, Loss: -0.953380, T=41s, eta=110s
Iteration 820, Loss: -0.953474, T=41s, eta=110s
Iteration 830, Loss: -0.953563, T=41s, eta=109s
Iteration 840, Loss: -0.953649, T=42s, eta=109s
Iteration 850, Loss: -0.953732, T=42s, eta=108s
Iteration 860, Loss: -0.953814, T=43s, eta=108s
Iteration 870, Loss: -0.953894, T=43s, eta=107s
Iteration 880, Loss: -0.953972, T=44s, eta=106s
Iteration 890, Loss: -0.954047, T=44s, eta=106s
Iteration 900, Loss: -0.954119, T=45s, eta=105s
Iteration 910, Loss: -0.954189, T=45s, eta=105s
Iteration 920, Loss: -0.954255, T=46s, eta=104s
Iteration 930, Loss: -0.954318, T=46s, eta=104s
Iteration 940, Loss: -0.954379, T=47s, eta=103s
Iteration 950, Loss: -0.954437, T=47s, eta=103s
Iteration 960, Loss: -0.954492, T=48s, eta=102s
Iteration 970, Loss: -0.954545, T=48s, eta=102s
Iteration 980, Loss: -0.954595, T=49s, eta=101s
Iteration 990, Loss: -0.954643, T=49s, eta=101s
Iteration 1000, Loss: -0.954688, T=50s, eta=100s
Iteration 1010, Loss: -0.954731, T=50s, eta=100s
Iteration 1020, Loss: -0.954772, T=51s, eta=99s
Iteration 1030, Loss: -0.954811, T=51s, eta=98s
Iteration 1040, Loss: -0.954847, T=52s, eta=98s
Iteration 1050, Loss: -0.954882, T=52s, eta=97s
Iteration 1060, Loss: -0.954915, T=53s, eta=97s
Iteration 1070, Loss: -0.954946, T=53s, eta=96s
Iteration 1080, Loss: -0.954976, T=54s, eta=96s
Iteration 1090, Loss: -0.955004, T=54s, eta=95s
Iteration 1100, Loss: -0.955031, T=55s, eta=95s
Iteration 1110, Loss: -0.955057, T=55s, eta=94s
Iteration 1120, Loss: -0.955082, T=56s, eta=94s
Iteration 1130, Loss: -0.955105, T=56s, eta=93s
Iteration 1140, Loss: -0.955128, T=57s, eta=93s
Iteration 1150, Loss: -0.955150, T=57s, eta=92s
Iteration 1160, Loss: -0.955172, T=58s, eta=92s
Iteration 1170, Loss: -0.955193, T=58s, eta=91s
Iteration 1180, Loss: -0.955213, T=58s, eta=90s
Iteration 1190, Loss: -0.955233, T=59s, eta=90s
Iteration 1200, Loss: -0.955253, T=59s, eta=89s
Iteration 1210, Loss: -0.955272, T=60s, eta=89s
Iteration 1220, Loss: -0.955291, T=60s, eta=88s
Iteration 1230, Loss: -0.955310, T=61s, eta=88s
Iteration 1240, Loss: -0.955328, T=62s, eta=88s
Iteration 1250, Loss: -0.955347, T=62s, eta=87s
Iteration 1260, Loss: -0.955365, T=63s, eta=87s
Iteration 1270, Loss: -0.955384, T=63s, eta=86s
Iteration 1280, Loss: -0.955402, T=64s, eta=86s
Iteration 1290, Loss: -0.955421, T=64s, eta=85s
Iteration 1300, Loss: -0.955439, T=65s, eta=85s
Iteration 1310, Loss: -0.955457, T=65s, eta=84s
Iteration 1320, Loss: -0.955475, T=66s, eta=84s
Iteration 1330, Loss: -0.955488, T=66s, eta=84s
Iteration 1340, Loss: -0.955507, T=67s, eta=83s
Iteration 1350, Loss: -0.955526, T=68s, eta=83s
Iteration 1360, Loss: -0.955550, T=68s, eta=82s
Iteration 1370, Loss: -0.955571, T=69s, eta=82s
Iteration 1380, Loss: -0.955591, T=69s, eta=81s
Iteration 1390, Loss: -0.955611, T=70s, eta=81s
Iteration 1400, Loss: -0.955631, T=70s, eta=80s
Iteration 1410, Loss: -0.955652, T=71s, eta=80s
Iteration 1420, Loss: -0.955674, T=72s, eta=80s
Iteration 1430, Loss: -0.955695, T=72s, eta=79s
Iteration 1440, Loss: -0.955718, T=73s, eta=79s
Iteration 1450, Loss: -0.955740, T=73s, eta=78s
Iteration 1460, Loss: -0.955764, T=74s, eta=78s
Iteration 1470, Loss: -0.955788, T=74s, eta=77s
Iteration 1480, Loss: -0.955812, T=75s, eta=77s
Iteration 1490, Loss: -0.955838, T=76s, eta=77s
Iteration 1500, Loss: -0.955863, T=76s, eta=76s
Iteration 1510, Loss: -0.955890, T=77s, eta=76s
Iteration 1520, Loss: -0.955918, T=77s, eta=75s
Iteration 1530, Loss: -0.955946, T=78s, eta=75s
Iteration 1540, Loss: -0.955975, T=79s, eta=75s
Iteration 1550, Loss: -0.956005, T=80s, eta=74s
Iteration 1560, Loss: -0.956036, T=80s, eta=74s
Iteration 1570, Loss: -0.956068, T=81s, eta=74s
Iteration 1580, Loss: -0.956102, T=81s, eta=73s
Iteration 1590, Loss: -0.956136, T=82s, eta=73s
Iteration 1600, Loss: -0.956171, T=83s, eta=72s
Iteration 1610, Loss: -0.956208, T=84s, eta=72s
Iteration 1620, Loss: -0.956246, T=84s, eta=72s
Iteration 1630, Loss: -0.956285, T=85s, eta=71s
Iteration 1640, Loss: -0.956318, T=85s, eta=71s
Iteration 1650, Loss: -0.956356, T=86s, eta=70s
Iteration 1660, Loss: -0.956409, T=87s, eta=70s
Iteration 1670, Loss: -0.956456, T=87s, eta=69s
Iteration 1680, Loss: -0.956504, T=88s, eta=69s
Iteration 1690, Loss: -0.956556, T=89s, eta=69s
Iteration 1700, Loss: -0.956609, T=89s, eta=68s
Iteration 1710, Loss: -0.956664, T=90s, eta=68s
Iteration 1720, Loss: -0.956722, T=91s, eta=67s
Iteration 1730, Loss: -0.956783, T=91s, eta=67s
Iteration 1740, Loss: -0.956847, T=92s, eta=66s
Iteration 1750, Loss: -0.956915, T=92s, eta=66s
Iteration 1760, Loss: -0.956986, T=93s, eta=65s
Iteration 1770, Loss: -0.957062, T=94s, eta=65s
Iteration 1780, Loss: -0.957143, T=94s, eta=65s
Iteration 1790, Loss: -0.957228, T=95s, eta=64s
Iteration 1800, Loss: -0.957318, T=96s, eta=64s
Iteration 1810, Loss: -0.957365, T=96s, eta=63s
Iteration 1820, Loss: -0.957490, T=97s, eta=63s
Iteration 1830, Loss: -0.957625, T=97s, eta=62s
Iteration 1840, Loss: -0.957737, T=98s, eta=62s
Iteration 1850, Loss: -0.957862, T=99s, eta=61s
Iteration 1860, Loss: -0.957993, T=99s, eta=61s
Iteration 1870, Loss: -0.958130, T=100s, eta=60s
Iteration 1880, Loss: -0.958275, T=100s, eta=60s
Iteration 1890, Loss: -0.958428, T=101s, eta=59s
Iteration 1900, Loss: -0.958588, T=102s, eta=59s
Iteration 1910, Loss: -0.958755, T=102s, eta=58s
Iteration 1920, Loss: -0.958928, T=103s, eta=58s
Iteration 1930, Loss: -0.959106, T=104s, eta=57s
Iteration 1940, Loss: -0.959289, T=104s, eta=57s
Iteration 1950, Loss: -0.959420, T=105s, eta=56s
Iteration 1960, Loss: -0.959639, T=106s, eta=56s
Iteration 1970, Loss: -0.959860, T=106s, eta=55s
Iteration 1980, Loss: -0.960062, T=107s, eta=55s
Iteration 1990, Loss: -0.960268, T=108s, eta=54s
Iteration 2000, Loss: -0.960483, T=108s, eta=54s
Iteration 2010, Loss: -0.960707, T=109s, eta=53s
Iteration 2020, Loss: -0.960945, T=109s, eta=53s
Iteration 2030, Loss: -0.961200, T=110s, eta=52s
Iteration 2040, Loss: -0.961478, T=110s, eta=52s
Iteration 2050, Loss: -0.961786, T=111s, eta=51s
Iteration 2060, Loss: -0.962134, T=112s, eta=51s
Iteration 2070, Loss: -0.962534, T=112s, eta=50s
Iteration 2080, Loss: -0.962968, T=113s, eta=50s
Iteration 2090, Loss: -0.963527, T=114s, eta=49s
Iteration 2100, Loss: -0.964281, T=114s, eta=49s
Iteration 2110, Loss: -0.965174, T=115s, eta=48s
Iteration 2120, Loss: -0.966338, T=115s, eta=48s
Iteration 2130, Loss: -0.967832, T=116s, eta=47s
Iteration 2140, Loss: -0.969651, T=116s, eta=47s
Iteration 2150, Loss: -0.971706, T=117s, eta=46s
Iteration 2160, Loss: -0.973925, T=118s, eta=45s
Iteration 2170, Loss: -0.976103, T=118s, eta=45s
Iteration 2180, Loss: -0.979339, T=119s, eta=44s
Iteration 2190, Loss: -0.983137, T=119s, eta=44s
Iteration 2200, Loss: -0.987161, T=120s, eta=43s
Iteration 2210, Loss: -0.990380, T=120s, eta=43s
Iteration 2220, Loss: -0.992954, T=121s, eta=42s
Iteration 2230, Loss: -0.994452, T=121s, eta=42s
Iteration 2240, Loss: -0.995231, T=122s, eta=41s
Iteration 2250, Loss: -0.995563, T=123s, eta=41s
Iteration 2260, Loss: -0.995764, T=123s, eta=40s
Iteration 2270, Loss: -0.995923, T=124s, eta=39s
Iteration 2280, Loss: -0.996060, T=124s, eta=39s
Iteration 2290, Loss: -0.996175, T=125s, eta=38s
Iteration 2300, Loss: -0.996272, T=126s, eta=38s
Iteration 2310, Loss: -0.996354, T=126s, eta=37s
Iteration 2320, Loss: -0.996425, T=127s, eta=37s
Iteration 2330, Loss: -0.996487, T=127s, eta=36s
Iteration 2340, Loss: -0.996542, T=128s, eta=36s
Iteration 2350, Loss: -0.996591, T=128s, eta=35s
Iteration 2360, Loss: -0.996635, T=129s, eta=35s
Iteration 2370, Loss: -0.996674, T=129s, eta=34s
Iteration 2380, Loss: -0.996701, T=130s, eta=34s
Iteration 2390, Loss: -0.996653, T=130s, eta=33s
Iteration 2400, Loss: -0.996765, T=131s, eta=32s
Iteration 2410, Loss: -0.996793, T=132s, eta=32s
Iteration 2420, Loss: -0.996819, T=132s, eta=31s
Iteration 2430, Loss: -0.996846, T=133s, eta=31s
Iteration 2440, Loss: -0.996867, T=133s, eta=30s
Iteration 2450, Loss: -0.996888, T=134s, eta=30s
Iteration 2460, Loss: -0.996907, T=135s, eta=29s
Iteration 2470, Loss: -0.996926, T=135s, eta=29s
Iteration 2480, Loss: -0.996943, T=136s, eta=28s
Iteration 2490, Loss: -0.996960, T=137s, eta=28s
Iteration 2500, Loss: -0.996976, T=137s, eta=27s
Iteration 2510, Loss: -0.996992, T=138s, eta=27s
Iteration 2520, Loss: -0.997007, T=138s, eta=26s
Iteration 2530, Loss: -0.997021, T=139s, eta=25s
Iteration 2540, Loss: -0.997035, T=140s, eta=25s
Iteration 2550, Loss: -0.997048, T=140s, eta=24s
Iteration 2560, Loss: -0.996958, T=141s, eta=24s
Iteration 2570, Loss: -0.996981, T=141s, eta=23s
Iteration 2580, Loss: -0.997077, T=142s, eta=23s
Iteration 2590, Loss: -0.997100, T=143s, eta=22s
Iteration 2600, Loss: -0.997109, T=144s, eta=22s
Iteration 2610, Loss: -0.997122, T=145s, eta=21s
Iteration 2620, Loss: -0.997134, T=145s, eta=21s
Iteration 2630, Loss: -0.997145, T=146s, eta=20s
Iteration 2640, Loss: -0.997156, T=146s, eta=20s
Iteration 2650, Loss: -0.997167, T=147s, eta=19s
Iteration 2660, Loss: -0.997177, T=147s, eta=18s
Iteration 2670, Loss: -0.997187, T=148s, eta=18s
Iteration 2680, Loss: -0.997197, T=149s, eta=17s
Iteration 2690, Loss: -0.997207, T=149s, eta=17s
Iteration 2700, Loss: -0.997216, T=150s, eta=16s
Iteration 2710, Loss: -0.997226, T=151s, eta=16s
Iteration 2720, Loss: -0.997235, T=151s, eta=15s
Iteration 2730, Loss: -0.997244, T=152s, eta=15s
Iteration 2740, Loss: -0.997252, T=152s, eta=14s
Iteration 2750, Loss: -0.997261, T=153s, eta=13s
Iteration 2760, Loss: -0.997269, T=153s, eta=13s
Iteration 2770, Loss: -0.997277, T=154s, eta=12s
Iteration 2780, Loss: -0.997285, T=154s, eta=12s
Iteration 2790, Loss: -0.997293, T=155s, eta=11s
Iteration 2800, Loss: -0.997300, T=156s, eta=11s
Iteration 2810, Loss: -0.997293, T=156s, eta=10s
Iteration 2820, Loss: -0.997250, T=157s, eta=10s
Iteration 2830, Loss: -0.997276, T=157s, eta=9s
Iteration 2840, Loss: -0.997324, T=158s, eta=8s
Iteration 2850, Loss: -0.997334, T=159s, eta=8s
Iteration 2860, Loss: -0.997339, T=159s, eta=7s
Iteration 2870, Loss: -0.997345, T=160s, eta=7s
Iteration 2880, Loss: -0.997352, T=160s, eta=6s
Iteration 2890, Loss: -0.997357, T=161s, eta=6s
Iteration 2900, Loss: -0.997363, T=161s, eta=5s
Iteration 2910, Loss: -0.997368, T=162s, eta=5s
Iteration 2920, Loss: -0.997373, T=162s, eta=4s
Iteration 2930, Loss: -0.997377, T=163s, eta=3s
Iteration 2940, Loss: -0.997382, T=164s, eta=3s
Iteration 2950, Loss: -0.997386, T=164s, eta=2s
Iteration 2960, Loss: -0.997390, T=165s, eta=2s
Iteration 2970, Loss: -0.997394, T=165s, eta=1s
Iteration 2980, Loss: -0.997398, T=166s, eta=1s
Iteration 2990, Loss: -0.997401, T=166s, eta=0s

result.final_fidelity

Array(0.99740477, 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.31525156e-06, -5.89273867e+00]], dtype=float64),
  Array([[-3.26988929]], dtype=float64)],
 [Array([[ 5.28977082e-06, -2.78282846e+00]], dtype=float64),
  Array([[2.16577799]], dtype=float64)],
 [Array([[ 2.19052949e-06, -3.49634372e+00]], dtype=float64),
  Array([[1.57935168]], dtype=float64)],
 [Array([[-5.27064663e-07, -2.99868057e+00]], dtype=float64),
  Array([[-1.98368342]], dtype=float64)],
 [Array([[ 1.42545146e-05, -9.48266314e+00]], dtype=float64),
  Array([[-1.58110898]], dtype=float64)],
 [Array([[-2.44811376e-06, -2.97317075e+00]], dtype=float64),
  Array([[-1.15033653]], dtype=float64)],
 [Array([[ 3.26297135e-06, -2.96861657e+00]], dtype=float64),
  Array([[-1.06996235]], dtype=float64)],
 [Array([[ 3.27621419e-07, -3.72471256e+00]], dtype=float64),
  Array([[-1.68242958]], dtype=float64)],
 [Array([[ 1.03427629e-06, -2.75121094e+00]], dtype=float64),
  Array([[-0.9894554]], dtype=float64)],
 [Array([[ 4.12795422e-06, -4.55488831e+00]], dtype=float64),
  Array([[-0.21148336]], dtype=float64)],
 [Array([[-1.14929512e-06, -4.47396234e+00]], dtype=float64),
  Array([[-1.43820629]], dtype=float64)],
 [Array([[ 1.77539958e-06, -3.52249274e+00]], dtype=float64),
  Array([[-1.07357422]], dtype=float64)],
 [Array([[ 2.45642025e-06, -4.85270580e+00]], dtype=float64),
  Array([[-1.79742222]], dtype=float64)],
 [Array([[-2.19583214e-05, -5.36539024e+00]], dtype=float64),
  Array([[-0.95500414]], dtype=float64)],
 [Array([[-3.28933134e-06, -2.74702593e+00]], dtype=float64),
  Array([[0.54108602]], dtype=float64)],
 [Array([[-3.94534672e-06, -3.40208410e+00]], dtype=float64),
  Array([[0.40112635]], dtype=float64)],
 [Array([[-6.20905287e-06, -4.68106789e+00]], dtype=float64),
  Array([[-1.39869067]], dtype=float64)],
 [Array([[ 1.01746731e-05, -3.91329288e+00]], dtype=float64),
  Array([[-0.87306763]], dtype=float64)],
 [Array([[-8.01372939e-06, -9.49427321e+00]], dtype=float64),
  Array([[-0.79292101]], dtype=float64)],
 [Array([[ 1.07440972e-03, -3.13650795e+00]], dtype=float64),
  Array([[0.00097306]], dtype=float64)]]

print(result.final_state)

[[[ 2.21775705e-02+7.61725082e-06j]
  [-9.42527381e-05+1.47265778e-07j]
  [ 6.25695969e-08+4.83372339e-06j]
  [ 1.30268004e-07-7.94846094e-05j]
  [ 8.37094166e-05+2.29733088e-08j]
  [-2.27302113e-04+1.55856699e-07j]
  [-9.37076162e-08-1.16420490e-04j]
  [-1.89259222e-07-2.12545332e-04j]
  [ 3.67964723e-01+1.26879193e-04j]
  [-9.17927584e-05+1.09820839e-07j]
  [ 2.69246147e-08-3.83959901e-05j]
  [ 1.00103664e-07-1.98415290e-04j]
  [ 6.05868693e-05-3.92709414e-10j]
  [-1.82161903e-04+1.21624383e-09j]
  [-1.25452245e-07-3.43227710e-05j]
  [-9.22556825e-08-6.03724801e-04j]
  [ 7.27258281e-01+2.50381102e-04j]
  [ 3.79799155e-04-5.46675642e-07j]
  [-8.39295796e-08+5.13440231e-04j]
  [ 6.32204594e-07+2.61714668e-04j]
  [-3.27414999e-04+3.69142201e-08j]
  [ 1.52502611e-03-6.86531660e-07j]
  [-3.47498641e-08+5.16989299e-04j]
  [ 3.58356674e-07+1.48995495e-03j]
  [ 5.39540966e-01+1.85789801e-04j]
  [-1.39793620e-04-8.29910695e-08j]
  [ 9.04292442e-08-3.75957735e-04j]
  [ 5.93473970e-07+2.75832237e-04j]
  [-2.00601486e-04-1.45497981e-07j]
  [-3.83016402e-03+1.28977332e-06j]
  [ 6.50168441e-07-1.15502956e-03j]
  [-5.69356143e-07-1.66583906e-03j]
  [ 2.09852474e-01+7.24481944e-05j]
  [-7.97365226e-04+8.90540785e-08j]
  [-5.57316578e-08+2.02331247e-05j]
  [-8.40683312e-08-1.09729652e-03j]
  [ 5.99716462e-03+2.07335590e-06j]
  [ 1.33348144e-05-9.70285964e-09j]
  [-6.16640670e-08+1.79892116e-04j]
  [ 4.39713158e-09-1.31756927e-05j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 3.91413137e-07+1.34170004e-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.99740509, 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()