Total Energy Minimization#
In this document, we start from the objective function of solid state DFT, starting from the abstract form and derive it until it is implementable. Solid state DFT solves the following energy minimization problem (we leave why the objective looks like this in the preliminaries of DFT and crystalgraphy)
The total energy functional is a functional of a set of wave functions \(\{\psi_{ik}\}\) and the occupation vectors \(f\). \(k\in K\) and \(i\in[1..I]\), we call \(I\) the number of bands and \(K\) the set of k points, and \(N\) is the number of electrons in the system. \(f\) specifies how the \(N\) electrons are distributed over the k points and bands. With the wave function and occupation vectors, we can calculate the electron density function as \(\rho(r)=\sum_i^I\sum_{k\in K} f_{ik}\psi^2_{ik}(r)\). As for all KS DFT calculations, the energy functional contains four terms, kinetic energy, external energy, hartree energy, and exchange-correlation energy.
And notice that except that the kinetic energy term is a direct functional of the wave functions and occupation, all other energies are direct functionals of the electron density (which in term depends on the wave function and occupation). We write below the definition of each energy term except the exchange-correlation functional. As the crystal system is periodic over the entire \(R^3\) space, we only calculate a fraction of the energy within a single unit cell \(\Omega\).
To this point, we have introduced DFT as an optimization problem with constraints in the function space.
The objective functional: \(E_\text{total}[\{\psi_{ik}\},f]\), which can be expanded into the above terms.
The parameter: \(f\) and \(\psi_{ik}\).
The constraints:
\(\psi_{ik}(r)=e^{ikr}u_{ik}(r)\), where \(u_{ik}\) is periodic over the unit cell, and \(\langle u_{ik}|u_{jk}\rangle=\delta_{ij}\).
\(f_{ik}\in [0,1]\) and \(\sum_k \sum_i f_{ik}=N\).
To make the problem computable, we just need to parameterize \(u_{ik}\) and \(f\) in a way that satisfy the constraints, plugging them back into the objective function and then perform the optimization in the parameter space. This is what we will do in the rest of this document.
Parameterizing \(u_{ik}(r)\) and \(f\)#
In planewave calculations, \(u_{ik}\) is parameterized as a linear combination over fourier components of different frequencies
where we limit the \(G\) in \(e^{iGr}\) to be frequency components that is periodic over the unit cell, making \(u_{ik}(r)\) satisfy the periodic constraint. At the same time, plugging the planewave back to \(\langle u_{ik}|u_{jk}\rangle=\delta_{ij}\), we translate the orthogonality constraint into \(\sum_G c^*_{ikG}c_{jkG}=I_{ij}\). A orthogonal matrix can be easily generated via reparameterization, for example, with the QR decomposition
key = jax.random.PRNGKey(0)
w = jax.random.normal(key, (num_K, num_G, I))
c, _ = jnp.linalg.qr(w)
The other parameter \(f\) needs to satisfy \(f_{ik}\in [0,1]\) and \(\sum_k \sum_i f_{ik}=N\). Similarly, we can use reparameterize it as
key = jax.random.PRNGKey(0)
v = jax.random.normal(key, (I*num_K, N*num_K))
Q, _ = jnp.linalg.qr(v) # Q has shape (I*num_K, N*num_K) and Q.T @ Q = I
f = jnp.diag(Q @ Q.T).reshape(I, num_K) # f has shape (I*num_K)
It is easily verifiable that \(\sum_{ik}f_{ik}=\Tr(QQ^\top)=\Tr(Q^\top Q)=N\) and \(f_{ik}=\|Q_{i*|K|+k}\|^2\in[0,1]\)
Casting into the parameter space#
Now that \(u_{ik}\) becomes a function parameterized by \(c\), we can substitute it back to the energy terms to cast the energy into a function of the finite dimensional parameters \(c\) and \(f\).
The Kinetic Energy#
Firstly, we apply the kinetic operator on the parameterized wave function
The kinetic energy is then reduced to the following using the property that \(\int_{\Omega} \psi^*_{ik}(r)\psi_{jk}(r) dr=\delta_{ij}\).
Reciprocal representation of the Coulombic potential#
The Coulombic potential generated from a charge \(\rho\) is
where \(\vb{R}\) is a Bravais lattice vector. Its reciprocal representation is given by (see Ewald Summation for derivation)
The External Energy#
The atomic point charge within the unit cell is
Therefore its reciprocal representation is given by
Now by Parseval’s theorem we have
where the \(\vb{G}=0\) term is removed due to neutral charge requirement (TODO: add doc on this). The reciprocal density \(\tilde{\rho} (\vb{G})\) can be calculated effciently using FFT (see jrystal.pw.density_grid_reciprocal()
).
The Hartree Energy#
Again using Parseval’s theorem we have