Density Matrix Examples#
CNOT Example#
# for testing purposes
import qutip as qt
from feedback_grape.utils.gates import cnot
import jax.numpy as jnp
from feedback_grape.utils.operators import identity, sigmax, sigmay, sigmaz
from feedback_grape.utils.tensor import tensor
from feedback_grape.grape import optimize_pulse, fidelity
from feedback_grape.utils.solver import sesolve
rho_target = cnot() @ cnot().conj().T
rho_initial = identity(4) @ identity(4).conj().T
H_drift = 0 * (tensor(sigmax(), sigmax()) + tensor(sigmay(), sigmay()))
H_ctrl = [
tensor(sigmax(), identity(2)),
tensor(sigmay(), identity(2)),
tensor(sigmaz(), identity(2)),
tensor(identity(2), sigmax()),
tensor(identity(2), sigmay()),
tensor(identity(2), sigmaz()),
tensor(sigmax(), sigmax()),
tensor(sigmay(), sigmay()),
tensor(sigmaz(), sigmaz()),
]
num_t_slots = 500
total_evo_time = 2 * jnp.pi
res = optimize_pulse(
H_drift,
H_ctrl,
rho_initial,
rho_target,
num_t_slots,
total_evo_time,
max_iter=100,
learning_rate=1e-2,
optimizer="adam",
evo_type="density",
)
print("final_fidelity: ", res.final_fidelity)
print("U_f \n", res.final_operator)
print("Converged after: ", res.iterations)
final_fidelity: 0.9999999999999998
U_f
[[ 1.00000000e+00+3.46944695e-18j 6.16284232e-17+3.68195058e-16j
-3.67039705e-16-7.04731412e-17j 3.19162015e-18-3.46944695e-18j]
[ 6.38230855e-17-3.40439482e-16j 1.00000000e+00+1.82145965e-17j
-1.10486978e-17+1.52655666e-16j 2.22094580e-16-4.29344060e-17j]
[-3.84005775e-16+5.85469173e-17j -1.11401773e-17-1.62630326e-16j
1.00000000e+00-1.12757026e-17j -6.15301674e-17-2.81458884e-16j]
[ 2.96122718e-18-8.67361738e-19j 2.18613274e-16+5.61616725e-17j
-5.52519591e-17+2.83627288e-16j 1.00000000e+00+9.97465999e-18j]]
Converged after: 2
qt.fidelity(
qt.Qobj(res.final_operator / jnp.trace(res.final_operator)),
qt.Qobj(rho_target / jnp.trace(rho_target)),
) # should be close to 1
np.float64(0.9999999999999993)
# Reconstucting hamiltonian
def build_ham_reconstructed(control_amplitudes):
"""
Build Hamiltonian for given (complex) e_qub and e_cav
"""
H_ctrl = (
control_amplitudes[0] * tensor(sigmax(), identity(2))
+ control_amplitudes[1] * tensor(sigmay(), identity(2))
+ control_amplitudes[2] * tensor(sigmaz(), identity(2))
+ control_amplitudes[3] * tensor(identity(2), sigmax())
+ control_amplitudes[4] * tensor(identity(2), sigmay())
+ control_amplitudes[5] * tensor(identity(2), sigmaz())
+ control_amplitudes[6] * tensor(sigmax(), sigmax())
+ control_amplitudes[7] * tensor(sigmay(), sigmay())
+ control_amplitudes[8] * tensor(sigmaz(), sigmaz())
)
H_drift = 0 * (tensor(sigmax(), sigmax()) + tensor(sigmay(), sigmay()))
H = H_drift + H_ctrl
return H
# Construct the Hamiltonian for each time slot
H_total = jnp.array(
[
build_ham_reconstructed(res.control_amplitudes[i, :])
for i in range(num_t_slots)
]
)
# Initial state
psi0_fg = identity(4) @ identity(4).conj().T
# Time intervals
delta_ts = jnp.repeat(total_evo_time / num_t_slots, num_t_slots).astype(
jnp.float32
)
# Solve the Schrödinger equation
psi_fg = sesolve(H_total, psi0_fg, delta_ts, evo_type="density")
# Target state
rho_target = cnot() @ cnot().conj().T
# Calculate fidelity
print(fidelity(C_target=rho_target, U_final=psi_fg, evo_type="density"))
1.0
Time Dependent with density Matrices#
## MAIN.py with time_dep example
import jax
import jax.numpy as jnp
from feedback_grape.grape import (
optimize_pulse,
fidelity,
)
from feedback_grape.utils.solver import sesolve
from feedback_grape.utils.operators import identity, destroy, sigmap, sigmaz
from feedback_grape.utils.tensor import tensor
from feedback_grape.utils.states import basis
# ruff: noqa
N_cav = 10
chi = 0.2385 * (2 * jnp.pi)
mu_qub = 4.0
mu_cav = 8.0
hconj = lambda a: jnp.swapaxes(a.conj(), -1, -2)
time_start = 0.0
time_end = 1.0
time_intervals_num = 5
N_cav = 10
# Eqivalant to delta_ts = jnp.repeat(0.2, time_intervals_num).astype(jnp.float32)
# However, it is implemented in this way to be more general and
# show that these are the differences between the time intervals
t_grid = jnp.linspace(time_start, time_end, time_intervals_num + 1)
delta_ts = t_grid[1:] - t_grid[:-1]
fake_random_key = jax.random.key(seed=0)
e_data = jax.random.uniform(
fake_random_key, shape=(4, len(delta_ts)), minval=-1, maxval=1
)
e_qub = e_data[0] + 1j * e_data[1]
e_cav = e_data[2] + 1j * e_data[3]
@jax.vmap
def build_ham(e_qub, e_cav):
"""
Build Hamiltonian for given (complex) e_qub and e_cav
"""
a = tensor(identity(2), destroy(N_cav))
adag = hconj(a)
n_phot = adag @ a
sigz = tensor(sigmaz(), identity(N_cav))
sigp = tensor(sigmap(), identity(N_cav))
one = tensor(identity(2), identity(N_cav))
H0 = +(chi / 2) * n_phot @ (sigz + one)
H_ctrl = mu_qub * sigp * e_qub + mu_cav * adag * e_cav
H_ctrl += hconj(H_ctrl)
# You just pass an array of the Hamiltonian matrices "Hs" corresponding to the time
# intervals "delta_ts" (that is, "Hs" is a 3D array).
return H0, H_ctrl
H0, H_ctrl = build_ham(e_qub, e_cav)
# Representation for time dependent Hamiltonian
def solve(Hs, delta_ts):
"""
Find evolution operator for piecewise Hs on time intervals delts_ts
"""
for i, (H, delta_t) in enumerate(zip(Hs, delta_ts)):
U_intv = jax.scipy.linalg.expm(-1j * H * delta_t)
U = U_intv if i == 0 else U_intv @ U
return U
U = solve(H0 + H_ctrl, delta_ts)
psi0 = tensor(basis(2), basis(N_cav)) @ tensor(basis(2), basis(N_cav)).conj().T
global psi_target_qt
psi_target_qt = psi_target = (U @ psi0) @ (U @ psi0).conj().T
def build_grape_format_ham():
"""
Build Hamiltonian for given (complex) e_qub and e_cav
"""
a = tensor(identity(2), destroy(N_cav))
adag = hconj(a)
n_phot = adag @ a
sigz = tensor(sigmaz(), identity(N_cav))
sigp = tensor(sigmap(), identity(N_cav))
one = tensor(identity(2), identity(N_cav))
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
def test_time_dep(optimizer="adam"):
H0_grape, H_ctrl_grape = build_grape_format_ham()
res = optimize_pulse(
H0_grape,
H_ctrl_grape,
psi0,
psi_target,
int(
(time_end - time_start) / delta_ts[0]
), # Ensure this is an integer
time_end - time_start,
max_iter=10000,
# when you decrease convergence threshold, it is more accurate
convergence_threshold=1e-3,
learning_rate=1e-2,
evo_type="density",
optimizer=optimizer,
)
return res
res = test_time_dep("l-bfgs")
print("final_fidelity: ", res.final_fidelity)
final_fidelity: 0.9976505826004491
### Define the time grid (same as defined)
time_start = 0.0
time_end = 1.0
time_intervals_num = 5
N_cav = 10
# Eqivalant to delta_ts = jnp.repeat(0.2, time_intervals_num).astype(jnp.float32)
# However, it is implemented in this way to be more general and
# show that these are the differences between the time intervals
t_grid = jnp.linspace(time_start, time_end, time_intervals_num + 1)
delta_ts = t_grid[1:] - t_grid[:-1]
### Build the Hamiltonian
def build_ham_reconstructed(u1, u2, u3, u4):
"""
Build Hamiltonian for given (complex) e_qub and e_cav
"""
a = tensor(identity(2), destroy(N_cav))
adag = hconj(a)
n_phot = adag @ a
sigz = tensor(sigmaz(), identity(N_cav))
sigp = tensor(sigmap(), identity(N_cav))
one = tensor(identity(2), identity(N_cav))
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)
# Apply control amplitudes
H_ctrl = (
u1 * H_ctrl_qub
+ u2 * H_ctrl_qub_dag
+ u3 * H_ctrl_cav
+ u4 * H_ctrl_cav_dag
)
H = H0 + H_ctrl
return H
u1 = res.control_amplitudes[:, 0]
u2 = res.control_amplitudes[:, 1]
u3 = res.control_amplitudes[:, 2]
u4 = res.control_amplitudes[:, 3]
u1
### Construct the Hamiltonian for each time step
H_total = jnp.array(
[
build_ham_reconstructed(u1[i], u2[i], u3[i], u4[i])
for i in range(len(u1))
]
)
H_total
H_total.shape
### Solve the Schrödinger Equation
psi0_fg = (
tensor(basis(2), basis(N_cav)) @ tensor(basis(2), basis(N_cav)).conj().T
)
psi_fg = sesolve(H_total, psi0_fg, delta_ts, evo_type="density")
### Calculate fidelity with target
print(fidelity(C_target=psi_target, U_final=psi_fg, evo_type="density"))
0.9976505915358532