common.py 42.8 KB
Newer Older
1
# Copyright 2011-2014 Kwant authors.
2
#
Christoph Groth's avatar
Christoph Groth committed
3
4
# This file is part of Kwant.  It is subject to the license terms in the file
# LICENSE.rst found in the top-level directory of this distribution and at
5
# http://kwant-project.org/license.  A list of Kwant authors can be found in
Christoph Groth's avatar
Christoph Groth committed
6
# the file AUTHORS.rst at the top-level directory of this distribution and at
7
8
# http://kwant-project.org/authors.

9
__all__ = ['SparseSolver', 'SMatrix', 'GreensFunction']
10
11

from collections import namedtuple
12
from itertools import product
13
import abc
14
from  numbers import Integral
15
16
import numpy as np
import scipy.sparse as sp
17
from .._common import ensure_isinstance, deprecate_args
18
from .. import system, physics
Joseph Weston's avatar
Joseph Weston committed
19
from functools import reduce
20

21
22
23
24
25
26
27
28
29
30
31
# Until v0.13.0, scipy.sparse did not support making block matrices out of
# matrices with one dimension being zero:
# https://github.com/scipy/scipy/issues/2127 Additionally, 'scipy.sparse.bmat'
# didn't support matrices with zero size until v0.18:
# https://github.com/scipy/scipy/issues/5976. For now we use NumPy dense
# matrices as a replacement.

# TODO: Once we depend on scipy >= 0.18, code for the special cases can be
# removed from _make_linear_sys, _solve_linear_sys and possibly other places
# marked by the line "See comment about zero-shaped sparse matrices at the top
# of common.py".
32

33
LinearSys = namedtuple('LinearSys', ['lhs', 'rhs', 'indices', 'num_orb'])
34

35

36
class SparseSolver(metaclass=abc.ABCMeta):
37
38
39
    """Solver class for computing physical quantities based on solving
    a liner system of equations.

40
41
    `SparseSolver` is an abstract base class.  It cannot be instantiated as it
    does not specify the actual linear solve step.  In order to be directly
Christoph Groth's avatar
Christoph Groth committed
42
    usable, a derived class must implement the methods
43

Christoph Groth's avatar
Christoph Groth committed
44
45
    - `_factorize` and `_solve_linear_sys`, that solve a linear system of
      equations
46

47
    and the following properties:
48

49
50
51
    - `lhsformat`, `rhsformat`: Sparse matrix format of left and right hand
       sides of the linear system, respectively.  Must be one of {'coo', 'csc',
       'csr'}.
52

53
54
55
56
    - `nrhs`: Number of right hand side vectors that should be solved at once
      in a call to `solve_linear_sys`, when the full solution is computed (i.e.
      kept_vars covers all entries in the solution). This should be not too big
      too avoid excessive memory usage, but for some solvers not too small for
Christoph Groth's avatar
Christoph Groth committed
57
      performance reasons.
58
59
60
    """

    @abc.abstractmethod
61
    def _factorized(self, a):
Christoph Groth's avatar
Christoph Groth committed
62
63
64
65
66
67
68
69
70
71
72
73
74
        """
        Return a preprocessed version of a matrix for the use with
        `solve_linear_sys`.

        Parameters
        ----------
        a : a scipy.sparse.coo_matrix sparse matrix.

        Returns
        -------
        factorized_a : object
            factorized lhs to be used with `_solve_linear_sys`.
        """
75
76
77
        pass

    @abc.abstractmethod
78
    def _solve_linear_sys(self, factorized_a, b, kept_vars):
79
80
81
82
        """
        Solve the linar system `a x = b`, returning the part of
        the result indicated in kept_vars.

Christoph Groth's avatar
Christoph Groth committed
83
84
85
86
        Parameters
        ----------
        factorized : object
            The result of calling `_factorized` for the matrix a.
87
        b : sparse matrix
Christoph Groth's avatar
Christoph Groth committed
88
            The right hand side. Its format must match `rhsformat`.
Christoph Groth's avatar
Christoph Groth committed
89
90
91
92
93
94
95
        kept_vars : slice object or sequence of integers
            A sequence of numbers of variables to keep in the solution

        Returns
        -------
        output : NumPy matrix
            Solution to the system of equations.
96
97
98
        """
        pass

99
    def _make_linear_sys(self, sys, in_leads, energy=0, args=(),
100
101
                         check_hermiticity=True, realspace=False,
                         *, params=None):
102
        """Make a sparse linear system of equations defining a scattering
103
104
105
106
107
        problem.

        Parameters
        ----------
        sys : kwant.system.FiniteSystem
Christoph Groth's avatar
Christoph Groth committed
108
            Low level system, containing the leads and the Hamiltonian of a
109
            scattering region.
Christoph Groth's avatar
Christoph Groth committed
110
111
        in_leads : sequence of integers
            Numbers of leads in which current or wave function is injected.
112
        energy : number
Christoph Groth's avatar
Christoph Groth committed
113
            Excitation energy at which to solve the scattering problem.
114
115
        args : tuple, defaults to empty
            Positional arguments to pass to the ``hamiltonian`` method.
116
            Deprecated in favor of 'params' (and mutually exclusive with it).
117
118
        check_hermiticity : bool
            Check if Hamiltonian matrices are in fact Hermitian.
119
            Enables deduction of missing transmission coefficients.
120
        realspace : bool
Christoph Groth's avatar
Christoph Groth committed
121
            Calculate Green's function between the outermost lead
122
123
            sites, instead of lead modes. This is almost always
            more computationally expensive and less stable.
124
125
126
        params : dict, optional
            Dictionary of parameter names and their values. Mutually exclusive
            with 'args'.
127
128
129

        Returns
        -------
130
        (lhs, rhs, indices, num_orb) : LinearSys
Anton Akhmerov's avatar
Anton Akhmerov committed
131
            `lhs` is a scipy.sparse.csc_matrix, containing the left hand side
132
133
134
135
136
            of the system of equations.  `rhs` is a list of matrices with the
            right hand side, with each matrix corresponding to one lead
            mentioned in `in_leads`. `indices` is a list of arrays of variables
            in the system of equations corresponding to the the outgoing modes
            in each lead, or the indices of variables, on which a lead defined
137
138
            via self-energy adds the self-energy. Finally `num_orb` is the
            total number of degrees of freedom in the scattering region.
139

140
        lead_info : list of objects
141
            Contains one entry for each lead.  If `realspace=False`, this is an
142
143
            instance of `~kwant.physics.PropagatingModes` for all leads that
            have modes, otherwise the lead self-energy matrix.
144
145
146

        Notes
        -----
147
148
        All the leads should implement a method `modes` if `realspace=False`
        and a method `selfenergy`.
149

150
151
        The system of equations that is created will be described in detail
        elsewhere.
152
        """
153
154
155

        syst = sys  # ensure consistent naming across function bodies
        ensure_isinstance(syst, system.System)
156

157
158
159
        splhsmat = getattr(sp, self.lhsformat + '_matrix')
        sprhsmat = getattr(sp, self.rhsformat + '_matrix')

160
        if not syst.lead_interfaces:
161
            raise ValueError('System contains no leads.')
162
        lhs, norb = syst.hamiltonian_submatrix(args, sparse=True,
163
164
                                               return_norb=True,
                                               params=params)[:2]
165
166
        lhs = getattr(lhs, 'to' + self.lhsformat)()
        lhs = lhs - energy * sp.identity(lhs.shape[0], format=self.lhsformat)
167
        num_orb = lhs.shape[0]
168

169
        if check_hermiticity and len(lhs.data):
170
171
172
173
            rtol = 1e-13
            atol = 1e-300
            tol = rtol * np.max(np.abs(lhs.data)) + atol
            if np.any(np.abs((lhs - lhs.T.conj()).data) > tol):
174
175
176
                raise ValueError('System Hamiltonian is not Hermitian. '
                                 'Use option `check_hermiticity=False` '
                                 'if this is intentional.')
177

178
179
        offsets = np.empty(norb.shape[0] + 1, int)
        offsets[0] = 0
180
181
182
183
184
        offsets[1 :] = np.cumsum(norb)

        # Process the leads, generate the eigenvector matrices and lambda
        # vectors. Then create blocks of the linear system and add them
        # step by step.
185
        indices = []
186
187
        rhs = []
        lead_info = []
188
189
        for leadnum, interface in enumerate(syst.lead_interfaces):
            lead = syst.leads[leadnum]
190
            if not realspace:
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
                if system.is_selfenergy_lead(lead):
                    # We make equivalent checks to this in 'smatrix',
                    # 'greens_function' and 'wave_function' (where we
                    # can give more informative error messages).
                    # Putting this check here again is just a sanity check,
                    assert leadnum not in in_leads
                    sigma = np.asarray(lead.selfenergy(energy, args, params=params))
                    rhs.append(None)
                    lead_info.append(sigma)
                    indices.append(None)
                    # add selfenergy to the LHS
                    coords = np.r_[tuple(slice(offsets[i], offsets[i + 1])
                                   for i in interface)]
                    if sigma.shape != 2 * coords.shape:
                        raise ValueError(
                           f"Self-energy dimension for lead {leadnum} "
                           "does not match the total number of orbitals "
                           "of the sites for which it is defined."
                        )
                    y, x = np.meshgrid(coords, coords)
                    sig_sparse = splhsmat((sigma.flat, [x.flat, y.flat]),
                                          lhs.shape)
                    lhs = lhs + sig_sparse
214
                else:
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
                    prop, stab = lead.modes(energy, args, params=params)
                    lead_info.append(prop)
                    u = stab.vecs
                    ulinv = stab.vecslmbdainv
                    nprop = stab.nmodes
                    svd_v = stab.sqrt_hop

                    if len(u) == 0:
                        rhs.append(None)
                        continue

                    indices.append(np.arange(lhs.shape[0], lhs.shape[0] + nprop))

                    u_out, ulinv_out = u[:, nprop:], ulinv[:, nprop:]
                    u_in, ulinv_in = u[:, :nprop], ulinv[:, :nprop]

                    # Construct a matrix of 1's that translates the
                    # inter-cell hopping to a proper hopping
                    # from the system to the lead.
                    iface_orbs = np.r_[tuple(slice(offsets[i], offsets[i + 1])
                                            for i in interface)]

                    n_lead_orbs = svd_v.shape[0]
                    if n_lead_orbs != len(iface_orbs):
                        msg = ('Lead {0} has hopping with dimensions '
                               'incompatible with its interface dimension.')
                        raise ValueError(msg.format(leadnum))

                    coords = np.r_[[np.arange(len(iface_orbs))], [iface_orbs]]
                    transf = sp.csc_matrix((np.ones(len(iface_orbs)), coords),
                                           shape=(iface_orbs.size, lhs.shape[0]))

                    v_sp = sp.csc_matrix(svd_v.T.conj()) * transf
                    vdaguout_sp = (transf.T *
                                   sp.csc_matrix(np.dot(svd_v, u_out)))
                    lead_mat = - ulinv_out

                    lhs = sp.bmat([[lhs, vdaguout_sp], [v_sp, lead_mat]],
                                  format=self.lhsformat)

                    if leadnum in in_leads and nprop > 0:
                        vdaguin_sp = transf.T * sp.csc_matrix(
                            -np.dot(svd_v, u_in))

                        # defer formation of the real matrix until the proper
                        # system size is known
                        rhs.append((vdaguin_sp, ulinv_in))
                    else:
                        rhs.append(None)
264
            else:
265
                sigma = np.asarray(lead.selfenergy(energy, args, params=params))
266
                lead_info.append(sigma)
267
                coords = np.r_[tuple(slice(offsets[i], offsets[i + 1])
268
                                      for i in interface)]
269

270
                if sigma.shape != 2 * coords.shape:
271
272
273
                    msg = ('Self-energy dimension for lead {0} does not '
                           'match the total number of orbitals of the '
                           'sites for which it is defined.')
274
275
                    raise ValueError(msg.format(leadnum))

276
                y, x = np.meshgrid(coords, coords)
277
278
                sig_sparse = splhsmat((sigma.flat, [x.flat, y.flat]),
                                      lhs.shape)
279
                lhs = lhs + sig_sparse  # __iadd__ is not implemented in v0.7
280
                indices.append(coords)
281
282
283
                if leadnum in in_leads:
                    # defer formation of true rhs until the proper system
                    # size is known
284
                    rhs.append((coords,))
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305

        # Resize the right-hand sides to be compatible with the full lhs
        for i, mats in enumerate(rhs):
            if isinstance(mats, tuple):
                if len(mats) == 1:
                    # self-energy lead
                    l = mats[0].shape[0]
                    rhs[i] = sprhsmat((-np.ones(l), [mats[0], np.arange(l)]),
                                      shape=(lhs.shape[0], l))
                else:
                    # lead with modes
                    zero_rows = (lhs.shape[0] - mats[0].shape[0] -
                                 mats[1].shape[0])

                    if zero_rows:
                        zero_mat = sprhsmat((zero_rows, mats[0].shape[1]))
                        bmat = [[mats[0]], [mats[1]], [zero_mat]]
                    else:
                        bmat = [[mats[0]], [mats[1]]]

                    rhs[i] = sp.bmat(bmat, format=self.rhsformat)
306
307
308
309
310
            elif mats is None:
                # A lead with no rhs.
                rhs[i] = np.zeros((lhs.shape[0], 0))
            else:
                raise RuntimeError('Unknown right-hand side format')
311

312
        return LinearSys(lhs, rhs, indices, num_orb), lead_info
313

314
    @deprecate_args
315
    def smatrix(self, sys, energy=0, args=(),
316
317
                out_leads=None, in_leads=None, check_hermiticity=True,
                *, params=None):
318
        """
319
        Compute the scattering matrix of a system.
320

321
322
        An alias exists for this common name: ``kwant.smatrix``.

323
324
325
        Parameters
        ----------
        sys : `kwant.system.FiniteSystem`
Christoph Groth's avatar
Christoph Groth committed
326
            Low level system, containing the leads and the Hamiltonian of a
327
328
            scattering region.
        energy : number
Christoph Groth's avatar
Christoph Groth committed
329
            Excitation energy at which to solve the scattering problem.
330
331
        args : tuple, defaults to empty
            Positional arguments to pass to the ``hamiltonian`` method.
332
            Deprecated in favor of 'params' (and mutually exclusive with it).
Christoph Groth's avatar
Christoph Groth committed
333
334
335
336
337
        out_leads : sequence of integers or ``None``
            Numbers of leads where current or wave function is extracted.  None
            is interpreted as all leads. Default is ``None`` and means "all
            leads".
        in_leads : sequence of integers or ``None``
Christoph Groth's avatar
Christoph Groth committed
338
            Numbers of leads in which current or wave function is injected.
Christoph Groth's avatar
Christoph Groth committed
339
340
341
            None is interpreted as all leads. Default is ``None`` and means
            "all leads".
        check_hermiticity : ``bool``
Christoph Groth's avatar
Christoph Groth committed
342
            Check if the Hamiltonian matrices are Hermitian.
343
            Enables deduction of missing transmission coefficients.
344
345
346
        params : dict, optional
            Dictionary of parameter names and their values. Mutually exclusive
            with 'args'.
347
348
349

        Returns
        -------
350
351
        output : `~kwant.solvers.common.SMatrix`
            See the notes below and `~kwant.solvers.common.SMatrix`
Christoph Groth's avatar
Christoph Groth committed
352
            documentation.
353
354
355

        Notes
        -----
Christoph Groth's avatar
Christoph Groth committed
356
357
        This function can be used to calculate the conductance and other
        transport properties of a system.  See the documentation for its output
358
        type, `~kwant.solvers.common.SMatrix`.
Christoph Groth's avatar
Christoph Groth committed
359

360
361
362
        The returned object contains the scattering matrix elements from the
        `in_leads` to the `out_leads` as well as information about the lead
        modes.
Christoph Groth's avatar
Christoph Groth committed
363
364

        Both `in_leads` and `out_leads` must be sorted and may only contain
365
366
367
        unique entries.
        """

368
369
370
371
        syst = sys  # ensure consistent naming across function bodies
        ensure_isinstance(syst, system.System)

        n = len(syst.lead_interfaces)
372
        if in_leads is None:
Joseph Weston's avatar
Joseph Weston committed
373
            in_leads = list(range(n))
374
375
        else:
            in_leads = list(in_leads)
376
        if out_leads is None:
Joseph Weston's avatar
Joseph Weston committed
377
            out_leads = list(range(n))
378
379
380
        else:
            out_leads = list(out_leads)
        if (np.any(np.diff(in_leads) <= 0) or np.any(np.diff(out_leads) <= 0)):
381
382
383
384
385
            raise ValueError("Lead lists must be sorted and "
                             "with unique entries.")
        if len(in_leads) == 0 or len(out_leads) == 0:
            raise ValueError("No output is requested.")

386
387
388
389
390
391
392
393
394
        for direction, leads in [("in", in_leads), ("out", out_leads)]:
            for lead in leads:
                if system.is_selfenergy_lead(syst.leads[lead]):
                    raise ValueError(
                        f"lead {lead} is only defined by a self-energy, "
                        "so we cannot calculate scattering matrix elements for it. "
                        f"Specify '{direction}_leads' without '{lead}'."
                    )

395
        linsys, lead_info = self._make_linear_sys(syst, in_leads, energy, args,
396
397
                                                  check_hermiticity, False,
                                                  params=params)
398

399
        kept_vars = np.concatenate([coords for i, coords in
400
401
402
                                    enumerate(linsys.indices) if i in
                                    out_leads])

403
404
        # Do not perform factorization if no calculation is to be done.
        len_rhs = sum(i.shape[1] for i in linsys.rhs)
405
        len_kv = len(kept_vars)
406
        if not(len_rhs and len_kv):
407
408
            return SMatrix(np.zeros((len_kv, len_rhs)), lead_info,
                           out_leads, in_leads, check_hermiticity)
409

410
411
412
        # See comment about zero-shaped sparse matrices at the top of common.py.
        rhs = sp.bmat([[i for i in linsys.rhs if i.shape[1]]],
                      format=self.rhsformat)
413
        flhs = self._factorized(linsys.lhs)
414
        data = self._solve_linear_sys(flhs, rhs, kept_vars)
415

416
        return SMatrix(data, lead_info, out_leads, in_leads, check_hermiticity)
417

418
    @deprecate_args
419
    def greens_function(self, sys, energy=0, args=(),
420
421
                        out_leads=None, in_leads=None, check_hermiticity=True,
                        *, params=None):
422
423
424
        """
        Compute the retarded Green's function of the system between its leads.

425
426
        An alias exists for this common name: ``kwant.greens_function``.

427
428
429
430
431
432
433
        Parameters
        ----------
        sys : `kwant.system.FiniteSystem`
            Low level system, containing the leads and the Hamiltonian of a
            scattering region.
        energy : number
            Excitation energy at which to solve the scattering problem.
434
435
        args : tuple, defaults to empty
            Positional arguments to pass to the ``hamiltonian`` method.
436
            Deprecated in favor of 'params' (and mutually exclusive with it).
437
438
439
440
441
442
443
444
445
446
        out_leads : sequence of integers or ``None``
            Numbers of leads where current or wave function is extracted.  None
            is interpreted as all leads. Default is ``None`` and means "all
            leads".
        in_leads : sequence of integers or ``None``
            Numbers of leads in which current or wave function is injected.
            None is interpreted as all leads. Default is ``None`` and means
            "all leads".
        check_hermiticity : ``bool``
            Check if the Hamiltonian matrices are Hermitian.
447
            Enables deduction of missing transmission coefficients.
448
449
450
        params : dict, optional
            Dictionary of parameter names and their values. Mutually exclusive
            with 'args'.
451
452
453
454
455
456
457
458
459
460
461
462

        Returns
        -------
        output : `~kwant.solvers.common.GreensFunction`
            See the notes below and `~kwant.solvers.common.GreensFunction`
            documentation.

        Notes
        -----
        This function can be used to calculate the conductance and other
        transport properties of a system.  It is often slower and less stable
        than the scattering matrix-based calculation executed by
463
        `~kwant.solvers.common.smatrix`, and is currently provided mostly for testing
464
465
466
467
468
469
470
471
472
473
474
        purposes and compatibility with RGF code.

        It returns an object encapsulating the Green's function elements
        between the system sites interfacing the leads in `in_leads` and those
        interfacing the leads in `out_leads`.  The returned object also
        contains a list with self-energies of the leads.

        Both `in_leads` and `out_leads` must be sorted and may only contain
        unique entries.
        """

475
476
477
478
        syst = sys  # ensure consistent naming across function bodies
        ensure_isinstance(syst, system.System)

        n = len(syst.lead_interfaces)
479
        if in_leads is None:
Joseph Weston's avatar
Joseph Weston committed
480
            in_leads = list(range(n))
481
482
        else:
            in_leads = list(in_leads)
483
        if out_leads is None:
Joseph Weston's avatar
Joseph Weston committed
484
            out_leads = list(range(n))
485
486
487
        else:
            out_leads = list(out_leads)
        if (np.any(np.diff(in_leads) <= 0) or np.any(np.diff(out_leads) <= 0)):
488
489
490
491
492
            raise ValueError("Lead lists must be sorted and "
                             "with unique entries.")
        if len(in_leads) == 0 or len(out_leads) == 0:
            raise ValueError("No output is requested.")

493
        linsys, lead_info = self._make_linear_sys(syst, in_leads, energy, args,
494
495
                                                  check_hermiticity, True,
                                                  params=params)
496

497
        kept_vars = np.concatenate([coords for i, coords in
498
499
500
501
502
503
504
                                    enumerate(linsys.indices) if i in
                                    out_leads])

        # Do not perform factorization if no calculation is to be done.
        len_rhs = sum(i.shape[1] for i in linsys.rhs)
        len_kv = len(kept_vars)
        if not(len_rhs and len_kv):
505
506
            return GreensFunction(np.zeros((len_kv, len_rhs)), lead_info,
                                  out_leads, in_leads, check_hermiticity)
507
508
509
510
511
512

        # See comment about zero-shaped sparse matrices at the top of common.py.
        rhs = sp.bmat([[i for i in linsys.rhs if i.shape[1]]],
                      format=self.rhsformat)
        flhs = self._factorized(linsys.lhs)
        data = self._solve_linear_sys(flhs, rhs, kept_vars)
513

514
515
        return GreensFunction(data, lead_info, out_leads, in_leads,
                              check_hermiticity)
516

517
    @deprecate_args
518
519
    def ldos(self, sys, energy=0, args=(), check_hermiticity=True,
             *, params=None):
520
521
522
        """
        Calculate the local density of states of a system at a given energy.

523
524
        An alias exists for this common name: ``kwant.ldos``.

525
526
527
        Parameters
        ----------
        sys : `kwant.system.FiniteSystem`
Christoph Groth's avatar
Christoph Groth committed
528
            Low level system, containing the leads and the Hamiltonian of the
529
530
            scattering region.
        energy : number
Christoph Groth's avatar
Christoph Groth committed
531
            Excitation energy at which to solve the scattering problem.
532
533
        args : tuple of arguments, or empty tuple
            Positional arguments to pass to the function(s) which
534
535
            evaluate the hamiltonian matrix elements.
            Deprecated in favor of 'params' (and mutually exclusive with it).
536
537
        check_hermiticity : ``bool``
            Check if the Hamiltonian matrices are Hermitian.
538
539
540
        params : dict, optional
            Dictionary of parameter names and their values. Mutually exclusive
            with 'args'.
541
542
543

        Returns
        -------
544
        ldos : a NumPy array
Christoph Groth's avatar
Christoph Groth committed
545
            Local density of states at each orbital of the system.
546
        """
547
548
549

        syst = sys  # ensure consistent naming across function bodies
        ensure_isinstance(syst, system.System)
550
551
552
553
        if not check_hermiticity:
            raise NotImplementedError("ldos for non-Hermitian Hamiltonians "
                                      "is not implemented yet.")

554
        for lead in syst.leads:
555
            if system.is_selfenergy_lead(lead):
556
                # TODO: fix this
557
                raise NotImplementedError("ldos for leads with only "
558
                                          "self-energy is not implemented yet.")
559

560
        linsys = self._make_linear_sys(syst, range(len(syst.leads)), energy,
561
562
                                       args, check_hermiticity,
                                       params=params)[0]
563

564
        ldos = np.zeros(linsys.num_orb, float)
565

566
        # Do not perform factorization if no further calculation is needed.
567
568
569
570
        if not sum(i.shape[1] for i in linsys.rhs):
            return ldos

        factored = self._factorized(linsys.lhs)
571

572
573
574
        # See comment about zero-shaped sparse matrices at the top of common.py.
        rhs = sp.bmat([[i for i in linsys.rhs if i.shape[1]]],
                      format=self.rhsformat)
Joseph Weston's avatar
Joseph Weston committed
575
        for j in range(0, rhs.shape[1], self.nrhs):
576
577
578
579
            jend = min(j + self.nrhs, rhs.shape[1])
            psi = self._solve_linear_sys(factored, rhs[:, j:jend],
                                         slice(linsys.num_orb))
            ldos += np.sum(np.square(abs(psi)), axis=1)
580
581
582

        return ldos * (0.5 / np.pi)

583
    @deprecate_args
584
585
    def wave_function(self, sys, energy=0, args=(), check_hermiticity=True,
                      *, params=None):
586
        r"""
587
588
589
        Return a callable object for the computation of the wave function
        inside the scattering region.

590
591
        An alias exists for this common name: ``kwant.wave_function``.

592
593
594
595
596
        Parameters
        ----------
        sys : `kwant.system.FiniteSystem`
            The low level system for which the wave functions are to be
            calculated.
597
598
        args : tuple of arguments, or empty tuple
            Positional arguments to pass to the function(s) which
599
600
            evaluate the hamiltonian matrix elements.
            Deprecated in favor of 'params' (and mutually exclusive with it).
601
602
        check_hermiticity : ``bool``
            Check if the Hamiltonian matrices are Hermitian.
603
604
605
        params : dict, optional
            Dictionary of parameter names and their values. Mutually exclusive
            with 'args'.
606
607
608
609

        Notes
        -----
        The returned object can be itself called like a function.  Given a lead
610
611
        number, it returns a 2d NumPy array that contains the wave function
        within the scattering region due to each incoming mode of the given
612
613
614
615
616
617
618
619
620
621
622
623
624
625
        lead.  Index 0 is the mode number, index 1 is the orbital number.

        The modes appear in the same order as the negative velocity modes in
        `kwant.physics.PropagatingModes`. In Kwant's convention leads are attached
        so that their translational symmetry points *away* from the scattering
        region::

             left lead    SR   right lead
             /---------\ /---\ /---------\
             ...-3-2-1-0-X-X-X-0-1-2-3-...

        This means that incoming modes (coming from infinity towards the
        scattering region) have *negative* velocity with respect to the
        lead's symmetry direction.
626
627
628

        Examples
        --------
629
        >>> wf = kwant.solvers.default.wave_function(some_syst, some_energy)
630
        >>> wfs_of_lead_2 = wf(2)
631

632
        """
633
        return WaveFunction(self, sys, energy, args, check_hermiticity, params)
634
635


636
class WaveFunction:
637
    def __init__(self, solver, sys, energy, args, check_hermiticity, params):
638
639
        syst = sys  # ensure consistent naming across function bodies
        ensure_isinstance(syst, system.System)
640
641
642
643
644
645
646

        modes_leads = [
            i for i, l in enumerate(syst.leads)
            if not system.is_selfenergy_lead(l)
        ]

        linsys = solver._make_linear_sys(syst, modes_leads, energy,
647
648
                                         args, check_hermiticity,
                                         params=params)[0]
649
        self.modes_leads = modes_leads
650
651
652
653
        self.solve = solver._solve_linear_sys
        self.rhs = linsys.rhs
        self.factorized_h = solver._factorized(linsys.lhs)
        self.num_orb = linsys.num_orb
654
655

    def __call__(self, lead):
656
657
658
659
660
        if lead not in self.modes_leads:
            msg = ('Scattering wave functions for leads with only '
                   'self-energy are undefined.')
            raise ValueError(msg)

661
662
        result = self.solve(self.factorized_h, self.rhs[lead],
                            slice(self.num_orb))
663
664
        return result.transpose()

665

666
class BlockResult(metaclass=abc.ABCMeta):
667
    """
668
    ABC for a linear system solution with variable grouping.
669
    """
670
671
672

    def __init__(self, data, lead_info, out_leads, in_leads, sizes,
                 current_conserving=False):
Michael Wimmer's avatar
Michael Wimmer committed
673
674
        self.data = data
        self.lead_info = lead_info
675
676
677
        self.out_leads = out_leads = list(out_leads)
        self.in_leads = in_leads = list(in_leads)
        self.sizes = sizes = np.array(sizes)
678
        self.current_conserving = current_conserving
679
680
        self.in_offsets = np.zeros(len(in_leads) + 1, int)
        self.in_offsets[1 :] = np.cumsum(self.sizes[in_leads])
681
        self.out_offsets = np.zeros(len(self.out_leads) + 1, int)
682
        self.out_offsets[1 :] = np.cumsum(self.sizes[out_leads])
Michael Wimmer's avatar
Michael Wimmer committed
683

684
685
686
687
    def block_coords(self, lead_out, lead_in):
        """
        Return slices corresponding to the block from lead_in to lead_out.
        """
Michael Wimmer's avatar
Michael Wimmer committed
688
689
690
        return self.out_block_coords(lead_out), self.in_block_coords(lead_in)

    def out_block_coords(self, lead_out):
691
        """Return a slice with the rows in the block corresponding to lead_out.
Michael Wimmer's avatar
Michael Wimmer committed
692
        """
693
        lead_out = self.out_leads.index(lead_out)
694
695
        return slice(self.out_offsets[lead_out],
                     self.out_offsets[lead_out + 1])
Michael Wimmer's avatar
Michael Wimmer committed
696
697

    def in_block_coords(self, lead_in):
698
699
        """
        Return a slice with the columns in the block corresponding to lead_in.
Michael Wimmer's avatar
Michael Wimmer committed
700
701
        """
        lead_in = self.in_leads.index(lead_in)
702
703
        return slice(self.in_offsets[lead_in],
                     self.in_offsets[lead_in + 1])
704
705
706
707
708

    def submatrix(self, lead_out, lead_in):
        """Return the matrix elements from lead_in to lead_out."""
        return self.data[self.block_coords(lead_out, lead_in)]

709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
    @abc.abstractmethod
    def num_propagating(self, lead):
        """Return the number of propagating modes in the lead."""
        pass

    @abc.abstractmethod
    def _transmission(self, lead_out, lead_in):
        """Return transmission from lead_in to lead_out.

        The calculation is expected to be done directly from the corresponding
        matrix elements.
        """
        pass

    def transmission(self, lead_out, lead_in):
        """Return transmission from lead_in to lead_out.

        If the option ``current_conserving`` has been enabled for this object,
        this method will deduce missing transmission values whenever possible.

        Current conservation is enabled by default for objects returned by
        `~kwant.solvers.default.smatrix` and
        `~kwant.solvers.default.greens_function` whenever the Hamiltonian has
        been verified to be Hermitian (option ``check_hermiticity``, enabled by
        default).

        """
        chosen = [lead_out, lead_in]
        present = self.out_leads, self.in_leads
        # OK means: chosen (outgoing, incoming) lead is present.
        ok = [int(c in p) for c, p in zip(chosen, present)]
        if all(ok):
            return self._transmission(lead_out, lead_in)

        if self.current_conserving:
            all_but_one = len(self.lead_info) - 1
            s_t = self._transmission
            if sum(ok) == 1:
                # Calculate the transmission element by summing over a single
                # row or column.
                sum_dim, ok_dim = ok
                if len(present[sum_dim]) == all_but_one:
                    return (self.num_propagating(chosen[ok_dim]) - sum(
                        s_t(*chosen) for chosen[sum_dim] in present[sum_dim]))
            elif all(len(p) == all_but_one for p in present):
                # Calculate the transmission element at the center of a single
                # missing row and a single missing column.
                return (sum(s_t(o, i) - self.num_propagating(o) if o == i else 0
                            for o, i in product(*present)))

        raise ValueError("Insufficient matrix elements to compute "
                         "transmission({0}, {1}).".format(*chosen))

    def conductance_matrix(self):
        """Return the conductance matrix.

        This is the matrix :math:`C` such that :math:`I = CV` where :math:`I`
        and :math:`V` are, respectively, the vectors of currents and voltages
        for each lead.

        This matrix is useful for calculating non-local resistances.  See
        Section 2.4 of the book by S. Datta.
        """
        n = len(self.lead_info)
Joseph Weston's avatar
Joseph Weston committed
773
        rn = range(n)
774
775
776
777
778
779
        result = np.array([[-self.transmission(i, j) if i != j else 0
                            for j in rn] for i in rn])
        # Set the diagonal elements such that the sums of columns are zero.
        result.flat[::n+1] = -result.sum(axis=0)
        return result

780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795

class SMatrix(BlockResult):
    """A scattering matrix.

    Transport properties can be easily accessed using the
    `~SMatrix.transmission` method (don't be fooled by the name,
    it can also compute reflection, which is just transmission from one
    lead back into the same lead.)

    `SMatrix` however also allows for a more direct access to the result: The
    data stored in `SMatrix` is a scattering matrix with respect to lead modes
    and these modes themselves. The details of this data can be directly
    accessed through the instance variables `data` and `lead_info`. Subblocks
    of data corresponding to particular leads are conveniently obtained by
    `~SMatrix.submatrix`.

796
797
798
799
800
801
802
803
    `SMatrix` also respects the conservation laws present in the lead, such as
    spin conservation, if they are declared during system construction. If
    queried with length-2 sequence the first number is the number of the lead,
    and the second number is the index of the corresponding conserved
    quantity. For example ``smatrix.transmission((1, 3), (0, 2))`` is
    transmission from block 2 of the conserved quantity in lead 0 to the block
    3 of the conserved quantity in lead 1.

804
805
    Attributes
    ----------
806
807
808
809
810
    data : NumPy array
        a matrix containing all the requested matrix elements of the scattering
        matrix.
    lead_info : list of data
        a list containing `kwant.physics.PropagatingModes` for each lead.
811
812
        If a lead is a selfenergy lead, then the corresponding entry
        in lead_info is a selfenergy.
813
    out_leads, in_leads : sequence of integers
814
815
816
817
818
        indices of the leads where current is extracted (out) or injected
        (in). Only those are listed for which SMatrix contains the
        calculated result.
    """

819
820
    def __init__(self, data, lead_info, out_leads, in_leads,
                 current_conserving=False):
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
        # An equivalent condition to this is checked in 'Solver.smatrix',
        # but we add it here again as a sanity check. If 'lead_info' is not
        # a 'PropagatingModes' that means that the corresponding lead is a
        # selfenergy lead. Scattering matrix elements to/from a selfenergy lead
        # are not defined.
        assert not any(
            not isinstance(info, physics.PropagatingModes)
            and leadnum in (out_leads + in_leads)
            for leadnum, info in enumerate(lead_info)
        )

        # 'getattr' in order to handle self-energy leads, for which
        # 'info' is just a matrix (so no 'momenta').
        sizes = [
            len(getattr(info, "momenta", [])) // 2 for info in lead_info
        ]
Anton Akhmerov's avatar
Anton Akhmerov committed
837
        super().__init__(data, lead_info, out_leads, in_leads,
838
                                      sizes, current_conserving)
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
        # in_offsets marks beginnings and ends of blocks in the scattering
        # matrix corresponding to the in leads. The end of the block
        # corresponding to lead i is taken as the beginning of the block of
        # i+1.  Same for out leads. For each lead block of the scattering
        # matrix, we want to pick out the computed blocks of the conservation
        # law.  The offsets of these symmetry blocks are stored in
        # block_offsets, for all leads.  List of lists containing the sizes of
        # symmetry blocks in all leads. leads_block_sizes[i][j] is the number
        # of left or right moving modes in symmetry block j of lead
        # i. len(leads_block_sizes[i]) is the number of blocks in lead i.
        leads_block_sizes = []
        for info in self.lead_info:
            # If a lead does not have block structure, append None.
            leads_block_sizes.append(getattr(info, 'block_nmodes', None))
        self.leads_block_sizes = leads_block_sizes
        block_offsets = []
        for lead_block_sizes in self.leads_block_sizes: # Cover all leads
            if lead_block_sizes is None:
                block_offsets.append(lead_block_sizes)
            else:
                block_offset = np.zeros(len(lead_block_sizes) + 1, int)
                block_offset[1:] = np.cumsum(lead_block_sizes)
                block_offsets.append(block_offset)
        # Symmetry block offsets for all leads - or None if lead does not have
        # blocks.
        self.block_offsets = block_offsets
        # Pick out symmetry block offsets for in and out leads
        self.in_block_offsets = \
                np.array(self.block_offsets)[list(self.in_leads)]
        self.out_block_offsets = \
                np.array(self.block_offsets)[list(self.out_leads)]
        # Block j of in lead i starts at in_block_offsets[i][j]
871

872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
    def out_block_coords(self, lead_out):
        """Return a slice with the rows in the block corresponding to lead_out.
        """
        if isinstance(lead_out, Integral):
            lead_out = self.out_leads.index(lead_out)
            return slice(self.out_offsets[lead_out],
                         self.out_offsets[lead_out + 1])
        else:
            lead_ind, block_ind = lead_out
            lead_ind = self.out_leads.index(lead_ind)
            return slice(self.out_offsets[lead_ind] +
                         self.out_block_offsets[lead_ind][block_ind],
                         self.out_offsets[lead_ind] +
                         self.out_block_offsets[lead_ind][block_ind + 1])

    def in_block_coords(self, lead_in):
        """
        Return a slice with the columns in the block corresponding to lead_in.
        """
        if isinstance(lead_in, Integral):
            lead_in = self.in_leads.index(lead_in)
            return slice(self.in_offsets[lead_in],
                     self.in_offsets[lead_in + 1])
        else:
            lead_ind, block_ind = lead_in
            lead_ind = self.in_leads.index(lead_ind)
            return slice(self.in_offsets[lead_ind] +
                         self.in_block_offsets[lead_ind][block_ind],
                         self.in_offsets[lead_ind] +
                         self.in_block_offsets[lead_ind][block_ind + 1])
902
903

    def _transmission(self, lead_out, lead_in):
904
905
        return np.linalg.norm(self.submatrix(lead_out, lead_in)) ** 2

906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
    def transmission(self, lead_out, lead_in):
        """Return transmission from lead_in to lead_out.

        If `lead_in` or `lead_out` are a length-2 sequence, the first number is
        the number of the lead, and the second number indexes the eigenvalue of
        the conserved quantity in that lead (e.g. spin) if it was specified.
        """
        # If transmission is between leads and not subblocks, we can use
        # unitarity (like BlockResult does).
        if isinstance(lead_in, Integral) and isinstance(lead_out, Integral):
            return super().transmission(lead_out, lead_in)
        else:
            return self._transmission(lead_out, lead_in)

    def num_propagating(self, lead):
        """Return the number of propagating modes in the lead."""
        return self.sizes[lead]

924
    def __repr__(self):
925
926
        return ("SMatrix(data=%r, lead_info=%r, out_leads=%r, in_leads=%r)" %
                (self.data, self.lead_info, self.out_leads, self.in_leads))
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944


class GreensFunction(BlockResult):
    """
    Retarded Green's function.

    Transport properties can be easily accessed using the
    `~GreensFunction.transmission` method (don't be fooled by the name, it can
    also compute reflection, which is just transmission from one lead back into
    the same lead).

    `GreensFunction` however also allows for a more direct access to the
    result: The data stored in `GreensFunction` is the real-space Green's
    function. The details of this data can be directly accessed through the
    instance variables `data` and `lead_info`. Subblocks of data corresponding
    to particular leads are conveniently obtained by
    `~GreensFunction.submatrix`.

945
946
    Attributes
    ----------
947
948
949
950
951
    data : NumPy array
        a matrix containing all the requested matrix elements of Green's
        function.
    lead_info : list of matrices
        a list with self-energies of each lead.
952
    out_leads, in_leads : sequence of integers
953
954
955
956
957
        indices of the leads where current is extracted (out) or injected
        (in). Only those are listed for which SMatrix contains the
        calculated result.
    """

958
959
    def __init__(self, data, lead_info, out_leads, in_leads,
                 current_conserving=False):
960
        sizes = [i.shape[0] for i in lead_info]
Anton Akhmerov's avatar
Anton Akhmerov committed
961
        super().__init__(
962
963
964
965
966
967
968
969
970
971
972
973
974
975
            data, lead_info, out_leads, in_leads, sizes, current_conserving)

    def num_propagating(self, lead):
        """Return the number of propagating modes in the lead."""
        gamma = 1j * (self.lead_info[lead] -
                      self.lead_info[lead].conj().T)

        # The number of channels is given by the number of
        # nonzero eigenvalues of Gamma
        # rationale behind the threshold from
        # Golub; van Loan, chapter 5.5.8
        eps = np.finfo(gamma.dtype).eps * 1000
        return np.sum(np.linalg.eigvalsh(gamma) >
                      eps * np.linalg.norm(gamma, np.inf))
976

977
    def _a_ttdagger_a_inv(self, lead_out, lead_in):
978
        """Return t * t^dagger in a certain basis."""
979
980
981
982
        gf = self.submatrix(lead_out, lead_in)
        factors = []
        for lead, gf2 in ((lead_out, gf), (lead_in, gf.conj().T)):
            possible_se = self.lead_info[lead]
983
            factors.append(1j * (possible_se - possible_se.conj().T))
984
985
986
            factors.append(gf2)
        return reduce(np.dot, factors)

987
    def _transmission(self, lead_out, lead_in):
988
        attdagainv = self._a_ttdagger_a_inv(lead_out, lead_in)
989
990
991
        result = np.trace(attdagainv).real
        if lead_out == lead_in:
            # For reflection we have to be more careful
992
            sigma = self.lead_info[lead_in]
993
            gamma = 1j * (sigma - sigma.conj().T)
994
995
996
997
998
999
            gf = self.submatrix(lead_out, lead_in)

            # The number of channels is given by the number of
            # nonzero eigenvalues of Gamma
            # rationale behind the threshold from
            # Golub; van Loan, chapter 5.5.8
1000
1001
            # We use ‖Σ‖, not ‖Γ‖, for the tolerance as ‖Γ‖~0 when there
            # are no open modes.
1002
            eps = (
1003
                1e6 * np.finfo(gamma.dtype).eps * np.linalg.norm(sigma, np.inf)
1004
1005
            )
            N = np.sum(np.linalg.eigvalsh(gamma) > eps)
1006
1007
1008
1009

            result += 2 * np.trace(np.dot(gamma, gf)).imag + N

        return result
Michael Wimmer's avatar
Michael Wimmer committed
1010
1011

    def __repr__(self):
1012
1013
1014
        return ("GreensFunction(data=%r, lead_info=%r, "
                "out_leads=%r, in_leads=%r)" %
                (self.data, self.lead_info, self.out_leads, self.in_leads))