Additional comment:
The quantity shown in Figure 5 of arXiv:1506.04707v1 is the square of Bloch states at the band edges.
Once you calculate the Bloch states at the band edges, you can do further manipulations by a short coding in order to the obtain the square of Bloch states.