1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402
|
# This file is part of QuTiP: Quantum Toolbox in Python.
#
# Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
#
# 3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
# of its contributors may be used to endorse or promote products derived
# from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################
"""
This module provides solvers for the unitary Schrodinger equation.
"""
__all__ = ['sesolve']
import os
import types
import numpy as np
import scipy.integrate
from scipy.linalg import norm as la_norm
from qutip.cy.stochastic import normalize_inplace
import qutip.settings as qset
from qutip.qobj import Qobj
from qutip.qobjevo import QobjEvo
from qutip.cy.spconvert import dense1D_to_fastcsr_ket, dense2D_to_fastcsr_fmode
from qutip.cy.spmatfuncs import (cy_expect_psi, cy_ode_psi_func_td,
cy_ode_psi_func_td_with_state)
from qutip.solver import Result, Options, config, solver_safe, SolverSystem
from qutip.superoperator import vec2mat
from qutip.ui.progressbar import (BaseProgressBar, TextProgressBar)
from qutip.cy.openmp.utilities import check_use_openmp, openmp_components
def sesolve(H, psi0, tlist, e_ops=None, args=None, options=None,
progress_bar=None, _safe_mode=True):
"""
Schrodinger equation evolution of a state vector or unitary matrix
for a given Hamiltonian.
Evolve the state vector (`psi0`) using a given
Hamiltonian (`H`), by integrating the set of ordinary differential
equations that define the system. Alternatively evolve a unitary matrix in
solving the Schrodinger operator equation.
The output is either the state vector or unitary matrix at arbitrary points
in time (`tlist`), or the expectation values of the supplied operators
(`e_ops`). If e_ops is a callback function, it is invoked for each
time in `tlist` with time and the state as arguments, and the function
does not use any return values. e_ops cannot be used in conjunction
with solving the Schrodinger operator equation
Parameters
----------
H : :class:`qutip.qobj`, :class:`qutip.qobjevo`, *list*, *callable*
system Hamiltonian as a Qobj, list of Qobj and coefficient, QobjEvo,
or a callback function for time-dependent Hamiltonians.
list format and options can be found in QobjEvo's description.
psi0 : :class:`qutip.qobj`
initial state vector (ket)
or initial unitary operator `psi0 = U`
tlist : *list* / *array*
list of times for :math:`t`.
e_ops : None / list of :class:`qutip.qobj` / callback function
single operator or list of operators for which to evaluate
expectation values.
For list operator evolution, the overlapse is computed:
tr(e_ops[i].dag()*op(t))
args : None / *dictionary*
dictionary of parameters for time-dependent Hamiltonians
options : None / :class:`qutip.Qdeoptions`
with options for the ODE solver.
progress_bar : None / BaseProgressBar
Optional instance of BaseProgressBar, or a subclass thereof, for
showing the progress of the simulation.
Returns
-------
output: :class:`qutip.solver`
An instance of the class :class:`qutip.solver`, which contains either
an *array* of expectation values for the times specified by `tlist`, or
an *array* or state vectors corresponding to the
times in `tlist` [if `e_ops` is an empty list], or
nothing if a callback function was given inplace of operators for
which to calculate the expectation values.
"""
if e_ops is None:
e_ops = []
if isinstance(e_ops, Qobj):
e_ops = [e_ops]
elif isinstance(e_ops, dict):
e_ops_dict = e_ops
e_ops = [e for e in e_ops.values()]
else:
e_ops_dict = None
if progress_bar is None:
progress_bar = BaseProgressBar()
if progress_bar is True:
progress_bar = TextProgressBar()
if not (psi0.isket or psi0.isunitary):
raise TypeError("The unitary solver requires psi0 to be"
" a ket as initial state"
" or a unitary as initial operator.")
if options is None:
options = Options()
if options.rhs_reuse and not isinstance(H, SolverSystem):
# TODO: deprecate when going to class based solver.
if "sesolve" in solver_safe:
# print(" ")
H = solver_safe["sesolve"]
else:
pass
# raise Exception("Could not find the Hamiltonian to reuse.")
if args is None:
args = {}
check_use_openmp(options)
if isinstance(H, SolverSystem):
ss = H
elif isinstance(H, (list, Qobj, QobjEvo)):
ss = _sesolve_QobjEvo(H, tlist, args, options)
elif callable(H):
ss = _sesolve_func_td(H, args, options)
else:
raise Exception("Invalid H type")
func, ode_args = ss.makefunc(ss, psi0, args, e_ops, options)
if _safe_mode:
v = psi0.full().ravel('F')
func(0., v, *ode_args) + v
res = _generic_ode_solve(func, ode_args, psi0, tlist, e_ops, options,
progress_bar, dims=psi0.dims)
if e_ops_dict:
res.expect = {e: res.expect[n]
for n, e in enumerate(e_ops_dict.keys())}
res.SolverSystem = ss
return res
# -----------------------------------------------------------------------------
# A time-dependent unitary wavefunction equation on the list-function format
#
def _sesolve_QobjEvo(H, tlist, args, opt):
"""
Prepare the system for the solver, H can be an QobjEvo.
"""
H_td = -1.0j * QobjEvo(H, args, tlist=tlist)
if opt.rhs_with_state:
H_td._check_old_with_state()
nthread = opt.openmp_threads if opt.use_openmp else 0
H_td.compile(omp=nthread)
ss = SolverSystem()
ss.H = H_td
ss.makefunc = _qobjevo_set
solver_safe["sesolve"] = ss
return ss
def _qobjevo_set(HS, psi, args, e_ops, opt):
"""
From the system, get the ode function and args
"""
H_td = HS.H
H_td.solver_set_args(args, psi, e_ops)
if psi.isunitary:
func = H_td.compiled_qobjevo.ode_mul_mat_f_vec
elif psi.isket:
func = H_td.compiled_qobjevo.mul_vec
else:
raise TypeError("The unitary solver requires psi0 to be"
" a ket as initial state"
" or a unitary as initial operator.")
return func, ()
# -----------------------------------------------------------------------------
# Wave function evolution using a ODE solver (unitary quantum evolution), for
# time dependent hamiltonians.
#
def _sesolve_func_td(H_func, args, opt):
"""
Prepare the system for the solver, H is a function.
"""
ss = SolverSystem()
ss.H = H_func
ss.makefunc = _Hfunc_set
solver_safe["sesolve"] = ss
return ss
def _Hfunc_set(HS, psi, args, e_ops, opt):
"""
From the system, get the ode function and args
"""
H_func = HS.H
if psi.isunitary:
if not opt.rhs_with_state:
func = _ode_oper_func_td
else:
func = _ode_oper_func_td_with_state
else:
if not opt.rhs_with_state:
func = cy_ode_psi_func_td
else:
func = cy_ode_psi_func_td_with_state
return func, (H_func, args)
# -----------------------------------------------------------------------------
# evaluate dU(t)/dt according to the schrodinger equation
#
def _ode_oper_func_td(t, y, H_func, args):
H = H_func(t, args).data * -1j
ym = vec2mat(y)
return (H * ym).ravel("F")
def _ode_oper_func_td_with_state(t, y, H_func, args):
H = H_func(t, y, args).data * -1j
ym = vec2mat(y)
return (H * ym).ravel("F")
# -----------------------------------------------------------------------------
# Solve an ODE for func.
# Calculate the required expectation values or invoke callback
# function at each time step.
def _generic_ode_solve(func, ode_args, psi0, tlist, e_ops, opt,
progress_bar, dims=None):
"""
Internal function for solving ODEs.
"""
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# This function is made similar to mesolve's one for futur merging in a
# solver class
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# prepare output array
n_tsteps = len(tlist)
output = Result()
output.solver = "sesolve"
output.times = tlist
if psi0.isunitary:
initial_vector = psi0.full().ravel('F')
oper_evo = True
size = psi0.shape[0]
# oper_n = dims[0][0]
# norm_dim_factor = np.sqrt(oper_n)
elif psi0.isket:
initial_vector = psi0.full().ravel()
oper_evo = False
# norm_dim_factor = 1.0
r = scipy.integrate.ode(func)
r.set_integrator('zvode', method=opt.method, order=opt.order,
atol=opt.atol, rtol=opt.rtol, nsteps=opt.nsteps,
first_step=opt.first_step, min_step=opt.min_step,
max_step=opt.max_step)
if ode_args:
r.set_f_params(*ode_args)
r.set_initial_value(initial_vector, tlist[0])
e_ops_data = []
output.expect = []
if callable(e_ops):
n_expt_op = 0
expt_callback = True
output.num_expect = 1
elif isinstance(e_ops, list):
n_expt_op = len(e_ops)
expt_callback = False
output.num_expect = n_expt_op
if n_expt_op == 0:
# fallback on storing states
opt.store_states = True
else:
for op in e_ops:
if op.isherm:
output.expect.append(np.zeros(n_tsteps))
else:
output.expect.append(np.zeros(n_tsteps, dtype=complex))
if oper_evo:
for e in e_ops:
e_ops_data.append(e.dag().data)
else:
for e in e_ops:
e_ops_data.append(e.data)
else:
raise TypeError("Expectation parameter must be a list or a function")
if opt.store_states:
output.states = []
if oper_evo:
def get_curr_state_data(r):
return vec2mat(r.y)
else:
def get_curr_state_data(r):
return r.y
#
# start evolution
#
dt = np.diff(tlist)
cdata = None
progress_bar.start(n_tsteps)
for t_idx, t in enumerate(tlist):
progress_bar.update(t_idx)
if not r.successful():
raise Exception("ODE integration error: Try to increase "
"the allowed number of substeps by increasing "
"the nsteps parameter in the Options class.")
# get the current state / oper data if needed
if opt.store_states or opt.normalize_output or n_expt_op > 0 or expt_callback:
cdata = get_curr_state_data(r)
if opt.normalize_output:
# normalize per column
if oper_evo:
cdata /= la_norm(cdata, axis=0)
#cdata *= norm_dim_factor / la_norm(cdata)
r.set_initial_value(cdata.ravel('F'), r.t)
else:
#cdata /= la_norm(cdata)
norm = normalize_inplace(cdata)
if norm > 1e-12:
# only reset the solver if state changed
r.set_initial_value(cdata, r.t)
else:
r._y = cdata
if opt.store_states:
if oper_evo:
fdata = dense2D_to_fastcsr_fmode(cdata, size, size)
output.states.append(Qobj(fdata, dims=dims))
else:
fdata = dense1D_to_fastcsr_ket(cdata)
output.states.append(Qobj(fdata, dims=dims, fast='mc'))
if expt_callback:
# use callback method
output.expect.append(e_ops(t, Qobj(cdata, dims=dims)))
if oper_evo:
for m in range(n_expt_op):
output.expect[m][t_idx] = (e_ops_data[m] * cdata).trace()
else:
for m in range(n_expt_op):
output.expect[m][t_idx] = cy_expect_psi(e_ops_data[m], cdata,
e_ops[m].isherm)
if t_idx < n_tsteps - 1:
r.integrate(r.t + dt[t_idx])
progress_bar.finished()
if opt.store_final_state:
cdata = get_curr_state_data(r)
if opt.normalize_output:
cdata /= la_norm(cdata, axis=0)
# cdata *= norm_dim_factor / la_norm(cdata)
output.final_state = Qobj(cdata, dims=dims)
return output
|