diff --git a/kwant/tests/test_builder.py b/kwant/tests/test_builder.py
index f41652a85c1306fbd25e6b1d5015dd1a550d50b2..f9832e3874f3bc5bb50948fbf6e739ce2301e91d 100644
--- a/kwant/tests/test_builder.py
+++ b/kwant/tests/test_builder.py
@@ -1,4 +1,4 @@
-# Copyright 2011-2018 Kwant authors.
+# Copyright 2011-2019 Kwant authors.
#
# 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
@@ -290,12 +290,28 @@ def random_hopping_integral(rng):
def check_onsite(fsyst, sites, subset=False, check_values=True):
+ vectorized = isinstance(fsyst, (system.FiniteVectorizedSystem, system.InfiniteVectorizedSystem))
+
+ if vectorized:
+ site_offsets = np.cumsum([0] + [len(s) for s in fsyst.site_arrays])
+
freq = {}
for node in range(fsyst.graph.num_nodes):
site = fsyst.sites[node].tag
freq[site] = freq.get(site, 0) + 1
if check_values and site in sites:
- assert fsyst.onsites[node][0] is sites[site]
+ 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]
if not subset:
# Check that all sites of `fsyst` are in `sites`.
for site in freq.keys():
@@ -306,24 +322,50 @@ def check_onsite(fsyst, sites, subset=False, check_values=True):
def check_hoppings(fsyst, hops):
+ vectorized = isinstance(fsyst, (system.FiniteVectorizedSystem, system.InfiniteVectorizedSystem))
+
+ if vectorized:
+ site_offsets = np.cumsum([0] + [len(s) for s in fsyst.site_arrays])
+
assert fsyst.graph.num_edges == 2 * len(hops)
for edge_id, edge in enumerate(fsyst.graph):
- tail, head = edge
- tail = fsyst.sites[tail].tag
- head = fsyst.sites[head].tag
- value = fsyst.hoppings[edge_id][0]
- if value is builder.Other:
- assert (head, tail) in hops
+ i, j = edge
+ tail = fsyst.sites[i].tag
+ head = fsyst.sites[j].tag
+
+ if vectorized:
+ term = fsyst._hopping_term_by_edge_id[edge_id]
+ 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])
else:
- assert (tail, head) in hops
- assert value is hops[tail, head]
+ 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]
def check_id_by_site(fsyst):
for i, site in enumerate(fsyst.sites):
assert fsyst.id_by_site[site] == i
-def test_finalization():
+@pytest.mark.parametrize("vectorize", [False, True])
+def test_finalization(vectorize):
"""Test the finalization of finite and infinite systems.
In order to exactly verify the finalization, low-level features of the
@@ -377,7 +419,7 @@ def test_finalization():
neighbors = sorted(neighbors)
# Build scattering region from blueprint and test it.
- syst = builder.Builder()
+ syst = builder.Builder(vectorize=vectorize)
fam = kwant.lattice.general(ta.identity(2), norbs=1)
for site, value in sr_sites.items():
syst[fam(*site)] = value
@@ -388,7 +430,7 @@ def test_finalization():
check_onsite(fsyst, sr_sites)
check_hoppings(fsyst, sr_hops)
# check that sites are sorted
- assert fsyst.sites == tuple(sorted(fam(*site) for site in sr_sites))
+ assert tuple(fsyst.sites) == tuple(sorted(fam(*site) for site in sr_sites))
# Build lead from blueprint and test it.
lead = builder.Builder(kwant.TranslationalSymmetry((size, 0)))
@@ -438,7 +480,7 @@ def test_finalization():
# site ordering independent of whether interface was specified
flead_order = builder.InfiniteSystem(lead, [system.Site(fam, n)
for n in neighbors])
- assert flead.sites == flead_order.sites
+ assert tuple(flead.sites) == tuple(flead_order.sites)
syst.leads[-1] = builder.BuilderLead(
@@ -457,47 +499,62 @@ def test_finalization():
raises(ValueError, lead.finalized)
-def test_site_ranges():
+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):
lat1a = kwant.lattice.chain(norbs=1, name='a')
lat1b = kwant.lattice.chain(norbs=1, name='b')
lat2 = kwant.lattice.chain(norbs=2)
- site_ranges = builder._site_ranges
# simple case -- single site family
for lat in (lat1a, lat2):
sites = list(map(lat, range(10)))
- ranges = site_ranges(sites)
+ syst = _make_system_from_sites(sites, vectorize)
+ ranges = syst.site_ranges
expected = [(0, lat.norbs, 0), (10, 0, 10 * lat.norbs)]
- assert ranges == expected
+ assert np.array_equal(ranges, expected)
# pair of site families
- sites = it.chain(map(lat1a, range(4)), map(lat1b, range(6)),
- map(lat1a, range(4)))
- expected = [(0, 1, 0), (4, 1, 4), (10, 1, 10), (14, 0, 14)]
- assert expected == site_ranges(tuple(sites))
+ 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)
sites = it.chain(map(lat2, range(4)), map(lat1a, range(6)),
map(lat1b, range(4)))
+ syst = _make_system_from_sites(sites, vectorize)
expected = [(0, 2, 0), (4, 1, 4*2), (10, 1, 4*2+6), (14, 0, 4*2+10)]
- assert expected == site_ranges(tuple(sites))
+ assert np.array_equal(expected, syst.site_ranges)
# test with an actual builder
for lat in (lat1a, lat2):
sites = list(map(lat, range(10)))
- syst = kwant.Builder()
+ syst = kwant.Builder(vectorize=vectorize)
syst[sites] = np.eye(lat.norbs)
ranges = syst.finalized().site_ranges
expected = [(0, lat.norbs, 0), (10, 0, 10 * lat.norbs)]
- assert ranges == expected
- # 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 == None
-
-
-def test_hamiltonian_evaluation():
+ 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):
def f_onsite(site):
return site.tag[0]
@@ -505,14 +562,26 @@ def test_hamiltonian_evaluation():
a, b = a.tag, b.tag
return complex(a[0] + b[0], a[1] - b[1])
+ 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])
+
tags = [(0, 0), (1, 1), (2, 2), (3, 3)]
edges = [(0, 1), (0, 2), (0, 3), (1, 2)]
- syst = builder.Builder()
+ syst = builder.Builder(vectorize=vectorize)
fam = builder.SimpleSiteFamily(norbs=1)
sites = [fam(*tag) for tag in tags]
- syst[(fam(*tag) for tag in tags)] = f_onsite
- syst[((fam(*tags[i]), fam(*tags[j])) for (i, j) in edges)] = f_hopping
+ 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
fsyst = syst.finalized()
assert fsyst.graph.num_nodes == len(tags)
@@ -521,12 +590,12 @@ def test_hamiltonian_evaluation():
for i in range(len(tags)):
site = fsyst.sites[i]
assert site in sites
- assert fsyst.hamiltonian(i, i) == syst[site](site)
+ assert fsyst.hamiltonian(i, i) == f_onsite(site)
for t, h in fsyst.graph:
tsite = fsyst.sites[t]
hsite = fsyst.sites[h]
- assert fsyst.hamiltonian(t, h) == syst[tsite, hsite](tsite, hsite)
+ assert fsyst.hamiltonian(t, h) == f_hopping(tsite, hsite)
# test when user-function raises errors
def onsite_raises(site):
@@ -555,7 +624,7 @@ def test_hamiltonian_evaluation():
syst[new_hop[0]] = onsite_raises
syst[new_hop] = hopping_raises
fsyst = syst.finalized()
- hop = tuple(map(fsyst.sites.index, new_hop))
+ hop = tuple(map(fsyst.id_by_site.__getitem__, new_hop))
test_raising(fsyst, hop)
# test with infinite system
@@ -563,7 +632,7 @@ def test_hamiltonian_evaluation():
for k, v in it.chain(syst.site_value_pairs(), syst.hopping_value_pairs()):
inf_syst[k] = v
inf_fsyst = inf_syst.finalized()
- hop = tuple(map(inf_fsyst.sites.index, new_hop))
+ hop = tuple(map(inf_fsyst.id_by_site.__getitem__, new_hop))
test_raising(inf_fsyst, hop)
@@ -1204,15 +1273,22 @@ def test_discrete_symmetries():
# We need to keep testing 'args', but we don't want to see
# all the deprecation warnings in the test logs
@pytest.mark.filterwarnings("ignore:.*'args' parameter")
-def test_argument_passing():
+@pytest.mark.parametrize("vectorize", [False, True])
+def test_argument_passing(vectorize):
chain = kwant.lattice.chain(norbs=1)
# Test for passing parameters to hamiltonian matrix elements
def onsite(site, p1, p2):
- return p1 + p2
+ if vectorize:
+ return (p1 + p2) * np.ones(len(site))
+ else:
+ return p1 + p2
def hopping(site1, site2, p1, p2):
- return p1 - p2
+ if vectorize:
+ return (p1 - p2) * np.ones(len(site1))
+ else:
+ return p1 - p2
def gen_fill_syst(onsite, hopping, syst):
syst[(chain(i) for i in range(3))] = onsite
@@ -1221,8 +1297,9 @@ def test_argument_passing():
fill_syst = ft.partial(gen_fill_syst, onsite, hopping)
- syst = fill_syst(kwant.Builder())
- inf_syst = fill_syst(kwant.Builder(kwant.TranslationalSymmetry((-3,))))
+ syst = fill_syst(kwant.Builder(vectorize=vectorize))
+ inf_syst = fill_syst(kwant.Builder(kwant.TranslationalSymmetry((-3,)),
+ vectorize=vectorize))
tests = (
syst.hamiltonian_submatrix,
@@ -1250,16 +1327,12 @@ def test_argument_passing():
# test that passing parameters without default values works, and that
# passing parameters with default values fails
- def onsite(site, p1, p2):
- return p1 + p2
-
- def hopping(site, site2, p1, p2):
- return p1 - p2
fill_syst = ft.partial(gen_fill_syst, onsite, hopping)
- syst = fill_syst(kwant.Builder())
- inf_syst = fill_syst(kwant.Builder(kwant.TranslationalSymmetry((-3,))))
+ syst = fill_syst(kwant.Builder(vectorize=vectorize))
+ inf_syst = fill_syst(kwant.Builder(kwant.TranslationalSymmetry((-3,)),
+ vectorize=vectorize))
tests = (
syst.hamiltonian_submatrix,
@@ -1275,12 +1348,18 @@ def test_argument_passing():
# Some common, some different args for value functions
def onsite2(site, a, b):
- return site.pos[0] + a + b
+ if vectorize:
+ return site.positions()[:, 0] + a + b
+ else:
+ return site.pos[0] + a + b
def hopping2(site1, site2, a, c, b):
- return a + b + c
+ if vectorize:
+ return (a + b + c) * np.ones(len(site1))
+ else:
+ return a + b + c
- syst = kwant.Builder()
+ syst = kwant.Builder(vectorize=vectorize)
syst[(chain(i) for i in range(3))] = onsite2
syst[((chain(i), chain(i + 1)) for i in range(2))] = hopping2
fsyst = syst.finalized()
@@ -1391,6 +1470,7 @@ def test_subs():
lead = lead.finalized()
assert lead.parameters == {'lead_a', 'lead_b', 'lead_c'}
+
def test_attach_stores_padding():
lat = kwant.lattice.chain(norbs=1)
syst = kwant.Builder()
diff --git a/kwant/tests/test_system.py b/kwant/tests/test_system.py
index 9bba152bc881013cd57bc949f46fca3244aa4153..3a4034776ee8492221910d0b00dd5a8e1aa97cf1 100644
--- a/kwant/tests/test_system.py
+++ b/kwant/tests/test_system.py
@@ -1,4 +1,4 @@
-# Copyright 2011-2016 Kwant authors.
+# Copyright 2011-2019 Kwant authors.
#
# 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
@@ -8,6 +8,7 @@
import pickle
import copy
+import pytest
from pytest import raises
import numpy as np
from scipy import sparse
@@ -15,9 +16,11 @@ import kwant
from kwant._common import ensure_rng
-def test_hamiltonian_submatrix():
- syst = kwant.Builder()
+@pytest.mark.parametrize("vectorize", [False, True])
+def test_hamiltonian_submatrix(vectorize):
+ syst = kwant.Builder(vectorize=vectorize)
chain = kwant.lattice.chain(norbs=1)
+ chain2 = kwant.lattice.chain(norbs=2)
for i in range(3):
syst[chain(i)] = 0.5 * i
for i in range(2):
@@ -27,7 +30,11 @@ def test_hamiltonian_submatrix():
mat = syst2.hamiltonian_submatrix()
assert mat.shape == (3, 3)
# Sorting is required due to unknown compression order of builder.
- perm = np.argsort([os[0] for os in syst2.onsites])
+ if vectorize:
+ _, (site_offsets, _) = syst2.subgraphs[0]
+ else:
+ site_offsets = [os[0] for os in syst2.onsites]
+ perm = np.argsort(site_offsets)
mat_should_be = np.array([[0, 1j, 0], [-1j, 0.5, 2j], [0, -2j, 1]])
mat = mat[perm, :]
@@ -41,19 +48,12 @@ def test_hamiltonian_submatrix():
mat = mat[:, perm]
np.testing.assert_array_equal(mat, mat_should_be)
- mat = syst2.hamiltonian_submatrix((), perm[[0, 1]], perm[[2]])
- np.testing.assert_array_equal(mat, mat_should_be[:2, 2:3])
-
- mat = syst2.hamiltonian_submatrix((), perm[[0, 1]], perm[[2]], sparse=True)
- mat = mat.toarray()
- np.testing.assert_array_equal(mat, mat_should_be[:2, 2:3])
-
# Test for correct treatment of matrix input.
- syst = kwant.Builder()
- syst[chain(0)] = np.array([[0, 1j], [-1j, 0]])
+ syst = kwant.Builder(vectorize=vectorize)
+ syst[chain2(0)] = np.array([[0, 1j], [-1j, 0]])
syst[chain(1)] = np.array([[1]])
syst[chain(2)] = np.array([[2]])
- syst[chain(1), chain(0)] = np.array([[1, 2j]])
+ syst[chain(1), chain2(0)] = np.array([[1, 2j]])
syst[chain(2), chain(1)] = np.array([[3j]])
syst2 = syst.finalized()
mat_dense = syst2.hamiltonian_submatrix()
@@ -62,9 +62,10 @@ def test_hamiltonian_submatrix():
# Test precalculation of modes.
rng = ensure_rng(5)
- lead = kwant.Builder(kwant.TranslationalSymmetry((-1,)))
- lead[chain(0)] = np.zeros((2, 2))
- lead[chain(0), chain(1)] = rng.randn(2, 2)
+ lead = kwant.Builder(kwant.TranslationalSymmetry((-1,)),
+ vectorize=vectorize)
+ lead[chain2(0)] = np.zeros((2, 2))
+ lead[chain2(0), chain2(1)] = rng.randn(2, 2)
syst.attach_lead(lead)
syst2 = syst.finalized()
smatrix = kwant.smatrix(syst2, .1).data
@@ -74,19 +75,26 @@ def test_hamiltonian_submatrix():
raises(ValueError, kwant.solvers.default.greens_function, syst3, 0.2)
# Test for shape errors.
- syst[chain(0), chain(2)] = np.array([[1, 2]])
+ syst[chain2(0), chain(2)] = np.array([[1, 2]])
syst2 = syst.finalized()
raises(ValueError, syst2.hamiltonian_submatrix)
raises(ValueError, syst2.hamiltonian_submatrix, sparse=True)
- syst[chain(0), chain(2)] = 1
+ syst[chain2(0), chain(2)] = 1
syst2 = syst.finalized()
raises(ValueError, syst2.hamiltonian_submatrix)
raises(ValueError, syst2.hamiltonian_submatrix, sparse=True)
+ if vectorize: # non-vectorized systems don't check this at finalization
+ # Add another hopping of the same type but with a different
+ # (and still incompatible) shape.
+ syst[chain2(0), chain(1)] = np.array([[1, 2]])
+ raises(ValueError, syst.finalized)
-def test_pickling():
- syst = kwant.Builder()
- lead = kwant.Builder(symmetry=kwant.TranslationalSymmetry([1.]))
+@pytest.mark.parametrize("vectorize", [False, True])
+def test_pickling(vectorize):
+ syst = kwant.Builder(vectorize=vectorize)
+ lead = kwant.Builder(symmetry=kwant.TranslationalSymmetry([1.]),
+ vectorize=vectorize)
lat = kwant.lattice.chain(norbs=1)
syst[lat(0)] = syst[lat(1)] = 0
syst[lat(0), lat(1)] = 1