Meta-optimization of maximally-localized Wannier functions

Abstract

Maximally-localized Wannier functions are quantum wavefunctions resembling atomic orbitals that are used to describe electrons in condensed matter. Since their introduction in 1997, these functions have become ubiquitous in ab initio materials simulations, including applications in linear-scaling methods, strongly-correlated electron systems, quantum transport, electron-phonon interactions, and topological materials. Despite their widespread adoption in a vast software ecosystem, Wannier functions have not yet attained their fullest potential in the presence of entangled bands, as their optimization remains challenging and labor-intensive. Here, we introduce a universal meta-optimization method that leverages workflow abstraction and machine learning techniques like differential evolution and Bayesian optimization to generate globally optimized Wannier functions without human intervention. We demonstrate this approach through three applications: (i) autonomous interpolation of entangled band structures with millielectronvolt accuracy starting from coarse Brillouin zone grids, (ii) thousand-fold acceleration of fully ab initio Boltzmann transport calculations via the use of minimal coarse Brillouin zone grids, and (iii) ultra-fast high-throughput calculations of high-precision Wannier functions for large materials libraries. This work brings calculations that previously required supercomputers within the reach of personal computers.

Figures (5)

• Figure 1: Meta-optimization of Wannier functions via global optimization. Schematic flowchart of the present method. The procedure starts with a standard DFT band structure calculation for a given crystal (in this work Quantum ESPRESSODFT2 is used). At Step 1, the DE or BO algorithm selects the hyperparameters needed to construct $U^0$. In this manuscript, the hyperparameters are the energies of the inner and outer disentanglement windows and the type of initial projections, which are required in input by the Wannier90 code.Wannier Steps 2:4 implement the mapping $U^0 \mapsto \mathcal{W}(U^0)$: in Step 2, the hyperparameters are used to build $U^0$ using data from the DFT calculation; in Step 3, local optimization of the Berry-phase spread functionalmarzari2012 is performed using Wannier90; in Step 4, the loss function is computed. This step uses MLWFs obtained at Step 3, as well as band structures from separate DFT calculations for the ground truth. In Step 5, the loss function is tested; if the test fails, the procedure is repeated for a new DE generation or BO iteration. The mapping $\mathcal{W}$ (yellow box) is encapsulated within a single Python function, and is optimized by calling the procedure differential_evolution of the SciPy library or BayesianOptimization of the bayes_opt package (green box). The complete workflow is implemented in the metaWann code provided in code ocean capsule. codeocean_capsule
• Figure 2: Wannier interpolation of entangled band structures with millielectronvolt accuracy. (a)-(f) 2D slices of the 5D hyper-parameter space explored by the DE meta-optimizer, for the conduction bands of silicon. The first, second, and third rows show data for DE generations 1:10, 100:110, and 140:150, respectively. Each disk corresponds to a single Wannierization, with red/blue indicating low/high loss function [colorbar in panel (e)]. We consider 285 types of initial projections, corresponding to hydrogenic orbitals with $s$, $p$, $d$, $f$ angular momenta, and a 6$\times$6$\times$6 coarse grid. We employ a DE population of 60 candidates, and evolve the system over 150 DE generations. (g) Meta-optimized MLWF of silicon, as obtained in (a)-(f). (h) Conduction bands of silicon: black lines are from Wannier interpolation using the MLWF shown in (g); disks are explicit DFT calculations. The underlying panel is the absolute value of the band interpolation error (red line). (i) Meta-optimized MLWFs for monolayer MoS$_2$: Comparison between DFT bands (disks) and Wannier-interpolated bands (black lines) near the conduction band bottom, and associated band interpolation error (red line). Such a low interpolation error refers to an energy window within 300 meV of the band bottom. A typical MLWF is shown in (j). Green and yellow atoms represent Mo and S, respectively. In this case we employ a 24$\times$24 coarse grid and 2 bands, a DE population of 30 candidates, and 50 DE generations. (k), (l) Meta-optimized MLWF of a (5,5) CNT and corresponding Wannier interpolation of band structures (disks: DFT, black lines: Wannier, red line: band interpolation error). In this case, the band interpolation error is evaluated over the two bands crossing the charge neutrality level (zero of the energy axis). For this case we employ a coarse one-dimensional BZ grid of 12 points, a DE population of 30 candidates, and 50 DE generations.
• Figure 3: Acceleration of aiBTE calculations of carrier transport using the evolutionary meta-optimizer. (a) Timing of aiBTE calculations of the electron mobility in silicon, for different choices of the coarse BZ grid $N$$\times$$N$$\times$$N$; measured times (blue) are compared to ideal scaling time (brown). The inset shows a ball-stick model of silicon. (b) Total electron-phonon interpolation time of aiBTE calculations of the electron mobility in monolayer MoS$_2$, for different choices of the number of Wannierized bands $N_{\rm b}$; measured times (blue) are compared to theoretical scaling times (brown), normalized to the values at $N_{\rm b}=10$. The inset shows a ball-stick model of monolayer MoS$_2$. The calculated mobilities from panels (a) and (b) are reported in Supplementary Table 2 for completeness. (c) Calculated temperature-dependent electron mobility of silicon, for different choices of the coarse BZ grid corresponding to the bar plot in panel (a), and relative error with respect to the densest grid ($N=16$). The linear grid size $N$ is indicated next to each curve. The corresponding band structures and Wannier functions are shown in Supplementary Figure 3. The square indicates the value from Ref. Ponce2021. (d) Calculated temperature-dependent electron mobility of monolayer MoS$_2$, for different choices of the number of bands corresponding to the bar plot in panel (b), and relative error with respect to the highest number of bands ($N_{\rm b}=10$). The corresponding band structures and Wannier functions are shown in Supplementary Figure 4. The square indicates the value from Ref. Ha2024. (e) Speedup in calculations of the electron mobility of silicon. We compare the present runtime on a PC (48 Cascade lake cores, blue) with the runtime of Ref. Ponce2021 on the MareNostrum 4 supercomputer (with 480 Skylake cores, gray). Timing data is normalized to a single core in both cases. (f) Computed electron mobilities of 56 2D materials from the 2D Materials Cloud Database, shown as histogram. All values are reported in Supplementary Table 3, and the optimized hyperparameters are reported in Supplementary Table 1. (g) Comparison of mobilities of 2D materials calculated with the present method (blue) with state-of-the-art aiBTE calculations for the 16 compounds reported in Ref. Ha2024 (gray). (h) Speedup in high-throughput calculations of the carrier mobility of 2D materials. The gray bar refers to the average runtime over the 16 compounds reported in (g), from Ref. Ha2024 (using the Frontera supercomputer, normalized to a single core). The blue bar is the average runtime over the 56 compounds in (f), using the present method on a PC.
• Figure 4: Building extreme tight-binding Hamiltonians with meta-optimization. (a) Wannier-interpolation of the band structure of monolayer MoS$_2$ over a wide energy range. Black lines are results from the present method, blue disks are explicit DFT calculations. For this calculation, we employ an energy window in the loss function that extends from 0.8 eV below the valence band maximum to 0.5 eV above the conduction band minimum, 14 MLWFs, and a coarse BZ grid with 6$\times$6 points. The key difference between this calculation and those in Fig.2 and Fig.S4 is that, here, we only retain the first and second nearest-neighbors in the Wannier representation. Specifically, in the interpolation of the smooth Bloch Hamiltonian starting from the Hamiltonian in the Wannier representation, $\tilde{H}_{mn}({\bf k}) = \sum_{\bf R} \exp(i{\bf k}\cdot{\bf R}) H_{mn}({\bf R})$, we only sum over ${\bf R}=0$ and the first and second shell of lattice vectors ${\bf R}\ne 0$. This truncated interpolation is used to evaluate the band interpolation error $\mathcal{E}$ in the loss function, and MLWFs are meta-optimized to minimize this error. (b) Spatial decay of the Hamiltonian in the Wannier representation, $|H_{mn}({\bf R})|/\mathrm{max}[|H_{mn}({\bf R})|]$. Blue disks correspond to the Hamiltonian obtained in Fig. S4 ("simple optimized"), while black disks correspond to the present nearest-neighbor optimization ("NN optimized"). The lines are guides to the eye. The shading indicates the range of lattice vectors included in the nearest-neighbor optimization. We observe that the present optimization leads to a faster decay and improved Wannier interpolation. (c) Schematic illustration of the nearest-neighbors retained in (a) and (b).
• Figure 5: Maximally-localized electron-phonon matrix elements. Spatial decay of the electron-phonon matrix element for the conduction bands of silicon, as a function of (a) the length of the electronic Wigner-Seitz vector ($|{\bf R}_{e}|$) and (b) the length of the phonon Wigner-Seitz vector ($|{\bf R}_{p}|$). Brown disks correspond to calculation without including the term $\xi_2$ in the loss function, while green circles are obtained by including this feature. Solid lines are the corresponding exponential fits. (c) Electron mobility at 300 K, obtained using electron-phonon-optimized Wannier functions (green) and standard Wannier functions (brown). The value on top of each bar is the relative error with respect to calculations using a dense coarse BZ grid.