onebody.py 20.5 KB
Newer Older
Kloss's avatar
Kloss committed
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's avatar
Kloss committed
12
13
14
from cmath import exp
import numpy as np
import scipy.sparse as sp
Kloss's avatar
Kloss committed
15
16
17
18

import kwant
import kwant_spectrum

Kloss's avatar
Kloss committed
19
from .. import leads, _common, _logging
20
from ..system import (extract_perturbation, hamiltonian_with_boundaries,
21
                      add_time_to_params)
Kloss's avatar
Kloss committed
22
23
from . import kernels, solvers

24
__all__ = ['WaveFunction', 'ScatteringStates', 'Task']
Kloss's avatar
Kloss committed
25
26


Kloss's avatar
Kloss committed
27
# set module logger
Kloss's avatar
Kloss committed
28
logger = _logging.make_logger(name=__name__)
Kloss's avatar
Kloss committed
29
30
log_func = _logging.log_func(logger)

Kloss's avatar
Kloss committed
31

Kloss's avatar
Kloss committed
32
# data formats
33

Kloss's avatar
Kloss committed
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
46
47
48
    momentum : float
        momentum of the onebody state
    weight : float or numpy float array
Kloss's avatar
Kloss committed
49
        weighting factor of the one-body state in the manybody sum
50
51
        weigth = math_weight * phys_weight, where math_weight is the weighting
        factor from the integration quadrature
52
53
    math_weight : float or numpy float array
        mathematical weighting factor from the numerical quadrature rule
54
55
    phys_weight : float
        physical weighting factor (fermi function, pi factors..)
Kloss's avatar
Kloss committed
56
    """
57
    def __init__(self, lead: int, mode: int, energy: float, weight: float,
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
63
        self.momentum = momentum
64
        self.weight = weight
65
        self.math_weight = math_weight
66
        self.phys_weight = phys_weight
Kloss's avatar
Kloss committed
67

Kloss's avatar
Kloss committed
68
69
70
71
    def __eq__(self, other) -> bool:
        if not isinstance(other, Task):
            return NotImplemented
        return (
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's avatar
Kloss committed
76
77

    def __str__(self):
78
79
        string = "onebody task: lead={lead}, mode={mode}, " \
                 "energy={energy}, momentum={momentum}, " \
80
81
                 "weight={weight}, math_weight={math_weight}, " \
                 "phys_weight={phys_weight}".format(**self.__dict__)
Kloss's avatar
Kloss committed
82
83
        return string

Kloss's avatar
Kloss committed
84

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's avatar
Kloss committed
90
class WaveFunction:
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's avatar
Kloss committed
110

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)
    <https://arxiv.org/abs/1510.05967>`_.
    """
    def __init__(self, H0, W, psi_init, energy=None, params=None,
                 solution_is_valid=None, time_is_valid=None,
Kloss's avatar
Kloss committed
124
                 time_start=_common.time_start, time_name=_common.time_name,
125
126
                 kernel_type=kernels.default, solver_type=solvers.default):
        r"""
Kloss's avatar
Kloss committed
127
128
        Parameters
        ----------
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's avatar
Kloss committed
134
        psi_init : array of complex
135
136
            The state :math:`\psi(t_0)` from which to start, defined over the
            central region.
Kloss's avatar
Kloss committed
137
138
        energy : float, optional
            If provided, then ``psi_init`` is assumed to be an eigenstate
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's avatar
Kloss committed
142
            on to the central part of the system.
143
        params : dict, optional
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's avatar
Kloss committed
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.
161
      """
Kloss's avatar
Kloss committed
162

163
        # add initial time to the params dict
164
165
        tparams = add_time_to_params(params, time_name=time_name,
                                     time=time_start, check_numeric_type=True)
166

Kloss's avatar
Kloss committed
167
168
169
        if energy is None:
            # starting from an arbitrary state, so we need
            # to be solving: H0 @ psi + W(t) @ psi
170
171
            kernel = kernel_type(H0, W, params=tparams)

Kloss's avatar
Kloss committed
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]),
177
                                 W, tparams, np.asarray(psi_init))
Kloss's avatar
Kloss committed
178
179
180
181

        # get the object that will actually do the time stepping
        self.solver = solver_type(kernel)

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's avatar
Kloss committed
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

199
200
201
202
        self._syst_size = syst_size
        self._solution_is_valid = solution_is_valid
        self._time_is_valid = time_is_valid

203
204
205
206
        self.time_name = time_name
        self.time_start = time_start

        self.time = time_start
207
        self.params = params
Kloss's avatar
Kloss committed
208
209
        self.energy = energy

210
    @classmethod
Kloss's avatar
Kloss committed
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,
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
273
274
275
            if boundaries is not None:
                raise ValueError('No boundary conditions must be provided '
                                 'for a system without leads.')
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's avatar
Kloss committed
285
                   time_start, time_name, kernel_type, solver_type)
286

Kloss's avatar
Kloss committed
287
    def psi(self):
288
        r"""Return the wavefunction :math:`\psi(t)` at the current time *t*.
Kloss's avatar
Kloss committed
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.
        """
295
        if self.energy is None:  # psi = psibar
296
            return self.psibar[:self._syst_size].copy()
297
        # else : psi = (psibar + psi_st) * exp(-i E (t - t0))
Kloss's avatar
Kloss committed
298
299
        return ((self.psibar[:self._syst_size] + self.psi_st) *
                exp(-1j * self.energy * (self.time - self.time_start)))
Kloss's avatar
Kloss committed
300

301
    def evolve(self, time, params=None):
302
303
304
        r"""
        Evolve the wavefunction :math:`\psi(t)` foreward in time up to
        :math:`t =` ``time``.
Kloss's avatar
Kloss committed
305
306
307
308
309

        Parameters
        ----------
        time : int or float
            time argument up to which the solver should be evolved
310
        params : dict, optional
311
            Extra arguments to pass to the Hamiltonian, excluding time.
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's avatar
Kloss committed
315
        """
316
317
318
319
        if params is not None:
            self.params = params
            self.solver.kernel.set_params(params)

Kloss's avatar
Kloss committed
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
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))
329

Kloss's avatar
Kloss committed
330
331
        # evolve forward in time
        next_psibar = self.solver(self.psibar, self.time, time)
332

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's avatar
Kloss committed
338
339
340
        # update internal state and return
        self.psibar, self.time = next_psibar, time

341
    def evaluate(self, observable):
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's avatar
Kloss committed
347
348
349
350

        Parameters
        ----------
        observable : callable or `kwant.operator`
351
            An operator :math:`\hat{O}` to evaluate the expectation value.
352
            Must have the calling signature of `kwant.operator`.
Kloss's avatar
Kloss committed
353
354
355
356

        Returns
        -------
        result : numpy array
357
            The expectation value :math:`O(t)` of ``observable``.
Kloss's avatar
Kloss committed
358
        """
359
360
361
        if _operator_bound(observable):
            raise ValueError("Operator must not use pre-bind values")

362
        tparams = add_time_to_params(self.params, time_name=self.time_name,
363
                                     time=self.time)
364
        return observable(self.psi(), params=tparams)
Kloss's avatar
Kloss committed
365
366


367
class ScatteringStates(collections.abc.Iterable):
Kloss's avatar
Kloss committed
368
369
    """Calculate time-dependent wavefunctions starting in an equilibrium scattering state"""

370
    def __init__(self, syst, energy, lead, tmax=None, params=None, spectra=None,
Kloss's avatar
Kloss committed
371
                 boundaries=None, equilibrium_solver=kwant.wave_function,
372
373
                 wavefunction_type=WaveFunction.from_kwant):
        """
Kloss's avatar
Kloss committed
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'.
387
        params : dict, optional
Kloss's avatar
Kloss committed
388
            Extra arguments to pass to the Hamiltonian of ``syst``, excluding time.
Kloss's avatar
Kloss committed
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's avatar
Kloss committed
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's avatar
Kloss committed
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's avatar
Kloss committed
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.
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's avatar
Kloss committed
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's avatar
Kloss committed
418
        if not _common.is_type(energy, 'real_number'):
Kloss's avatar
Kloss committed
419
            raise TypeError('energy must be a real number')
Kloss's avatar
Kloss committed
420
        if not _common.is_type(lead, 'integer'):
Kloss's avatar
Kloss committed
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's avatar
Kloss committed
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))
438

439
        # add initial time to the params dict
440
441
        tparams = add_time_to_params(params, time_name=time_name,
                                     time=time_start, check_numeric_type=True)
Kloss's avatar
Kloss committed
442
443
444

        if boundaries is None:
            if spectra is None:
445
                spectra = kwant_spectrum.spectra(syst.leads, params=tparams)
Kloss's avatar
Kloss committed
446
447
            boundaries = leads.automatic_boundary(spectra, tmax)

448
        scattering_states = equilibrium_solver(syst, energy=energy, params=tparams)
Kloss's avatar
Kloss committed
449
        self._psi_st = scattering_states(lead)
450
451
        self._wavefunction = functools.partial(wavefunction_type, syst=syst,
                                               boundaries=boundaries, params=params)
Kloss's avatar
Kloss committed
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's avatar
Kloss committed
469
        if not _common.is_type(mode, 'integer'):
Kloss's avatar
Kloss committed
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)))

476
        return self._wavefunction(psi_init=self._psi_st[mode], energy=self.energy)
Kloss's avatar
Kloss committed
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)))