From af5eea742ad9319cb539365bab649f363705cdd8 Mon Sep 17 00:00:00 2001
From: Joseph Weston <joseph.weston08@gmail.com>
Date: Tue, 24 May 2016 10:42:42 +0200
Subject: [PATCH] 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.
---
kwant/builder.py | 27 +++++++++++++++++++++++----
kwant/tests/test_builder.py | 6 ++++++
2 files changed, 29 insertions(+), 4 deletions(-)
diff --git a/kwant/builder.py b/kwant/builder.py
index c789c732..b6712a2e 100644
--- a/kwant/builder.py
+++ b/kwant/builder.py
@@ -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:
diff --git a/kwant/tests/test_builder.py b/kwant/tests/test_builder.py
index 0bce96fe..6b4456c6 100644
--- a/kwant/tests/test_builder.py
+++ b/kwant/tests/test_builder.py
@@ -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)
--
GitLab