test_builder.py 62.2 KB
Newer Older
1
# Copyright 2011-2019 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
# https://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
# https://kwant-project.org/authors.
8

9
import warnings
Anton Akhmerov's avatar
Anton Akhmerov committed
10
import pickle
11
import itertools as it
12
13
14
15
16
import functools as ft
from random import Random

import numpy as np
import tinyarray as ta
17
import pytest
Christoph Groth's avatar
Christoph Groth committed
18
from pytest import raises, warns
19
from numpy.testing import assert_almost_equal
20

21
import kwant
22
from kwant import builder, system
23
from kwant._common import ensure_rng, KwantDeprecationWarning
24
25


26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
@ft.total_ordering
class SimpleSiteFamily(system.SiteFamily):
    """A site family used as an example and for testing.

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

    It exists to provide a basic site family that can be used for testing the
    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
    `SimpleSiteFamily` when `kwant.lattice.Monatomic` would also work.
    """

    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)

    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


57
58
59
def test_bad_keys():

    def setitem(key):
60
        syst[key] = None
61

62
    fam = SimpleSiteFamily(norbs=1)
63
    syst = builder.Builder()
64
65
66

    failures = [
        # Invalid single keys
67
        ([syst.__contains__, syst.__getitem__, setitem, syst.__delitem__],
68
69
70
71
72
73
74
75
76
77
         [(TypeError, [123,
                       (0, 1),
                       (fam(0), 123),
                       (123, (fam(0)))]),
          (IndexError, [(fam(0),),
                        (fam(0), fam(1), fam(2))]),
          (ValueError, [(fam(0), fam(0)),
                        (fam(2), fam(2))])]),

        # Hoppings that contain sites that do not belong to the system
78
        ([syst.__getitem__, setitem, syst.__delitem__],
79
80
81
82
         [(KeyError, [(fam(0), fam(3)), (fam(2), fam(1)),
                      (fam(2), fam(3))])]),

        # Sequences containing a bad key.
83
        ([setitem, syst.__delitem__],
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
         [(TypeError, [[fam(0), fam(1), 123],
                       [fam(0), (fam(1),)],
                       [fam(0), (fam(1), fam(2))],
                       [(fam(0), fam(1)), (0, 1)],
                       [(fam(0), fam(1)), (fam(0), 123)],
                       [(fam(0), fam(1)), (123, fam(0))],
                       [(fam(0), fam(1)), fam(2)]]),
          (IndexError, [[(fam(0), fam(1)), (fam(2),)]]),
          (ValueError, [[(fam(0), fam(1)), (fam(2), fam(2))],
                        [(fam(0), fam(0)), (fam(1), fam(0))]]),
          (KeyError, [[(fam(0), fam(1)), (fam(0), fam(3))],
                      [(fam(0), fam(1)), (fam(2), fam(1))],
                      [(fam(1), fam(2)), (fam(0), fam(1))]])]),

        # Sites that do not belong to the system, also as part of a
        # sequence
100
        ([syst.__delitem__],
101
102
103
104
105
         [(KeyError, [fam(123),
                      [fam(0), fam(123)],
                      [fam(123), fam(1)]])]),

        # Various things that are not sites present in the system.
106
        ([syst.degree, lambda site: list(syst.neighbors(site))],
107
108
109
110
111
112
113
114
115
116
117
118
119
         [(TypeError, [123,
                       [0, 1, 2],
                       (0, 1),
                       (fam(0), fam(1)),
                       [fam(0), fam(1)],
                       [fam(1), fam(2)],
                       [fam(3), fam(0)]]),
          (KeyError, [fam(123)])])]

    for funcs, errors in failures:
        for error, keys in errors:
            for key in keys:
                for func in funcs:
120
121
                    syst[[fam(0), fam(1)]] = None
                    syst[fam(0), fam(1)] = None
122
                    try:
Anton Akhmerov's avatar
Anton Akhmerov committed
123
                        raises(error, func, key)
124
                    except AssertionError:
Joseph Weston's avatar
Joseph Weston committed
125
                        print(func, error, key)
126
127
128
                        raise


129
def test_site_families():
130
    syst = builder.Builder()
131
132
133
    fam = SimpleSiteFamily(norbs=1)
    ofam = SimpleSiteFamily(norbs=1)
    yafam = SimpleSiteFamily('another_name', norbs=1)
134

135
    syst[fam(0)] = 7
136
    assert syst[fam(0)] == 7
137

138
    assert len(set([fam, ofam, fam('a'), ofam('a'), yafam])) == 3
139
    syst[fam(1)] = 123
140
141
    assert syst[fam(1)] == 123
    assert syst[ofam(1)] == 123
Anton Akhmerov's avatar
Anton Akhmerov committed
142
    raises(KeyError, syst.__getitem__, yafam(1))
143

144
    # test site families compare equal/not-equal
145
    assert fam == ofam
Anton Akhmerov's avatar
Anton Akhmerov committed
146
    assert fam != yafam
147
    assert fam is not None
Anton Akhmerov's avatar
Anton Akhmerov committed
148
    assert fam != 'a'
149

150
    # test site families sorting
151
152
    fam1 = SimpleSiteFamily(norbs=1)
    fam2 = SimpleSiteFamily(norbs=2)
Anton Akhmerov's avatar
Anton Akhmerov committed
153
    assert fam1 < fam2  # string '1' is lexicographically less than '2'
154
155


156
class VerySimpleSymmetry(system.Symmetry):
157
158
    def __init__(self, period):
        self.period = period
159

160
161
162
    @property
    def num_directions(self):
        return 1
163

164
    def has_subgroup(self, other):
165
        if isinstance(other, system.NoSymmetry):
166
167
168
169
170
171
172
173
174
175
176
177
178
179
            return True
        elif isinstance(other, VerySimpleSymmetry):
            return not other.period % self.period
        else:
            return False

    def subgroup(self, *generators):
        generators = ta.array(generators)
        assert generators.shape == (1, 1)
        if generators.dtype != int:
            raise ValueError('Generators must be sequences of integers.')
        g = generators[0, 0]
        return VerySimpleSymmetry(g * self.period)

180
181
182
183
    def which(self, site):
        return ta.array((site.tag[0] // self.period,), int)

    def act(self, element, a, b=None):
184
        shifted = lambda site, delta: site.family(*ta.add(site.tag, delta))
185
186
        delta = (self.period * element[0],) + (len(a.tag) - 1) * (0,)
        if b is None:
Anton Akhmerov's avatar
Anton Akhmerov committed
187
            return shifted(a, delta)
188
        else:
Anton Akhmerov's avatar
Anton Akhmerov committed
189
            return shifted(a, delta), shifted(b, delta)
190
191
192
193
194


# The hoppings have to form a ring.  Some other implicit assumptions are also
# made.
def check_construction_and_indexing(sites, sites_fd, hoppings, hoppings_fd,
195
                                    unknown_hoppings, sym=None):
196
    fam = SimpleSiteFamily(norbs=1)
197
    syst = builder.Builder(sym)
198
    t, V = 1.0j, 0.0
199
    syst[sites] = V
200
    for site in sites:
201
202
        syst[site] = V
    syst[hoppings] = t
203
    for hopping in hoppings:
204
        syst[hopping] = t
205

206
    for hopping in unknown_hoppings:
Anton Akhmerov's avatar
Anton Akhmerov committed
207
        raises(KeyError, syst.__setitem__, hopping, t)
208

209
210
211
    assert (fam(5), fam(123)) not in syst
    assert (sites[0], fam(5, 123)) not in syst
    assert (fam(7, 8), sites[0]) not in syst
212
    for site in sites:
213
        assert site in syst
Anton Akhmerov's avatar
Anton Akhmerov committed
214
        assert syst[site] == V
215
216
    for hop in hoppings:
        rev_hop = hop[1], hop[0]
217
218
        assert hop in syst
        assert rev_hop in syst
219
220
        assert syst[hop] == t
        assert syst[rev_hop] == t.conjugate()
221

222
223
224
    assert syst.degree(sites[0]) == 2
    assert (sorted(s for s in syst.neighbors(sites[0])) ==
            sorted([sites[1], sites[-1]]))
225

226
    del syst[hoppings]
227
    assert list(syst.hoppings()) == []
228
    syst[hoppings] = t
229

230
    del syst[sites[0]]
231
232
233
    assert sorted(tuple(s) for s in syst.sites()) == sorted(sites_fd[1:])
    assert (sorted((a, b) for a, b in syst.hoppings()) ==
            sorted(hoppings_fd[1:-1]))
234

235
236
237
238
239
240
241
    assert (sorted((tuple(site.tag), value)
                        for site, value in syst.site_value_pairs()) ==
            sorted((tuple(site.tag), syst[site]) for site in syst.sites()))
    assert (sorted((tuple(a.tag), tuple(b.tag), value)
                   for (a, b), value in syst.hopping_value_pairs()) ==
            sorted((tuple(a.tag), tuple(b.tag), syst[a, b])
                   for a, b in syst.hoppings()))
242

Anton Akhmerov's avatar
Anton Akhmerov committed
243

244
245
def test_construction_and_indexing():
    # Without symmetry
246
    fam = SimpleSiteFamily(norbs=1)
247
248
249
250
    sites = [fam(0, 0), fam(0, 1), fam(1, 0)]
    hoppings = [(fam(0, 0), fam(0, 1)),
                (fam(0, 1), fam(1, 0)),
                (fam(1, 0), fam(0, 0))]
251
    unknown_hoppings = [(fam(0, 1), fam(7, 8)),
252
                        (fam(12, 14), fam(0, 1))]
253
    check_construction_and_indexing(sites, sites, hoppings, hoppings,
254
                                    unknown_hoppings)
255
256

    # With symmetry
257
258
259
260
261
262
263
264
265
266
    sites = [fam(0, 0), fam(1, 1), fam(2, 1), fam(4, 2)]
    sites_fd = [fam(0, 0), fam(1, 1), fam(0, 1), fam(0, 2)]
    hoppings = [(fam(0, 0), fam(1, 1)),
                (fam(1, 1), fam(2, 1)),
                (fam(2, 1), fam(4, 2)),
                (fam(4, 2), fam(0, 0))]
    hoppings_fd = [(fam(0, 0), fam(1, 1)),
                   (fam(1, 1), fam(2, 1)),
                   (fam(0, 1), fam(2, 2)),
                   (fam(0, 2), fam(-4, 0))]
267
    unknown_hoppings = [(fam(0, 0), fam(0, 3)), (fam(0, 4), fam(0, 0)),
268
269
                        (fam(0, 0), fam(2, 3)), (fam(2, 4), fam(0, 0)),
                        (fam(4, 2), fam(6, 3)), (fam(6, 4), fam(4, 2))]
270
271
    sym = VerySimpleSymmetry(2)
    check_construction_and_indexing(sites, sites_fd, hoppings, hoppings_fd,
272
                                    unknown_hoppings, sym)
273
274
275


def test_hermitian_conjugation():
276
    def f(i, j, arg):
277
        i, j = i.tag, j.tag
278
        if j[0] == i[0] + 1:
279
            return arg * ta.array([[1, 2j], [3 + 1j, 4j]])
280
281
282
        else:
            raise ValueError

283
    syst = builder.Builder()
284
    fam = SimpleSiteFamily(norbs=1)
285
    syst[fam(0)] = syst[fam(1)] = ta.identity(2)
286

287
288
289
    syst[fam(0), fam(1)] = f
    assert syst[fam(0), fam(1)] is f
    assert isinstance(syst[fam(1), fam(0)], builder.HermConjOfFunc)
290
291
    assert (syst[fam(1), fam(0)](fam(1), fam(0), 2) ==
            syst[fam(0), fam(1)](fam(0), fam(1), 2).conjugate().transpose())
292
293
294
    syst[fam(0), fam(1)] = syst[fam(1), fam(0)]
    assert isinstance(syst[fam(0), fam(1)], builder.HermConjOfFunc)
    assert syst[fam(1), fam(0)] is f
295
296
297


def test_value_equality_and_identity():
298
    m = ta.array([[1, 2], [3j, 4j]])
299
    syst = builder.Builder()
300
    fam = SimpleSiteFamily(norbs=1)
301

302
303
304
    syst[fam(0)] = m
    syst[fam(1)] = m
    assert syst[fam(1)] is m
305

306
    syst[fam(0), fam(1)] = m
307
    assert syst[fam(1), fam(0)] == m.transpose().conjugate()
308
    assert syst[fam(0), fam(1)] is m
309

310
    syst[fam(1), fam(0)] = m
311
    assert syst[fam(0), fam(1)] == m.transpose().conjugate()
312
    assert syst[fam(1), fam(0)] is m
313
314
315
316
317


def random_onsite_hamiltonian(rng):
    return 2 * rng.random() - 1

Anton Akhmerov's avatar
Anton Akhmerov committed
318

319
320
321
def random_hopping_integral(rng):
    return complex(2 * rng.random() - 1, 2 * rng.random() - 1)

Anton Akhmerov's avatar
Anton Akhmerov committed
322

323
def check_onsite(fsyst, sites, subset=False, check_values=True):
324
    vectorized = system.is_vectorized(fsyst)
325
326

    if vectorized:
327
        site_offsets, _, _ = fsyst.site_ranges.transpose()
328

329
    freq = {}
330
331
    for node in range(fsyst.graph.num_nodes):
        site = fsyst.sites[node].tag
332
333
        freq[site] = freq.get(site, 0) + 1
        if check_values and site in sites:
334
335
336
337
338
339
340
341
342
343
344
345
            if vectorized:
                term = fsyst._onsite_term_by_site_id[node]
                value = fsyst._term_values[term]
                if callable(value):
                    assert value is sites[site]
                else:
                    (w, _), (off, _) = fsyst.subgraphs[fsyst.terms[term].subgraph]
                    node_off = node - site_offsets[w]
                    selector = np.searchsorted(off, node_off)
                    assert np.allclose(value[selector], sites[site])
            else:
                assert fsyst.onsites[node][0] is sites[site]
346
    if not subset:
347
        # Check that all sites of `fsyst` are in `sites`.
Joseph Weston's avatar
Joseph Weston committed
348
        for site in freq.keys():
349
            assert site in sites
350
    # Check that all sites of `sites` are in `fsyst`.
351
    for site in sites:
352
        assert freq[site] == 1
353

Anton Akhmerov's avatar
Anton Akhmerov committed
354

355
def check_hoppings(fsyst, hops):
356
    vectorized = system.is_vectorized(fsyst)
357
358

    if vectorized:
359
        site_offsets, _, _ = fsyst.site_ranges.transpose()
360

361
    assert fsyst.graph.num_edges == 2 * len(hops)
362
    for edge_id, edge in enumerate(fsyst.graph):
363
364
365
366
367
368
        i, j = edge
        tail = fsyst.sites[i].tag
        head = fsyst.sites[j].tag

        if vectorized:
            term = fsyst._hopping_term_by_edge_id[edge_id]
369
370
371
372
373
374
            # Map sites in previous cell to fundamental domain; vectorized
            # infinite systems only store sites in the FD. We already know
            # that this hopping is between unit cells because this is encoded
            # in the edge_id and term id.
            if system.is_infinite(fsyst):
                i, j = i % fsyst.cell_size, j % fsyst.cell_size
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
            if term < 0:  # Hermitian conjugate
                assert (head, tail) in hops
            else:
                value = fsyst._term_values[term]
                assert (tail, head) in hops
                if callable(value):
                    assert value is hops[tail, head]
                else:
                    dtype = np.dtype([('f0', int), ('f1', int)])
                    subgraph = fsyst.terms[term].subgraph
                    (to_w, from_w), hoppings = fsyst.subgraphs[subgraph]
                    hop = (i - site_offsets[to_w], j - site_offsets[from_w])
                    hop = np.array(hop, dtype=dtype)
                    hoppings = hoppings.transpose().view(dtype).reshape(-1)
                    selector = np.recarray.searchsorted(hoppings, hop)
                    assert np.allclose(value[selector], hops[tail, head])
391
        else:
392
393
394
395
396
397
            value = fsyst.hoppings[edge_id][0]
            if value is builder.Other:
                assert (head, tail) in hops
            else:
                assert (tail, head) in hops
                assert value is hops[tail, head]
398

399
400
def check_id_by_site(fsyst):
    for i, site in enumerate(fsyst.sites):
401
        assert fsyst.id_by_site[site] == i
402

Anton Akhmerov's avatar
Anton Akhmerov committed
403

404
405
@pytest.mark.parametrize("vectorize", [False, True])
def test_finalization(vectorize):
406
407
408
409
410
411
412
413
    """Test the finalization of finite and infinite systems.

    In order to exactly verify the finalization, low-level features of the
    build module are used directly.  This is not the way one would use a
    finalized system in normal code.
    """
    def set_sites(dest):
        while len(dest) < n_sites:
414
            site = ta.array([rng.randrange(size), rng.randrange(size)])
415
416
417
418
419
            if site not in dest:
                dest[site] = random_onsite_hamiltonian(rng)

    def set_hops(dest, sites):
        while len(dest) < n_hops:
420
            a, b = rng.sample(list(sites), 2)
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
            if (a, b) not in dest and (b, a) not in dest:
                dest[a, b] = random_hopping_integral(rng)

    rng = Random(123)
    size = 20
    n_sites = 120
    n_hops = 500

    # Make scattering region blueprint.
    sr_sites = {}
    set_sites(sr_sites)
    sr_hops = {}
    set_hops(sr_hops, sr_sites)

    # Make lead blueprint.
    possible_neighbors = rng.sample(list(sr_sites), n_sites // 2)
    lead_sites = {}
    for pn in possible_neighbors:
        lead_sites[pn] = random_hopping_integral(rng)
    set_sites(lead_sites)
    lead_hops = {}        # Hoppings within a single lead unit cell
    set_hops(lead_hops, lead_sites)
    lead_sites_list = list(lead_sites)
    neighbors = set()
Joseph Weston's avatar
Joseph Weston committed
445
    for i in range(n_hops):
446
447
448
449
        while True:
            a = rng.choice(lead_sites_list)
            b = rng.choice(possible_neighbors)
            neighbors.add(b)
450
            b = ta.array([b[0] - size, b[1]])
451
452
453
454
455
456
457
458
            if rng.randrange(2):
                a, b = b, a
            if (a, b) not in lead_hops and (b, a) not in lead_hops:
                break
        lead_hops[a, b] = random_hopping_integral(rng)
    neighbors = sorted(neighbors)

    # Build scattering region from blueprint and test it.
459
    syst = builder.Builder(vectorize=vectorize)
460
    fam = kwant.lattice.general(ta.identity(2), norbs=1)
Joseph Weston's avatar
Joseph Weston committed
461
    for site, value in sr_sites.items():
462
        syst[fam(*site)] = value
Joseph Weston's avatar
Joseph Weston committed
463
    for hop, value in sr_hops.items():
464
465
        syst[fam(*hop[0]), fam(*hop[1])] = value
    fsyst = syst.finalized()
466
    check_id_by_site(fsyst)
467
468
    check_onsite(fsyst, sr_sites)
    check_hoppings(fsyst, sr_hops)
469
    # check that sites are sorted
470
    assert tuple(fsyst.sites) == tuple(sorted(fam(*site) for site in sr_sites))
471
472

    # Build lead from blueprint and test it.
473
474
    lead = builder.Builder(kwant.TranslationalSymmetry((size, 0)),
                           vectorize=vectorize)
Joseph Weston's avatar
Joseph Weston committed
475
    for site, value in lead_sites.items():
476
477
        shift = rng.randrange(-5, 6) * size
        site = site[0] + shift, site[1]
478
        lead[fam(*site)] = value
479
480
481
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        lead.finalized()        # Trigger the warning.
482
        assert len(w) == 1
483
484
        assert issubclass(w[0].category, RuntimeWarning)
        assert "disconnected" in str(w[0].message)
Joseph Weston's avatar
Joseph Weston committed
485
    for (a, b), value in lead_hops.items():
486
487
488
        shift = rng.randrange(-5, 6) * size
        a = a[0] + shift, a[1]
        b = b[0] + shift, b[1]
489
        lead[fam(*a), fam(*b)] = value
490
491
    flead = lead.finalized()
    all_sites = list(lead_sites)
492
    all_sites.extend(ta.array([x - size, y]) for (x, y) in neighbors)
493
    check_id_by_site(fsyst)
494
495
496
497
    check_onsite(flead, all_sites, check_values=False)
    check_onsite(flead, lead_sites, subset=True)
    check_hoppings(flead, lead_hops)

498
    # Attach lead to system with empty interface.
499
    syst.leads.append(builder.BuilderLead(lead, ()))
Anton Akhmerov's avatar
Anton Akhmerov committed
500
    raises(ValueError, syst.finalized)
501
502

    # Attach lead with improper interface.
503
    syst.leads[-1] = builder.BuilderLead(
504
        lead, 2 * tuple(system.Site(fam, n) for n in neighbors))
Anton Akhmerov's avatar
Anton Akhmerov committed
505
    raises(ValueError, syst.finalized)
506
507

    # Attach lead properly.
508
    syst.leads[-1] = builder.BuilderLead(
509
        lead, (system.Site(fam, n) for n in neighbors))
510
    fsyst = syst.finalized()
511
512
513
    assert len(fsyst.lead_interfaces) == 1
    assert ([fsyst.sites[i].tag for i in fsyst.lead_interfaces[0]] ==
            neighbors)
514
515

    # test that we cannot finalize a system with a badly sorted interface order
516
    raises(ValueError, builder.InfiniteSystem, lead,
517
           [system.Site(fam, n) for n in reversed(neighbors)])
518
    # site ordering independent of whether interface was specified
519
    flead_order = builder.InfiniteSystem(lead, [system.Site(fam, n)
520
                                                for n in neighbors])
521
    assert tuple(flead.sites) == tuple(flead_order.sites)
522
523


524
    syst.leads[-1] = builder.BuilderLead(
525
        lead, (system.Site(fam, n) for n in neighbors))
526
    fsyst = syst.finalized()
527
528
529
    assert len(fsyst.lead_interfaces) == 1
    assert ([fsyst.sites[i].tag for i in fsyst.lead_interfaces[0]] ==
            neighbors)
530

531
    # Add a hopping to the lead which couples two next-nearest cells and check
532
533
534
535
    # whether this leads to an error.
    a = rng.choice(lead_sites_list)
    b = rng.choice(possible_neighbors)
    b = b[0] + 2 * size, b[1]
536
    lead[fam(*a), fam(*b)] = random_hopping_integral(rng)
Anton Akhmerov's avatar
Anton Akhmerov committed
537
    raises(ValueError, lead.finalized)
538
539


540
541
542
543
544
545
546
547
548
549
550
551
552
def _make_system_from_sites(sites, vectorize):
    syst = builder.Builder(vectorize=vectorize)
    for s in sites:
        norbs = s.family.norbs
        if norbs:
            syst[s] = np.eye(s.family.norbs)
        else:
            syst[s] = None
    return syst.finalized()


@pytest.mark.parametrize("vectorize", [False, True])
def test_site_ranges(vectorize):
553
554
555
556
557
558
559
    lat1a = kwant.lattice.chain(norbs=1, name='a')
    lat1b = kwant.lattice.chain(norbs=1, name='b')
    lat2 = kwant.lattice.chain(norbs=2)

    # simple case -- single site family
    for lat in (lat1a, lat2):
        sites = list(map(lat, range(10)))
560
561
        syst = _make_system_from_sites(sites, vectorize)
        ranges = syst.site_ranges
562
        expected = [(0, lat.norbs, 0), (10, 0, 10 * lat.norbs)]
563
        assert np.array_equal(ranges, expected)
564
565

    # pair of site families
566
567
568
569
    sites = it.chain(map(lat1a, range(4)), map(lat1b, range(6)))
    syst = _make_system_from_sites(sites, vectorize)
    expected = [(0, 1, 0), (4, 1, 4), (10, 0, 10)]
    assert np.array_equal(expected, syst.site_ranges)
570
571
    sites = it.chain(map(lat2, range(4)), map(lat1a, range(6)),
                     map(lat1b, range(4)))
572
    syst = _make_system_from_sites(sites, vectorize)
573
    expected = [(0, 2, 0), (4, 1, 4*2), (10, 1, 4*2+6), (14, 0, 4*2+10)]
574
    assert np.array_equal(expected, syst.site_ranges)
575
576
577
578

    # test with an actual builder
    for lat in (lat1a, lat2):
        sites = list(map(lat, range(10)))
579
        syst = kwant.Builder(vectorize=vectorize)
580
581
582
        syst[sites] = np.eye(lat.norbs)
        ranges = syst.finalized().site_ranges
        expected = [(0, lat.norbs, 0), (10, 0, 10 * lat.norbs)]
583
584
585
586
587
588
589
590
591
592
593
594
595
        assert np.array_equal(ranges, expected)
        if not vectorize:  # vectorized systems *must* have all norbs
            # poison system with a single site with no norbs defined.
            # Also catch the deprecation warning.
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                syst[kwant.lattice.chain(norbs=None)(0)] = 1
            ranges = syst.finalized().site_ranges
            assert ranges is None


@pytest.mark.parametrize("vectorize", [False, True])
def test_hamiltonian_evaluation(vectorize):
596
597
598
599
600
601
602
    def f_onsite(site):
        return site.tag[0]

    def f_hopping(a, b):
        a, b = a.tag, b.tag
        return complex(a[0] + b[0], a[1] - b[1])

603
604
605
606
607
608
609
    def f_onsite_vectorized(sites):
        return sites.tags[:, 0]

    def f_hopping_vectorized(a, b):
        a, b = a.tags, b.tags
        return a[:, 0] + b[:, 0] + 1j * (a[:, 1] - b[:, 1])

610
611
612
    tags = [(0, 0), (1, 1), (2, 2), (3, 3)]
    edges = [(0, 1), (0, 2), (0, 3), (1, 2)]

613
    syst = builder.Builder(vectorize=vectorize)
614
    fam = SimpleSiteFamily(norbs=1)
615
    sites = [fam(*tag) for tag in tags]
616
617
618
619
620
621
622
    hoppings = [(sites[i], sites[j]) for i, j in edges]
    if vectorize:
        syst[sites] = f_onsite_vectorized
        syst[hoppings] = f_hopping_vectorized
    else:
        syst[sites] = f_onsite
        syst[hoppings] = f_hopping
623
    fsyst = syst.finalized()
624

625
626
    assert fsyst.graph.num_nodes == len(tags)
    assert fsyst.graph.num_edges == 2 * len(edges)
627
628

    for i in range(len(tags)):
629
        site = fsyst.sites[i]
630
        assert site in sites
631
632
633
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=KwantDeprecationWarning)
            assert fsyst.hamiltonian(i, i) == f_onsite(site)
634

635
636
637
    for t, h in fsyst.graph:
        tsite = fsyst.sites[t]
        hsite = fsyst.sites[h]
638
639
640
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=KwantDeprecationWarning)
            assert fsyst.hamiltonian(t, h) == f_hopping(tsite, hsite)
641

642
643
644
645
646
647
648
649
650
651
    # test when user-function raises errors
    def onsite_raises(site):
        raise ValueError()

    def hopping_raises(a, b):
        raise ValueError('error message')

    def test_raising(fsyst, hop):
        a, b = hop
        # exceptions are converted to kwant.UserCodeError and we add our message
652
        with raises(kwant.UserCodeError) as exc_info:
653
654
655
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", category=KwantDeprecationWarning)
                fsyst.hamiltonian(a, a)
656
        msg = 'Error occurred in user-supplied value function "onsite_raises"'
657
        assert msg in str(exc_info.value)
658
659

        for hop in [(a, b), (b, a)]:
660
            with raises(kwant.UserCodeError) as exc_info:
661
662
663
664
                with warnings.catch_warnings():
                    # Ignore deprecation warnings for 'hamiltonian'
                    warnings.filterwarnings("ignore", category=KwantDeprecationWarning)
                    fsyst.hamiltonian(*hop)
665
666
            msg = ('Error occurred in user-supplied '
                   'value function "hopping_raises"')
667
            assert msg in str(exc_info.value)
668
669
670
671
672
673

    # test with finite system
    new_hop = (fam(-1, 0), fam(0, 0))
    syst[new_hop[0]] = onsite_raises
    syst[new_hop] = hopping_raises
    fsyst = syst.finalized()
674
    hop = tuple(map(fsyst.id_by_site.__getitem__, new_hop))
675
676
677
678
    test_raising(fsyst, hop)

    # test with infinite system
    inf_syst = kwant.Builder(VerySimpleSymmetry(2))
679
680
    for k, v in it.chain(syst.site_value_pairs(), syst.hopping_value_pairs()):
        inf_syst[k] = v
681
    inf_fsyst = inf_syst.finalized()
682
    hop = tuple(map(inf_fsyst.id_by_site.__getitem__, new_hop))
683
684
    test_raising(inf_fsyst, hop)

685

686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
def test_vectorized_hamiltonian_evaluation():

    def onsite(site):
        return site.tag[0]

    def vectorized_onsite(sites):
        return sites.tags[:, 0]

    def hopping(to_site, from_site):
        a, b = to_site.tag, from_site.tag
        return a[0] + b[0] + 1j * (a[1] - b[1])

    def vectorized_hopping(to_sites, from_sites):
        a, b = to_sites.tags, from_sites.tags
        return a[:, 0] + b[:, 0] + 1j * (a[:, 1] - b[:, 1])

    tags = [(0, 0), (1, 1), (2, 2), (3, 3)]
    edges = [(0, 1), (0, 2), (0, 3), (1, 2)]

705
    fam = SimpleSiteFamily(norbs=1)
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
    sites = [fam(*tag) for tag in tags]
    hops = [(fam(*tags[i]), fam(*tags[j])) for (i, j) in edges]

    syst_simple = builder.Builder(vectorize=False)
    syst_simple[sites] = onsite
    syst_simple[hops] = hopping
    fsyst_simple = syst_simple.finalized()

    syst_vectorized = builder.Builder(vectorize=True)
    syst_vectorized[sites] = vectorized_onsite
    syst_vectorized[hops] = vectorized_hopping
    fsyst_vectorized = syst_vectorized.finalized()

    assert fsyst_vectorized.graph.num_nodes == len(tags)
    assert fsyst_vectorized.graph.num_edges == 2 * len(edges)
    assert len(fsyst_vectorized.site_arrays) == 1
    assert fsyst_vectorized.site_arrays[0] == system.SiteArray(fam, tags)

    assert np.allclose(
        fsyst_simple.hamiltonian_submatrix(),
        fsyst_vectorized.hamiltonian_submatrix(),
    )

    for i in range(len(tags)):
        site = fsyst_vectorized.sites[i]
        assert site in sites
732
733
734
735
736
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=KwantDeprecationWarning)
            assert (
                fsyst_vectorized.hamiltonian(i, i)
                == fsyst_simple.hamiltonian(i, i))
737
738

    for t, h in fsyst_vectorized.graph:
739
740
741
742
743
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=KwantDeprecationWarning)
            assert (
                fsyst_vectorized.hamiltonian(t, h)
                == fsyst_simple.hamiltonian(t, h))
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

    # Test infinite system, including hoppings that go both ways into
    # the next cell
    lat = kwant.lattice.square(norbs=1)

    syst_vectorized = builder.Builder(kwant.TranslationalSymmetry((-1, 0)),
                                      vectorize=True)
    syst_vectorized[lat(0, 0)] = 4
    syst_vectorized[lat(0, 1)] = 5
    syst_vectorized[lat(0, 2)] = vectorized_onsite
    syst_vectorized[(lat(1, 0), lat(0, 0))] = 1j
    syst_vectorized[(lat(2, 1), lat(1, 1))] = vectorized_hopping
    fsyst_vectorized = syst_vectorized.finalized()

    syst_simple = builder.Builder(kwant.TranslationalSymmetry((-1, 0)),
                                      vectorize=False)
    syst_simple[lat(0, 0)] = 4
    syst_simple[lat(0, 1)] = 5
    syst_simple[lat(0, 2)] = onsite
    syst_simple[(lat(1, 0), lat(0, 0))] = 1j
    syst_simple[(lat(2, 1), lat(1, 1))] = hopping
    fsyst_simple = syst_simple.finalized()

    assert np.allclose(
        fsyst_vectorized.cell_hamiltonian(),
        fsyst_simple.cell_hamiltonian(),
    )
    assert np.allclose(
        fsyst_vectorized.inter_cell_hopping(),
        fsyst_simple.inter_cell_hopping(),
    )


777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
def test_vectorized_value_normalization():
    # Here we test whether all legal shapes for values for vectorized Builders
    # are accepted:
    # + single scalars are interpreted as 1x1 matrices, and broadcast into an
    #   (N, 1, 1) array
    # + single (n, m) matrices are broadcast into an (N, n, m) array
    # + a shape (N,) array is interpreted as an (N, 1, 1) array

    sz = np.array([[1, 0], [0, -1]])

    scalars = [
        2,
        lambda s: 2,
        lambda s: [2] * len(s)
    ]
    matrices = [
        sz,
        lambda s: sz,
        lambda s: [sz] * len(s)
    ]
    inter_lattice_matrices =[
        [[1, 0]],
        lambda a, b: [[1, 0]],
        lambda a, b: [[[1, 0]]] * len(a)
    ]

    lat_a = kwant.lattice.chain(norbs=1)
    lat_b = kwant.lattice.chain(norbs=2)

    hams = []
    for s, m, v in it.product(scalars, matrices, inter_lattice_matrices):
        syst = builder.Builder(vectorize=True)
        syst[map(lat_a, range(10))] = s
        syst[map(lat_b, range(10))] = m
        syst[((lat_a(i), lat_b(i)) for i in range(10))] = v
        h = syst.finalized().hamiltonian_submatrix()
        hams.append(h)
    assert np.all(hams == hams[0])


817
@pytest.mark.parametrize("sym", [
818
    system.NoSymmetry(),
819
820
821
    kwant.TranslationalSymmetry([-1]),
])
def test_vectorized_requires_norbs(sym):
822
823
824
825

    # Catch deprecation warning for lack of norbs
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
826
        lat = kwant.lattice.chain()
827

828
    syst = builder.Builder(sym, vectorize=True)
829
830
    syst[lat(0)] = syst[lat(1)] = 2
    syst[lat(1), lat(0)] = -1
831
832
833
834

    raises(ValueError, syst.finalized)


835
836
837
838
839
def test_dangling():
    def make_system():
        #        1
        #       / \
        #    3-0---2-4-5  6-7  8
840
        syst = builder.Builder()
841
        fam = SimpleSiteFamily(norbs=1)
842
843
844
845
846
847
848
        syst[(fam(i) for i in range(9))] = None
        syst[[(fam(0), fam(1)), (fam(1), fam(2)), (fam(2), fam(0))]] = None
        syst[[(fam(0), fam(3)), (fam(2), fam(4)), (fam(4), fam(5))]] = None
        syst[fam(6), fam(7)] = None
        return syst

    syst0 = make_system()
849
850
    assert (sorted(site.tag for site in syst0.dangling()) ==
            sorted([(3,), (5,), (6,), (7,), (8,)]))
851
    syst0.eradicate_dangling()
852

853
    syst1 = make_system()
854
    while True:
855
        dangling = list(syst1.dangling())
Anton Akhmerov's avatar
Anton Akhmerov committed
856
857
        if not dangling:
            break
858
        del syst1[dangling]
859

860
861
862
863
    assert (sorted(site.tag for site in syst0.sites()) ==
            sorted([(0,), (1,), (2,)]))
    assert (sorted(site.tag for site in syst0.sites()) ==
            sorted(site.tag for site in syst1.sites()))
864

865
def test_dangling_with_symmetry():
866
867
868
869
870
871
872
873
874
875
876
877
878
879
    # lenght = 3 is the special case that gives all dangling bonds
    # lenght = 4 is the standard case of dangling bonds across the symmetry
    for length in [3, 4]:
        symm = kwant.TranslationalSymmetry((length, 0))
        lat = kwant.lattice.square(norbs=1)
        syst = kwant.Builder(symmetry=symm)
        for x in range(length):
            syst[lat(x, 0)] = 0
            syst[lat(x, 1)] = 0
        syst[lat.neighbors()] = -1
        # remove neighbors of site at (lenght-1, 0), dangling across symm
        del syst[lat(length - 2, 0)]
        del syst[lat(length - 1, 1)]
        syst.eradicate_dangling()
880
881

def test_builder_with_symmetry():
882
    g = kwant.lattice.general(ta.identity(3), norbs=1)
883
    sym = kwant.TranslationalSymmetry((0, 0, 3), (0, 2, 0))
884
    syst = builder.Builder(sym)
885
886

    t, V = 1.0j, 0.0
887
888
889
890
891
892
893
894
895
896
897
    hoppings = [(g(5, 0, 0), g(0, 5, 0)),
                (g(0, 5, 0), g(0, 0, 5)),
                (g(0, 0, 5), g(5, 0, 0)),
                (g(0, 3, 0), g(0, 0, 5)),
                (g(0, 7, -6), g(5, 6, -6))]
    hoppings_fd = [(g(5, 0, 0), g(0, 5, 0)),
                   (g(0, 1, 0), g(0, -4, 5)),
                   (g(0, 0, 2), g(5, 0, -3)),
                   (g(0, 1, 0), g(0, -2, 5)),
                   (g(0, 1, 0), g(5, 0, 0))]

898
899
    syst[(a for a, b in hoppings)] = V
    syst[hoppings] = t
900

901
    # TODO: Once Tinyarray supports "<" the conversion to tuple can be removed.
902
903
    assert (sorted(tuple(site.tag) for site in syst.sites()) ==
            sorted(set(tuple(hop[0].tag) for hop in hoppings_fd)))
904
905
    for sites in hoppings_fd:
        for site in sites:
906
            assert site in syst
907
            assert syst[site] == V
908

909
    # TODO: Once Tinyarray supports "<" the conversion to tuple can be removed.
910
911
    assert (sorted((tuple(a.tag), tuple(b.tag)) for a, b in syst.hoppings()) ==
            sorted((tuple(a.tag), tuple(b.tag)) for a, b in hoppings_fd))
912
913
    for hop in hoppings_fd:
        rhop = hop[1], hop[0]
914
915
        assert hop in syst
        assert rhop in syst
916
917
        assert syst[hop] == t
        assert syst[rhop] == t.conjugate()
918

919
920
    del syst[g(0, 6, -4), g(0, 11, -9)]
    assert (g(0, 1, 0), g(0, -4, 5)) not in syst
921

922
    del syst[g(0, 3, -3)]
923
924
    assert (list((a.tag, b.tag) for a, b in syst.hoppings()) == [((0, 0, 2),
                                                                  (5, 0, -3))])
925
926


Anton Akhmerov's avatar
Anton Akhmerov committed
927
def test_fill():
928
    g = kwant.lattice.square(norbs=1)
Anton Akhmerov's avatar
Anton Akhmerov committed
929
930
931
932
    sym_x = kwant.TranslationalSymmetry((-1, 0))
    sym_xy = kwant.TranslationalSymmetry((-1, 0), (0, 1))

    template_1d = builder.Builder(sym_x)
933
934
    template_1d[g(0, 0)] = None
    template_1d[g.neighbors()] = None
Anton Akhmerov's avatar
Anton Akhmerov committed
935
936
937
938

    def line_200(site):
        return -100 <= site.pos[0] < 100

939
    ## Test that copying a builder by "fill" preserves everything.
940
941
    for sym, func in [(kwant.TranslationalSymmetry(*np.diag([3, 4, 5])),
                       lambda pos: True),
942
                      (system.NoSymmetry(),
943
                       lambda pos: ta.dot(pos, pos) < 17)]:
944
        cubic = kwant.lattice.general(ta.identity(3), norbs=1)
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969

        # Make a weird system.
        orig = kwant.Builder(sym)
        sites = cubic.shape(func, (0, 0, 0))
        for i, site in enumerate(orig.expand(sites)):
            if i % 7 == 0:
                continue
            orig[site] = i
        for i, hopp in enumerate(orig.expand(cubic.neighbors(1))):
            if i % 11 == 0:
                continue
            orig[hopp] = i * 1.2345
        for i, hopp in enumerate(orig.expand(cubic.neighbors(2))):
            if i % 13 == 0:
                continue
            orig[hopp] = i * 1j

        # Clone the original using fill.
        clone = kwant.Builder(sym)
        clone.fill(orig, lambda s: True, (0, 0, 0))

        # Verify that both are identical.
        assert set(clone.site_value_pairs()) == set(orig.site_value_pairs())
        assert (set(clone.hopping_value_pairs())
                == set(orig.hopping_value_pairs()))
970

971
    ## Test for warning when "start" is outside the filling shape.
972
    target = builder.Builder()
Christoph Groth's avatar
Christoph Groth committed
973
974
975
    for start in [(-101, 0), (101, 0)]:
        with warns(RuntimeWarning):
            target.fill(template_1d, line_200, start)
976

Christoph Groth's avatar
Christoph Groth committed
977
978
979
    ## Test filling of infinite builder.
    for n in [1, 2, 4]:
        sym_n = kwant.TranslationalSymmetry((n, 0))
980
981
982
983
984
985
986
        for start in [g(0, 0), g(20, 0)]:
            target = builder.Builder(sym_n)
            sites = target.fill(template_1d, lambda s: True, start,
                                max_sites=10)
            assert len(sites) == n
            assert len(list(target.hoppings())) == n
            assert set(sym_n.to_fd(s) for s in sites) == set(target.sites())
Christoph Groth's avatar
Christoph Groth committed
987

Anton Akhmerov's avatar
Anton Akhmerov committed
988
989
990
991
    ## test max_sites
    target = builder.Builder()
    for max_sites in (-1, 0):
        with raises(ValueError):
992
993
            target.fill(template_1d, lambda site: True, g(0, 0),
                        max_sites=max_sites)
994
        assert len(list(target.sites())) == 0
Anton Akhmerov's avatar
Anton Akhmerov committed
995
996
    target = builder.Builder()
    with raises(RuntimeError):
997
        target.fill(template_1d, line_200, g(0, 0) , max_sites=10)
Anton Akhmerov's avatar
Anton Akhmerov committed
998
999
    ## test filling
    target = builder.Builder()
1000
    added_sites = target.fill(template_1d, line_200, g(0, 0))
Anton Akhmerov's avatar
Anton Akhmerov committed
1001
    assert len(added_sites) == 200
1002
    # raise warning if target already contains all starting sites