How to get a Wannier function of tight-binding model numerically (i.e. Python)

I have a question on construction of a Wannier function for tight-binding model. Let's say we consider the tight-binding model of 1D chain with two atoms( site A and B in a unit cell). In k-space we can write the single-particle Hamiltonian in following way:

$$$$H = \sum_{k} \Psi_{k}^\dagger h(k) \Psi_k~~,~~ \Psi_k = ( c_{A,k} ~~ c_{B,k}) ~~,~~ h(k)= \begin{pmatrix} \epsilon(k) & \Delta(k) \\ \Delta(k)^\dagger & \epsilon(k) \end{pmatrix}$$$$

By diagonalizing the single-particle Hamiltonian $$h(k)$$, we get the Bloch wavefunction $$| \psi_{nk} \rangle$$ for each $$k$$, where the Bloch wavefunctions are 2-dimensional. To construct the Wannier function, we can think of it as the inverse Fourier function of Bloch waves: $$$$| R n \rangle = \frac{V}{(2 \pi)^D} \int d^{D}k ~~ e^{-i k\cdot R} |\psi_{nk} \rangle$$$$ where $$|R n \rangle$$ is the Wannier function, $$n,R$$ are the band index and real space lattice vector respectively.

My question is how to construct the Wannier function numerically in Python. The picture in my mind is that the resulted Wannier function should be centred at the home unit cell. However, $$| \psi_{nk} \rangle$$ is a 2D array, it means that the Wannier function that I compute should be a vector rather than a function. So, how can I interpret the elements of the resulted vector after Wannierization? After Wannierization, what things should I plot in order to see the localized function in real space? I appreciate all comments.

This is my rough and trial code on Wannierization process in python

import numpy as np
import scipy.linalg as la

# Suppose I solved the Hamiltonian and get a set of states
psi_k = np.array([ la.eigh(H(k))[0] for k in k_list ]) #[0] means returns the eigvecs

# Select the lower band ( 1D chain with 2 atoms > 2 bands in k-space)
low_band = psi_k[:,:,0]

# Wannier function
def Wannier(psi_k,r_list, k_list):
WF = 0
for r in r_list:
for i in range(len(k_list)):
WF += np.exp(-j*k*R) * psi_k[i,:,0]
return WF*len(k_list)**-1
$$$$
`
• Would Computational Science be a better home for this question? We can migrate it if you like.
– rob
Feb 23 at 17:17
• Sure! Thank you for your suggestion Feb 23 at 17:18
• Matter Modelling (mattermodeling.stackexchange.com) might be even better Feb 23 at 18:16

$$|{\psi_{nk}}>$$ is a set of states parameterized by two indices, n and k. To get the real space plot of this function, we would compute $$<{\vec r}|{\psi_{nk}}> \equiv \psi_{nk}(\vec r)$$.
Similarly, $$|{Rn}>$$ is the Wannier function centered in unit cell $$R$$ of orbital $$n$$. You can plot it spatially by again taking the inner product of this state with the position vector, $$ \equiv \phi_{Rn}(\vec r)$$. You'll note there are again two numbers parameterizing these states, R and n.
Depending on the code you are using, $$|\psi_{nk}>$$ may be stored either in position space, momentum space ($$\psi_{nk}(\vec q)$$), or some other basis. I assume you know what basis it's being stored in, so that you can take the inner products above.
The only question now is - what are the sizes of these parameters? Well, Wannier functions are a 1-1 map from the Bloch states, so naturally the number of R has to equal the number of k. The easiest way to see this is the case is to review Bloch's theorem: You get discretized $$k$$ points by assuming periodicity of your electronic states in some supercell. It turns out this supercell has $$N_R$$ cells inside of it, and so $$\vec R$$ tracks which cell you're talking about in this supercell.
• Thank you for your comment! Suppose that I solve my 2nd quantized Hamiltonian in Momentum space, then I get a set of Bloch states $| \psi_{nk} \rangle$. How can I find a real space position state $| r \rangle$ to project the Bloch state to this basis(i.e. $\langle r | \psi_{nk} \rangle = \psi_{nk}(r)$? Mar 2 at 17:02
• Suppose you have $|𝜓_{𝑛𝑘}>$ stored in a basis, such that you know $<{b_i}|{\psi_{nk}}>$ for each basis vector $<{b_i}|$. Then you need to perform a change of variables to have it in position space. Or, in the bra-ket notation, insert the identity. $<{r}|{\psi_{nk}}> = \sum_{b_i} <{r}|{b_i}><{b_i}|{\psi_{nk}}>$. Therefore, you need to know the basis you're storing the Bloch wavefunction in, as well as its overlaps with the position basis. As an example, $<r|q> = e^{i{\vec q \cdot \vec r}}$ for momentum basis. Mar 2 at 21:02