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

Attachment: signature.asc
Description: OpenPGP digital signature

Reply via email to