builder.py 86.5 KB
Newer Older
Anton Akhmerov's avatar
Anton Akhmerov 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.

Anton Akhmerov's avatar
Anton Akhmerov committed
9
import abc
10
import warnings
11
import operator
12
import collections
13
import copy
14
from functools import total_ordering, wraps, update_wrapper
Joseph Weston's avatar
Joseph Weston committed
15
from itertools import islice, chain
16
import numbers
17
import inspect
18
import tinyarray as ta
19
import numpy as np
20
from scipy import sparse
21
from . import system, graph, KwantDeprecationWarning, UserCodeError
Christoph Groth's avatar
Christoph Groth committed
22
from .linalg import lll
23
24
from .operator import Density
from .physics import DiscreteSymmetry
25
from ._common import (ensure_isinstance, get_parameters, reraise_warnings,
26
                      interleave, deprecate_args)
27

28

29
30
31
__all__ = ['Builder', 'Site', 'SiteFamily', 'SimpleSiteFamily', 'Symmetry',
           'HoppingKind', 'Lead', 'BuilderLead', 'SelfEnergyLead', 'ModesLead']

Christoph Groth's avatar
Christoph Groth committed
32

33
################ Sites and site families
34

35
class Site(tuple):
36
    """A site, member of a `SiteFamily`.
37
38
39
40

    Sites are the vertices of the graph which describes the tight binding
    system in a `Builder`.

41
    A site is uniquely identified by its family and its tag.
42
43
44

    Parameters
    ----------
45
    family : an instance of `SiteFamily`
Christoph Groth's avatar
Christoph Groth committed
46
        The 'type' of the site.
47
    tag : a hashable python object
48
        The unique identifier of the site within the site family, typically a
Christoph Groth's avatar
Christoph Groth committed
49
        vector of integers.
50
51
52
53

    Raises
    ------
    ValueError
54
        If `tag` is not a proper tag for `family`.
55
56
57

    Notes
    -----
58
    For convenience, ``family(*tag)`` can be used instead of ``Site(family,
59
60
61
62
    tag)`` to create a site.

    The parameters of the constructor (see above) are stored as instance
    variables under the same names.  Given a site ``site``, common things to
63
    query are thus ``site.family``, ``site.tag``, and ``site.pos``.
64
    """
65
    __slots__ = ()
66

67
68
69
70
    family = property(operator.itemgetter(0),
                      doc="The site family to which the site belongs.")
    tag = property(operator.itemgetter(1), doc="The tag of the site.")

71

72
    def __new__(cls, family, tag, _i_know_what_i_do=False):
Christoph Groth's avatar
Christoph Groth committed
73
        if _i_know_what_i_do:
74
            return tuple.__new__(cls, (family, tag))
75
        try:
76
            tag = family.normalize_tag(tag)
77
        except (TypeError, ValueError) as e:
78
            msg = 'Tag {0} is not allowed for site family {1}: {2}'
79
            raise type(e)(msg.format(repr(tag), repr(family), e.args[0]))
80
        return tuple.__new__(cls, (family, tag))
81
82

    def __repr__(self):
83
        return 'Site({0}, {1})'.format(repr(self.family), repr(self.tag))
84

85
    def __str__(self):
86
        sf = self.family
87
        return '<Site {0} of {1}>'.format(self.tag, sf.name if sf.name else sf)
88

Anton Akhmerov's avatar
Anton Akhmerov committed
89
90
91
    def __getnewargs__(self):
        return (self.family, self.tag, True)

92
93
    @property
    def pos(self):
94
95
96
97
        """Real space position of the site.

        This relies on ``family`` having a ``pos`` method (see `SiteFamily`).
        """
98
        return self.family.pos(self.tag)
99
100


101
@total_ordering
102
class SiteFamily(metaclass=abc.ABCMeta):
103
104
105
106
107
    """Abstract base class for site families.

    Site families are the 'type' of `Site` objects.  Within a family, individual
    sites are uniquely identified by tags.  Valid tags must be hashable Python
    objects, further details are up to the family.
108

109
110
111
112
113
114
    Site families must be immutable and fully defined by their initial
    arguments.  They must inherit from this abstract base class and call its
    __init__ function providing it with two arguments: a canonical
    representation and a name.  The canonical representation will be returned as
    the objects representation and must uniquely identify the site family
    instance.  The name is a string used to distinguish otherwise identical site
115
116
117
118
    families.  It may be empty. ``norbs`` defines the number of orbitals
    on sites associated with this site family; it may be `None`, in which case
    the number of orbitals is not specified.

119

120
    All site families must define the method `normalize_tag` which brings a tag
121
    to the standard format for this site family.
122

123
124
125
126
    Site families that are intended for use with plotting should also provide a
    method `pos(tag)`, which returns a vector with real-space coordinates of the
    site belonging to this family with a given tag.

127
128
129
130
131
132
133
134
135
    If the ``norbs`` of a site family are provided, and sites of this family
    are used to populate a `~kwant.builder.Builder`, then the associated
    Hamiltonian values must have the correct shape. That is, if a site family
    has ``norbs = 2``, then any on-site terms for sites belonging to this
    family should be 2x2 matrices. Similarly, any hoppings to/from sites
    belonging to this family must have a matrix structure where there are two
    rows/columns. This condition applies equally to Hamiltonian values that
    are given by functions. If this condition is not satisfied, an error will
    be raised.
136
137
    """

138
    def __init__(self, canonical_repr, name, norbs):
139
140
141
        self.canonical_repr = canonical_repr
        self.hash = hash(canonical_repr)
        self.name = name
142
143
144
145
146
        if norbs is not None:
            if int(norbs) != norbs or norbs <= 0:
                raise ValueError('The norbs parameter must be an integer > 0.')
            norbs = int(norbs)
        self.norbs = norbs
147

148
    def __repr__(self):
149
150
151
        return self.canonical_repr

    def __str__(self):
152
        if self.name:
153
            msg = '<{0} site family {1}{2}>'
154
        else:
155
156
157
            msg = '<unnamed {0} site family{2}>'
        orbs = ' with {0} orbitals'.format(self.norbs) if self.norbs else ''
        return msg.format(self.__class__.__name__, self.name, orbs)
158

159
    def __hash__(self):
160
        return self.hash
161
162
163
164
165
166
167

    def __eq__(self, other):
        try:
            return self.canonical_repr == other.canonical_repr
        except AttributeError:
            return False

168
    def __ne__(self, other):
169
170
171
        try:
            return self.canonical_repr != other.canonical_repr
        except AttributeError:
172
            return True
173

174
    def __lt__(self, other):
175
176
177
        # If this raises an AttributeError, we were trying
        # to compare it to something non-comparable anyway.
        return self.canonical_repr < other.canonical_repr
178

179
    @abc.abstractmethod
180
181
182
183
184
    def normalize_tag(self, tag):
        """Return a normalized version of the tag.

        Raises TypeError or ValueError if the tag is not acceptable.
        """
185
186
187
188
189
190
        pass

    def __call__(self, *tag):
        """
        A convenience function.

191
        This function allows to write fam(1, 2) instead of Site(fam, (1, 2)).
192
193
194
        """
        # Catch a likely and difficult to find mistake.
        if tag and isinstance(tag[0], tuple):
195
196
            raise ValueError('Use site_family(1, 2) instead of '
                             'site_family((1, 2))!')
197
198
199
        return Site(self, tag)


200
201
class SimpleSiteFamily(SiteFamily):
    """A site family used as an example and for testing.
202

203
    A family of sites tagged by any python objects where object satisfied
204
205
    condition ``object == eval(repr(object))``.

206
    It exists to provide a basic site family that can be used for testing the
207
208
209
210
    builder module without other dependencies.  It can be also used to tag
    sites with non-numeric objects like strings should this every be useful.

    Due to its low storage efficiency for numbers it is not recommended to use
211
    `SimpleSiteFamily` when `kwant.lattice.Monatomic` would also work.
212
    """
213

214
215
216
217
    def __init__(self, name=None, norbs=None):
        canonical_repr = '{0}({1}, {2})'.format(self.__class__, repr(name),
                                                repr(norbs))
        super().__init__(canonical_repr, name, norbs)
218

219
220
221
222
223
224
225
226
227
    def normalize_tag(self, tag):
        tag = tuple(tag)
        try:
            if eval(repr(tag)) != tag:
                raise RuntimeError()
        except:
            raise TypeError('It must be possible to recreate the tag from '
                            'its representation.')
        return tag
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

def validate_hopping(hopping):
    """Verify that the argument is a valid hopping."""

    # This check is essential to maintaining the requirement that hoppings must
    # be tuples.  Without it, list "hoppings" would work in some cases
    # (e.g. with Builder.__contains__).
    if not isinstance(hopping, tuple):
        raise TypeError("{0} object is not a valid key.".format(
            type(hopping).__name__))

    # The following check is not strictly necessary (there would be an error
    # anyway), but the error message would be confusing.
    try:
        a, b = hopping
    except:
        raise IndexError("Only length 2 tuples (=hoppings) are valid keys.")

    # This check is essential for Builder.__contains__ - without it a builder
    # would simply "not contain" such invalid hoppings.  In other cases it
    # provides a nicer error message.
    for site in hopping:
        if not isinstance(site, Site):
            raise TypeError("Hopping elements must be Site objects, not {0}."
                            .format(type(site).__name__))

    # This again is an essential check.  Without it, builders would accept loop
    # hoppings.
    if a == b:
        raise ValueError("A hopping connects the following site to itself:\n"
                         "{0}".format(a))


Christoph Groth's avatar
Christoph Groth committed
262

263
264
################ Symmetries

265
class Symmetry(metaclass=abc.ABCMeta):
266
267
    """Abstract base class for spatial symmetries.

268
269
    Many physical systems possess a discrete spatial symmetry, which results in
    special properties of these systems.  This class is the standard way to
270
    describe discrete spatial symmetries in Kwant.  An instance of this class
271
272
273
    can be passed to a `Builder` instance at its creation.  The most important
    kind of symmetry is translational symmetry, used to define scattering
    leads.
274
275
276

    Each symmetry has a fundamental domain -- a set of sites and hoppings,
    generating all the possible sites and hoppings upon action of symmetry
277
278
279
280
281
282
283
284
285
    group elements.  A class derived from `Symmetry` has to implement mapping
    of any site or hopping into the fundamental domain, applying a symmetry
    group element to a site or a hopping, and a method `which` to determine the
    group element bringing some site from the fundamental domain to the
    requested one.  Additionally, it has to have a property `num_directions`
    returning the number of independent symmetry group generators (number of
    elementary periods for translational symmetry).

    A ``ValueError`` must be raised by the symmetry class whenever a symmetry
286
    is used together with sites whose site family is not compatible with it.  A
287
288
    typical example of this is when the vector defining a translational
    symmetry is not a lattice vector.
Christoph Groth's avatar
Christoph Groth committed
289
290
291
292
293

    The type of the domain objects as handled by the methods of this class is
    not specified.  The only requirement is that it must support the unary
    minus operation.  The reference implementation of `to_fd()` is hence
    `self.act(-self.which(a), a, b)`.
294
295
296
297
298
299
300
301
302
303
304
305
    """

    @abc.abstractproperty
    def num_directions(self):
        """Number of elementary periods of the symmetry."""
        pass

    @abc.abstractmethod
    def which(self, site):
        """Calculate the domain of the site.

        Return the group element whose action on a certain site from the
306
        fundamental domain will result in the given ``site``.
307
308
309
310
311
312
313
314
315
316
317
        """
        pass

    @abc.abstractmethod
    def act(self, element, a, b=None):
        """Act with a symmetry group element on a site or hopping."""
        pass

    def to_fd(self, a, b=None):
        """Map a site or hopping to the fundamental domain.

Christoph Groth's avatar
Christoph Groth committed
318
319
320
        If ``b`` is None, return a site equivalent to ``a`` within the
        fundamental domain.  Otherwise, return a hopping equivalent to ``(a,
        b)`` but where the first element belongs to the fundamental domain.
321

Christoph Groth's avatar
Christoph Groth committed
322
        Equivalent to `self.act(-self.which(a), a, b)`.
323
        """
324
        return self.act(-self.which(a), a, b)
325
326

    def in_fd(self, site):
327
        """Tell whether ``site`` lies within the fundamental domain."""
328
329
330
331
332
        for d in self.which(site):
            if d != 0:
                return False
        return True

333
334
335
336
337
338
    @abc.abstractmethod
    def subgroup(self, *generators):
        """Return the subgroup generated by a sequence of group elements."""
        pass

    @abc.abstractmethod
339
340
    def has_subgroup(self, other):
        """Test whether `self` has the subgroup `other`...
Christoph Groth's avatar
Christoph Groth committed
341
342

        or, in other words, whether `other` is a subgroup of `self`.  The
343
        reason why this is the abstract method (and not `is_subgroup`) is that
Christoph Groth's avatar
Christoph Groth committed
344
345
346
        in general it's not possible for a subgroup to know its supergroups.

        """
347
348
        pass

349
350
351
352

class NoSymmetry(Symmetry):
    """A symmetry with a trivial symmetry group."""

Daniel Jaschke's avatar
Daniel Jaschke committed
353
354
355
356
357
358
    def __eq__(self, other):
        return isinstance(other, NoSymmetry)

    def __ne__(self, other):
        return not self.__eq__(other)

359
360
361
362
363
364
365
    def __repr__(self):
        return 'NoSymmetry()'

    @property
    def num_directions(self):
        return 0

366
367
    periods = ()

368
369
    _empty_array = ta.array((), int)

370
    def which(self, site):
371
        return self._empty_array
372
373
374
375

    def act(self, element, a, b=None):
        if element:
            raise ValueError('`element` must be empty for NoSymmetry.')
376
        return a if b is None else (a, b)
377
378
379
380
381
382
383

    def to_fd(self, a, b=None):
        return a if b is None else (a, b)

    def in_fd(self, site):
        return True

384
385
386
387
388
    def subgroup(self, *generators):
        if any(generators):
            raise ValueError('Generators must be empty for NoSymmetry.')
        return NoSymmetry(generators)

389
390
    def has_subgroup(self, other):
        return isinstance(other, NoSymmetry)
391
392


Christoph Groth's avatar
Christoph Groth committed
393

394
395
################ Hopping kinds

396
class HoppingKind(tuple):
397
398
    """A pattern for matching hoppings.

399
400
    An alias exists for this common name: ``kwant.HoppingKind``.

401
402
    A hopping ``(a, b)`` matches precisely when the site family of ``a`` equals
    `family_a` and that of ``b`` equals `family_b` and ``(a.tag - b.tag)`` is
403
    equal to `delta`.  In other words, the matching hoppings have the form:
404
    ``(family_a(x + delta), family_b(x))``
405
406
407
408
409

    Parameters
    ----------
    delta : Sequence of integers
        The sequence is interpreted as a vector with integer elements.
410
411
412
    family_a : `~kwant.builder.SiteFamily`
    family_b : `~kwant.builder.SiteFamily` or ``None`` (default)
        The default value means: use the same family as `family_a`.
413
414
415

    Notes
    -----
Christoph Groth's avatar
Christoph Groth committed
416
417
418
419
420
421
422
    A ``HoppingKind`` is a callable object: When called with a
    `~kwant.builder.Builder` as sole argument, an instance of this class will
    return an iterator over all possible matching hoppings whose sites are
    already present in the system.  The hoppings do *not* have to be already
    present in the system.  For example::

        kind = kwant.builder.HoppingKind((1, 0), lat)
423
        syst[kind(syst)] = 1
Christoph Groth's avatar
Christoph Groth committed
424
425
426
427
428
429
430

    Because a `~kwant.builder.Builder` can be indexed with functions or
    iterables of functions, ``HoppingKind`` instances (or any non-tuple
    iterables of them, e.g. a list) can be used directly as "wildcards" when
    setting or deleting hoppings::

        kinds = [kwant.builder.HoppingKind(v, lat) for v in [(1, 0), (0, 1)]]
431
        syst[kinds] = 1
432
    """
433
434
435
436
437
438
439
440
    __slots__ = ()

    delta = property(operator.itemgetter(0),
                     doc="The difference between the tags of the hopping's sites")
    family_a = property(operator.itemgetter(1),
                        doc="The family of the first site in the hopping")
    family_b = property(operator.itemgetter(2),
                        doc="The family of the second site in the hopping")
441

442
443
    def __new__(cls, delta, family_a, family_b=None):
        delta = ta.array(delta, int)
444
445
        ensure_isinstance(family_a, SiteFamily)
        if family_b is None:
446
            family_b = family_a
447
448
        else:
            ensure_isinstance(family_b, SiteFamily)
449
            family_b = family_b
450
451
452
453
454
455
456
457
458
459
460

        try:
            Site(family_b, family_a.normalize_tag(delta) - delta)
        except Exception as e:
            same_fams = family_b is family_a
            msg = (str(family_a),
                   'and {} are'.format(family_b) if not same_fams else ' is',
                   'not compatible with delta={}'.format(delta),
                  )
            raise ValueError(' '.join(msg)) from e

461
        return tuple.__new__(cls, (delta, family_a, family_b))
462

463
    def __call__(self, builder):
464
        delta = self.delta
465
466
        family_a = self.family_a
        family_b = self.family_b
467
468
469
470
        H = builder.H
        symtofd = builder.symmetry.to_fd

        for a in H:
471
            if a.family != family_a:
472
                continue
473
            b = Site(family_b, a.tag - delta, True)
474
475
476
477
            if symtofd(b) in H:
                yield a, b

    def __repr__(self):
Christoph Groth's avatar
Christoph Groth committed
478
        return '{0}({1}, {2}{3})'.format(
479
            self.__class__.__name__, repr(tuple(self.delta)),
480
481
            repr(self.family_a),
            ', ' + repr(self.family_b) if self.family_a != self.family_b else '')
482

483
484
    def __str__(self):
        return '{0}({1}, {2}{3})'.format(
485
486
            self.__class__.__name__, tuple(self.delta),
            self.family_a,
487
488
489
            ', ' + str(self.family_b) if self.family_a != self.family_b else '')


Christoph Groth's avatar
Christoph Groth committed
490

491
492
493
494
495
496
497
498
499
500
501
502
################ Support for Hermitian conjugation

def herm_conj(value):
    """
    Calculate the hermitian conjugate of a python object.

    If the object is neither a complex number nor a matrix, the original value
    is returned.  In the context of this module, this is the correct behavior
    for functions.
    """
    if hasattr(value, 'conjugate'):
        value = value.conjugate()
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
    elif not isinstance(value, numbers.Number):
        # Fallback for the case of array-like: see issue 406
        # https://gitlab.kwant-project.org/kwant/kwant/-/issues/406
        # If we get a list of lists or another array-like, conjugate()
        # attribute won't be present. To be able to conjugate them, we try to
        # convert them to tinyarrays
        try:
            value = ta.array(value).conjugate()
        except Exception:
            pass

    if hasattr(value, 'shape'):
        if len(value.shape) > 2:
            is_ta = isinstance(value, (
                ta.ndarray_int, ta.ndarray_float, ta.ndarray_complex))
            value = np.swapaxes(value, -1, -2)
            if is_ta:
                value = ta.array(value)
        else:
522
            value = value.transpose()
523

524
525
526
    return value


527
class HermConjOfFunc:
528
529
    """Proxy returning the hermitian conjugate of the original result."""
    __slots__ = ('function')
Anton Akhmerov's avatar
Anton Akhmerov committed
530

531
532
533
    def __init__(self, function):
        self.function = function

534
535
    def __call__(self, i, j, *args, **kwargs):
        return herm_conj(self.function(j, i, *args, **kwargs))
536

537
538
539
540
    @property
    def __signature__(self):
        return inspect.signature(self.function)

Christoph Groth's avatar
Christoph Groth committed
541

542
543
################ Leads

544
class Lead(metaclass=abc.ABCMeta):
545
546
    """Abstract base class for leads that can be attached to a `Builder`.

547
548
549
550
    To attach a lead to a builder, append it to the builder's `~Builder.leads`
    instance variable.  See the documentation of `kwant.builder` for the
    concrete classes of leads derived from this one.

551
552
    Attributes
    ----------
553
    interface : sequence of sites
554

555
556
557
    """

    @abc.abstractmethod
558
    def finalized(self):
559
560
561
562
563
564
565
566
        """Return a finalized version of the lead.

        Returns
        -------
        finalized_lead

        Notes
        -----
567
568
569
570
        The finalized lead must be an object that can be used as a lead
        in a `kwant.system.FiniteSystem`, i.e. an instance of
        `kwant.system.InfiniteSystem`.  Typically it will be a finalized
        builder: `kwant.builder.InfiniteSystem`.
571

Christoph Groth's avatar
Christoph Groth committed
572
573
        The order of sites for the finalized lead must be the one specified in
        `interface`.
574
575
576
577
578
579
580
581
582
583
584
585
        """
        pass


class BuilderLead(Lead):
    """A lead made from a `Builder` with a spatial symmetry.

    Parameters
    ----------
    builder : `Builder`
        The tight-binding system of a lead. It has to possess appropriate
        symmetry, and it may not contain hoppings between further than
Christoph Groth's avatar
Christoph Groth committed
586
        neighboring images of the fundamental domain.
587
    interface : sequence of `Site` instances
588
589
590
        Sequence of sites in the scattering region to which the lead is
        attached.

591
592
593
594
595
596
    Attributes
    ----------
    builder : `Builder`
        The tight-binding system of a lead.
    interface : list of `Site` instances
        A sorted list of interface sites.
597
598
599
600
    padding : list of `Site` instances
        A sorted list of sites that originate from the lead, have the same
        onsite Hamiltonian, and are connected by the same hoppings as the lead
        sites.
601

602
603
    Notes
    -----
604
605
606
607
    The hopping from the scattering region to the lead is assumed to be equal to
    the hopping from a lead unit cell to the next one in the direction of the
    symmetry vector (i.e. the lead is 'leaving' the system and starts with a
    hopping).
608
609
610
611
612

    Every system has an attribute `leads`, which stores a list of
    `BuilderLead` objects with all the information about the leads that are
    attached.
    """
613
    def __init__(self, builder, interface, padding=None):
614
        self.builder = builder
615
616
        self.interface = sorted(interface)
        self.padding = sorted(padding) if padding is not None else []
617
618

    def finalized(self):
619
        """Return a `kwant.builder.InfiniteSystem` corresponding to the
620
621
        compressed lead.

622
        The order of interface sites is kept during finalization.
623
        """
624
        return InfiniteSystem(self.builder, self.interface)
625
626


627
def _ensure_signature(func):
628
629
630
631
632
633
634
635
636
637
638
639
    """
    Ensure that a modes/selfenergy function has a keyword-only parameter
    ``params``, or takes ``**kwargs`` by potentially wrapping it.
    """
    parameters = inspect.signature(func).parameters
    has_params = bool(parameters.get('params'))
    has_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD
                     for p in parameters.values())
    if has_params or has_kwargs:
        return func

    # function conforming to old API: needs wrapping
640
    @deprecate_args
641
642
643
644
    def wrapper(energy, args=(), *, params=None):
        return func(energy, args)

    return wrapper
645
646


647
class SelfEnergyLead(Lead):
648
649
650
651
    """A general lead defined by its self energy.

    Parameters
    ----------
652
    selfenergy_func : function
653
654
        Has the same signature as `selfenergy` (without the ``self``
        parameter) and returns the self energy matrix for the interface sites.
655
    interface : sequence of `Site` instances
656
657
    parameters : sequence of strings
        The parameters on which the lead depends.
658
    """
659
    def __init__(self, selfenergy_func, interface, parameters):
660
        self.interface = tuple(interface)
661
662
663
664
        # we changed the API of 'selfenergy_func' to have a keyword-only
        # parameter 'params', but we still need to support the old API
        # XXX: remove this when releasing Kwant 2.0
        self.selfenergy_func = _ensure_signature(selfenergy_func)
665
        self.parameters = frozenset(parameters)
666
667
668
669
670

    def finalized(self):
        """Trivial finalization: the object is returned itself."""
        return self

671
    @deprecate_args
672
673
    def selfenergy(self, energy, args=(), *, params=None):
        return self.selfenergy_func(energy, args, params=params)
674
675
676
677
678
679
680
681


class ModesLead(Lead):
    """A general lead defined by its modes wave functions.

    Parameters
    ----------
    modes_func : function
682
683
684
        Has the same signature as `modes` (without the ``self`` parameter)
        and returns the modes of the lead as a tuple of
        `~kwant.physics.PropagatingModes` and `~kwant.physics.StabilizedModes`.
685
686
    interface :
        sequence of `Site` instances
687
688
    parameters : sequence of strings
        The parameters on which the lead depends.
689
    """
690
    def __init__(self, modes_func, interface, parameters):
691
        self.interface = tuple(interface)
692
693
694
695
        # we changed the API of 'selfenergy_func' to have a keyword-only
        # parameter 'params', but we still need to support the old API
        # XXX: remove this when releasing Kwant 2.0
        self.modes_func = _ensure_signature(modes_func)
696
        self.parameters = frozenset(parameters)
697
698
699
700
701

    def finalized(self):
        """Trivial finalization: the object is returned itself."""
        return self

702
    @deprecate_args
703
704
    def modes(self, energy, args=(), *, params=None):
        return self.modes_func(energy, args, params=params)
705

706
    @deprecate_args
707
708
    def selfenergy(self, energy, args=(), *, params=None):
        stabilized = self.modes(energy, args, params=params)[1]
709
        return stabilized.selfenergy()
Christoph Groth's avatar
Christoph Groth committed
710

711

Christoph Groth's avatar
Christoph Groth committed
712

713
714
################ Builder class

715
# A marker, meaning for hopping (i, j): this value is given by the Hermitian
716
# conjugate the value of the hopping (j, i).  Used by Builder and System.
717
class Other:
Christoph Groth's avatar
Christoph Groth committed
718
    pass
719

720

721
def edges(seq):
722
723
    result = interleave(seq)
    next(result)  # Skip the special loop edge.
724
    return result
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
def _site_ranges(sites):
    """Return a sequence of ranges for ``sites``.

    Here, "ranges" are defined as sequences of sites that have the same site
    family. Because site families now have a fixed number of orbitals,
    this coincides with the definition given in `~kwant.system.System`.
    """
    # we shall start a new range of different `SiteFamily`s separately,
    # even if they happen to contain the same number of orbitals.
    total_norbs = 0
    current_fam = None
    site_ranges = []
    for idx, fam in enumerate(s.family for s in sites):
        if not fam.norbs:
            # can't provide site_ranges if norbs not given
            return None
        if fam != current_fam:  # start a new run
            current_fam = fam
            current_norbs = fam.norbs
            site_ranges.append((idx, current_norbs, total_norbs))
        total_norbs += current_norbs
    # add sentinel to the end
    site_ranges.append((len(sites), 0, total_norbs))
    return site_ranges


753
754
class _Substituted:
    """Proxy that renames function parameters."""
755

756
757
758
759
    def __init__(self, func, params):
        self.func = func
        self.params = params
        update_wrapper(self, func)
760

761
762
763
764
    def __eq__(self, other):
        if not isinstance(other, _Substituted):
            return False
        return (self.func == other.func and self.params == other.params)
765

766
767
    def __hash__(self):
        return hash((self.func, self.params))
768

769
770
771
772
773
    @property
    def __signature__(self):
        return inspect.Signature(
            [inspect.Parameter(name, inspect.Parameter.POSITIONAL_ONLY)
             for name in self.params])
774

775
776
    def __call__(self, *args):
        return self.func(*args)
777
778


779
780
781
def _substitute_params(func, subs):
    """Substitute 'params' from 'subs' into 'func'."""
    assert callable(func)
782

783
784
785
786
787
788
    if isinstance(func, _Substituted):
        old_params = func.params
        old_func = func.func
    else:
        old_params = get_parameters(func)
        old_func = func
789

790
    params = tuple(subs.get(p, p) for p in old_params)
791

792
793
794
795
796
797
798
    duplicates = [p for p, count in collections.Counter(params).items()
                  if count > 1]
    if duplicates:
        msg = ('Cannot rename parameters ',
               ','.join('"{}"'.format(d) for d in duplicates),
               ': parameters with the same name exist')
        raise ValueError(''.join(msg))
799

800
801
802
803
    if params == old_params:
        return func
    else:
        return _Substituted(old_func, params)
804
805


806
class Builder:
807
808
    """A tight binding system defined on a graph.

809
810
    An alias exists for this common name: ``kwant.Builder``.

811
    This is one of the central types in Kwant.  It is used to construct tight
812
813
814
815
    binding systems in a flexible way.

    The nodes of the graph are `Site` instances.  The edges, i.e. the hoppings,
    are pairs (2-tuples) of sites.  Each node and each edge has a value
Christoph Groth's avatar
Christoph Groth committed
816
    associated with it.  The values associated with nodes are interpreted as
817
818
    on-site Hamiltonians, the ones associated with edges as hopping integrals.

Christoph Groth's avatar
Christoph Groth committed
819
    To make the graph accessible in a way that is natural within the Python
820
    language it is exposed as a *mapping* (much like a built-in Python
Christoph Groth's avatar
Christoph Groth committed
821
    dictionary).  Keys are sites or hoppings.  Values are 2d arrays
822
    (e.g. NumPy or Tinyarray) or numbers (interpreted as 1 by 1 matrices).
823
824
825
826

    Parameters
    ----------
    symmetry : `Symmetry` or `None`
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
        The spatial symmetry of the system.
    conservation_law : 2D array, dictionary, function, or `None`
        An onsite operator with integer eigenvalues that commutes with the
        Hamiltonian.  The ascending order of eigenvalues corresponds to the
        selected ordering of the Hamiltonian subblocks.  If a dict is
        given, it maps from site families to such matrices. If a function is
        given it must take the same arguments as the onsite Hamiltonian
        functions of the system and return the onsite matrix.
    time_reversal : scalar, 2D array, dictionary, function, or `None`
        The unitary part of the onsite time-reversal symmetry operator.
        Same format as that of `conservation_law`.
    particle_hole : scalar, 2D array, dictionary, function, or `None`
        The unitary part of the onsite particle-hole symmetry operator.
        Same format as that of `conservation_law`.
    chiral : 2D array, dictionary, function or `None`
        The unitary part of the onsite chiral symmetry operator.
        Same format as that of `conservation_law`.
844
845
846

    Notes
    -----
847

Christoph Groth's avatar
Christoph Groth committed
848
849
850
851
852
853
    Values can be also functions that receive the site or the hopping (passed
    to the function as two sites) and possibly additional arguments and are
    expected to return a valid value.  This allows to define systems quickly,
    to modify them without reconstructing, and to save memory for many-orbital
    models.

854
855
    In addition to simple keys (single sites and hoppings) more powerful keys
    are possible as well that allow to manipulate multiple sites/hoppings in a
Christoph Groth's avatar
Christoph Groth committed
856
857
    single operation.  Such keys are internally expanded into a sequence of
    simple keys by using the method `Builder.expand`.  For example,
858
    ``syst[general_key] = value`` is equivalent to ::
859

860
861
        for simple_key in syst.expand(general_key):
            syst[simple_key] = value
Christoph Groth's avatar
Christoph Groth committed
862

863
864
865
866
    Builder instances automatically ensure that every hopping is Hermitian, so
    that if ``builder[a, b]`` has been set, there is no need to set
    ``builder[b, a]``.

Christoph Groth's avatar
Christoph Groth committed
867
868
    Builder instances can be made to automatically respect a `Symmetry` that is
    passed to them during creation.  The behavior of builders with a symmetry
Anton Akhmerov's avatar
Anton Akhmerov committed
869
870
871
    is slightly more sophisticated: all keys are mapped to the fundamental
    domain of the symmetry before storing them.  This may produce confusing
    results when neighbors of a site are queried.
872

Christoph Groth's avatar
Christoph Groth committed
873
874
875
    The method `attach_lead` *works* only if the sites affected by them have
    tags which are sequences of integers.  It *makes sense* only when these
    sites live on a regular lattice, like the ones provided by `kwant.lattice`.
876

877
    Attaching a lead manually (without the use of `~Builder.attach_lead`)
Christoph Groth's avatar
Christoph Groth committed
878
879
    amounts to creating a `Lead` object and appending it to the list of leads
    accessbile as the `~Builder.leads` attribute.
880

881
882
883
884
    `conservation_law`, `time_reversal`, `particle_hole`, and `chiral`
    affect the basis in which scattering modes derived from the builder
    are expressed - see `~kwant.physics.DiscreteSymmetry` for details.

885
886
887
888
889
890
    .. warning::

        If functions are used to set values in a builder with a symmetry, then
        they must satisfy the same symmetry.  There is (currently) no check and
        wrong results will be the consequence of a misbehaving function.

891
892
893
894
895
    Attributes
    ----------
    leads : list of `Lead` instances
        The leads that are attached to the system.

896
897
898
899
900
901
902
903
    Examples
    --------
    Define a site.

    >>> builder[site] = value

    Print the value of a site.

Joseph Weston's avatar
Joseph Weston committed
904
    >>> print(builder[site])
905
906
907
908
909
910
911
912
913

    Define a hopping.

    >>> builder[site1, site2] = value

    Delete a site.

    >>> del builder[site3]

914
915
916
917
918
    Detach the last lead.  (This does not remove the sites that were added to
    the scattering region by `~Builder.attach_lead`.)

    >>> del builder.leads[-1]

919
    """
920

921
922
    def __init__(self, symmetry=None, *, conservation_law=None, time_reversal=None,
                 particle_hole=None, chiral=None):
923
924
        if symmetry is None:
            symmetry = NoSymmetry()
925
926
        else:
            ensure_isinstance(symmetry, Symmetry)
927
        self.symmetry = symmetry
928
929
930
931
        self.conservation_law = conservation_law
        self.time_reversal = time_reversal
        self.particle_hole = particle_hole
        self.chiral = chiral
932
        self.leads = []
933
934
935
936
937
938
939
940
941
942
        self.H = {}

    #### Note on H ####
    #
    # This dictionary stores a directed graph optimized for efficient querying
    # and modification.  The nodes are instances of `Site`.
    #
    # Each edge, specified by a ``(tail, head)`` pair of nodes, holds an object
    # as a value.  Likewise, each tail which occurs in the graph also holds a
    # value.  (Nodes which only occur as heads are not required to have
943
944
    # values.) Every tail node has to be in the fundamental domain of the
    # builder's symmetry.
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
    #
    # For a given `tail` site, H[tail] is a list alternately storing
    # heads and values.  (The heads occupy even locations followed by the
    # values at odd locations.)  Each pair of entries thus describes a single
    # directed edge of the graph.
    #
    # The first pair of entries in each list is special: it always
    # corresponds to a loop edge.  (The head is equal to the tail.)  This
    # special edge has two purposes: It is used to store the value
    # associated with the tail node itself, and it is necessary for the
    # method getkey_tail which helps to conserve memory by storing equal
    # node label only once.

    def _get_edge(self, tail, head):
        for h, value in edges(self.H[tail]):
            if h == head:
                return value
962
963
964
965
966
967
968
969

        # (tail, head) is not present in the system, but tail is.
        if head in self.H:
            raise KeyError((tail, head))
        else:
            # If already head is missing, we only report this.  This way the
            # behavior is symmetric with regard to tail and head.
            raise KeyError(head)
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984

    def _set_edge(self, tail, head, value):
        hvhv = self.H[tail]
        heads = hvhv[2::2]
        if head in heads:
            i = 2 + 2 * heads.index(head)
            hvhv[i] = head
            hvhv[i + 1] = value
        else:
            hvhv.append(head)
            hvhv.append(value)

    def _del_edge(self, tail, head):
        hvhv = self.H[tail]
        heads = hvhv[2::2]
985
986
987
988
989
990
991
992
993
994
995
996

        try:
            i = 2 + 2 * heads.index(head)
        except ValueError:
            # (tail, head) is not present in the system, but tail is.
            if head in self.H:
                raise KeyError((tail, head))
            else:
                # If already head is missing, we only report this.  This way the
                # behavior is symmetric with regard to tail and head.
                raise KeyError(head)

997
998
999
        del hvhv[i : i + 2]

    def _out_neighbors(self, tail):
1000
        hvhv = self.H[tail]
1001
1002
1003
        return islice(hvhv, 2, None, 2)

    def _out_degree(self, tail):
1004
        hvhv = self.H[tail]
1005
        return len(hvhv) // 2 - 1
1006

1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
    def __copy__(self):
        """Shallow copy"""
        result = object.__new__(Builder)
        result.symmetry = self.symmetry
        result.conservation_law = self.conservation_law
        result.time_reversal = self.time_reversal
        result.particle_hole = self.particle_hole
        result.chiral = self.chiral
        result.leads = self.leads
        result.H = self.H
        return result

1019
    # TODO: write a test for this method.
1020
1021
1022
1023
1024
1025
1026
1027
1028
    def reversed(self):
        """Return a shallow copy of the builder with the symmetry reversed.

        This method can be used to attach the same infinite system as lead from
        two opposite sides.  It requires a builder to which an infinite
        symmetry is associated.
        """
        if self.leads:
            raise ValueError('System to be reversed may not have leads.')
1029
1030
1031
1032
        result = copy.copy(self)
        # if we don't assign a new list we will inadvertantly add leads to
        # the reversed system if we add leads to *this* system
        # (because we only shallow copy)
1033
        result.leads = []
1034
        result.symmetry = self.symmetry.reversed()
1035
1036
        return result

Joseph Weston's avatar
Joseph Weston committed
1037
    def __bool__(self):
1038
        return bool(self.H)
1039

1040
1041
    def expand(self, key):
        """
Christoph Groth's avatar
Christoph Groth committed
1042
        Expand a general key into an iterator over simple keys.
1043

Christoph Groth's avatar
Christoph Groth committed
1044
1045
1046
1047
        Parameters
        ----------
        key : builder key (see notes)
            The key to be expanded
1048

Christoph Groth's avatar
Christoph Groth committed
1049
1050
        Notes
        -----
1051
1052
1053
1054
1055
1056
1057
        Keys are (recursively):
            * Simple keys: sites or 2-tuples of sites (=hoppings).
            * Any (non-tuple) iterable of keys, e.g. a list or a generator
              expression.
            * Any function that returns a key when passed a builder as sole
              argument, e.g. a `HoppingKind` instance or the function returned
              by `~kwant.lattice.Polyatomic.shape`.
Christoph Groth's avatar
Christoph Groth committed
1058
1059

        This method is internally used to expand the keys when getting or
1060
1061
        deleting items of a builder (i.e. ``syst[key] = value`` or ``del
        syst[key]``).
Christoph Groth's avatar
Christoph Groth committed
1062

1063
1064
1065
1066
1067
        """
        itr = iter((key,))
        iter_stack = [None]
        while iter_stack:
            for key in itr:
1068
                while callable(key):
Christoph Groth's avatar
Christoph Groth committed
1069
                    key = key(self)
1070
1071
                if isinstance(key, tuple):
                    # Site instances are also tuples.
1072
                    yield key
1073
                else:
1074
                    iter_stack.append(itr)
1075
1076
1077
1078
1079
                    try:
                        itr = iter(key)
                    except TypeError:
                        raise TypeError("{0} object is not a valid key."
                                        .format(type(key).__name__))
1080
1081
1082
                    break
            else:
                itr = iter_stack.pop()
1083

1084
1085
1086
1087
    def __getitem__(self, key):
        """Get the value of a single site or hopping."""
        if isinstance(key, Site):
            site = self.symmetry.to_fd(key)
1088
            return self.H[site][1]
1089
1090

        sym = self.symmetry
1091
1092
1093
        validate_hopping(key)
        a, b = sym.to_fd(*key)
        value = self._get_edge(a, b)
1094
        if value is Other:
1095
1096
1097
            if not sym.in_fd(b):
                b, a = sym.to_fd(b, a)
                assert not sym.in_fd(a)
1098
            value = self._get_edge(b, a)
1099
            if callable(value):
1100
1101
1102
1103
1104
1105
1106
1107
                assert not isinstance(value, HermConjOfFunc)
                value = HermConjOfFunc(value)
            else:
                value = herm_conj(value)
        return value

    def __contains__(self, key):
        """Tell whether the system contains a site or hopping."""
1108
        if isinstance(key, Site):
1109
1110
1111
1112
1113
1114
1115
            key = self.symmetry.to_fd(key)
            return key in self.H

        validate_hopping(key)
        a, b = self.symmetry.to_fd(*key)
        hvhv = self.H.get(a, ())
        return b in islice(hvhv, 2, None, 2)
1116

1117
    def _set_site(self, site, value):
1118
        """Set a single site."""
1119
1120
        if not isinstance(site, Site):
            raise TypeError('Expecting a site, got {0} instead.'.format(type(site).__name__))
1121
        site = self.symmetry.to_fd(site)
1122
1123
1124
1125
1126
        hvhv = self.H.setdefault(site, [])
        if hvhv:
            hvhv[1] = value
        else:
            hvhv[:] = [site, value]
1127

1128
    def _set_hopping(self, hopping, value):
1129
1130
        """Set a single hopping."""
        sym = self.symmetry
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
        validate_hopping(hopping)
        a, b = sym.to_fd(*hopping)

        if sym.in_fd(b):
            # Make sure that we do not waste space by storing multiple instances
            # of identical sites.
            a2 = a = self.H[a][0]
            b2 = b = self.H[b][0]
        else:
            b2, a2 = sym.to_fd(b, a)
            assert not sym.in_fd(a2)