diff --git a/kwant/builder.py b/kwant/builder.py
index 24e2c9d6d5d1868ba46d296e4fbe074cce5a6d57..40047f27d63266047aff6f0e330dc533f6c6c1ab 100644
--- a/kwant/builder.py
+++ b/kwant/builder.py
@@ -593,6 +593,32 @@ def edges(seq):
return result
+def _site_ranges(sites):
+ """Return a sequence of ranges for ``sites``.
+
+ Here, "ranges" are defined as sequences of sites that have the same site
+ family. Because site families now have a fixed number of orbitals,
+ this coincides with the definition given in `~kwant.system.System`.
+ """
+ # we shall start a new range of different `SiteFamily`s separately,
+ # even if they happen to contain the same number of orbitals.
+ total_norbs = 0
+ current_fam = None
+ site_ranges = []
+ for idx, fam in enumerate(s.family for s in sites):
+ if not fam.norbs:
+ # can't provide site_ranges if norbs not given
+ return None
+ if fam != current_fam: # start a new run
+ current_fam = fam
+ current_norbs = fam.norbs
+ site_ranges.append((idx, current_norbs, total_norbs))
+ total_norbs += current_norbs
+ # add sentinel to the end
+ site_ranges.append((len(sites), 0, total_norbs))
+ return site_ranges
+
+
class Builder:
"""A tight binding system defined on a graph.
@@ -1261,6 +1287,7 @@ class Builder:
result = FiniteSystem()
result.graph = g
result.sites = sites
+ result.site_ranges = _site_ranges(sites)
result.id_by_site = id_by_site
result.leads = finalized_leads
result.hoppings = [self._get_edge(sites[tail], sites[head])
@@ -1388,6 +1415,7 @@ class Builder:
result.cell_size = cell_size
result.sites = sites
result.id_by_site = id_by_site
+ result.site_ranges = _site_ranges(sites)
result.graph = g
result.hoppings = hoppings
result.onsite_hamiltonians = onsite_hamiltonians
diff --git a/kwant/system.py b/kwant/system.py
index 2096026376835cfc9f9a7ff05dfd51555090a9b4..afdfa5137312cbfc7783ef8cd114435bfb5e9711 100644
--- a/kwant/system.py
+++ b/kwant/system.py
@@ -23,6 +23,14 @@ class System(metaclass=abc.ABCMeta):
----------
graph : kwant.graph.CGraph
The system graph.
+ site_ranges : None or sorted sequence of triples of integers
+ If provided, encodes ranges of sites that have the same number of
+ orbitals. Each triple consists of ``(first_site, norbs, orb_offset)``:
+ the first site in the range, the number of orbitals on each site in the
+ range, and the offset of the first orbital of the first site in the
+ range. In addition, the final triple should have the form
+ ``(len(graph.num_nodes), tot_norbs, 0)`` where ``tot_norbs`` is the
+ total number of orbitals in the system.
Notes
-----
@@ -31,6 +39,13 @@ class System(metaclass=abc.ABCMeta):
Optionally, a class derived from `System` can provide a method `pos` which
is assumed to return the real-space position of a site given its index.
+
+ Due to the ordering semantics of sequences, and the fact that a given
+ ``first_site`` can only appear *at most once* in ``site_ranges``,
+ ``site_ranges`` is ordered according to ``first_site``.
+
+ Consecutive elements in ``site_ranges`` are not required to have different
+ numbers of orbitals.
"""
@abc.abstractmethod
diff --git a/kwant/tests/test_builder.py b/kwant/tests/test_builder.py
index 6b4456c62c323fe9b0313672d5c1fe2dea9fbb69..71ef0e73a2ac718458f93e84458d1ea24f12df6d 100644
--- a/kwant/tests/test_builder.py
+++ b/kwant/tests/test_builder.py
@@ -8,9 +8,11 @@
import warnings
from random import Random
+import itertools as it
from nose.tools import assert_raises, assert_true, assert_not_equal
from numpy.testing import assert_equal, assert_almost_equal
import tinyarray as ta
+import numpy as np
import kwant
from kwant import builder
@@ -412,6 +414,43 @@ def test_finalization():
assert_raises(ValueError, lead.finalized)
+def test_site_ranges():
+ 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)
+ expected = [(0, lat.norbs, 0), (10, 0, 10 * lat.norbs)]
+ assert_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_equal(expected, site_ranges(tuple(sites)))
+ sites = it.chain(map(lat2, range(4)), map(lat1a, range(6)),
+ map(lat1b, range(4)))
+ expected = [(0, 2, 0), (4, 1, 4*2), (10, 1, 4*2+6), (14, 0, 4*2+10)]
+ assert_equal(expected, site_ranges(tuple(sites)))
+
+ # test with an actual builder
+ for lat in (lat1a, lat2):
+ sites = list(map(lat, range(10)))
+ syst = kwant.Builder()
+ syst[sites] = np.eye(lat.norbs)
+ ranges = syst.finalized().site_ranges
+ expected = [(0, lat.norbs, 0), (10, 0, 10 * lat.norbs)]
+ assert_equal(ranges, expected)
+ # poison system with a single site with no norbs defined
+ syst[kwant.lattice.chain()(0)] = 1
+ ranges = syst.finalized().site_ranges
+ assert_equal(ranges, None)
+
+
def test_hamiltonian_evaluation():
def f_onsite(site):
return site.tag[0]