Commit af5eea74 authored by Joseph Weston's avatar Joseph Weston
Browse files

add inverse of `sites` mapping to finalized builders

The inverse map is created during finalization anyway, now
we just save it along with the `sites` map.
parent bf1d3177
......@@ -1242,6 +1242,7 @@ class Builder:
result = FiniteSystem()
result.graph = g
result.sites = sites
result.id_by_site = id_by_site
result.leads = finalized_leads
result.hoppings = [self._get_edge(sites[tail], sites[head])
for tail, head in g]
......@@ -1350,7 +1351,6 @@ class Builder:
# correspond to one left out just above.
g.add_edge(head_id, tail_id)
g.add_edge(tail_id, head_id)
del id_by_site
g = g.compressed()
#### Extract hoppings.
......@@ -1368,6 +1368,7 @@ class Builder:
result = InfiniteSystem()
result.cell_size = cell_size
result.sites = sites
result.id_by_site = id_by_site
result.graph = g
result.hoppings = hoppings
result.onsite_hamiltonians = onsite_hamiltonians
......@@ -1384,11 +1385,19 @@ def _raise_user_error(exc, func):
class FiniteSystem(system.FiniteSystem):
"""
Finalized `Builder` with leads.
"""Finalized `Builder` with leads.
Usable as input for the solvers in `kwant.solvers`.
Attributes
----------
sites : sequence
``sites[i]`` is the `~kwant.builder.Site` instance that corresponds
to the integer-labeled site ``i`` of the low-level system.
id_by_site : dict
The inverse of ``sites``; maps from ``i`` to ``sites[i]``.
"""
def hamiltonian(self, i, j, *args):
if i == j:
value = self.onsite_hamiltonians[i]
......@@ -1426,7 +1435,17 @@ class FiniteSystem(system.FiniteSystem):
class InfiniteSystem(system.InfiniteSystem):
"""Finalized infinite system, extracted from a `Builder`."""
"""Finalized infinite system, extracted from a `Builder`.
Attributes
----------
sites : sequence
Mapping from site IDs (integers) to the associated
`~kwant.builder.Site` objects.
id_by_site : dict
The inverse of the ``sites`` mapping.
"""
def hamiltonian(self, i, j, *args):
if i == j:
if i >= self.cell_size:
......
......@@ -291,6 +291,10 @@ def check_hoppings(fsyst, hops):
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_equal(fsyst.id_by_site[site], i)
def test_finalization():
"""Test the finalization of finite and infinite systems.
......@@ -353,6 +357,7 @@ def test_finalization():
for hop, value in sr_hops.items():
syst[fam(*hop[0]), fam(*hop[1])] = value
fsyst = syst.finalized()
check_id_by_site(fsyst)
check_onsite(fsyst, sr_sites)
check_hoppings(fsyst, sr_hops)
......@@ -376,6 +381,7 @@ def test_finalization():
flead = lead.finalized()
all_sites = list(lead_sites)
all_sites.extend((x - size, y) for (x, y) in neighbors)
check_id_by_site(fsyst)
check_onsite(flead, all_sites, check_values=False)
check_onsite(flead, lead_sites, subset=True)
check_hoppings(flead, lead_hops)
......
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