# Numerical calculation of Wannier function in optical lattice

I am working out some optical lattice band structures (example here). I have no issue with setting up the eigenvalue equation: $$H_{jj'}c_{j'}=Ec_{j'}$$ Where $$H$$ is the tri-diagonal matrix that equations the different Fourier components. The Bloch wave is given by

$$\psi =e^{iqx} u_n^q = e^{iqx} \sum_j a_c e^{2ikxj}$$

And the Wannier function is $$\mathcal{W}=\int dq \ u_n^q \ e^{-iqx}$$

So my pseudocode is:

1. For each $$x$$:

2. For each $$q$$:

3. Find the eigenvectors for $$c_j$$. Multiply that $$c_j$$ by its Fourier component $$\exp(2ik j)$$

4. Sum those terms and multiply by $$\exp(iq_i x_q)$$

5. Repeat for $$q$$, then add all terms

6. Repeat for all $$x$$

My python code is below. When I plot x,|w|^2, I do not get anything resembling the gaussian approximation. Also, I figure this can be vectorized, but I am struggling to get it working in the loop form.

x = np.linspace(-2,2,101)
lmax = 10
l=np.arange(-lmax,lmax+1)
V0 = 5
wavelength = np.pi
k_lattice = 2*np.pi/wavelength
qx = np.linspace(-1,1,101) #This is qx/k_lattice
wave_dict = {q:None for q in qx}

for q in qx:
diags = [(q+2*k)**2 for k in l]
Hmat = np.diag(diags)
evals, evecs = LA.eigh(Hmat)

coefs = evecs[:,0][:,None]
planewaves = coefs*np.exp(1j*(2*k_lattice)*np.outer(l,x)) # exp(2ik_l x)
psi = planewaves.sum(axis=0)
wave_dict[q] = psi

w = np.zeros(x.shape,dtype = 'complex128')
for q in qx:

w+= wave_dict[q]* np.exp(1j*x*q*k_lattice)

w/=len(qx)
plt.plot(x,np.abs(w)**2)
plt.plot(x,np.sin(k_lattice*x)**2)
plt.show()

• Could you post your full code? What is lmax for instance? I've also understood that you have taken $k=1$? Maybe you should indicate it somewhere. Also, not that I think it matters, but at the last step, you're multiplying by $\exp(-iqx)$ when you said you wanted to multiply by $\exp(iqx)$. At last, it seems that you're taking the highest of the $n$ eigenvectors (evecs[:,np.argsort(evals)[-1]]). Why not take the eigenvector with the lowest energy instead? It will definitely show more localisation. – QuantumApple May 3 at 11:36

As I've said in a comment, you should not take the eigenvector with the highest energy, as the highest bands correspond typically to very delocalized wavefunctions, so the corresponding Wannier function will not surely not be Gaussian.

You should also write explicitely your choice of $$k$$ as a parameter, as it can help preventing mistakes. I think that apart from the choice of the last eigenvector, you've made 2 major mistakes:

• The Fourier series development writes $$u_n^q(x) = \sum_j c_j e^{i 2 \pi x j/a}$$, where $$a$$ is the period of your lattice. But because you've taken $$k = 1$$, and $$V(x) = \frac{V_0}{2} \cos(2 k x)$$, the period of the lattice here is $$a=\pi$$. So the Fourier series in your case should be $$u_n^q(x) = \sum_j c_j e^{i 2 x j}$$, and not $$u_n^q(x) = \sum_j c_j e^{i 2 \pi x j}$$ as you have written it. This is why it is especially important to explicitely name your variables if you don't wan't to be confused.
• Wannier wavefunctions are supposed to be taken at the minima of potential. But looking at your choice of matrix Hmat, it appears that you have taken $$V(x) = \frac{V_0}{2} \cos(2 k x)$$. You can see that $$V(x)$$ is maximal at $$x=0$$, and you're trying to look at the Wannier function at $$X = 0$$, which does not make sense. Instead, if you want to have a nice Wannier function at $$X=0$$, you should take $$V(x) = -\frac{V_0}{2} \cos(2 k x)$$ instead (switching the sign of the off-diagonal terms).

With these corrections, the code should work better. I suggest also a few improvements:

• Work with arrays. It will make your life easier. For instance in your code, you're diagonalizing the whole matrix for each position of $$x$$ and for each $$q$$, whereas it would be enough to diagonlize it once per $$q$$ only (the values do not depend on $$x$$).
• Use scipy.linalg.eigh to minimize your matrix instead. It works with Hermitian/real symmetric matrices, and gives you real eigenvalues, ordered from the lowest to the highest. It is more efficient and also will save you the trouble of ordering the eigenvalues/eigenvectors yourself.
• Try to test your algorithm against something that you know. For instance, before thinking of computing the Wannier function, check that your program gives the right form for the energy bands. Once this is done, try to get the Wannier function and compare them against the Gaussian approximation for a harmonic trap.

You can find the code with the corrected mistakes and a few improvements below (I've only partially vectorized the loop, so it could still be slightly optimized, but at least you're not diagonalizing the same matrix 100 times):


import numpy as np
import scipy.linalg as LA
import matplotlib.pyplot as plt

Ei = []
phii = []
w = []

k_light = 1 #wavevector of the light beam creating the lattice
a_lattice = np.pi/k_light #period of the lattice
lmax = 20
l=np.arange(-lmax,lmax+1)
V0 = 5 #V0/E_recoil
V0 *= k_light**2 #"true" V0
x= np.linspace(-a_lattice/2*1.5,a_lattice/2*1.5,4000)
dx = x-x
qx = np.linspace(-k_light,k_light,100, endpoint=False)
for q in qx:
u = 0
diags = [(q+2*k*k_light)**2 for k in l]
Hmat = np.diag(diags)
Hmat += -V0/4 *(np.diag(np.ones(len(l)-1),1)+np.diag(np.ones(len(l)-1),-1))
evals, evecs = LA.eigh(Hmat)
Ei.append(evals)
phii.append(evecs)
phii = np.array(phii)
Ei = np.array(Ei)

for xi in x:
b = 0
a = np.exp(1j*2*np.pi*xi*l/a_lattice)
for p in range(len(qx)):
b += np.sum(a*phii[p, :, 0])*np.exp(1j*qx[p]*xi)
w.append(b/(len(qx)))
w = np.array(w)
plt.plot(x,abs(w)**2/np.sum(abs(w)**2)/dx, label='Wannier wavefunction')

X_harmonic_sq = 1/(k_light*V0**0.5)

plt.plot(x, np.exp(-x**2/X_harmonic_sq)/(np.sqrt(np.pi*X_harmonic_sq)), label='Harmonic approximation')

plt.legend()

#plt.plot(qx, Ei[:, 0])
#plt.plot(qx, Ei[:, 1])
#plt.plot(qx, Ei[:, 2])



Plot of the Wannier wavefunction for $$k=1$$, and $$V_0=5$$, compared to the harmonic approximation: Tell me if you have further questions.

• Thanks for this! I had the  in the lowest for my eigenvalues, I just was playing around with other eigenvalues and didn't catch that I uploaded this with [-1]. My energy bands are the correct form with the folded quadratic at V=0 and the gaps opening with increasing V0. I was also using a sin^2 potential, and dropped a (-) sign, so my matrix should have -V0. I think I got something that works, I edited my code to reflect it. Thanks for the guidance! – yankeefan11 May 4 at 14:19
• You're welcome! :) Yes, $\sin^2(kx)$ potential would also work (it is actually the same as $\cos(2kx)$ up to a constant). I also remember playing with Wannier functions and I think I got the same problem at some point because this was familiar. I find it a bit counter-intuitive that changing a sign in the Hamiltonian matrix (which does not depend explicitely on $x$) changes the actual position of the sites (here from $a/2$ to $a$), but that's actually the case so you need to be very careful about the sign. – QuantumApple May 4 at 16:59