leads.py 52.1 KB
Newer Older
Tómas's avatar
Tómas committed
1
# Copyright 2011-2016 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.

Joseph Weston's avatar
Joseph Weston committed
9

10
from math import sin, cos, sqrt, pi, copysign
11
from collections import namedtuple
Joseph Weston's avatar
Joseph Weston committed
12

Tómas's avatar
Tómas committed
13
from itertools import combinations_with_replacement
14
15
16
import numpy as np
import numpy.linalg as npl
import scipy.linalg as la
17
from .. import linalg as kla
Tómas's avatar
Tómas committed
18
19
20
from scipy.linalg import block_diag
from scipy.sparse import (identity as sp_identity, hstack as sp_hstack,
                          csr_matrix)
21
22
23

dot = np.dot

24
__all__ = ['selfenergy', 'modes', 'PropagatingModes', 'StabilizedModes']
25

Anton Akhmerov's avatar
Anton Akhmerov committed
26

Tómas's avatar
Tómas committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
# TODO: Use scipy block_diag once we depend on scipy>=0.19
try:
    # Throws ValueError, but if fixed ensure that works as intended
    bdiag_broken = block_diag(np.zeros((1,1)), np.zeros((2,0))).shape != (3, 1)
except ValueError:  # skip coverage
    bdiag_broken = True
if bdiag_broken:  # skip coverage
    def block_diag(*matrices):
        """Construct a block diagonal matrix out of the input matrices.

        Like scipy.linalg.block_diag, but works for zero size matrices."""
        rows, cols = np.sum([mat.shape for mat in matrices], axis=0)
        b_mat = np.zeros((rows,cols), dtype='complex')
        rows, cols = 0, 0
        for mat in matrices:
            new_rows = rows + mat.shape[0]
            new_cols = cols + mat.shape[1]
            b_mat[rows:new_rows, cols:new_cols] = mat
            rows, cols = new_rows, new_cols
        return b_mat


49
# TODO: Remove the workaround once we depend on scipy >= 1.0
Tómas's avatar
Tómas committed
50
51
52
53
def lstsq(a, b):
    """Least squares version that works also with 0-shaped matrices."""
    if a.shape[1] == 0:
        return np.empty((0, 0), dtype=np.common_type(a, b))
54
    return la.lstsq(a, b)[0]
Tómas's avatar
Tómas committed
55
56


Tómas's avatar
Tómas committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def nonzero_symm_projection(matrix):
    """Check whether a discrete symmetry relation between two blocks of the
    Hamiltonian vanishes or not.

    For a discrete symmetry S, the projection that maps from the block i to
    block j of the Hamiltonian with projectors P_i and P_j is
    S_ji = P_j^+ S P_i.
    This function determines whether S_ji vanishes or not, i. e. whether a
    symmetry relation exists between blocks i and j.

    Here, discrete symmetries and projectors are in canonical form, such that
    S maps each block at most either to itself or to a single other block.
    In other words, for each j, the block S_ji is nonzero at most for one i,
    while all other blocks vanish. If a nonzero block exists, it is unitary
    and hence has norm 1. To identify nonzero blocks without worrying about
    numerical errors, we thus check that its norm is larger than 0.5.
    """
    if not isinstance(matrix, np.ndarray):
        matrix = matrix.data
    return np.linalg.norm(matrix) > 0.5


def group_halves(arr_list):
    """Split and rearrange a list of arrays.

    Combine a list of arrays into a single array where the first halves
    of each array appear first:
    `[a b], [c d], [e f] -> [a c e b d f]`
    """
86
    list_ = [np.split(arr, 2) for arr in arr_list]
Tómas's avatar
Tómas committed
87
88
89
90
    lefts, rights = zip(*list_)
    return np.r_[tuple(lefts + rights)]


91
92
93
94
# Container classes
Linsys = namedtuple('Linsys', ['eigenproblem', 'v', 'extract'])


95
class PropagatingModes:
96
    """The calculated propagating modes of a lead.
97

98
99
    Attributes
    ----------
100
    wave_functions : numpy array
101
        The wave functions of the propagating modes.
102
103
104
105
    momenta : numpy array
        Momenta of the modes.
    velocities : numpy array
        Velocities of the modes.
Tómas's avatar
Tómas committed
106
107
108
    block_nmodes: list of integers
        Number of left or right moving propagating modes
        per conservation law block of the Hamiltonian.
109
110
111
112

    Notes
    =====
    The sort order of all the three arrays is identical. The first half of the
Joseph Weston's avatar
Joseph Weston committed
113
114
115
116
117
118
    modes have negative velocity, the second half have positive velocity.
    Within these halves the modes are ordered by the eigenvalues of any
    declared conservation laws. Within blocks with the same conservation law
    eigenvalue the modes with negative velocity are ordered by increasing
    momentum, and the modes with positive velocity are ordered by decreasing
    momentum. Finally, modes are ordered by the magnitude of their velocity.
119
120
121
122
    To summarize, the modes are ordered according to the key
    `(sign(v), conserved_quantity, sign(v) * k , abs(v))` where `v` is
    velocity, `k` is momentum and `conserved_quantity` is the conservation
    law eigenvalue.
123

124
125
126
    In the above, the positive velocity and momentum directions are defined
    with respect to the translational symmetry direction of the system.

Christoph Groth's avatar
Christoph Groth committed
127
128
129
130
131
    The first dimension of `wave_functions` corresponds to the orbitals of all
    the sites in a unit cell, the second one to the number of the mode.  Each
    mode is normalized to carry unit current. If several modes have the same
    momentum and velocity, an arbitrary orthonormal basis in the subspace of
    these modes is chosen.
Tómas's avatar
Tómas committed
132

133
134
135
136
    If a conservation law is specified to block diagonalize the Hamiltonian,
    then `block_nmodes[i]` is the number of left or right moving propagating
    modes in conservation law block `i`. The ordering of blocks is the same as
    the ordering of the projectors used to block diagonalize the Hamiltonian.
137
138
139
140
141
142
143
    """
    def __init__(self, wave_functions, velocities, momenta):
        kwargs = locals()
        kwargs.pop('self')
        self.__dict__.update(kwargs)


144
class StabilizedModes:
145
146
147
    """Stabilized eigendecomposition of the translation operator.

    Due to the lack of Hermiticity of the translation operator, its
148
    eigendecomposition is frequently poorly conditioned. Solvers in Kwant use
149
150
151
152
153
154
155
156
157
158
159
160
    this stabilized decomposition of the propagating and evanescent modes in
    the leads. If the hopping between the unit cells of an infinite system is
    invertible, the translation eigenproblem is written in the basis `psi_n,
    h_hop^+ * psi_(n+1)`, with ``h_hop`` the hopping between unit cells.  If
    `h_hop` is not invertible, and has the singular value decomposition `u s
    v^+`, then the eigenproblem is written in the basis `sqrt(s) v^+ psi_n,
    sqrt(s) u^+ psi_(n+1)`. In this basis we calculate the eigenvectors of the
    propagating modes, and the Schur vectors (an orthogonal basis) in the space
    of evanescent modes.

    `vecs` and `vecslmbdainv` are the first and the second halves of the wave
    functions.  The first `nmodes` are eigenmodes moving in the negative
161
    direction (hence they are incoming into the system in Kwant convention),
162
163
164
165
166
167
    the second `nmodes` are eigenmodes moving in the positive direction. The
    remaining modes are the Schur vectors of the modes evanescent in the
    positive direction. Propagating modes with the same eigenvalue are
    orthogonalized, and all the propagating modes are normalized to carry unit
    current. Finally the `sqrt_hop` attribute is `v sqrt(s)`.

168
169
    Attributes
    ----------
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
    vecs : numpy array
        Translation eigenvectors.
    vecslmbdainv : numpy array
        Translation eigenvectors divided by the corresponding eigenvalues.
    nmodes : int
        Number of left-moving (or right-moving) modes.
    sqrt_hop : numpy array or None
        Part of the SVD of `h_hop`, or None if the latter is invertible.
    """

    def __init__(self, vecs, vecslmbdainv, nmodes, sqrt_hop=None):
        kwargs = locals()
        kwargs.pop('self')
        self.__dict__.update(kwargs)

185
186
187
188
189
190
191
192
193
194
195
196
197
198
    def selfenergy(self):
        """
        Compute the self-energy generated by lead modes.

        Returns
        -------
        Sigma : numpy array, real or complex, shape (M,M)
            The computed self-energy. Note that even if `h_cell` and `h_hop` are
            both real, `Sigma` will typically be complex. (More precisely, if
            there is a propagating mode, `Sigma` will definitely be complex.)
        """
        v = self.sqrt_hop
        vecs = self.vecs[:, self.nmodes:]
        vecslmbdainv = self.vecslmbdainv[:, self.nmodes:]
Tómas's avatar
Tómas committed
199
        return dot(v, dot(vecs, la.solve(vecslmbdainv, v.T.conj())))
200

201
202

# Auxiliary functions that perform different parts of the calculation.
203
204
def setup_linsys(h_cell, h_hop, tol=1e6, stabilization=None):
    """Make an eigenvalue problem for eigenvectors of translation operator.
205
206
207

    Parameters
    ----------
208
209
    h_cell : numpy array with shape (n, n)
        Hamiltonian of a single lead unit cell.
210
    h_hop : numpy array with shape (n, m), m <= n
211
        Hopping Hamiltonian from a cell to the next one.
212
213
214
    tol : float
        Numbers are considered zero when they are smaller than `tol` times
        the machine precision.
215
    stabilization : sequence of 2 booleans or None
Christoph Groth's avatar
Christoph Groth committed
216
        Which steps of the eigenvalue problem stabilization to perform. If the
217
        value is `None`, then Kwant chooses the fastest (and least stable)
218
        algorithm that is expected to be sufficient.  For any other value,
219
220
        Kwant forms the eigenvalue problem in the basis of the hopping singular
        values.  The first element set to `True` forces Kwant to add an
221
222
223
        anti-Hermitian term to the cell Hamiltonian before inverting. If it is
        set to `False`, the extra term will only be added if the cell
        Hamiltonian isn't invertible. The second element set to `True` forces
224
        Kwant to solve a generalized eigenvalue problem, and not to reduce it
225
        to the regular one.  If it is `False`, reduction to a regular problem
226
        is performed if possible.
227
228
229

    Returns
    -------
230
    linsys : namedtuple
231
232
233
234
235
        A named tuple containing `matrices` a matrix pencil defining
        the eigenproblem, `v` a hermitian conjugate of the last matrix in
        the hopping singular value decomposition, and functions for
        extracting the wave function in the unit cell from the wave function
        in the basis of the nonzero singular exponents of the hopping.
236
237
238

    Notes
    -----
239
240
    The lead problem with degenerate hopping is rather complicated, and the
    details of the algorithm will be published elsewhere.
241

242
    """
243
    n = h_cell.shape[0]
244
    m = h_hop.shape[1]
245
246
    if stabilization is not None:
        stabilization = list(stabilization)
247

248
    if not np.any(h_hop):  # skip coverage
249
        # Inter-cell hopping is zero.  The current algorithm is not suited to
250
        # treat this extremely singular case.
251
        raise ValueError("Inter-cell hopping is exactly zero.")
252

253
    # If both h and t are real, it may be possible to use the real eigenproblem.
254
    if (not np.any(h_hop.imag)) and (not np.any(h_cell.imag)):
255
        h_hop = h_hop.real
256
        h_cell = h_cell.real
257

258
    eps = np.finfo(np.common_type(h_cell, h_hop)).eps * tol
259

260
261
262
    # First check if the hopping matrix has singular values close to 0.
    # (Close to zero is defined here as |x| < eps * tol * s[0] , where
    # s[0] is the largest singular value.)
263

264
    u, s, vh = la.svd(h_hop)
265
    assert m == vh.shape[1], "Corrupt output of svd."
266
    n_nonsing = np.sum(s > eps * s[0])
267

268
    if (n_nonsing == n and stabilization is None):
269
        # The hopping matrix is well-conditioned and can be safely inverted.
270
        # Hence the regular transfer matrix may be used.
271
272
273
274
275
        h_hop_sqrt = sqrt(np.linalg.norm(h_hop))
        A = h_hop / h_hop_sqrt
        B = h_hop_sqrt
        B_H_inv = 1.0 / B     # just a real scalar here
        A_inv = la.inv(A)
276

277
278
279
280
        lhs = np.zeros((2*n, 2*n), dtype=np.common_type(h_cell, h_hop))
        lhs[:n, :n] = -dot(A_inv, h_cell) * B_H_inv
        lhs[:n, n:] = -A_inv * B
        lhs[n:, :n] = A.T.conj() * B_H_inv
281
282

        def extract_wf(psi, lmbdainv):
283
            return B_H_inv * np.copy(psi[:n])
284

285
286
        matrices = (lhs, None)
        v_out = h_hop_sqrt * np.eye(n)
287
    else:
288
289
        if stabilization is None:
            stabilization = [None, False]
290

291
292
293
294
295
        # The hopping matrix has eigenvalues close to 0 - those
        # need to be eliminated.

        # Recast the svd of h_hop = u s v^dagger such that
        # u, v are matrices with shape n x n_nonsing.
Anton Akhmerov's avatar
Anton Akhmerov committed
296
297
        u = u[:, :n_nonsing]
        s = s[:n_nonsing]
298
        u = u * np.sqrt(s)
299
300
        # pad v with zeros if necessary
        v = np.zeros((n, n_nonsing), dtype=vh.dtype)
301
        v[:vh.shape[1]] = vh[:n_nonsing].T.conj()
302
        v = v * np.sqrt(s)
303

304
305
306
307
308
309
310
311
312
313
314
        # Eliminating the zero eigenvalues requires inverting the on-site
        # Hamiltonian, possibly including a self-energy-like term.  The
        # self-energy-like term stabilizes the inversion, but the most stable
        # choice is inherently complex. This can be disadvantageous if the
        # Hamiltonian is real, since staying in real arithmetics can be
        # significantly faster.  The strategy here is to add a complex
        # self-energy-like term always if the original Hamiltonian is complex,
        # and check for invertibility first if it is real

        matrices_real = issubclass(np.common_type(h_cell, h_hop), np.floating)
        add_imaginary = stabilization[0] or ((stabilization[0] is None) and
315
                                             not matrices_real)
316
317
318
319
        # Check if there is a chance we will not need to add an imaginary term.
        if not add_imaginary:
            h = h_cell
            sol = kla.lu_factor(h)
320
321
            rcond = kla.rcond_from_lu(sol, npl.norm(h, 1))

322
323
324
            if rcond < eps:
                need_to_stabilize = True
            else:
325
                need_to_stabilize = False
326

327
        if add_imaginary or need_to_stabilize:
328
            need_to_stabilize = True
329
            # Matrices are complex or need self-energy-like term to be
330
            # stabilized.
331
            temp = dot(u, u.T.conj()) + dot(v, v.T.conj())
332
            h = h_cell + 1j * temp
333
334
335
336
337
338

            sol = kla.lu_factor(h)
            rcond = kla.rcond_from_lu(sol, npl.norm(h, 1))

            # If the condition number of the stabilized h is
            # still bad, there is nothing we can do.
339
            if rcond < eps:
340
341
342
                raise RuntimeError("Flat band encountered at the requested "
                                   "energy, result is badly defined.")

343
344
        # The function that can extract the full wave function psi from
        # the projected one (v^dagger psi lambda^-1, s u^dagger psi).
345
346

        def extract_wf(psi, lmbdainv):
347
            wf = -dot(u, psi[: n_nonsing] * lmbdainv) - dot(v, psi[n_nonsing:])
348
349
350
351
            if need_to_stabilize:
                wf += 1j * (dot(v, psi[: n_nonsing]) +
                            dot(u, psi[n_nonsing:] * lmbdainv))
            return kla.lu_solve(sol, wf)
352
353
354

        # Setup the generalized eigenvalue problem.

355
356
        A = np.zeros((2 * n_nonsing, 2 * n_nonsing), np.common_type(h, h_hop))
        B = np.zeros((2 * n_nonsing, 2 * n_nonsing), np.common_type(h, h_hop))
357

358
359
        begin, end = slice(n_nonsing), slice(n_nonsing, None)

Anton Akhmerov's avatar
Anton Akhmerov committed
360
        A[end, begin] = np.identity(n_nonsing)
361
        temp = kla.lu_solve(sol, v)
362
        temp2 = dot(u.T.conj(), temp)
363
        if need_to_stabilize:
Anton Akhmerov's avatar
Anton Akhmerov committed
364
365
            A[begin, begin] = -1j * temp2
        A[begin, end] = temp2
366
        temp2 = dot(v.T.conj(), temp)
367
        if need_to_stabilize:
Anton Akhmerov's avatar
Anton Akhmerov committed
368
369
            A[end, begin] -= 1j *temp2
        A[end, end] = temp2
370

Anton Akhmerov's avatar
Anton Akhmerov committed
371
        B[begin, end] = -np.identity(n_nonsing)
372
373
        temp = kla.lu_solve(sol, u)
        temp2 = dot(u.T.conj(), temp)
Anton Akhmerov's avatar
Anton Akhmerov committed
374
        B[begin, begin] = -temp2
375
        if need_to_stabilize:
Anton Akhmerov's avatar
Anton Akhmerov committed
376
            B[begin, end] += 1j * temp2
377
        temp2 = dot(v.T.conj(), temp)
Anton Akhmerov's avatar
Anton Akhmerov committed
378
        B[end, begin] = -temp2
379
        if need_to_stabilize:
Anton Akhmerov's avatar
Anton Akhmerov committed
380
            B[end, end] = 1j * temp2
381
382

        v_out = v[:m]
383
384
385
386
387
388
389
390
391
392

        # Solving a generalized eigenproblem is about twice as expensive
        # as solving a regular eigenvalue problem.
        # Computing the LU factorization is negligible compared to both
        # (approximately 1/30th of a regular eigenvalue problem).
        # Because of this, it makes sense to try to reduce
        # the generalized eigenvalue problem to a regular one, provided
        # the matrix B can be safely inverted.

        lu_b = kla.lu_factor(B)
393
        if not stabilization[1]:
394
395
396
            rcond = kla.rcond_from_lu(lu_b, npl.norm(B, 1))
            # A more stringent condition is used here since errors can
            # accumulate from here to the eigenvalue calculation later.
397
            stabilization[1] = rcond > eps * tol
398

399
        if stabilization[1]:
400
            matrices = (kla.lu_solve(lu_b, A), None)
401
        else:
402
            matrices = (A, B)
403
404
405
406
407
408
409
410
    return Linsys(matrices, v_out, extract_wf)


def unified_eigenproblem(a, b=None, tol=1e6):
    """A helper routine for modes(), that wraps eigenproblems.

    This routine wraps the regular and general eigenproblems that can arise
    in a unfied way.
411

412
413
414
415
416
417
418
419
420
421
    Parameters
    ----------
    a : numpy array
        The matrix on the left hand side of a regular or generalized eigenvalue
        problem.
    b : numpy array or None
        The matrix on the right hand side of the generalized eigenvalue problem.
    tol : float
        The tolerance for separating eigenvalues with absolute value 1 from the
        rest.
422

423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
    Returns
    -------
    ev : numpy array
        An array of eigenvalues (can contain NaNs and Infs, but those
        are not accessed in `modes()`) The number of eigenvalues equals
        twice the number of nonzero singular values of
        `h_hop` (so `2*h_cell.shape[0]` if `h_hop` is invertible).
    evanselect : numpy integer array
        Index array of right-decaying modes.
    propselect : numpy integer array
        Index array of propagating modes (both left and right).
    vec_gen(select) : function
        A function that computes the eigenvectors chosen by the array select.
    ord_schur(select) : function
        A function that computes the unitary matrix (corresponding to the right
        eigenvector space) of the (general) Schur decomposition reordered such
        that the eigenvalues chosen by the array select are in the top left
        block.
441
    """
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
    if b is None:
        eps = np.finfo(a.dtype).eps * tol
        t, z, ev = kla.schur(a)

        # Right-decaying modes.
        select = abs(ev) > 1 + eps
        # Propagating modes.
        propselect = abs(abs(ev) - 1) < eps

        vec_gen = lambda x: kla.evecs_from_schur(t, z, select=x)
        ord_schur = lambda x: kla.order_schur(x, t, z, calc_ev=False)[1]

    else:
        eps = np.finfo(np.common_type(a, b)).eps * tol
        s, t, z, alpha, beta = kla.gen_schur(a, b, calc_q=False)

        # Right-decaying modes.
        select = abs(alpha) > (1 + eps) * abs(beta)
        # Propagating modes.
        propselect = (abs(abs(alpha) - abs(beta)) < eps * abs(beta))

463
464
        with np.errstate(divide='ignore', invalid='ignore'):
            ev = alpha / beta
465
466
467
468
469
470
471
472
473
474
        # Note: the division is OK here, since we later only access
        #       eigenvalues close to the unit circle

        vec_gen = lambda x: kla.evecs_from_gen_schur(s, t, z=z, select=x)
        ord_schur = lambda x: kla.order_gen_schur(x, s, t, z=z,
                                                  calc_ev=False)[2]

    return ev, select, propselect, vec_gen, ord_schur


475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
def phs_symmetrization(wfs, particle_hole):
    """Makes the wave functions that have the same velocity at a time-reversal
    invariant momentum (TRIM) particle-hole symmetric.

    If P is the particle-hole operator and P^2 = 1, then a particle-hole
    symmetric wave function at a TRIM is an eigenstate of P with eigenvalue 1.
    If P^2 = -1, wave functions with the same velocity at a TRIM come in pairs.
    Such a pair is particle-hole symmetric if the wave functions are related by
    P, i. e. the pair can be expressed as [psi_n, P psi_n] where psi_n is a wave
    function.

    To ensure proper ordering of modes, this function also returns an array
    of indices which ensures that particle-hole partners are properly ordered
    in this subspace of modes. These are later used with np.lexsort to ensure
    proper ordering.

    Parameters
    ----------
    wfs : numpy array
        A matrix of propagating wave functions at a TRIM that all have the same
495
        velocity. The orthonormal wave functions form the columns of this matrix.
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
    particle_hole : numpy array
        The matrix representation of the unitary part of the particle-hole
        operator, expressed in the tight binding basis.

    Returns
    -------
    new_wfs : numpy array
        The matrix of particle-hole symmetric wave functions.
    TRIM_sort: numpy integer array
        Index array that stores the proper sort order of particle-hole symmetric
        wave functions in this subspace.
    """

    def Pdot(mat):
        """Apply the particle-hole operator to an array. """
        return particle_hole.dot(mat.conj())

513
514
515
516
517
518
519
520
521
522
523
524
525
    # Take P in the subspace of W = wfs: U = W^+ @ P @ W^*.
    U = wfs.T.conj().dot(Pdot(wfs))
    # Check that wfs are orthonormal and the space spanned
    # by them is closed under ph, meaning U is unitary.
    if not np.allclose(U.dot(U.T.conj()), np.eye(U.shape[0])):
        raise ValueError('wfs are not orthonormal or not closed under particle_hole.')
    P_squared = U.dot(U.conj())
    if np.allclose(P_squared, np.eye(U.shape[0])):
        P_squared = 1
    elif np.allclose(P_squared, -np.eye(U.shape[0])):
        P_squared = -1
    else:
        raise ValueError('particle_hole must square to +1 or -1. P_squared = {}'.format(P_squared))
526
527

    if P_squared == 1:
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
        # Use the matrix square root method from
        # Applied Mathematics and Computation 234 (2014) 380-384.
        assert np.allclose(U, U.T)
        # Schur decomposition guarantees that vecs are orthonormal.
        vals, vecs = la.schur(U)
        # U should be normal, so vals is diagonal.
        assert np.allclose(np.diag(np.diag(vals)), vals)
        vals = np.diag(vals)
        # Need to take safe square root of U, the branch cut should not go
        # through any eigenvalues. Otherwise the square root may not be symmetric.
        # Find largest gap between eigenvalues
        phases = np.sort(np.angle(vals))
        dph = np.append(np.diff(phases), phases[0] + 2*np.pi - phases[-1])
        i = np.argmax(dph)
        shift = -np.pi - (phases[i] + dph[i]/2)
        # Take matrix square root with branch cut in largest gap
        vals = np.sqrt(vals * np.exp(1j * shift)) * np.exp(-0.5j * shift)
        sqrtU = vecs.dot(np.diag(vals)).dot(vecs.T.conj())
        # For symmetric U sqrt(U) is also symmetric.
        assert np.allclose(sqrtU, sqrtU.T)
        # We want a new basis W_new such that W_new^+ @ P @ W_new^* = 1.
        # This is achieved by W_new = W @ sqrt(U).
        new_wfs = wfs.dot(sqrtU)
551
552
        # If P^2 = 1, there is no need to sort the modes further.
        TRIM_sort = np.zeros((wfs.shape[1],), dtype=int)
Tómas's avatar
Tómas committed
553
    else:
Tómas's avatar
Tómas committed
554
555
        # P^2 = -1.
        # Iterate over wave functions to construct
556
        # particle-hole partners.
Tómas's avatar
Tómas committed
557
558
559
560
561
562
        new_wfs = []
        # The number of modes. This is always an even number >=2.
        N_modes = wfs.shape[1]
        # If there are only two modes in this subspace, they are orthogonal
        # so we replace the second one with the P applied to the first one.
        if N_modes == 2:
563
            wf = wfs[:, 0]
Tómas's avatar
Tómas committed
564
565
566
567
568
569
570
571
572
573
574
575
            # Store psi_n and P psi_n.
            new_wfs.append(wf)
            new_wfs.append(Pdot(wf))
        # If there are more than two modes, iterate over wave functions
        # and construct their particle-hole partners one by one.
        else:
            # We construct pairs of modes that are particle-hole partners.
            # Need to iterate over all pairs except the final one.
            iterations = range((N_modes-2)//2)
            for i in iterations:
                # Take a mode psi_n from the basis - the first column
                # of the matrix of remaining modes.
576
                wf = wfs[:, 0]
577
578
                # Store psi_n and P psi_n.
                new_wfs.append(wf)
Tómas's avatar
Tómas committed
579
580
581
582
                P_wf = Pdot(wf)
                new_wfs.append(P_wf)
                # Remove psi_n and P psi_n from the basis matrix of modes.
                # First remove psi_n.
583
                wfs = wfs[:, 1:]
Tómas's avatar
Tómas committed
584
                # Now we project the remaining modes onto the orthogonal
585
586
                # complement of P psi_n. projector:
                projector = wfs.dot(wfs.T.conj()) - \
Tómas's avatar
Tómas committed
587
588
589
590
                            np.outer(P_wf, P_wf.T.conj())
                # After the projection, the mode matrix is rank deficient -
                # the span of the column space has dimension one less than
                # the number of columns.
591
                wfs = projector.dot(wfs)
Tómas's avatar
Tómas committed
592
593
594
595
596
597
598
599
                wfs = la.qr(wfs, mode='economic', pivoting=True)[0]
                # Remove the redundant column.
                wfs = wfs[:, :-1]
                # If this is the final iteration, we only have two modes
                # left and can construct particle-hole partners without
                # the projection.
                if i == iterations[-1]:
                    assert wfs.shape[1] == 2
600
                    wf = wfs[:, 0]
601
602
                    # Store psi_n and P psi_n.
                    new_wfs.append(wf)
603
                    new_wfs.append(Pdot(wf))
Tómas's avatar
Tómas committed
604
                assert np.allclose(wfs.T.conj().dot(wfs),
605
                                   np.eye(wfs.shape[1]))
Tómas's avatar
Tómas committed
606
607
608
609
610
        new_wfs = np.hstack([col.reshape(len(col), 1)/npl.norm(col) for
                             col in new_wfs])
        assert np.allclose(new_wfs[:, 1::2], Pdot(new_wfs[:, ::2]))
        # Store sort ordering in this subspace of modes
        TRIM_sort = np.arange(new_wfs.shape[1])
611
612
613
614
    assert np.allclose(new_wfs.T.conj().dot(new_wfs), np.eye(new_wfs.shape[1]))
    return new_wfs, TRIM_sort


Tómas's avatar
Tómas committed
615
616
def make_proper_modes(lmbdainv, psi, extract, tol, particle_hole,
                      time_reversal, chiral):
617
618
    """
    Find, normalize and sort the propagating eigenmodes.
619
620
621
622
623

    Special care is taken of the case of degenerate k-values, where the
    numerically computed modes are typically a superposition of the real
    modes. In this case, also the proper (orthogonal) modes are computed.
    """
624
    vel_eps = np.finfo(psi.dtype).eps * tol
625
626

    nmodes = psi.shape[1]
627
    n = len(psi) // 2
628
629

    # Array for the velocities.
630
    velocities = np.empty(nmodes, dtype=float)
631

632
633
634
    # Array of indices to sort modes at a TRIM by PHS.
    TRIM_PHS_sort = np.zeros(nmodes, dtype=int)

635
636
    # Calculate the full wave function in real space.
    full_psi = extract(psi, lmbdainv)
637

638
639
640
641
642
643
644
    # Cast the types if any of the symmetry operators is complex
    for symm in time_reversal, particle_hole, chiral:
        if symm is None:
            continue
        full_psi = full_psi.astype(np.common_type(symm, full_psi))
        psi = psi.astype(np.common_type(symm, psi))

645
646
647
    # Find clusters of nearby eigenvalues. Since the eigenvalues occupy the
    # unit circle, special care has to be taken to not introduce a cut at
    # lambda = -1.
648
649
650
651
652
    eps = np.finfo(lmbdainv.dtype).eps * tol
    angles = np.angle(lmbdainv)
    sort_order = np.resize(np.argsort(angles), (2 * len(angles,)))
    boundaries = np.argwhere(np.abs(np.diff(lmbdainv[sort_order]))
                             > eps).flatten() + 1
653

654
655
656
657
    # Detect the singular case of all eigenvalues equal.
    if boundaries.shape == (0,) and len(angles):
        boundaries = np.array([0, len(angles)])

Joseph Weston's avatar
Joseph Weston committed
658
    for interval in zip(boundaries[:-1], boundaries[1:]):
659
660
        if interval[1] > boundaries[0] + len(angles):
            break
661

662
        indx = sort_order[interval[0] : interval[1]]
663

664
665
666
667
668
669
        # If there is a degenerate eigenvalue with several different
        # eigenvectors, the numerical routines return some arbitrary
        # overlap of the real, physical solutions. In order
        # to figure out the correct wave function, we need to
        # have the full, not the projected wave functions
        # (at least to our current knowledge).
670

671
        # Finding the true modes is done in two steps:
672

673
674
675
676
677
678
679
680
        # 1. The true transversal modes should be orthogonal to each other, as
        # they share the same Bloch momentum (note that transversal modes with
        # different Bloch momenta k1 and k2 need not be orthogonal, the full
        # modes are orthogonal because of the longitudinal dependence e^{i k1
        # x} and e^{i k2 x}).  The modes with the same k are therefore
        # orthogonalized. Moreover for the velocity to have a proper value the
        # modes should also be normalized.

681
        q, r = la.qr(full_psi[:, indx], mode='economic')
682

683
684
685
686
687
688
        # If the eigenvectors were purely real up to this stage,
        # they will typically become complex after the rotation.
        if psi.dtype != np.common_type(psi, r):
            psi = psi.astype(np.common_type(psi, r))
        if full_psi.dtype != np.common_type(full_psi, q):
            full_psi = full_psi.astype(np.common_type(psi, q))
689

690
691
        full_psi[:, indx] = q
        psi[:, indx] = la.solve(r.T, psi[:, indx].T).T
692

693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
        # 2. Moving infinitesimally away from the degeneracy
        # point, the modes should diagonalize the velocity
        # operator (i.e. when they are non-degenerate any more)
        # The modes are therefore rotated properly such that they
        # diagonalize the velocity operator.
        # Note that step 2. does not give a unique result if there are
        # two modes with the same velocity, or if the modes stay
        # degenerate even for a range of Bloch momenta (and hence
        # must have the same velocity). However, this does not matter,
        # since we are happy with any superposition in this case.

        vel_op = -1j * dot(psi[n:, indx].T.conj(), psi[:n, indx])
        vel_op = vel_op + vel_op.T.conj()
        vel_vals, rot = la.eigh(vel_op)

        # If the eigenvectors were purely real up to this stage,
        # they will typically become complex after the rotation.

        if psi.dtype != np.common_type(psi, rot):
            psi = psi.astype(np.common_type(psi, rot))
713
714
715
        if full_psi.dtype != np.common_type(full_psi, rot):
            full_psi = full_psi.astype(np.common_type(psi, rot))

716
        psi[:, indx] = dot(psi[:, indx], rot)
717
        full_psi[:, indx] = dot(full_psi[:, indx], rot)
718
719
        velocities[indx] = vel_vals

720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
        # With particle-hole symmetry, treat TRIMs individually.
        # Particle-hole conserves velocity.
        # If P^2 = 1, we can pick modes at a TRIM as particle-hole eigenstates.
        # If P^2 = -1, a mode at a TRIM and its particle-hole partner are
        # orthogonal, and we pick modes such that they are related by
        # particle-hole symmetry.

        # At a TRIM, propagating translation eigenvalues are +1 or -1.
        if (particle_hole is not None and
            (np.abs(np.abs(lmbdainv[indx].real) - 1) < eps).all()):
            assert not len(indx) % 2
            # Set the eigenvalues to the exact TRIM values.
            if (np.abs(lmbdainv[indx].real - 1) < eps).all():
                lmbdainv[indx] = 1
            else:
Tómas's avatar
Tómas committed
735
736
737
738
                # Momenta are the negative arguments of the translation
                # eigenvalues, as computed below using np.angle. np.angle of -1
                # is pi, so this assigns k = -pi to modes with translation
                # eigenvalue -1.
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
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
                lmbdainv[indx] = -1

            # Original wave functions
            orig_wf = full_psi[:, indx]

            # Modes are already sorted by velocity in ascending order, as
            # returned by la.eigh above. The first half is thus incident,
            # and the second half outgoing.
            # Here we work within a subspace of modes with a fixed velocity.
            # Mostly, this is done to ensure that modes of different velocities
            # are not mixed when particle-hole partners are constructed for
            # P^2 = -1. First, we identify which modes have the same velocity.
            # In each such subspace of modes, we construct wave functions that
            # are particle-hole partners.
            vels = velocities[indx]
            # Velocities are sorted in ascending order. Find the indices of the
            # last instance of each unique velocity.
            inds = [ind+1 for ind, vel in enumerate(vels[:-1])
                    if np.abs(vel-vels[ind+1])>vel_eps]
            inds = [0] + inds + [len(vels)]
            inds = zip(inds[:-1], inds[1:])
            # Now possess an iterator over tuples, where each tuple (i,j)
            # contains the starting and final indices i and j of a submatrix
            # of the modes matrix, such that all modes in the submatrix
            # have the same velocity.

            # Iterate over all submatrices of modes with the same velocity.
            new_wf = []
            TRIM_sorts = []
            for ind_tuple in inds:
                # Pick out wave functions that have a given velocity
                wfs = orig_wf[:, slice(*ind_tuple)]
                # Make particle-hole symmetric modes
                new_modes, TRIM_sort = phs_symmetrization(wfs, particle_hole)
                new_wf.append(new_modes)
                # Store sorting indices of the TRIM modes with the given
                # velocity.
                TRIM_sorts.append(TRIM_sort)
            # Gather into a matrix of modes
            new_wf = np.hstack(new_wf)
            # Store the sort order of all modes at the TRIM.
            # Used later with np.lexsort when the ordering
            # of modes is done.
            TRIM_PHS_sort[indx] = np.hstack(TRIM_sorts)
            # Replace the old modes.
            full_psi[:, indx] = new_wf
            # For both cases P^2 = +-1, must rotate wave functions in the
            # singular value basis. Find the rotation from new basis to old.
            rot = new_wf.T.conj().dot(orig_wf)
            # Rotate the wave functions in the singular value basis
            psi[:, indx] = psi[:, indx].dot(rot.T.conj())

        # Ensure proper usage of chiral symmetry.
        if chiral is not None and time_reversal is None:
            out_orig = full_psi[:, indx[len(indx)//2:]]
            out = chiral.dot(full_psi[:, indx[:len(indx)//2]])
Tómas's avatar
Tómas committed
795
            # No least squares below because the modes should be orthogonal.
796
797
798
799
            rot = out_orig.T.conj().dot(out)
            full_psi[:, indx[len(indx)//2:]] = out
            psi[:, indx[len(indx)//2:]] = psi[:, indx[len(indx)//2:]].dot(rot)

800
801
    if np.any(abs(velocities) < vel_eps):
        raise RuntimeError("Found a mode with zero or close to zero velocity.")
802
    if 2 * np.sum(velocities < 0) != len(velocities):
803
        raise RuntimeError("Numbers of left- and right-propagating "
804
805
                           "modes differ, possibly due to a numerical "
                           "instability.")
806

807
    momenta = -np.angle(lmbdainv)
808
809
810
811
    # Sort the modes. The modes are sorted first by velocity and momentum,
    # and finally TRIM modes are properly ordered.
    order = np.lexsort([TRIM_PHS_sort, velocities,
                        -np.sign(velocities) * momenta, np.sign(velocities)])
812

813
814
    velocities = velocities[order]
    momenta = momenta[order]
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
    full_psi = full_psi[:, order]
    psi = psi[:, order]

    # Use particle-hole symmetry to relate modes that propagate in the
    # same direction but at opposite momenta.
    # Modes are sorted by velocity (first incident, then outgoing).
    # Incident modes are then sorted by momentum in ascending order,
    # and outgoing modes in descending order.
    # Adopted convention is to get modes with negative k (both in and out)
    # by applying particle-hole operator to modes with positive k.
    if particle_hole is not None:
        N = nmodes//2  # Number of incident or outgoing modes.
        # With particle-hole symmetry, N must be an even number.
        # Incident modes
        positive_k = (np.pi - eps > momenta[:N]) * (momenta[:N] > eps)
        # Original wave functions with negative values of k
        orig_neg_k = full_psi[:, :N][:, positive_k[::-1]]
        # For incident modes, ordering of wfs by momentum as returned by kwant
        # is [-k2, -k1, k1, k2], if k2, k1 > 0 and k2 > k1.
        # To maintain this ordering with ki and -ki as particle-hole partners,
        # reverse the order of the product at the end.
Tómas's avatar
Tómas committed
836
837
        wf_neg_k = particle_hole.dot(
                (full_psi[:, :N][:, positive_k]).conj())[:, ::-1]
Tómas's avatar
Tómas committed
838
        rot = lstsq(orig_neg_k, wf_neg_k)
839
        full_psi[:, :N][:, positive_k[::-1]] = wf_neg_k
Tómas's avatar
Tómas committed
840
841
        psi[:, :N][:, positive_k[::-1]] = \
                psi[:, :N][:, positive_k[::-1]].dot(rot)
842
843
844
845
846
847
848

        # Outgoing modes
        positive_k = (np.pi - eps > momenta[N:]) * (momenta[N:] > eps)
        # Original wave functions with negative values of k
        orig_neg_k = full_psi[:, N:][:, positive_k[::-1]]
        # For outgoing modes, ordering of wfs by momentum as returned by kwant
        # is like [k2, k1, -k1, -k2], if k2, k1 > 0 and k2 > k1.
Tómas's avatar
Tómas committed
849
850
851
852

        # Reverse order at the end to match momenta of opposite sign.
        wf_neg_k = particle_hole.dot(
                full_psi[:, N:][:, positive_k].conj())[:, ::-1]
Tómas's avatar
Tómas committed
853
        rot = lstsq(orig_neg_k, wf_neg_k)
854
        full_psi[:, N:][:, positive_k[::-1]] = wf_neg_k
Tómas's avatar
Tómas committed
855
856
        psi[:, N:][:, positive_k[::-1]] = \
                psi[:, N:][:, positive_k[::-1]].dot(rot)
857
858
859
860
861
862
863
864

    # Modes are ordered by velocity.
    # Use time-reversal symmetry to relate modes of opposite velocity.
    if time_reversal is not None:
        # Note: within this function, nmodes refers to the total number
        # of propagating modes, not either left or right movers.
        out_orig = full_psi[:, nmodes//2:]
        out = time_reversal.dot(full_psi[:, :nmodes//2].conj())
Tómas's avatar
Tómas committed
865
        rot = lstsq(out_orig, out)
866
867
868
869
870
871
        full_psi[:, nmodes//2:] = out
        psi[:, nmodes//2:] = psi[:, nmodes//2:].dot(rot)

    norm = np.sqrt(abs(velocities))
    full_psi = full_psi / norm
    psi = psi / norm
872
873

    return psi, PropagatingModes(full_psi, velocities, momenta)
874

Anton Akhmerov's avatar
Anton Akhmerov committed
875

Tómas's avatar
Tómas committed
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
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
def compute_block_modes(h_cell, h_hop, tol, stabilization,
                        time_reversal, particle_hole, chiral):
    """Calculate modes corresponding to a single projector. """
    n, m = h_hop.shape

    # Defer most of the calculation to helper routines.
    matrices, v, extract = setup_linsys(h_cell, h_hop, tol, stabilization)
    ev, evanselect, propselect, vec_gen, ord_schur = unified_eigenproblem(
        *(matrices + (tol,)))

    # v is never None.
    # h_hop.shape[0] and v.shape[1] not always the same.
    # Adding this makes the difference for some tests
    n = v.shape[1]

    nrightmovers = np.sum(propselect) // 2
    nevan = n - nrightmovers
    evan_vecs = ord_schur(evanselect)[:, :nevan]

    # Compute the propagating eigenvectors.
    prop_vecs = vec_gen(propselect)
    # Compute their velocity, and, if necessary, rotate them.

    # prop_vecs here is 'psi' in make_proper_modes, i.e. the wf in the SVD
    # basis.  It is in turn used to construct vecs and vecslmbdainv (the
    # propagating parts).  The evanescent parts of vecs and vecslmbdainv come
    # from evan_vecs above.
    prop_vecs, real_space_data = make_proper_modes(ev[propselect], prop_vecs,
                                                   extract, tol, particle_hole,
                                                   time_reversal, chiral)

    vecs = np.c_[prop_vecs[n:], evan_vecs[n:]]
    vecslmbdainv = np.c_[prop_vecs[:n], evan_vecs[:n]]

    # Prepare output for a single block
    wave_functions = real_space_data.wave_functions
    momenta = real_space_data.momenta
    velocities = real_space_data.velocities

    return (wave_functions, momenta, velocities, vecs, vecslmbdainv, v)


def transform_modes(modes_data, unitary=None, time_reversal=None,
                    particle_hole=None, chiral=None):
    """Transform the modes data for a given block of the Hamiltonian using a
    discrete symmetry (see arguments). The symmetry operator can also be
    specified as can also be identity, for the case when blocks are identical.

    Assume that modes_data has the form returned by compute_block_modes, i.e. a
    tuple (wave_functions, momenta, velocities, nmodes, vecs, vecslmbdainv, v)
    containing the block modes data.

    Assume that the symmetry operator is written in the proper basis (block
    basis, not full TB).

    """
    wave_functions, momenta, velocities, vecs, vecslmbdainv, v = modes_data
Tómas's avatar
Tómas committed
933
934
935
936
937
938
939
940
941

    # Copy to not overwrite modes from previous blocks
    wave_functions = wave_functions.copy()
    momenta = momenta.copy()
    velocities = velocities.copy()
    vecs = vecs.copy()
    vecslmbdainv = vecslmbdainv.copy()
    v = v.copy()

Tómas's avatar
Tómas committed
942
943
944
945
946
    nmodes = wave_functions.shape[1] // 2

    if unitary is not None:
        perm = np.arange(2*nmodes)
        conj = False
Tómas's avatar
Tómas committed
947
        flip_energy = False
Tómas's avatar
Tómas committed
948
949
950
951
    elif time_reversal is not None:
        unitary = time_reversal
        perm = np.arange(2*nmodes)[::-1]
        conj = True
Tómas's avatar
Tómas committed
952
        flip_energy = False
Tómas's avatar
Tómas committed
953
954
955
956
957
    elif particle_hole is not None:
        unitary = particle_hole
        perm = ((-1-np.arange(2*nmodes)) % nmodes +
                nmodes * (np.arange(2*nmodes) // nmodes))
        conj = True
Tómas's avatar
Tómas committed
958
        flip_energy = True
Tómas's avatar
Tómas committed
959
960
961
962
963
    elif chiral is not None:
        unitary = chiral
        perm = (np.arange(2*nmodes) % nmodes +
                nmodes * (np.arange(2*nmodes) < nmodes))
        conj = False
Tómas's avatar
Tómas committed
964
        flip_energy = True
Tómas's avatar
Tómas committed
965
966
967
968
969
970
971
972
973
974
    else:  # skip coverage
        raise ValueError("No relation between blocks was provided.")

    if conj:
        momenta *= -1
        vecs = vecs.conj()
        vecslmbdainv = vecslmbdainv.conj()
        wave_functions = wave_functions.conj()
        v = v.conj()

Tómas's avatar
Tómas committed
975
976
977
978
979
980
    if flip_energy:
        vecslmbdainv *= -1

    if flip_energy != conj:
        velocities *= -1

Tómas's avatar
Tómas committed
981
982
983
984
985
986
987
988
989
    wave_functions = unitary.dot(wave_functions)[:, perm]
    v = unitary.dot(v)
    vecs[:, :2*nmodes] = vecs[:, perm]
    vecslmbdainv[:, :2*nmodes] = vecslmbdainv[:, perm]
    velocities = velocities[perm]
    momenta = momenta[perm]
    return (wave_functions, momenta, velocities, vecs, vecslmbdainv, v)


990
991
992
def modes(h_cell, h_hop, tol=1e6, stabilization=None, *,
          discrete_symmetry=None, projectors=None, time_reversal=None,
          particle_hole=None, chiral=None):
993
    """Compute the eigendecomposition of a translation operator of a lead.
994
995
996

    Parameters
    ----------
997
998
    h_cell : numpy array, real or complex, shape (N,N) The unit cell
        Hamiltonian of the lead unit cell.
999
    h_hop : numpy array, real or complex, shape (N,M)
1000
        The hopping matrix from a lead cell to the one on which self-energy
For faster browsing, not all history is shown. View entire blame