onebody.py 20.5 KB
 Kloss committed May 29, 2019 1 2 3 4 5 6 7 8 9 10 11 # Copyright 2016-2019 tkwant authors. # # This file is part of tkwant. It is subject to the license terms in the file # LICENSE.rst found in the top-level directory of this distribution and at # http://kwant-project.org/license. A list of tkwant authors can be found in # the file AUTHORS.rst at the top-level directory of this distribution and at # http://kwant-project.org/authors. """Tools for solving the one-body time-dependent Schrödinger equation.""" import collections.abc import functools  Kloss committed Aug 23, 2019 12 13 14 from cmath import exp import numpy as np import scipy.sparse as sp  Kloss committed May 29, 2019 15 16 17 18  import kwant import kwant_spectrum  Kloss committed Mar 22, 2020 19 from .. import leads, _common, _logging  Kloss committed Jul 12, 2019 20 from ..system import (extract_perturbation, hamiltonian_with_boundaries,  Kloss committed Jul 31, 2019 21  add_time_to_params)  Kloss committed May 29, 2019 22 23 from . import kernels, solvers  Kloss committed Nov 11, 2019 24 __all__ = ['WaveFunction', 'ScatteringStates', 'Task']  Kloss committed May 29, 2019 25 26   Kloss committed Mar 22, 2020 27 # set module logger  Kloss committed Mar 22, 2020 28 logger = _logging.make_logger(name=__name__)  Kloss committed Mar 22, 2020 29 30 log_func = _logging.log_func(logger)  Kloss committed May 29, 2019 31   Kloss committed Mar 22, 2020 32 # data formats  Kloss committed Jul 12, 2019 33   Kloss committed May 29, 2019 34 35 36 37 38 39 40 41 42 43 44 45 class Task: """Data format to store the set of quantum numbers that uniquely indentifies a onebody state and the weight of that state in the manybody sum. Attributes ---------- lead : int lead index mode : int scattering mode index energy : float energy of the onebody state  Kloss committed Jan 16, 2020 46 47 48  momentum : float momentum of the onebody state weight : float or numpy float array  Kloss committed May 29, 2019 49  weighting factor of the one-body state in the manybody sum  Kloss committed Oct 16, 2019 50 51  weigth = math_weight * phys_weight, where math_weight is the weighting factor from the integration quadrature  Kloss committed Jan 16, 2020 52 53  math_weight : float or numpy float array mathematical weighting factor from the numerical quadrature rule  Kloss committed Oct 16, 2019 54 55  phys_weight : float physical weighting factor (fermi function, pi factors..)  Kloss committed May 29, 2019 56  """  Kloss committed Oct 16, 2019 57  def __init__(self, lead: int, mode: int, energy: float, weight: float,  Kloss committed Jan 16, 2020 58 59  math_weight: float = None, phys_weight: float = None, momentum: float = None):  60 61 62  self.lead = lead self.mode = mode self.energy = energy  Kloss committed Jan 16, 2020 63  self.momentum = momentum  64  self.weight = weight  Kloss committed Oct 16, 2019 65  self.math_weight = math_weight  Kloss committed Jan 16, 2020 66  self.phys_weight = phys_weight  Kloss committed May 29, 2019 67   Kloss committed Aug 23, 2019 68 69 70 71  def __eq__(self, other) -> bool: if not isinstance(other, Task): return NotImplemented return (  Kloss committed Jan 16, 2020 72 73 74 75  (self.lead, self.mode, self.energy, self.momentum, self.weight, self.math_weight, self.phys_weight) == (other.lead, other.mode, other.energy, other.momentum, other.weight, other.math_weight, other.phys_weight))  Kloss committed Aug 23, 2019 76 77  def __str__(self):  Kloss committed Jan 16, 2020 78 79  string = "onebody task: lead={lead}, mode={mode}, " \ "energy={energy}, momentum={momentum}, " \  Kloss committed Oct 16, 2019 80 81  "weight={weight}, math_weight={math_weight}, " \ "phys_weight={phys_weight}".format(**self.__dict__)  Kloss committed Aug 23, 2019 82 83  return string  Kloss committed May 29, 2019 84   Kloss committed Jul 12, 2019 85 86 87 88 89 def _operator_bound(operator): """Return True iff operator is bound""" return bool(operator._bound_onsite or operator._bound_hamiltonian)  Kloss committed May 29, 2019 90 class WaveFunction:  Kloss committed Nov 11, 2019 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109  r"""A class to solve the time-dependent single-particle wavefunction. The time-dependent single-particle Schrödinger equation is :math:i \partial_t \psi = (H_0 + W(t)) \psi, where the total Hamiltonian :math:H(t) = H_0 + W(t) has been splitted into a static part :math:H_0 and a time-dependent perturbation :math:W(t). Moreover, the initial condition is :math:\psi(t_0) and :math:W(t) is expected to be absent before the initial time :math:t_0: :math:W(t) = 0 \,\, \text{for} \,\, t \leq t_0. If an energy :math:E is provided, this routine expects that the initial condition represents the scattering state :math:\psi_{st} that solves the time-independent Schrödinger equation :math:H_0 \psi_{st} = E \psi_{st}, \, \psi_{st} = \psi(t_0).  Kloss committed May 29, 2019 110   Kloss committed Nov 11, 2019 111 112 113 114 115 116 117 118 119 120 121 122 123  For numerical reasons, the evolution is then performed in the variable :math:i \partial_t \bar{\psi} = (H_0 + W(t) - E) \bar{\psi} + W(t) \psi_{st}, where :math:\psi = (\bar{\psi} - \psi_{st}) e^{-i E (t - t_0)}. See J. Weston and X. Waintal, Phys. Rev. B 93, 134506 (2016) _. """ def __init__(self, H0, W, psi_init, energy=None, params=None, solution_is_valid=None, time_is_valid=None,  Kloss committed Mar 22, 2020 124  time_start=_common.time_start, time_name=_common.time_name,  Kloss committed Nov 11, 2019 125 126  kernel_type=kernels.default, solver_type=solvers.default): r"""  Kloss committed May 29, 2019 127 128  Parameters ----------  Kloss committed Nov 11, 2019 129 130 131 132 133  H0 : array-like The static part of the Hamiltonian matrix, :math:H_0. W : callable or None Time-dependent part of the Hamiltonian matrix, :math:W(t). Typically the object returned by tkwant.system.extract_perturbation.  Kloss committed May 29, 2019 134  psi_init : array of complex  Kloss committed Nov 11, 2019 135 136  The state :math:\psi(t_0) from which to start, defined over the central region.  Kloss committed May 29, 2019 137 138  energy : float, optional If provided, then psi_init is assumed to be an eigenstate  Kloss committed Nov 11, 2019 139 140 141  of energy :math:E. If the Hamiltonian represents an open quantum system with leads, then psi_init is assumed to be the projection of a scattering state at energy :math:E  Kloss committed May 29, 2019 142  on to the central part of the system.  Kloss committed Jul 12, 2019 143  params : dict, optional  Kloss committed Nov 11, 2019 144 145 146 147 148 149 150 151 152 153 154 155  Extra arguments to pass to the time-dependent Hamiltonian function :math:W(t), excluding time. solution_is_valid : callable, optional Function to detect spurious reflections for a system with leads. See tkwant.leads.EvaluatedBoundary. time_is_valid : callable, optional Function to check if boundary conditions are valid at a given time. See tkwant.leads.EvaluatedBoundary. time_start : float, optional The initial time :math:t_0. Default value is zero. time_name : str, optional The name of the time argument :math:t. Default name: *time*.  Kloss committed May 29, 2019 156 157 158 159 160  kernel_type : tkwant.onebody.solvers.default, optional The kernel to calculate the right-hand-site of the Schrödinger equation. solver_type : tkwant.onebody.solvers.default, optional The solver used to evolve the wavefunction forward in time.  Kloss committed Nov 11, 2019 161  """  Kloss committed May 29, 2019 162   Kloss committed Jul 12, 2019 163  # add initial time to the params dict  Kloss committed Jul 31, 2019 164 165  tparams = add_time_to_params(params, time_name=time_name, time=time_start, check_numeric_type=True)  Kloss committed Jul 12, 2019 166   Kloss committed May 29, 2019 167 168 169  if energy is None: # starting from an arbitrary state, so we need # to be solving: H0 @ psi + W(t) @ psi  Kloss committed Jul 12, 2019 170 171  kernel = kernel_type(H0, W, params=tparams)  Kloss committed May 29, 2019 172 173 174 175 176  else: # we are starting from an eigenstate, so we need to # be solving: (H0 - E) @ psibar + W(t) @ (psibar + psi_st) # and psi = (psibar + psi_st) * exp(-1j * energy * time) kernel = kernel_type(H0 - energy * sp.eye(H0.shape[0]),  Kloss committed Jul 12, 2019 177  W, tparams, np.asarray(psi_init))  Kloss committed May 29, 2019 178 179 180 181  # get the object that will actually do the time stepping self.solver = solver_type(kernel)  Kloss committed Nov 11, 2019 182 183 184 185 186 187 188 189  # The size of the central scattering region. The central scattering # region can be smaller as the total size (kernel.size) of the system # in case of boundary conditions. syst_size = psi_init.size # transform initial psi to psibar and psi_st # note that the dgl is always solved in the variable psibar psibar = np.zeros((kernel.size,), complex)  Kloss committed May 29, 2019 190 191 192 193 194 195 196 197 198  if energy is not None: psi_st = np.array(psi_init, complex) else: psi_st = None psibar[:syst_size] = psi_init self.psibar = psibar self.psi_st = psi_st  Kloss committed Nov 11, 2019 199 200 201 202  self._syst_size = syst_size self._solution_is_valid = solution_is_valid self._time_is_valid = time_is_valid  Kloss committed Jul 31, 2019 203 204 205 206  self.time_name = time_name self.time_start = time_start self.time = time_start  Kloss committed Jul 12, 2019 207  self.params = params  Kloss committed May 29, 2019 208 209  self.energy = energy  Kloss committed Nov 11, 2019 210  @classmethod  Kloss committed Mar 22, 2020 211 212  def from_kwant(cls, syst, psi_init, boundaries=None, energy=None, params=None, time_start=_common.time_start, time_name=_common.time_name,  Kloss committed Nov 11, 2019 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  kernel_type=kernels.default, solver_type=solvers.default): """Set up a time-dependent onebody wavefunction from a kwant system. Parameters ---------- syst : kwant.builder.FiniteSystem The low level system for which the wave functions are to be calculated. psi_init : array of complex The state from which to start, defined over the central region. boundaries : sequence of ~tkwant.leads.BoundaryBase, optional The boundary conditions for each lead attached to syst. Must be provided for a system with leads. energy : float, optional If provided, then psi_init is assumed to be an eigenstate of energy *E*. If syst has leads, then psi_init is assumed to be the projection of a scattering state at energy *E* on to the central part of the system. params : dict, optional Extra arguments to pass to the Hamiltonian of syst, excluding time. time_start : float, optional The initial time :math:t_0. Default value is zero. time_name : str, optional The name of the time argument :math:t. Default name: *time*. kernel_type : tkwant.onebody.solvers.default, optional The kernel to calculate the right-hand-site of the Schrödinger equation. solver_type : tkwant.onebody.solvers.default, optional The solver used to evolve the wavefunction forward in time. Returns ------- wave_function : tkwant.onebody.WaveFunction A time-dependent onebody wavefunction at the initial time. """ syst_size = syst.site_ranges[-1][2] if not psi_init.size == syst_size: raise ValueError('Size of the initial condition={} does not match ' 'the total number of orbitals in the central ' 'system ={}'.format(psi_init.size, syst_size)) # add initial time to the params dict tparams = add_time_to_params(params, time_name=time_name, time=time_start, check_numeric_type=True) if syst.leads: # need to add boundary conditions if not boundaries: raise ValueError('Must provide boundary conditions for systems ' 'with leads.') # get static part of Hamiltonian for central system + boundary conditions # TODO: Make this work with arbitrary kwant.system.FiniteSystem. extended_system = hamiltonian_with_boundaries(syst, boundaries, params=tparams) H0 = extended_system.hamiltonian solution_is_valid = extended_system.solution_is_valid time_is_valid = extended_system.time_is_valid else: # true finite systems (no leads) so no need for boundary conditions  Kloss committed Apr 01, 2020 273 274 275  if boundaries is not None: raise ValueError('No boundary conditions must be provided ' 'for a system without leads.')  Kloss committed Nov 11, 2019 276 277 278 279 280 281 282 283 284  H0 = syst.hamiltonian_submatrix(params=tparams, sparse=True) solution_is_valid = None time_is_valid = None # TODO: Make this work with arbitrary kwant.system.FiniteSystem. W = extract_perturbation(syst, time_name, time_start, params=tparams) return cls(H0, W, psi_init, energy, params, solution_is_valid, time_is_valid,  Kloss committed Mar 13, 2020 285  time_start, time_name, kernel_type, solver_type)  Kloss committed Nov 11, 2019 286   Kloss committed May 29, 2019 287  def psi(self):  Kloss committed Nov 11, 2019 288  r"""Return the wavefunction :math:\psi(t) at the current time *t*.  Kloss committed May 29, 2019 289 290 291 292 293 294  psi : ~numpy.ndarray The wavefunction at time. If this wavefunction is for a system with leads, then the wavefunction projected onto the central region is returned. """  Kloss committed Nov 11, 2019 295  if self.energy is None: # psi = psibar  Kloss committed Jul 12, 2019 296  return self.psibar[:self._syst_size].copy()  Kloss committed Nov 11, 2019 297  # else : psi = (psibar + psi_st) * exp(-i E (t - t0))  Kloss committed Aug 23, 2019 298 299  return ((self.psibar[:self._syst_size] + self.psi_st) * exp(-1j * self.energy * (self.time - self.time_start)))  Kloss committed May 29, 2019 300   Kloss committed Jul 12, 2019 301  def evolve(self, time, params=None):  Kloss committed Nov 11, 2019 302 303 304  r""" Evolve the wavefunction :math:\psi(t) foreward in time up to :math:t = time.  Kloss committed May 29, 2019 305 306 307 308 309  Parameters ---------- time : int or float time argument up to which the solver should be evolved  Kloss committed Jul 12, 2019 310  params : dict, optional  Kloss committed Nov 11, 2019 311  Extra arguments to pass to the Hamiltonian, excluding time.  Kloss committed Jul 12, 2019 312 313 314  If present, the public params attribute is updated to the new value of params. By default, the params from initialization are taken.  Kloss committed May 29, 2019 315  """  Kloss committed Jul 12, 2019 316 317 318 319  if params is not None: self.params = params self.solver.kernel.set_params(params)  Kloss committed May 29, 2019 320 321 322 323 324  # time corresponds to the future time, to which we will evolve if time < self.time: raise ValueError('Cannot evolve backwards in time') if time == self.time: return  Kloss committed Nov 11, 2019 325 326 327 328  if self._time_is_valid is not None: if not self._time_is_valid(time): raise RuntimeError('Cannot evolve up to {} with the given ' 'boundary conditions'.format(time))  Kloss committed Jul 12, 2019 329   Kloss committed May 29, 2019 330 331  # evolve forward in time next_psibar = self.solver(self.psibar, self.time, time)  Kloss committed Jul 12, 2019 332   Kloss committed Nov 11, 2019 333 334 335 336 337  if self._solution_is_valid is not None: if not self._solution_is_valid(next_psibar): raise RuntimeError('Evolving between {} and {} resulted in an ' 'unphysical result due to the boundary ' 'conditions'.format(self.time, next_time=time))  Kloss committed May 29, 2019 338 339 340  # update internal state and return self.psibar, self.time = next_psibar, time  Kloss committed Jul 12, 2019 341  def evaluate(self, observable):  Kloss committed Nov 11, 2019 342 343 344 345 346  r""" Evaluate the expectation value of an operator at the current time *t*. For an operator :math:\hat{O} the expectation value is :math:O(t) = \langle \psi(t) | \hat{O} |\psi(t) \rangle.  Kloss committed May 29, 2019 347 348 349 350  Parameters ---------- observable : callable or kwant.operator  Kloss committed Nov 11, 2019 351  An operator :math:\hat{O} to evaluate the expectation value.  Kloss committed Jul 12, 2019 352  Must have the calling signature of kwant.operator.  Kloss committed May 29, 2019 353 354 355 356  Returns ------- result : numpy array  Kloss committed Nov 11, 2019 357  The expectation value :math:O(t) of observable.  Kloss committed May 29, 2019 358  """  Kloss committed Jul 12, 2019 359 360 361  if _operator_bound(observable): raise ValueError("Operator must not use pre-bind values")  Kloss committed Jul 31, 2019 362  tparams = add_time_to_params(self.params, time_name=self.time_name,  Kloss committed Jul 12, 2019 363  time=self.time)  Kloss committed Nov 11, 2019 364  return observable(self.psi(), params=tparams)  Kloss committed May 29, 2019 365 366   Kloss committed May 29, 2019 367 class ScatteringStates(collections.abc.Iterable):  Kloss committed May 29, 2019 368 369  """Calculate time-dependent wavefunctions starting in an equilibrium scattering state"""  Kloss committed Jul 12, 2019 370  def __init__(self, syst, energy, lead, tmax=None, params=None, spectra=None,  Kloss committed May 29, 2019 371  boundaries=None, equilibrium_solver=kwant.wave_function,  Kloss committed Nov 11, 2019 372 373  wavefunction_type=WaveFunction.from_kwant): """  Kloss committed May 29, 2019 374 375 376 377 378 379 380 381 382 383 384 385 386  Parameters ---------- syst : kwant.builder.FiniteSystem The low level system for which the wave functions are to be calculated. energy : float Energy of the scattering eigenstate. lead : int Lead index to construct the scattering state. tmax : float, optional The maximum time up to which to simulate. Sets the boundary conditions such that they are guaranteed to be correct up to 'tmax'. Mutually exclusive with 'boundaries'.  Kloss committed Jul 12, 2019 387  params : dict, optional  Kloss committed May 29, 2019 388  Extra arguments to pass to the Hamiltonian of syst, excluding time.  Kloss committed Mar 22, 2020 389 390 391 392 393  spectra : sequence of kwant_spectrum.spectrum, optional Energy dispersion :math:E_n(k) for the leads. Must have the same length as syst.leads. Required only if no boundaries are provided. If needed but not present, it will be calculated on the fly from syst.leads.  Kloss committed May 29, 2019 394 395 396 397 398  boundaries : sequence of ~tkwant.leads.BoundaryBase, optional The boundary conditions for each lead attached to syst. Mutually exclusive with 'tmax'. equilibrium_solver : kwant.wave_function, optional Solver for initial equilibrium scattering problem.  Kloss committed Mar 22, 2020 399 400 401 402  wavefunction_type : WaveFunction, optional One-body time-dependent wave function. Name of the time argument and initial time are taken from this class. If this is not possible, default values are used as a fallback.  Kloss committed May 29, 2019 403 404 405 406 407  Notes ----- An instance of this class behaves like a sequence of WaveFunction instances. The index in the sequence corresponds to the scattering mode.  Kloss committed Jul 31, 2019 408 409 410 411  The name of the time argument (time_name) and the initial time of the evolution (time_start) are taken from the default values of the WaveFunction.__init__ method. Changing the default values by partial prebind (e.g. via functools) is possible.  Kloss committed May 29, 2019 412 413 414 415 416 417  """ if not syst.leads: raise AttributeError("system has no leads.") if tmax is not None and boundaries is not None: raise ValueError("'boundaries' and 'tmax' are mutually exclusive.")  Kloss committed Aug 23, 2019 418  if not _common.is_type(energy, 'real_number'):  Kloss committed May 29, 2019 419  raise TypeError('energy must be a real number')  Kloss committed Aug 23, 2019 420  if not _common.is_type(lead, 'integer'):  Kloss committed May 29, 2019 421 422 423 424  raise TypeError('lead index must be an integer') if lead >= len(syst.leads): raise ValueError("lead index must be smaller than {}.".format(len(syst.leads)))  Kloss committed Mar 22, 2020 425 426 427 428 429 430 431 432 433 434 435 436 437  # get initial time and time argument name from the onebody wavefunction try: time_name = _common.get_default_function_argument(wavefunction_type, 'time_name') time_start = _common.get_default_function_argument(wavefunction_type, 'time_start') except Exception: time_name = _common.time_name time_start = _common.time_start logger.warning('retrieving initial time and time argument name from', 'the onebody wavefunction failed, use default values: ', '"time_name"={}, "time_start"={}'.format(time_name, time_start))  Kloss committed Jul 31, 2019 438   Kloss committed Jul 12, 2019 439  # add initial time to the params dict  Kloss committed Jul 31, 2019 440 441  tparams = add_time_to_params(params, time_name=time_name, time=time_start, check_numeric_type=True)  Kloss committed May 29, 2019 442 443 444  if boundaries is None: if spectra is None:  Kloss committed Jul 12, 2019 445  spectra = kwant_spectrum.spectra(syst.leads, params=tparams)  Kloss committed May 29, 2019 446 447  boundaries = leads.automatic_boundary(spectra, tmax)  Kloss committed Jul 12, 2019 448  scattering_states = equilibrium_solver(syst, energy=energy, params=tparams)  Kloss committed May 29, 2019 449  self._psi_st = scattering_states(lead)  Kloss committed Nov 11, 2019 450 451  self._wavefunction = functools.partial(wavefunction_type, syst=syst, boundaries=boundaries, params=params)  Kloss committed May 29, 2019 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468  self.energy = energy self.lead = lead def __getitem__(self, mode): """Return a time-dependent wavefunction for the corresponding scattering mode. Parameters ---------- mode : int Mode index of the scattering wavefunction. Returns ------- wave_function : WaveFunction The time-dependent wave function starting in the scattering mode. """  Kloss committed Aug 23, 2019 469  if not _common.is_type(mode, 'integer'):  Kloss committed May 29, 2019 470 471 472 473 474 475  raise TypeError('mode index must be an integer') if mode >= len(self._psi_st): raise KeyError('no open scattering mode={}; ' 'number of open modes={} for energy={} and lead={}' .format(self.lead, mode, self.energy, len(self._psi_st)))  Kloss committed Nov 11, 2019 476  return self._wavefunction(psi_init=self._psi_st[mode], energy=self.energy)  Kloss committed May 29, 2019 477 478 479 480 481 482 483 484  def __len__(self): """Return the number of open scattering modes.""" return len(self._psi_st) def __iter__(self): """Return an iterable sequence of all open modes.""" return (self[mode] for mode in range(len(self)))