[Kwant] Wavefunction at each site

Dear Kwant users and developers, I have a question about plotting wavefunction. This question is based on http://kwant-project.org/doc/1.0/tutorial/tutorial6. If I want to print the amplitude of wavefunction at the A site / B site only at the edge, I suppose I have to get the index of the A/B atom at the edge first. Is there anyway by which we can get the index of them? Thank you very much. Best, Johnny