diff --git a/kwant/builder.py b/kwant/builder.py index c789c732c3d5471afe8d0c87d3fa453b5ebc2547..b6712a2ed260febef7251846e3cda4e2f7988bc3 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 0bce96fe97d4c4cf8269b0a9449b700adc54b042..6b4456c62c323fe9b0313672d5c1fe2dea9fbb69 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)