builder.py 46.8 KB
Newer Older
1
# Copyright 2011-2013 Kwant authors.
2
#
3
# This file is part of Kwant.  It is subject to the license terms in the
4
# LICENSE file 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
6
7
8
# the AUTHORS file at the top-level directory of this distribution and at
# http://kwant-project.org/authors.

9
10
from __future__ import division

11
__all__ = ['Builder', 'Site', 'SiteFamily', 'SimpleSiteFamily', 'Symmetry',
12
           'HoppingKind', 'Lead', 'BuilderLead', 'SelfEnergyLead', 'ModesLead']
13

Anton Akhmerov's avatar
Anton Akhmerov committed
14
15
import abc
import sys
16
import warnings
17
import operator
18
from itertools import izip, islice, chain
19
import tinyarray as ta
20
import numpy as np
21
from . import system, graph, physics
22

Christoph Groth's avatar
Christoph Groth committed
23

24
################ Sites and site families
25

26
class Site(tuple):
27
    """A site, member of a `SiteFamily`.
28
29
30
31

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

32
    A site is uniquely identified by its family and its tag.
33
34
35

    Parameters
    ----------
36
    family : an instance of `SiteFamily`
Christoph Groth's avatar
Christoph Groth committed
37
        The 'type' of the site.
38
    tag : a hashable python object
39
        The unique identifier of the site within the site family, typically a
Christoph Groth's avatar
Christoph Groth committed
40
        vector of integers.
41
42
43
44

    Raises
    ------
    ValueError
45
        If `tag` is not a proper tag for `family`.
46
47
48

    Notes
    -----
49
    For convenience, ``family(*tag)`` can be used instead of ``Site(family,
50
51
52
53
    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
54
    query are thus ``site.family``, ``site.tag``, and ``site.pos``.
55
    """
56
    __slots__ = ()
57

58
59
60
61
    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.")

62

63
    def __new__(cls, family, tag, _i_know_what_i_do=False):
Christoph Groth's avatar
Christoph Groth committed
64
        if _i_know_what_i_do:
65
            return tuple.__new__(cls, (family, tag))
66
        try:
67
            tag = family.normalize_tag(tag)
68
69
        except (TypeError, ValueError):
            t, v, tb = sys.exc_info()
70
71
72
            msg = 'Tag {0} is not allowed for site family {1}: {2}'
            raise t(msg.format(repr(tag), repr(family), v))
        return tuple.__new__(cls, (family, tag))
73
74

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

77
    def __str__(self):
78
        return '{1} of {0}'.format(str(self.family), str(self.tag))
79

80
81
82
    @property
    def pos(self):
        """Real space position of the site."""
83
        return self.family.pos(self.tag)
84
85


86
class SiteFamily(object):
87
    """
88
    Abstract base class for site families.
89

90
91
92
93
94
95
96
    A site family is a 'type' of sites.  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 families.  It may be empty.
97

98
    All site families must define the method `normalize_tag` which brings a tag
99
    to the standard format for this site family.
100

101
    Site families which are intended for use with plotting should also provide a
102
    method `pos(tag)`, which returns a vector with real-space coordinates of
103
    the site belonging to this family with a given tag.
104
105
106
    """
    __metaclass__ = abc.ABCMeta

107
108
109
110
111
    def __init__(self, canonical_repr, name):
        self.canonical_repr = canonical_repr
        self.hash = hash(canonical_repr)
        self.name = name

112
    def __repr__(self):
113
114
115
116
        return self.canonical_repr

    def __str__(self):
        return '{0} object {1}'.format(self.__class__, self.name)
117

118
    def __hash__(self):
119
        return self.hash
120
121
122
123
124
125
126
127
128
129
130
131

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

    def __neq__(self, other):
        try:
            return self.canonical_repr != other.canonical_repr
        except AttributeError:
            return False
132

133
    @abc.abstractmethod
134
135
136
137
138
    def normalize_tag(self, tag):
        """Return a normalized version of the tag.

        Raises TypeError or ValueError if the tag is not acceptable.
        """
139
140
141
142
143
144
145
146
147
148
        pass

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

        This function allows to write sg(1, 2) instead of Site(sg, (1, 2)).
        """
        # Catch a likely and difficult to find mistake.
        if tag and isinstance(tag[0], tuple):
149
150
            raise ValueError('Use site_family(1, 2) instead of '
                             'site_family((1, 2))!')
151
152
153
        return Site(self, tag)


154
155
class SimpleSiteFamily(SiteFamily):
    """A site family used as an example and for testing.
156

157
    A family of sites tagged by any python objects where object satisfied
158
159
    condition ``object == eval(repr(object))``.

160
    It exists to provide a basic site family that can be used for testing the
161
162
163
164
    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
165
    `SimpleSiteFamily` when `kwant.lattice.Monatomic` would also work.
166
    """
167
168

    def __init__(self, name=None):
169
170
        canonical_repr = '{0}({1})'.format(self.__class__, repr(name))
        super(SimpleSiteFamily, self).__init__(canonical_repr, name)
171

172
173
174
175
176
177
178
179
180
    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
181

Christoph Groth's avatar
Christoph Groth committed
182

183
184
185
186
187
################ Symmetries

class Symmetry(object):
    """Abstract base class for spatial symmetries.

188
189
    Many physical systems possess a discrete spatial symmetry, which results in
    special properties of these systems.  This class is the standard way to
190
    describe discrete spatial symmetries in Kwant.  An instance of this class
191
192
193
    can be passed to a `Builder` instance at its creation.  The most important
    kind of symmetry is translational symmetry, used to define scattering
    leads.
194
195
196

    Each symmetry has a fundamental domain -- a set of sites and hoppings,
    generating all the possible sites and hoppings upon action of symmetry
197
198
199
200
201
202
203
204
205
    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
206
    is used together with sites whose site family is not compatible with it.  A
207
208
    typical example of this is when the vector defining a translational
    symmetry is not a lattice vector.
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
    """
    __metaclass__ = abc.ABCMeta

    @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
        fundamental domain will result in the given `site`.
        """
        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.

        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.

        This default implementation works but may be not efficient.
        """
240
        return self.act(-self.which(a), a, b)
241
242
243
244
245
246
247
248
249
250
251
252

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


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

Daniel Jaschke's avatar
Daniel Jaschke committed
253
254
255
256
257
258
    def __eq__(self, other):
        return isinstance(other, NoSymmetry)

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

259
260
261
262
263
264
265
    def __repr__(self):
        return 'NoSymmetry()'

    @property
    def num_directions(self):
        return 0

266
267
    _empty_array = ta.array((), int)

268
    def which(self, site):
269
        return self._empty_array
270
271
272
273

    def act(self, element, a, b=None):
        if element:
            raise ValueError('`element` must be empty for NoSymmetry.')
274
        return a if b is None else (a, b)
275
276
277
278
279
280
281

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

    def in_fd(self, site):
        return True

Christoph Groth's avatar
Christoph Groth committed
282

283
284
285
286
287
################ Hopping kinds

class HoppingKind(object):
    """A pattern for matching hoppings.

288
289
    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
290
    equal to `delta`.  In other words, the matching hoppings have the form:
291
    ``(family_a(x + delta), family_b(x))``
292
293
294
295
296

    Parameters
    ----------
    delta : Sequence of integers
        The sequence is interpreted as a vector with integer elements.
297
298
299
    family_a : `~kwant.builder.SiteFamily`
    family_b : `~kwant.builder.SiteFamily` or ``None`` (default)
        The default value means: use the same family as `family_a`.
300
301
302

    Notes
    -----
Christoph Groth's avatar
Christoph Groth committed
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
    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)
        sys[kind(sys)] = 1

    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)]]
        sys[kinds] = 1
319
    """
320
    __slots__ = ('delta', 'family_a', 'family_b')
321

322
    def __init__(self, delta, family_a, family_b=None):
323
        self.delta = ta.array(delta, int)
324
325
        self.family_a = family_a
        self.family_b = family_b if family_b is not None else family_a
326

327
    def __call__(self, builder):
328
        delta = self.delta
329
330
        family_a = self.family_a
        family_b = self.family_b
331
332
333
334
        H = builder.H
        symtofd = builder.symmetry.to_fd

        for a in H:
335
            if a.family != family_a:
336
                continue
337
            b = Site(family_b, a.tag - delta, True)
338
339
340
341
            if symtofd(b) in H:
                yield a, b

    def __repr__(self):
Christoph Groth's avatar
Christoph Groth committed
342
        return '{0}({1}, {2}{3})'.format(
343
            self.__class__.__name__, repr(tuple(self.delta)),
344
345
            repr(self.family_a),
            ', ' + repr(self.family_b) if self.family_a != self.family_b else '')
346

Christoph Groth's avatar
Christoph Groth committed
347

348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
################ 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()
        if hasattr(value, 'transpose'):
            value = value.transpose()
    return value


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

369
370
371
372
373
374
    def __init__(self, function):
        self.function = function

    def __call__(self, i, j):
        return herm_conj(self.function(j, i))

Christoph Groth's avatar
Christoph Groth committed
375

376
377
378
379
380
################ Leads

class Lead(object):
    """Abstract base class for leads that can be attached to a `Builder`.

381
382
    Attributes
    ----------
383
    interface : sequence of sites
384
385
386
387
388
389
390
391
392
393
394
395
396
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def finalized():
        """Return a finalized version of the lead.

        Returns
        -------
        finalized_lead

        Notes
        -----
Christoph Groth's avatar
Christoph Groth committed
397
398
399
        The finalized lead must be an object that can be used as a lead in a
        `kwant.system.FiniteSystem`.  It could be an instance of
        `kwant.system.InfiniteSystem` for example.
400

Christoph Groth's avatar
Christoph Groth committed
401
402
        The order of sites for the finalized lead must be the one specified in
        `interface`.
403
404
405
406
407
408
409
410
411
412
413
414
        """
        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
415
        neighboring images of the fundamental domain.
416
    interface : sequence of `Site` instances
417
418
419
420
421
        Sequence of sites in the scattering region to which the lead is
        attached.

    Notes
    -----
422
423
424
425
    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).
426

427
    The given order of interface sites is preserved throughout finalization.
428
429
430
431
432

    Every system has an attribute `leads`, which stores a list of
    `BuilderLead` objects with all the information about the leads that are
    attached.
    """
433
    def __init__(self, builder, interface):
434
        self.builder = builder
435
        self.interface = tuple(interface)
436
437
438
439
440

    def finalized(self):
        """Return a `kwant.system.InfiniteSystem` corresponding to the
        compressed lead.

441
        The order of interface sites is kept during finalization.
442
        """
443
        return self.builder._finalized_infinite(self.interface)
444
445


446
class SelfEnergyLead(Lead):
447
448
449
450
    """A general lead defined by its self energy.

    Parameters
    ----------
451
    selfenergy_func : function
452
        Function which returns the self energy matrix for the interface sites
Anton Akhmerov's avatar
Anton Akhmerov committed
453
        given the energy and optionally a list of extra arguments.
454
    interface : sequence of `Site` instances
455
    """
456
457
    def __init__(self, selfenergy_func, interface):
        self._selfenergy_func = selfenergy_func
458
        self.interface = tuple(interface)
459
460
461
462
463

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

464
465
    def selfenergy(self, energy, args=()):
        return self._selfenergy_func(energy, args)
466
467
468
469
470
471
472
473


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

    Parameters
    ----------
    modes_func : function
474
475
476
477
478
479
        Function which returns the modes of the lead as a tuple of
        `~kwant.physics.PropagatingModes` and `~kwant.physics.StabilizedModes`
        given the energy and optionally a list of extra arguments.
    interface :
        sequence of `Site` instances

480
481
    """
    def __init__(self, modes_func, interface):
482
        self.modes_func = modes_func
483
484
485
486
487
488
489
        self.interface = tuple(interface)

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

    def modes(self, energy, args=()):
490
        return self.modes_func(energy, args)
491

492
    def selfenergy(self, energy, args=()):
493
494
        propagating, stabilized = self.modes(energy, args)
        return stabilized.selfenergy()
Christoph Groth's avatar
Christoph Groth committed
495

496

Christoph Groth's avatar
Christoph Groth committed
497

498
499
################ Builder class

500
# A marker, meaning for hopping (i, j): this value is given by the Hermitian
501
# conjugate the value of the hopping (j, i).  Used by Builder and System.
502
Other = type('Other', (object,), {'__repr__': lambda s: 'Other'})()
503

504

505
506
507
508
509
510
511
def edges(seq):
    # izip, when given the same iterator twice, turns a sequence into a
    # sequence of pairs.
    seq_iter = iter(seq)
    result = izip(seq_iter, seq_iter)
    next(result)                # Skip the special loop edge.
    return result
512

513

514
515
516
class Builder(object):
    """A tight binding system defined on a graph.

517
    This is one of the central types in Kwant.  It is used to construct tight
518
519
520
521
    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
522
    associated with it.  The values associated with nodes are interpreted as
523
524
    on-site Hamiltonians, the ones associated with edges as hopping integrals.

Christoph Groth's avatar
Christoph Groth committed
525
    To make the graph accessible in a way that is natural within the Python
526
    language it is exposed as a *mapping* (much like a built-in Python
Christoph Groth's avatar
Christoph Groth committed
527
    dictionary).  Keys are sites or hoppings.  Values are 2d arrays
528
    (e.g. NumPy or Tinyarray) or numbers (interpreted as 1 by 1 matrices).
529
530
531
532
533
534
535
536

    Parameters
    ----------
    symmetry : `Symmetry` or `None`
        The symmetry of the system.

    Notes
    -----
Christoph Groth's avatar
Christoph Groth committed
537
538
539
540
541
542
    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.

543
544
    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
545
546
547
    single operation.  Such keys are internally expanded into a sequence of
    simple keys by using the method `Builder.expand`.  For example,
    ``sys[general_key] = value`` is equivalent to ::
548

Christoph Groth's avatar
Christoph Groth committed
549
        for simple_key in sys.expand(general_key):
550
            sys[simple_key] = value
Christoph Groth's avatar
Christoph Groth committed
551

552
553
554
555
    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
556
557
558
    Builder instances can be made to automatically respect a `Symmetry` that is
    passed to them during creation.  The behavior of builders with a symmetry
    is slightly more sophisticated.  First of all, it is implicitly assumed
559
    throughout Kwant that **every** function assigned as a value to a builder
Christoph Groth's avatar
Christoph Groth committed
560
561
562
    with a symmetry possesses the same symmetry.  Secondly, 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.
563

Christoph Groth's avatar
Christoph Groth committed
564
565
566
    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`.
567

Christoph Groth's avatar
Christoph Groth committed
568
569
570
571
    ``builder0 += builder1`` adds all the sites, hoppings, and leads of
    ``builder1`` to ``builder0``.  Sites and hoppings present in both systems
    are overwritten by those in ``builder1``.  The leads of ``builder1`` are
    appended to the leads of the system being extended.
572

573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
    .. 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.

    Examples
    --------
    Define a site.

    >>> builder[site] = value

    Print the value of a site.

    >>> print builder[site]

    Define a hopping.

    >>> builder[site1, site2] = value

    Delete a site.

    >>> del builder[site3]

    """
598

599
600
601
602
603
    def __init__(self, symmetry=None):
        if symmetry is None:
            symmetry = NoSymmetry()
        self.symmetry = symmetry
        self.leads = []
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
        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
    # values.)
    #
    # 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
Christoph Groth's avatar
Christoph Groth committed
632
        raise ValueError(tail, head)
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647

    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]
Christoph Groth's avatar
Christoph Groth committed
648
        i = 2 + 2 * heads.index(head)
649
650
651
652
653
654
655
656
657
        del hvhv[i : i + 2]

    def _out_neighbors(self, tail):
        hvhv = self.H.get(tail, ())
        return islice(hvhv, 2, None, 2)

    def _out_degree(self, tail):
        hvhv = self.H.get(tail, ())
        return len(hvhv) // 2 - 1
658

659
    # TODO: write a test for this method.
660
661
662
663
664
665
666
667
668
669
670
671
    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.
        """
        result = object.__new__(Builder)
        result.symmetry = self.symmetry.reversed()
        if self.leads:
            raise ValueError('System to be reversed may not have leads.')
        result.leads = []
672
        result.H = self.H
673
674
675
        return result

    def __nonzero__(self):
676
        return bool(self.H)
677

678
679
    def expand(self, key):
        """
Christoph Groth's avatar
Christoph Groth committed
680
        Expand a general key into an iterator over simple keys.
681

Christoph Groth's avatar
Christoph Groth committed
682
683
684
685
        Parameters
        ----------
        key : builder key (see notes)
            The key to be expanded
686

Christoph Groth's avatar
Christoph Groth committed
687
688
        Notes
        -----
689
690
691
692
693
694
695
        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
696
697
698
699
700

        This method is internally used to expand the keys when getting or
        deleting items of a builder (i.e. ``sys[key] = value`` or ``del
        sys[key]``).

701
702
703
704
705
        """
        itr = iter((key,))
        iter_stack = [None]
        while iter_stack:
            for key in itr:
Christoph Groth's avatar
Christoph Groth committed
706
707
                while callable(key):
                    key = key(self)
708
709
                if isinstance(key, Site) or isinstance(key, tuple):
                    yield key
710
                else:
711
712
713
714
715
                    iter_stack.append(itr)
                    itr = iter(key)
                    break
            else:
                itr = iter_stack.pop()
716

717
718
    def _get_site(self, site):
        site = self.symmetry.to_fd(site)
719
        try:
720
            return self.H[site][1]
721
        except KeyError:
722
            raise KeyError(site)
723

724
    def _get_hopping(self, hopping):
725
726
        sym = self.symmetry
        try:
727
            a, b = hopping
728
        except:
729
            raise KeyError(hopping)
730
        try:
731
            a, b = sym.to_fd(a, b)
732
            value = self._get_edge(a, b)
Christoph Groth's avatar
Christoph Groth committed
733
        except ValueError:
734
            raise KeyError(hopping)
735
        if value is Other:
736
737
738
            if not sym.in_fd(b):
                b, a = sym.to_fd(b, a)
                assert not sym.in_fd(a)
739
            value = self._get_edge(b, a)
Christoph Groth's avatar
Christoph Groth committed
740
            if callable(value):
741
742
743
744
745
746
747
748
                assert not isinstance(value, HermConjOfFunc)
                value = HermConjOfFunc(value)
            else:
                value = herm_conj(value)
        return value

    def __getitem__(self, key):
        """Get the value of a single site or hopping."""
749
        if isinstance(key, Site):
750
            return self._get_site(key)
751
        elif isinstance(key, tuple):
752
            return self._get_hopping(key)
753
754
        else:
            raise KeyError(key)
755
756
757

    def __contains__(self, key):
        """Tell whether the system contains a site or hopping."""
758
759
        if isinstance(key, Site):
            site = self.symmetry.to_fd(key)
760
            return site in self.H
761
        elif isinstance(key, tuple):
762
            a, b = key
763
            a, b = self.symmetry.to_fd(a, b)
764
765
            hvhv = self.H.get(a, ())
            return b in islice(hvhv, 2, None, 2)
766
767
        else:
            raise KeyError(key)
768

769
    def _set_site(self, site, value):
770
        """Set a single site."""
771
        site = self.symmetry.to_fd(site)
772
773
774
775
776
        hvhv = self.H.setdefault(site, [])
        if hvhv:
            hvhv[1] = value
        else:
            hvhv[:] = [site, value]
777

778
    def _set_hopping(self, hopping, value):
779
780
781
        """Set a single hopping."""
        # Avoid nested HermConjOfFunc instances.
        try:
782
            a, b = hopping
783
        except:
784
            raise KeyError(hopping)
785
786
        if a == b:
            raise KeyError(hopping)
787
788
789
790
791
792
793
        if isinstance(value, HermConjOfFunc):
            a, b = b, a
            value = value.function

        sym = self.symmetry

        try:
794
            a, b = sym.to_fd(a, b)
795
            if sym.in_fd(b):
796
                # These two following lines make sure we do not waste space by
797
798
                # storing different instances of identical sites.  They also
                # verify that sites a and b already belong to the system.
Anton Akhmerov's avatar
Anton Akhmerov committed
799
800
                a = self.H[a][0]                 # Might fail.
                b = self.H[b][0]                 # Might fail.
801
                self._set_edge(a, b, value)      # Will work.
802
                self._set_edge(b, a, Other)      # Will work.
803
            else:
804
                b2, a2 = sym.to_fd(b, a)
805
                if b2 not in self.H:
806
807
                    raise KeyError()
                assert not sym.in_fd(a2)
Anton Akhmerov's avatar
Anton Akhmerov committed
808
                self._set_edge(a, b, value)      # Might fail.
809
                self._set_edge(b2, a2, Other)    # Will work.
810
        except KeyError:
811
            raise KeyError(hopping)
812
813

    def __setitem__(self, key, value):
814
815
816
817
818
819
820
821
822
823
        """Set a single site/hopping or a bunch of them."""
        # TODO: Once we can take Python 3 for granted, get rid of the if-clause
        # inside the loop by defining a special func that rebinds itself upon
        # the first call.
        func = None
        for sh in self.expand(key):
            if func is None:
                func = self._set_site if isinstance(sh, Site) \
                    else self._set_hopping
            func(sh, value)
824

825
    def _del_site(self, site):
826
827
        """Delete a single site and all associated hoppings."""
        tfd = self.symmetry.to_fd
828
        site = tfd(site)
829
        try:
830
831
832
            for neighbor in self._out_neighbors(site):
                if neighbor in self.H:
                    self._del_edge(neighbor, site)
833
834
835
                else:
                    assert not self.symmetry.in_fd(neighbor)
                    a, b = tfd(neighbor, site)
836
                    self._del_edge(a, b)
Christoph Groth's avatar
Christoph Groth committed
837
        except ValueError:
838
            raise KeyError(site)
839
        del self.H[site]
840

841
    def _del_hopping(self, hopping):
842
843
844
845
        """Delete a single hopping."""
        sym = self.symmetry

        try:
846
            a, b = hopping
847
        except:
848
            raise KeyError(hopping)
849
        try:
850
            a, b = sym.to_fd(a, b)
851
            if sym.in_fd(b):
852
853
                self._del_edge(a, b)
                self._del_edge(b, a)
854
            else:
855
                self._del_edge(a, b)
856
857
                b, a = sym.to_fd(b, a)
                assert not sym.in_fd(a)
858
                self._del_edge(b, a)
Christoph Groth's avatar
Christoph Groth committed
859
        except ValueError:
860
            raise KeyError(hopping)
861
862

    def __delitem__(self, key):
863
864
865
866
867
868
869
870
871
872
        """Delete a single site/hopping or bunch of them."""
        # TODO: Once we can take Python 3 for granted, get rid of the if-clause
        # inside the loop by defining a special func that rebinds itself upon
        # the first call.
        func = None
        for sh in self.expand(key):
            if func is None:
                func = self._del_site if isinstance(sh, Site) \
                    else self._del_hopping
            func(sh)
873
874
875

    def eradicate_dangling(self):
        """Keep deleting dangling sites until none are left."""
876
877
        sites = list(site for site in self.H
                     if self._out_degree(site) < 2)
878
        for site in sites:
Anton Akhmerov's avatar
Anton Akhmerov committed
879
880
            if site not in self.H:
                continue
881
            while site:
882
883
884
885
886
887
888
                neighbors = tuple(self._out_neighbors(site))
                if neighbors:
                    assert len(neighbors) == 1
                    neighbor = neighbors[0]
                    self._del_edge(neighbor, site)
                    if self._out_degree(neighbor) > 1:
                        neighbor = False
889
                else:
890
                    neighbor = False
891
                del self.H[site]
892
                site = neighbor
893
894
895

    def __iter__(self):
        """Return an iterator over all sites and hoppings."""
896
        return chain(self.H, self.hoppings())
897
898

    def sites(self):
Anton Akhmerov's avatar
Anton Akhmerov committed
899
900
901
902
903
904
        """Return a read-only set over all sites.

        The sites that are returned belong to the fundamental domain of the
        `Builder` symmetry, and are not necessarily the ones that were set
        initially (but always the equivalent ones).
        """
905
906
907
908
        try:
            return self.H.viewkeys()
        except AttributeError:
            return frozenset(self.H)
909
910
911

    def site_value_pairs(self):
        """Return an iterator over all (site, value) pairs."""
912
913
        for site, hvhv in self.H.iteritems():
            yield site, hvhv[1]
914
915

    def hoppings(self):
Anton Akhmerov's avatar
Anton Akhmerov committed
916
917
918
919
920
921
        """Return an iterator over all Builder hoppings.

        The hoppings that are returned belong to the fundamental domain of the
        `Builder` symmetry, and are not necessarily the ones that were set
        initially (but always the equivalent ones).
        """
922
923
        for tail, hvhv in self.H.iteritems():
            for head, value in edges(hvhv):
924
                if value is Other:
Anton Akhmerov's avatar
Anton Akhmerov committed
925
                    continue
926
                yield (tail, head)
927
928
929

    def hopping_value_pairs(self):
        """Return an iterator over all (hopping, value) pairs."""
930
931
        for tail, hvhv in self.H.iteritems():
            for head, value in edges(hvhv):
932
                if value is Other:
Anton Akhmerov's avatar
Anton Akhmerov committed
933
                    continue
934
                yield (tail, head), value
935
936
937

    def dangling(self):
        """Return an iterator over all dangling sites."""
938
939
        for site in self.H:
            if self._out_degree(site) < 2:
940
                yield site
941

942
    def degree(self, site):
943
        """Return the number of neighbors of a site."""
944
        site = self.symmetry.to_fd(site)
945
        return self._out_degree(site)
946

947
    def neighbors(self, site):
948
        """Return an iterator over all neighbors of a site."""
949
        a = self.symmetry.to_fd(site)
950
        return self._out_neighbors(a)
951
952

    def __iadd__(self, other_sys):
Daniel Jaschke's avatar
Daniel Jaschke committed
953
954
955
956
957
958
        for site, value in other_sys.site_value_pairs():
            self[site] = value
        for hop, value in other_sys.hopping_value_pairs():
            self[hop] = value
        self.leads.extend(other_sys.leads)
        return self
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975

    def attach_lead(self, lead_builder, origin=None):
        """Attach a lead to the builder, possibly adding missing sites.

        Parameters
        ----------
        lead_builder : `Builder` with 1D translational symmetry
            Builder of the lead which has to be attached.
        origin : `Site`
            The site which should belong to a domain where the lead should
            begin. It is used to attach a lead inside the system, e.g. to an
            inner radius of a ring.

        Raises
        ------
        ValueError
            If `lead_builder` does not have proper symmetry, has hoppings with
976
            range of more than one lead unit cell, or if it is not completely
977
978
979
980
981
982
983
984
            interrupted by the system.

        Notes
        -----
        This method is not fool-proof, i.e. if it returns an error, there is
        no guarantee that the system stayed unaltered.
        """
        sym = lead_builder.symmetry
985
        H = lead_builder.H
986

987
988
989
990
        if sym.num_directions != 1:
            raise ValueError('Only builders with a 1D symmetry are allowed.')
        for hopping in lead_builder.hoppings():
            if not -1 <= sym.which(hopping[1])[0] <= 1:
991
992
                msg = 'Hopping {0} connects non-neighboring lead unit cells.' +\
                      'Only nearest-cell hoppings are allowed ' +\
993
994
                      '(consider increasing the lead period).'
                raise ValueError(msg.format(hopping))
995
        if not H:
996
997
            raise ValueError('Lead to be attached contains no sites.')

998
        # Check if site families of the lead are present in the system (catches
999
        # a common and a hard to find bug).
1000
        families = set(site.family for site in H)
For faster browsing, not all history is shown. View entire blame