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

Question

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:

\begin{equation} 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} \end{equation}

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: \begin{equation} | R n \rangle = \frac{V}{(2 \pi)^D} \int d^{D}k ~~ e^{-i k\cdot R} |\psi_{nk} \rangle \end{equation} 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 
```

Answer 1

$|{\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, $<r|Rn> \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.