Hi,
I think I misunderstood what you are trying to do. You talked about "filtering" only the sites which were in a plane, so I thought you wanted to produce a projection of the wavefunction onto the sites in the plane. If you just want to be able to access wavefunction elements "by their site" you could do something like the following
def site_indexer(wavefunction, sites, n_orbs=4): site_idx_map = {site: i for i, site in enumerate(sites)} def indexer(site): orbital = site_idx_map[site] * n_orbs return wavefunction[orbital: orbital + n_orbs]
return indexer
lat = kwant.lattice.square() sys = make_system() fsys = sys.finalized() wf = kwant.wavefunction(...)
better_wf = site_indexer(wf, fsys.sites)
# we get an array of length 4, for the 4 orbitals on the site print better_wf(lat(0, 2))
You could then use `better_wf` to access the wavefunction on the orbitals belonging to different sites.
Joe