grape_2level.py

import sys
import random
import numpy as np
import matplotlib.pyplot as plt
import datetime
import qutip as qt
import qutip.logging_utils as logging
logger = logging.get_logger()
#QuTiP control modules
import qutip.control.pulseoptim as cpo

example_name = '2leveldiss'
log_level = logging.INFO

# ****************************************************************
# Define the physics of the problem
stateA = qt.basis(2,0)
stateB = qt.basis(2,1)

H1 = -1j*qt.sigmax()
H2 = -1j*qt.sigmaz()


gamma1=0.1
gamma2=0
H_d = qt.Qobj(np.diagflat([-0.5*gamma1,-0.5*gamma2]))

H_c = [H1,H2]
# Number of ctrls
n_ctrls = len(H_c)

psi_0 = stateA
psi_T = stateB

# ***** Define time evolution parameters *****
# Number of time slots
n_ts = 1000
# Time allowed for the evolution
evo_time = 15
ubound=1
lbound=0
# Fidelity error target
fid_err_targ = 0.0001

# Maximum iterations for the optisation algorithm
max_iter = 500
# Maximum (elapsed) time allowed in seconds
max_wall_time = 30


# Initial pulse type
# pulse type alternatives: RND|ZERO|LIN|SINE|SQUARE|SAW|TRIANGLE|
p_type = 'LIN'
# *************************************************************
# File extension for output files

f_ext = "{}_n_ts{}_ptype{}.txt".format(example_name, n_ts, p_type)

# Run the optimisation
print("\n***********************************")
print("Starting pulse optimisation")
result = cpo.optimize_pulse(H_d, H_c, psi_0, psi_T, n_ts, evo_time, 
                fid_err_targ=fid_err_targ,amp_ubound=ubound,amp_lbound=lbound,
                max_iter=max_iter, max_wall_time=max_wall_time, 
#                dyn_params={'oper_dtype':Qobj},
                # comment in/out these next three lines for CRAB/GRAPE
#                alg='CRAB',
#                alg_params={'init_coeff_scaling':5.0, 'num_coeffs':5, 'guess_pulse_type':None},
#                method_params={'xtol':1e-3},
               # fid_params={'phase_option':'PSU'},
                out_file_ext=f_ext, init_pulse_type=p_type,#dyn_type='GEN_MAT',prop_type='FRECHET',
                log_level=log_level, gen_stats=True)


print("\n***********************************")
print("Optimising complete. Stats follow:")
result.stats.report()
print("\nFinal evolution\n{}\n".format(result.evo_full_final))

print("********* Summary *****************")
print("Initial fidelity error {}".format(result.initial_fid_err))
print("Final fidelity error {}".format(result.fid_err))
print("Final gradient normal {}".format(result.grad_norm_final))
print("Terminated due to {}".format(result.termination_reason))
print("Number of iterations {}".format(result.num_iter))
#print("wall time: ", result.wall_time
print("Completed in {} HH:MM:SS.US".\
        format(datetime.timedelta(seconds=result.wall_time)))
print("***********************************")

# Plot the initial and final amplitudes
fig1 = plt.figure()
ax1 = fig1.add_subplot(2, 1, 1)
ax1.set_title("Initial control amps")
ax1.set_xlabel("Time")
ax1.set_ylabel("Control amplitude")
for j in range(n_ctrls):
    ax1.step(result.time, 
             np.hstack((result.initial_amps[:, j], result.initial_amps[-1, j])), 
             where='post')
             
ax2 = fig1.add_subplot(2, 1, 2)
ax2.set_title("Optimised Control Sequences")
ax2.set_xlabel("Time")
ax2.set_ylabel("Control amplitude")
for j in range(n_ctrls):
    ax2.step(result.time, 
             np.hstack((result.final_amps[:, j], result.final_amps[-1, j])), 
             where='post')
plt.show()