Commit 91b58d4a authored by Joseph Weston's avatar Joseph Weston
Browse files

add ``site-ranges`` to low-level systems and Builder finalization

parent d4c1cc6c
......@@ -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
......
......@@ -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
......
......@@ -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]
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment