def test_parallel_map():
"parallel_map"
args = (1, 2, 3)
kwargs = {'d': 4, 'e': 5, 'f': 6}
x = np.arange(10)
y1 = list(map(_func1, x))
y1 = [_func2(xx, *args, **kwargs) for xx in x]
y2 = parallel_map(_func2, x, args, kwargs, num_cpus=1)
assert_((np.array(y1) == np.array(y2)).all())
y2 = parallel_map(_func2, x, args, kwargs, num_cpus=2)
assert_((np.array(y1) == np.array(y2)).all())
def test_parallel_map():
"parallel_map"
args = (1, 2, 3)
kwargs = {'d': 4, 'e': 5, 'f': 6}
x = np.arange(10)
y1 = list(map(_func1, x))
y1 = [_func2(xx, *args, **kwargs)for xx in x]
y2 = parallel_map(_func2, x, args, kwargs, num_cpus=1)
assert_((np.array(y1) == np.array(y2)).all())
y2 = parallel_map(_func2, x, args, kwargs, num_cpus=2)
assert_((np.array(y1) == np.array(y2)).all())
def propagator(H, t, c_op_list=[], args={}, options=None,
unitary_mode='batch', parallel=False,
progress_bar=None, **kwargs):
"""
Calculate the propagator U(t) for the density matrix or wave function such
that :math:`\psi(t) = U(t)\psi(0)` or
:math:`\\rho_{\mathrm vec}(t) = U(t) \\rho_{\mathrm vec}(0)`
where :math:`\\rho_{\mathrm vec}` is the vector representation of the
density matrix.
Parameters
----------
H : qobj or list
Hamiltonian as a Qobj instance of a nested list of Qobjs and
coefficients in the list-string or list-function format for
time-dependent Hamiltonians (see description in :func:`qutip.mesolve`).
t : float or array-like
Time or list of times for which to evaluate the propagator.
c_op_list : list
List of qobj collapse operators.
args : list/array/dictionary
Parameters to callback functions for time-dependent Hamiltonians and
collapse operators.
options : :class:`qutip.Options`
with options for the ODE solver.
unitary_mode = str ('batch', 'single')
Solve all basis vectors simulaneously ('batch') or individually
('single').
parallel : bool {False, True}
Run the propagator in parallel mode. This will override the
unitary_mode settings if set to True.
progress_bar: BaseProgressBar
Optional instance of BaseProgressBar, or a subclass thereof, for
showing the progress of the simulation. By default no progress bar
is used, and if set to True a TextProgressBar will be used.
Returns
-------
a : qobj
Instance representing the propagator :math:`U(t)`.
"""
kw = _default_kwargs()
if 'num_cpus' in kwargs:
num_cpus = kwargs['num_cpus']
else:
num_cpus = kw['num_cpus']
if progress_bar is None:
progress_bar = BaseProgressBar()
elif progress_bar is True:
progress_bar = TextProgressBar()
if options is None:
options = Options()
options.rhs_reuse = True
rhs_clear()
if isinstance(t, (int, float, np.integer, np.floating)):
tlist = [0, t]
else:
tlist = t
td_type = _td_format_check(H, c_op_list, solver='me')
if isinstance(H, (types.FunctionType, types.BuiltinFunctionType,
functools.partial)):
H0 = H(0.0, args)
elif isinstance(H, list):
H0 = H[0][0] if isinstance(H[0], list) else H[0]
else:
H0 = H
if len(c_op_list) == 0 and H0.isoper:
# calculate propagator for the wave function
N = H0.shape[0]
dims = H0.dims
if parallel:
unitary_mode = 'single'
u = np.zeros([N, N, len(tlist)], dtype=complex)
output = parallel_map(_parallel_sesolve,range(N),
task_args=(N,H, tlist,args,options),
progress_bar=progress_bar, num_cpus=num_cpus)
for n in range(N):
for k, t in enumerate(tlist):
u[:, n, k] = output[n].states[k].full().T
else:
if unitary_mode == 'single':
u = np.zeros([N, N, len(tlist)], dtype=complex)
progress_bar.start(N)
for n in range(0, N):
progress_bar.update(n)
psi0 = basis(N, n)
output = sesolve(H, psi0, tlist, [], args, options, _safe_mode=False)
for k, t in enumerate(tlist):
u[:, n, k] = output.states[k].full().T
progress_bar.finished()
elif unitary_mode =='batch':
u = np.zeros(len(tlist), dtype=object)
_rows = np.array([(N+1)*m for m in range(N)])
_cols = np.zeros_like(_rows)
_data = np.ones_like(_rows,dtype=complex)
psi0 = Qobj(sp.coo_matrix((_data,(_rows,_cols))).tocsr())
if td_type[1] > 0 or td_type[2] > 0:
H2 = []
for k in range(len(H)):
if isinstance(H[k], list):
H2.append([tensor(qeye(N), H[k][0]), H[k][1]])
else:
H2.append(tensor(qeye(N), H[k]))
else:
H2 = tensor(qeye(N), H)
output = sesolve(H2, psi0, tlist, [] , args = args, _safe_mode=False,
options=Options(normalize_output=False))
for k, t in enumerate(tlist):
u[k] = sp_reshape(output.states[k].data, (N, N))
unit_row_norm(u[k].data, u[k].indptr, u[k].shape[0])
u[k] = u[k].T.tocsr()
else:
raise Exception('Invalid unitary mode.')
elif len(c_op_list) == 0 and H0.issuper:
# calculate the propagator for the vector representation of the
# density matrix (a superoperator propagator)
unitary_mode = 'single'
N = H0.shape[0]
sqrt_N = int(np.sqrt(N))
dims = H0.dims
u = np.zeros([N, N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_mesolve,range(N * N),
task_args=(sqrt_N,H,tlist,c_op_list,args,options),
progress_bar=progress_bar, num_cpus=num_cpus)
for n in range(N * N):
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output[n].states[k].full()).T
else:
progress_bar.start(N)
for n in range(0, N):
progress_bar.update(n)
col_idx, row_idx = np.unravel_index(n,(sqrt_N,sqrt_N))
rho0 = Qobj(sp.csr_matrix(([1],([row_idx],[col_idx])), shape=(sqrt_N,sqrt_N), dtype=complex))
output = mesolve(H, rho0, tlist, [], [], args, options, _safe_mode=False)
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output.states[k].full()).T
progress_bar.finished()
else:
# calculate the propagator for the vector representation of the
# density matrix (a superoperator propagator)
unitary_mode = 'single'
N = H0.shape[0]
dims = [H0.dims, H0.dims]
u = np.zeros([N * N, N * N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_mesolve,range(N * N),
task_args=(N,H,tlist,c_op_list,args,options),
progress_bar=progress_bar, num_cpus=num_cpus)
for n in range(N * N):
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output[n].states[k].full()).T
else:
progress_bar.start(N * N)
for n in range(N * N):
progress_bar.update(n)
col_idx, row_idx = np.unravel_index(n,(N,N))
rho0 = Qobj(sp.csr_matrix(([1],([row_idx],[col_idx])), shape=(N,N), dtype=complex))
output = mesolve(H, rho0, tlist, c_op_list, [], args, options, _safe_mode=False)
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output.states[k].full()).T
progress_bar.finished()
if len(tlist) == 2:
if unitary_mode == 'batch':
return Qobj(u[-1], dims=dims)
else:
return Qobj(u[:, :, 1], dims=dims)
else:
if unitary_mode == 'batch':
return np.array([Qobj(u[k], dims=dims) for k in range(len(tlist))], dtype=object)
else:
return np.array([Qobj(u[:, :, k], dims=dims) for k in range(len(tlist))], dtype=object)
def propagator(H,
t,
c_op_list=[],
args={},
options=None,
unitary_mode='batch',
parallel=False,
progress_bar=None,
_safe_mode=True,
**kwargs):
r"""
Calculate the propagator U(t) for the density matrix or wave function such
that :math:`\psi(t) = U(t)\psi(0)` or
:math:`\rho_{\mathrm vec}(t) = U(t) \rho_{\mathrm vec}(0)`
where :math:`\rho_{\mathrm vec}` is the vector representation of the
density matrix.
Parameters
----------
H : qobj or list
Hamiltonian as a Qobj instance of a nested list of Qobjs and
coefficients in the list-string or list-function format for
time-dependent Hamiltonians (see description in :func:`qutip.mesolve`).
t : float or array-like
Time or list of times for which to evaluate the propagator.
c_op_list : list
List of qobj collapse operators.
args : list/array/dictionary
Parameters to callback functions for time-dependent Hamiltonians and
collapse operators.
options : :class:`qutip.Options`
with options for the ODE solver.
unitary_mode = str ('batch', 'single')
Solve all basis vectors simulaneously ('batch') or individually
('single').
parallel : bool {False, True}
Run the propagator in parallel mode. This will override the
unitary_mode settings if set to True.
progress_bar: BaseProgressBar
Optional instance of BaseProgressBar, or a subclass thereof, for
showing the progress of the simulation. By default no progress bar
is used, and if set to True a TextProgressBar will be used.
Returns
-------
a : qobj
Instance representing the propagator :math:`U(t)`.
"""
kw = _default_kwargs()
if 'num_cpus' in kwargs:
num_cpus = kwargs['num_cpus']
else:
num_cpus = kw['num_cpus']
if progress_bar is None:
progress_bar = BaseProgressBar()
elif progress_bar is True:
progress_bar = TextProgressBar()
if options is None:
options = Options()
options.rhs_reuse = True
rhs_clear()
if isinstance(t, (int, float, np.integer, np.floating)):
tlist = [0, t]
else:
tlist = t
if _safe_mode:
_solver_safety_check(H, None, c_ops=c_op_list, e_ops=[], args=args)
td_type = _td_format_check(H, c_op_list, solver='me')
if isinstance(
H,
(types.FunctionType, types.BuiltinFunctionType, functools.partial)):
H0 = H(0.0, args)
if unitary_mode == 'batch':
# batch don't work with function Hamiltonian
unitary_mode = 'single'
elif isinstance(H, list):
H0 = H[0][0] if isinstance(H[0], list) else H[0]
else:
H0 = H
if len(c_op_list) == 0 and H0.isoper:
# calculate propagator for the wave function
N = H0.shape[0]
dims = H0.dims
if parallel:
unitary_mode = 'single'
u = np.zeros([N, N, len(tlist)], dtype=complex)
output = parallel_map(_parallel_sesolve,
range(N),
task_args=(N, H, tlist, args, options),
progress_bar=progress_bar,
num_cpus=num_cpus)
for n in range(N):
for k, t in enumerate(tlist):
u[:, n, k] = output[n].states[k].full().T
else:
if unitary_mode == 'single':
output = sesolve(H,
qeye(dims[0]),
tlist, [],
args,
options,
_safe_mode=False)
if len(tlist) == 2:
return output.states[-1]
else:
return output.states
elif unitary_mode == 'batch':
u = np.zeros(len(tlist), dtype=object)
_rows = np.array([(N + 1) * m for m in range(N)])
_cols = np.zeros_like(_rows)
_data = np.ones_like(_rows, dtype=complex)
psi0 = Qobj(sp.coo_matrix((_data, (_rows, _cols))).tocsr())
if td_type[1] > 0 or td_type[2] > 0:
H2 = []
for k in range(len(H)):
if isinstance(H[k], list):
H2.append([tensor(qeye(N), H[k][0]), H[k][1]])
else:
H2.append(tensor(qeye(N), H[k]))
else:
H2 = tensor(qeye(N), H)
options.normalize_output = False
output = sesolve(H2,
psi0,
tlist, [],
args=args,
options=options,
_safe_mode=False)
for k, t in enumerate(tlist):
u[k] = sp_reshape(output.states[k].data, (N, N))
unit_row_norm(u[k].data, u[k].indptr, u[k].shape[0])
u[k] = u[k].T.tocsr()
else:
raise Exception('Invalid unitary mode.')
elif len(c_op_list) == 0 and H0.issuper:
# calculate the propagator for the vector representation of the
# density matrix (a superoperator propagator)
unitary_mode = 'single'
N = H0.shape[0]
sqrt_N = int(np.sqrt(N))
dims = H0.dims
u = np.zeros([N, N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_mesolve,
range(N * N),
task_args=(sqrt_N, H, tlist, c_op_list, args,
options),
progress_bar=progress_bar,
num_cpus=num_cpus)
for n in range(N * N):
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output[n].states[k].full()).T
else:
rho0 = qeye(N, N)
rho0.dims = [[sqrt_N, sqrt_N], [sqrt_N, sqrt_N]]
output = mesolve(H,
psi0,
tlist, [],
args,
options,
_safe_mode=False)
if len(tlist) == 2:
return output.states[-1]
else:
return output.states
else:
# calculate the propagator for the vector representation of the
# density matrix (a superoperator propagator)
unitary_mode = 'single'
N = H0.shape[0]
dims = [H0.dims, H0.dims]
u = np.zeros([N * N, N * N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_mesolve,
range(N * N),
task_args=(N, H, tlist, c_op_list, args,
options),
progress_bar=progress_bar,
num_cpus=num_cpus)
for n in range(N * N):
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output[n].states[k].full()).T
else:
progress_bar.start(N * N)
for n in range(N * N):
progress_bar.update(n)
col_idx, row_idx = np.unravel_index(n, (N, N))
rho0 = Qobj(
sp.csr_matrix(([1], ([row_idx], [col_idx])),
shape=(N, N),
dtype=complex))
output = mesolve(H,
rho0,
tlist,
c_op_list, [],
args,
options,
_safe_mode=False)
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output.states[k].full()).T
progress_bar.finished()
if len(tlist) == 2:
if unitary_mode == 'batch':
return Qobj(u[-1], dims=dims)
else:
return Qobj(u[:, :, 1], dims=dims)
else:
if unitary_mode == 'batch':
return np.array([Qobj(u[k], dims=dims) for k in range(len(tlist))],
dtype=object)
else:
return np.array(
[Qobj(u[:, :, k], dims=dims) for k in range(len(tlist))],
dtype=object)
args = {'w': omega}
freezing_pts = jn_zeros(0, 4)
H0 = -2.0 * q**2 + h0 * (1. - 4. * q**2).sqrtm() * p.cosm()
def floquet_esys(ampl):
H1 = ampl * (1. - 4. * q**2).sqrtm() * p.cosm()
H = [H0, [H1, lambda t, args: np.cos(args["w"] * t)]]
f_modes_0, f_energies = floquet_modes(H, time_period, args, True)
return f_energies
ampls = np.linspace(1.0, 250.0, 300)
quasi_evals = parallel_map(floquet_esys, ampls, progress_bar=True)
quasi_evals = np.array(quasi_evals)
np.savetxt("N_%d_w_%1.3lf.csv" % (N, omega),
np.vstack((ampls, quasi_evals.T)).T,
delimiter=',')
for i in range(N):
# plt.scatter(ampls/args["w"], quasi_evals[:, i], s=0.1, c='b')
plt.plot(4.0 * ampls / args["w"], quasi_evals[:, i], c='b')
for pt in freezing_pts:
plt.axvline(x=pt, color='gray', linestyle="--")
plt.xlabel(r'$4h\;N/\omega$')
def propagator(H, t, c_op_list=[], args={}, options=None,
parallel=False, progress_bar=None, **kwargs):
"""
Calculate the propagator U(t) for the density matrix or wave function such
that :math:`\psi(t) = U(t)\psi(0)` or
:math:`\\rho_{\mathrm vec}(t) = U(t) \\rho_{\mathrm vec}(0)`
where :math:`\\rho_{\mathrm vec}` is the vector representation of the
density matrix.
Parameters
----------
H : qobj or list
Hamiltonian as a Qobj instance of a nested list of Qobjs and
coefficients in the list-string or list-function format for
time-dependent Hamiltonians (see description in :func:`qutip.mesolve`).
t : float or array-like
Time or list of times for which to evaluate the propagator.
c_op_list : list
List of qobj collapse operators.
args : list/array/dictionary
Parameters to callback functions for time-dependent Hamiltonians and
collapse operators.
options : :class:`qutip.Options`
with options for the ODE solver.
parallel : bool {False, True}
Run the propagator in parallel mode.
progress_bar: BaseProgressBar
Optional instance of BaseProgressBar, or a subclass thereof, for
showing the progress of the simulation. By default no progress bar
is used, and if set to True a TextProgressBar will be used.
Returns
-------
a : qobj
Instance representing the propagator :math:`U(t)`.
"""
kw = _default_kwargs()
if 'num_cpus' in kwargs:
num_cpus = kwargs['num_cpus']
else:
num_cpus = kw['num_cpus']
if progress_bar is None:
progress_bar = BaseProgressBar()
elif progress_bar is True:
progress_bar = TextProgressBar()
if options is None:
options = Options()
options.rhs_reuse = True
rhs_clear()
if isinstance(t, (int, float, np.integer, np.floating)):
tlist = [0, t]
else:
tlist = t
td_type = _td_format_check(H, c_op_list, solver='me')[2]
if td_type > 0:
rhs_generate(H, c_op_list, args=args, options=options)
if isinstance(H, (types.FunctionType, types.BuiltinFunctionType,
functools.partial)):
H0 = H(0.0, args)
elif isinstance(H, list):
H0 = H[0][0] if isinstance(H[0], list) else H[0]
else:
H0 = H
if len(c_op_list) == 0 and H0.isoper:
# calculate propagator for the wave function
N = H0.shape[0]
dims = H0.dims
u = np.zeros([N, N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_sesolve,range(N),
task_args=(N,H,tlist,args,options),
progress_bar=progress_bar, num_cpus=num_cpus)
for n in range(N):
for k, t in enumerate(tlist):
u[:, n, k] = output[n].states[k].full().T
else:
progress_bar.start(N)
for n in range(0, N):
progress_bar.update(n)
psi0 = basis(N, n)
output = sesolve(H, psi0, tlist, [], args, options)
for k, t in enumerate(tlist):
u[:, n, k] = output.states[k].full().T
progress_bar.finished()
# todo: evolving a batch of wave functions:
# psi_0_list = [basis(N, n) for n in range(N)]
# psi_t_list = mesolve(H, psi_0_list, [0, t], [], [], args, options)
# for n in range(0, N):
# u[:,n] = psi_t_list[n][1].full().T
elif len(c_op_list) == 0 and H0.issuper:
# calculate the propagator for the vector representation of the
# density matrix (a superoperator propagator)
N = H0.shape[0]
sqrt_N = int(np.sqrt(N))
dims = H0.dims
u = np.zeros([N, N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_mesolve,range(N * N),
task_args=(sqrt_N,H,tlist,c_op_list,args,options),
progress_bar=progress_bar, num_cpus=num_cpus)
for n in range(N * N):
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output[n].states[k].full()).T
else:
progress_bar.start(N)
for n in range(0, N):
progress_bar.update(n)
col_idx, row_idx = np.unravel_index(n,(sqrt_N,sqrt_N))
rho0 = Qobj(sp.csr_matrix(([1],([row_idx],[col_idx])), shape=(sqrt_N,sqrt_N), dtype=complex))
output = mesolve(H, rho0, tlist, [], [], args, options)
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output.states[k].full()).T
progress_bar.finished()
else:
# calculate the propagator for the vector representation of the
# density matrix (a superoperator propagator)
N = H0.shape[0]
dims = [H0.dims, H0.dims]
u = np.zeros([N * N, N * N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_mesolve,range(N * N),
task_args=(N,H,tlist,c_op_list,args,options),
progress_bar=progress_bar, num_cpus=num_cpus)
for n in range(N * N):
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output[n].states[k].full()).T
else:
progress_bar.start(N * N)
for n in range(N * N):
progress_bar.update(n)
col_idx, row_idx = np.unravel_index(n,(N,N))
rho0 = Qobj(sp.csr_matrix(([1],([row_idx],[col_idx])), shape=(N,N), dtype=complex))
output = mesolve(H, rho0, tlist, c_op_list, [], args, options)
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output.states[k].full()).T
progress_bar.finished()
if len(tlist) == 2:
return Qobj(u[:, :, 1], dims=dims)
else:
return np.array([Qobj(u[:, :, k], dims=dims) for k in range(len(tlist))], dtype=object)
def propagator(H,
t,
c_op_list=[],
args={},
options=None,
parallel=False,
progress_bar=None,
**kwargs):
"""
Calculate the propagator U(t) for the density matrix or wave function such
that :math:`\psi(t) = U(t)\psi(0)` or
:math:`\\rho_{\mathrm vec}(t) = U(t) \\rho_{\mathrm vec}(0)`
where :math:`\\rho_{\mathrm vec}` is the vector representation of the
density matrix.
Parameters
----------
H : qobj or list
Hamiltonian as a Qobj instance of a nested list of Qobjs and
coefficients in the list-string or list-function format for
time-dependent Hamiltonians (see description in :func:`qutip.mesolve`).
t : float or array-like
Time or list of times for which to evaluate the propagator.
c_op_list : list
List of qobj collapse operators.
args : list/array/dictionary
Parameters to callback functions for time-dependent Hamiltonians and
collapse operators.
options : :class:`qutip.Options`
with options for the ODE solver.
parallel : bool {False, True}
Run the propagator in parallel mode.
progress_bar: BaseProgressBar
Optional instance of BaseProgressBar, or a subclass thereof, for
showing the progress of the simulation. By default no progress bar
is used, and if set to True a TextProgressBar will be used.
Returns
-------
a : qobj
Instance representing the propagator :math:`U(t)`.
"""
kw = _default_kwargs()
if 'num_cpus' in kwargs:
num_cpus = kwargs['num_cpus']
else:
num_cpus = kw['num_cpus']
if progress_bar is None:
progress_bar = BaseProgressBar()
elif progress_bar is True:
progress_bar = TextProgressBar()
if options is None:
options = Options()
options.rhs_reuse = True
rhs_clear()
if isinstance(t, (int, float, np.integer, np.floating)):
tlist = [0, t]
else:
tlist = t
td_type = _td_format_check(H, c_op_list, solver='me')[2]
if td_type > 0:
rhs_generate(H, c_op_list, args=args, options=options)
if isinstance(
H,
(types.FunctionType, types.BuiltinFunctionType, functools.partial)):
H0 = H(0.0, args)
elif isinstance(H, list):
H0 = H[0][0] if isinstance(H[0], list) else H[0]
else:
H0 = H
if len(c_op_list) == 0 and H0.isoper:
# calculate propagator for the wave function
N = H0.shape[0]
dims = H0.dims
u = np.zeros([N, N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_sesolve,
range(N),
task_args=(N, H, tlist, args, options),
progress_bar=progress_bar,
num_cpus=num_cpus)
for n in range(N):
for k, t in enumerate(tlist):
u[:, n, k] = output[n].states[k].full().T
else:
progress_bar.start(N)
for n in range(0, N):
progress_bar.update(n)
psi0 = basis(N, n)
output = sesolve(H,
psi0,
tlist, [],
args,
options,
_safe_mode=False)
for k, t in enumerate(tlist):
u[:, n, k] = output.states[k].full().T
progress_bar.finished()
# todo: evolving a batch of wave functions:
# psi_0_list = [basis(N, n) for n in range(N)]
# psi_t_list = mesolve(H, psi_0_list, [0, t], [], [], args, options)
# for n in range(0, N):
# u[:,n] = psi_t_list[n][1].full().T
elif len(c_op_list) == 0 and H0.issuper:
# calculate the propagator for the vector representation of the
# density matrix (a superoperator propagator)
N = H0.shape[0]
sqrt_N = int(np.sqrt(N))
dims = H0.dims
u = np.zeros([N, N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_mesolve,
range(N * N),
task_args=(sqrt_N, H, tlist, c_op_list, args,
options),
progress_bar=progress_bar,
num_cpus=num_cpus)
for n in range(N * N):
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output[n].states[k].full()).T
else:
progress_bar.start(N)
for n in range(0, N):
progress_bar.update(n)
col_idx, row_idx = np.unravel_index(n, (sqrt_N, sqrt_N))
rho0 = Qobj(
sp.csr_matrix(([1], ([row_idx], [col_idx])),
shape=(sqrt_N, sqrt_N),
dtype=complex))
output = mesolve(H,
rho0,
tlist, [], [],
args,
options,
_safe_mode=False)
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output.states[k].full()).T
progress_bar.finished()
else:
# calculate the propagator for the vector representation of the
# density matrix (a superoperator propagator)
N = H0.shape[0]
dims = [H0.dims, H0.dims]
u = np.zeros([N * N, N * N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_mesolve,
range(N * N),
task_args=(N, H, tlist, c_op_list, args,
options),
progress_bar=progress_bar,
num_cpus=num_cpus)
for n in range(N * N):
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output[n].states[k].full()).T
else:
progress_bar.start(N * N)
for n in range(N * N):
progress_bar.update(n)
col_idx, row_idx = np.unravel_index(n, (N, N))
rho0 = Qobj(
sp.csr_matrix(([1], ([row_idx], [col_idx])),
shape=(N, N),
dtype=complex))
output = mesolve(H,
rho0,
tlist,
c_op_list, [],
args,
options,
_safe_mode=False)
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output.states[k].full()).T
progress_bar.finished()
if len(tlist) == 2:
return Qobj(u[:, :, 1], dims=dims)
else:
return np.array(
[Qobj(u[:, :, k], dims=dims) for k in range(len(tlist))],
dtype=object)
def tran_spectrum(
Mol: molSystem,
delta_max: float,
delta_step: float,
sampletlist: List[float],
method="theory",
calculation="parallel",
save_data=False,
):
"""
input:
an instance of molSystem (neglect the detuning setting)
delta_max: Unit: MHz. positive maximum detuning frequency.
delta_step: Unit: MHz. Increasing step of detuning frequency from zero
sampletlist: The list of sampling time to draw spectrum
output:
transmission spectrum data
"""
if Mol.tlist[-1] < sampletlist[-1]:
raise ValueError("tlist range does not cover sampletlist!")
if method == "simulation" or method == "theory":
print("The way to derive transmission spectrum is " + method)
else:
raise ValueError("method value should be 'simulation' or 'theory'")
if calculation == "serial" or calculation == "parallel":
print("The way to derive transmission spectrum is " + calculation)
else:
raise ValueError("calculation value should be 'serial' or 'parallel'")
deltalist = arange(0, delta_max, delta_step)
if calculation == "serial":
# serial version
spectrum_data = []
for delta in deltalist:
spectrum_data.append(
_helper(2 * pi * delta, Mol, sampletlist, method))
elif calculation == "parallel":
# parallel version
spectrum_data = parallel_map(
_helper,
2 * pi * array(deltalist), # convert the unit to Mega rad/s
task_kwargs=dict(Mol_original=Mol,
sampletlist=sampletlist,
method=method),
progress_bar=True,
)
if save_data:
save(Mol._path + "tran_spectrum.npy", spectrum_data)
save(Mol._path + "deltalist.npy", array(deltalist))
save(Mol._path + "sampletlist.npy", array(sampletlist))
# draw graph
delta_minus = -copy(deltalist)
delta_data = concatenate([flip(delta_minus), deltalist])
spec_minus = flip(copy(spectrum_data), axis=0)
spec_data = transpose(concatenate([spec_minus, spectrum_data]))
fig, axes = plt.subplots(1, 1, figsize=(10, 8))
for t_idx, spec in enumerate(spec_data):
axes.plot(delta_data, spec, label=str(sampletlist[t_idx]) + r"$\mu s$")
axes.legend(loc=0)
if not exists(Mol._path + "figures"): # detect figures folder
makedirs(Mol._path + "figures")
plt.savefig(Mol._path + "figures/Transmission_Spectrum.png")