From dae719f6c21492676f64efb457bd9f8a28f88c05 Mon Sep 17 00:00:00 2001
From: Joseph Weston <joseph@weston.cloud>
Date: Tue, 26 Nov 2019 14:56:59 +0100
Subject: [PATCH] replace '_count_norbs' with 'site_ranges'
The removed logic is an artifact from when vectorized and
unvectorized systems shared the same logic for 'hamiltonian_submatrix'.
This was accidentally left in when we reverted to the old logic
for unvectorized systems.
---
kwant/_system.pyx | 85 ++---------------------------------------------
1 file changed, 3 insertions(+), 82 deletions(-)
diff --git a/kwant/_system.pyx b/kwant/_system.pyx
index 2f58dd4f..5bc6c235 100644
--- a/kwant/_system.pyx
+++ b/kwant/_system.pyx
@@ -453,68 +453,6 @@ def _vectorized_make_dense(subgraphs, hams, long [:] norbs, long [:] orb_offsets
return mat
-def _count_norbs(syst, site_offsets, hams, args=(), params=None):
- """Return the norbs and orbital offset of each site family in 'syst'
-
- Parameters
- ----------
- syst : `kwant.system.System`
- site_offsets : sequence of int
- Contains the index of the first site for each site array
- in 'syst.site_arrays'.
- hams : sequence of arrays or 'None'
- The Hamiltonian for each term in 'syst.terms'. 'None'
- if the corresponding term has not been evaluated.
- args, params : positional and keyword arguments to the system Hamiltonian
- """
- site_ranges = syst.site_ranges
- if site_ranges is None:
- # NOTE: Remove this branch once we can rely on
- # site families storing the norb information.
-
- site_arrays = syst.site_arrays
- family_norbs = {s.family: None for s in site_arrays}
-
- # Compute the norbs per site family using already evaluated
- # Hamiltonian terms where possible
- for ham, t in zip(hams, syst.terms):
- (to_w, from_w), _ = syst.subgraphs[t.subgraph]
- if ham is not None:
- family_norbs[site_arrays[to_w].family] = ham.shape[1]
- family_norbs[site_arrays[from_w].family] = ham.shape[2]
-
- # Evaluate Hamiltonian terms where we do not already have them
- for n, t in enumerate(syst.terms):
- (to_w, from_w), _ = syst.subgraphs[t.subgraph]
- to_fam = site_arrays[to_w].family
- from_fam = site_arrays[from_w].family
- if family_norbs[to_fam] is None or family_norbs[from_fam] is None:
- ham = syst.hamiltonian_term(n, args=args, params=params)
- family_norbs[to_fam] = ham.shape[1]
- family_norbs[from_fam] = ham.shape[2]
-
- # This case is technically allowed by the format (some sites present
- # but no hamiltonian terms that touch them) but is very unlikely.
- if any(norbs is None for norbs in family_norbs.values()):
- raise ValueError("Cannot determine the number of orbitals for "
- "some site families.")
-
- orb_offset = 0
- site_ranges = []
- for start, site_array in zip(site_offsets, syst.site_arrays):
- norbs = family_norbs[site_array.family]
- site_ranges.append((start, norbs, orb_offset))
- orb_offset += len(site_array) * norbs
- site_ranges.append((site_offsets[-1], 0, orb_offset))
- site_ranges = np.array(site_ranges)
-
- _, norbs, orb_offsets = site_ranges.transpose()
- # The final (extra) element in orb_offsets is the total number of orbitals
- assert len(orb_offsets) == len(syst.site_arrays) + 1
-
- return norbs, orb_offsets
-
-
def _expand_norbs(compressed_norbs, site_offsets):
"Return norbs per site, given norbs per site array."
norbs = np.empty(site_offsets[-1], int)
@@ -586,8 +524,7 @@ def vectorized_hamiltonian_submatrix(self, args=(), sparse=False,
if t.hermitian
]
- norbs, orb_offsets = _count_norbs(self, site_offsets, hams,
- args=args, params=params)
+ _, norbs, orb_offsets = self.site_ranges.transpose()
func = _vectorized_make_sparse if sparse else _vectorized_make_dense
mat = func(subgraphs, hams, norbs, orb_offsets, site_offsets)
@@ -642,13 +579,7 @@ def vectorized_cell_hamiltonian(self, args=(), sparse=False, *, params=None):
if self.terms[n].hermitian
]
- # _count_norbs requires passing hamiltonians for all terms, or 'None'
- # if it has not been evaluated
- all_hams = [None] * len(self.terms)
- for n, ham in zip(cell_terms, hams):
- all_hams[n] = ham
- norbs, orb_offsets = _count_norbs(self, site_offsets, all_hams,
- args=args, params=params)
+ _, norbs, orb_offsets = self.site_ranges.transpose()
shape = (orb_offsets[next_cell], orb_offsets[next_cell])
func = _vectorized_make_sparse if sparse else _vectorized_make_dense
@@ -719,17 +650,7 @@ def vectorized_inter_cell_hopping(self, args=(), sparse=False, *, params=None):
for (to_sa, from_sa), (to_off, from_off) in subgraphs
]
- # _count_norbs requires passing hamiltonians for all terms, or 'None'
- # if it has not been evaluated
- all_hams = [None] * len(self.terms)
- for n, ham in zip(inter_cell_hopping_terms, inter_cell_hams):
- all_hams[n] = ham
- for n, ham in zip(reversed_inter_cell_hopping_terms,
- reversed_inter_cell_hams):
- # Transpose to get back correct shape wrt. original term
- all_hams[n] = ham.transpose(0, 2, 1)
- norbs, orb_offsets = _count_norbs(self, site_offsets, all_hams,
- args=args, params=params)
+ _, norbs, orb_offsets = self.site_ranges.transpose()
shape = (orb_offsets[next_cell],
orb_offsets[len(self.site_arrays) - next_cell])
--
GitLab