GRAPE for systems with Dissipation example

import jax.numpy as jnp
from feedback_grape.grape import *
from feedback_grape.utils.gates import *
from feedback_grape.utils.operators import *
from feedback_grape.utils.states import *
from feedback_grape.utils.tensor import *

Sx = sigmax()
Sz = sigmaz()
Sm = sigmam()
# Hamiltonian
Del = 0.1  # Tunnelling term
wq = 1.0  # Energy of the 2-level system.
H0 = 0.5 * wq * Sz + 0.5 * Del * Sx

# Amplitude damping#
# Damping rate:
gamma = 0.01
l_ops = [jnp.sqrt(gamma) * Sm]
# Kraus operators


# Drift
drift = H0
# Controls - different combinations can be tried
ctrls = [Sz, Sx]
# Number of ctrls
n_ctrls = len(ctrls)

# Number of time slots
n_ts = 10
# Time allowed for the evolution
evo_time = 2

from feedback_grape.utils.states import basis, coherent

psi0 = basis(2)  # Initial state
psi_target = coherent(2, 1.5)  # Target state

result = optimize_pulse(
    H_drift=drift,
    H_control=ctrls,
    U_0=psi0 @ psi0.conj().T,  # Initial state as a density matrix
    C_target=psi_target
    @ psi_target.conj().T,  # Target state as a density matrix
    c_ops=l_ops,
    num_t_slots=n_ts,
    total_evo_time=evo_time,
    evo_type="density",
    optimizer="adam",
    convergence_threshold=1e-16,
    max_iter=1000,
    learning_rate=0.01,
)

result.final_fidelity

Array(0.99033529, dtype=float64)

from feedback_grape.utils.solver import mesolve
from feedback_grape.utils.states import coherent

gamma = 0.3
jump_ops = [jnp.sqrt(gamma) * Sm, jnp.sqrt(gamma * 2) * Sm.T]

H0 = 0.5 * wq * sigmaz() + 0.5 * Del * sigmax()
H_ctrl = [Sz, Sx]

# tsave should be the time points at which the state is saved, not a list of delta t
tsave = jnp.linspace(0, 1, 1)

psi = coherent(2, alpha=1.0)
rho0 = psi @ psi.conj().T

print(rho0.shape)

mesolve_result_2 = mesolve(
    H=[H0] + H_ctrl, jump_ops=jump_ops, rho0=rho0, tsave=tsave
)

(2, 2)

print(len([H0] + H_ctrl))

3

print(rho0)

[[0.36787944+0.j 0.36787944+0.j]
 [0.36787944+0.j 0.36787944+0.j]]

print(mesolve_result_2.shape)

(2, 2)

print(mesolve_result_2)

[[0.36787944+0.j 0.36787944+0.j]
 [0.36787944+0.j 0.36787944+0.j]]

print(tsave)

[0.]

for item in zip([H0] + H_ctrl, tsave):
    print(item)

(Array([[ 0.5 +0.j,  0.05+0.j],
       [ 0.05+0.j, -0.5 +0.j]], dtype=complex128), Array(0., dtype=float64))

c_ops = (
    [  # c_ops for each decay index
        [
            tensor(identity(30), jnp.sqrt(0.01) * sigmam()),
        ],
        ["hi"],
    ],
)

print(c_ops[0][0])

[Array([[0. +0.j, 0. +0.j, 0. +0.j, ..., 0. +0.j, 0. +0.j, 0. +0.j],
       [0.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.1+0.j, 0. +0.j]],      dtype=complex128)]