Rewrite Wannier interface

Project.toml

@@ -104,7 +104,9 @@ QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
+Wannier = "2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9"
+WannierIO = "cb1bc77f-5443-4951-af9f-05b616a3e422"
 wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 
 [targets]
-test = ["Test", "ASEconvert", "Aqua", "AtomsIO", "AtomsIOPython", "CUDA", "ComponentArrays", "DoubleFloats", "FiniteDiff", "FiniteDifferences", "GenericLinearAlgebra", "IntervalArithmetic", "JLD2", "Logging", "Plots", "QuadGK", "Random", "KrylovKit", "WriteVTK", "wannier90_jll"]
+test = ["Test", "ASEconvert", "Aqua", "AtomsIO", "AtomsIOPython", "CUDA", "ComponentArrays", "DoubleFloats", "FiniteDiff", "FiniteDifferences", "GenericLinearAlgebra", "IntervalArithmetic", "JLD2", "Logging", "Plots", "QuadGK", "Random", "KrylovKit", "WriteVTK", "Wannier", "WannierIO", "wannier90_jll"]

examples/wannier.jl

@@ -0,0 +1,92 @@
+# # Wannierization using Wannier.jl
+#
+# DFTK features an interface with the program
+# [Wannier.jl](https://www.wannierjl.org/),
+# in order to compute maximally-localized Wannier functions (MLWFs)
+# from an initial self consistent field calculation.
+# All processes are handled by calling the routine `run_wannier`.
+#
+# !!! warning "No guarantees on Wannier.jl interface"
+#     This code is at an early stage and has so far not been fully tested.
+#     Bugs are likely and we welcome issues in case you find any!
+#
+# This example shows how to obtain the MLWFs corresponding
+# to the first five bands of graphene. Since the bands 2 to 11 are entangled,
+# 15 bands are first computed to obtain 5 MLWFs by a disantanglement procedure.
+
+using DFTK
+using Unitful
+using UnitfulAtomic
+
+d = 10u"Å"
+a = 2.641u"Å"  # Graphene Lattice constant
+lattice = [a  -a/2    0;
+           0  √3*a/2  0;
+           0     0    d]
+
+C = ElementPsp(:C, psp=load_psp("hgh/pbe/c-q4"))
+atoms     = [C, C]
+positions = [[0.0, 0.0, 0.0], [1//3, 2//3, 0.0]]
+model  = model_PBE(lattice, atoms, positions)
+basis  = PlaneWaveBasis(model; Ecut=15, kgrid=[5, 5, 1])
+nbandsalg = AdaptiveBands(basis.model; n_bands_converge=15)
+scfres = self_consistent_field(basis; nbandsalg, tol=1e-5);
+
+# Plot bandstructure of the system
+using Plots
+
+plot_bandstructure(scfres; kline_density=10)
+
+# Now we use the `run_wannier` routine to generate all files needed by
+# wannier90 and to perform the Wannierization procedure using Wannier.jl.
+# In Wannier90's convention, all files are named with the same prefix and only differ by
+# their extensions. By default all generated input and output files are stored
+# in the subfolder "wannier" under the prefix "wannier" (i.e. "wannier/wannier.win",
+# "wannier/wannier.wout", etc.). A different file prefix can be given with the
+# keyword argument `fileprefix` as shown below.
+#
+# We now solve for 5 MLWF using Wannier.jl:
+
+using Wannier  # Needed to make run_wannier available
+
+# The Wannier.jl interface is very similar to the Wannier90 example,
+# except that the function `run_wannier` is used instead of `run_wannier90`.
+# To further explore the functionalities of the MLWF interface, in this example
+# we use SCDM to generate a better initial guess for the MLWFs
+# (by default, `run_wannier` will use random initial guess which is not good).
+# We need to first unfold the `scfres` to a MP kgrid for Wannierization,
+# and remove the possibly unconverged bands (bands above `scfres.n_bands_converge`)
+exclude_bands = DFTK._default_exclude_bands(scfres)
+basis, ψ, eigenvalues = DFTK.unfold_scfres_wannier(scfres, exclude_bands)
+
+# Then compute the SCDM initial guess with [`compute_amn_scdm`](@ref)
+# Since this is an entangled case, we need a weight factor, using `erfc` function.
+# Note that the unit of `μ` and `σ` are `Ha`, as in the DFTK convention.
+# Here we choose these numbers by inspecting the band structure, you can also
+# test with different values to see how it affects the result.
+μ, σ = 0.0, 0.01
+f = DFTK.scdm_f_erfc(basis, eigenvalues, μ, σ)
+# and construct 5 MLWFs
+n_wann = 5
+A = DFTK.compute_amn_scdm(basis, ψ, n_wann, f)
+# we pass the `A` matrix to `run_wannier`, so it will skip the `compute_amn` step
+wann_model, = run_wannier(
+    scfres;
+    fileprefix="wannier/graphene",
+    n_wann,
+    A,
+    dis_froz_max=0.1,
+);
+# Note we unwrap the returned objects since in `:collinear` case, the
+# `run_wannier` will return two `Wannier.Model` objects.
+# As can be observed standard optional arguments for the disentanglement
+# can be passed directly to `run_wannier` as keyword arguments.
+
+# The MLWF centers and spreads can be obtained from
+Wannier.omega(wann_model)
+
+# Please refer to the [Wannier.jl documentation](https://wannierjl.org/)
+# for more advanced usage of the Wannier function interface.
+
+# (Delete temporary files.)
+rm("wannier", recursive=true)

examples/wannier90.jl

@@ -33,23 +33,26 @@ nbandsalg = AdaptiveBands(basis.model; n_bands_converge=15)
 scfres = self_consistent_field(basis; nbandsalg, tol=1e-5);
 
 # Plot bandstructure of the system
-
+using Plots
 plot_bandstructure(scfres; kline_density=10)
 
 # Now we use the `run_wannier90` routine to generate all files needed by
-# wannier90 and to perform the wannierization procedure.
+# wannier90 and to perform the Wannierization procedure.
 # In Wannier90's convention, all files are named with the same prefix and only differ by
 # their extensions. By default all generated input and output files are stored
-# in the subfolder "wannier90" under the prefix "wannier" (i.e. "wannier90/wannier.win",
-# "wannier90/wannier.wout", etc.). A different file prefix can be given with the
+# in the subfolder "wannier" under the prefix "wannier" (i.e. "wannier/wannier.win",
+# "wannier/wannier.wout", etc.). A different file prefix can be given with the
 # keyword argument `fileprefix` as shown below.
 #
 # We now solve for 5 MLWF using wannier90:
 
+using WannierIO  # Needed to save Wannier90 files, e.g., amn, mmn, etc.
 using wannier90_jll  # Needed to make run_wannier90 available
+
 run_wannier90(scfres;
               fileprefix="wannier/graphene",
-              n_wannier=5,
+              n_wann=5,  # Number of MLWFs
+              # following are inputs to wannier90
               num_print_cycles=25,
               num_iter=200,
               ##
@@ -58,6 +61,8 @@ run_wannier90(scfres;
               dis_num_iter=300,
               dis_mix_ratio=1.0,
               ##
+              bands_plot=true,
+              ##
               wannier_plot=true,
               wannier_plot_format="cube",
               wannier_plot_supercell=5,

src/DFTK.jl

@@ -183,7 +183,7 @@ include("pseudo/attach_psp.jl")
 
 export DFTKPotential
 export atomic_system, periodic_system  # Reexport from AtomsBase
-export run_wannier90
+export run_wannier90, run_wannier
 include("external/atomsbase.jl")
 include("external/interatomicpotentials.jl")
 include("external/stubs.jl")  # Function stubs for conditionally defined methods
@@ -233,8 +233,14 @@ function __init__()
     @require Plots="91a5bcdd-55d7-5caf-9e0b-520d859cae80" include("plotting.jl")
     @require JLD2="033835bb-8acc-5ee8-8aae-3f567f8a3819"  include("external/jld2io.jl")
     @require WriteVTK="64499a7a-5c06-52f2-abe2-ccb03c286192" include("external/vtkio.jl")
-    @require wannier90_jll="c5400fa0-8d08-52c2-913f-1e3f656c1ce9" begin
-        include("external/wannier90.jl")
+    @require WannierIO="cb1bc77f-5443-4951-af9f-05b616a3e422" begin
+        include("external/wannierio.jl")
+        @require wannier90_jll="c5400fa0-8d08-52c2-913f-1e3f656c1ce9" begin
+            include("external/wannier90.jl")
+        end
+        @require Wannier="2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9" begin
+            include("external/wannier.jl")
+        end
     end
     @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba"  begin
         include("workarounds/cuda_arrays.jl")

src/common/ortho.jl

@@ -7,3 +7,8 @@
         return x[:, 1:size(φk, 2)]
     end
 end
+
+function ortho_lowdin(A::AbstractMatrix)
+    F = svd(A)
+    return F.U * F.Vt
+end

src/external/stubs.jl

@@ -1,17 +1 @@
 # Stubs for conditionally defined functions
-
-"""
-Wannerize the obtained bands using wannier90. By default all converged
-bands from the `scfres` are employed (change with `n_bands` kwargs)
-and `n_wannier = n_bands` wannier functions are computed using
-random Gaussians as guesses. All keyword arguments supported by
-Wannier90 for the disentanglement may be added as keyword arguments.
-The function returns the `fileprefix`.
-
-!!! warning "Experimental feature"
-    Currently this is an experimental feature, which has not yet been tested
-    to full depth. The interface is considered unstable and may change
-    incompatibly in the future. Use at your own risk and please report bugs
-    in case you encounter any.
-"""
-function run_wannier90 end

src/external/wannier.jl

@@ -0,0 +1,69 @@
+"""
+Call `Wannier.jl` to Wannierize the results.
+
+# Arguments
+- `scfres`: the result of scf.
+
+`kwargs` will be passed to [`save_wannier`](@ref).
+"""
+function run_wannier(
+    scfres::NamedTuple;
+    fileprefix::AbstractString="wannier/wannier",
+    exclude_bands::AbstractArray{<:Integer}=_default_exclude_bands(scfres),
+    kwargs...,
+)
+    basis, ψ, eigenvalues = unfold_scfres_wannier(scfres, exclude_bands)
+
+    n_bands = length(eigenvalues[1])
+    @assert haskey(kwargs, :n_wann) "Must specify `n_wann` in `kwargs`"
+    n_wann = kwargs[:n_wann]
+
+    # Although with Wannier.jl we can pass matrices in-memory,
+    # however I still save the files so that we can restart later.
+    if basis.model.spin_polarization == :collinear
+        prefixes = ["$(fileprefix)_up", "$(fileprefix)_dn"]
+    else
+        prefixes = [fileprefix]
+    end
+
+    @timing "Compute b-vectors" begin
+        win = get_wannier90_win(basis; num_wann=n_wann, num_bands=n_bands)
+        # Note I am not using directly `basis.model.recip_lattice` here since
+        # that one is in Bohr unit, while `win.unit_cell_cart` is in Angstrom.
+        recip_lattice = compute_recip_lattice(win.unit_cell_cart)
+        bvectors = Wannier.generate_kspace_stencil(recip_lattice, win.mp_grid, win.kpoints)
+        # TODO I am naming kpb_G as kpb_b in WannierIO.jl for the moment,
+        # probably going to rename it in the future.
+        kpb_k, kpb_G = bvectors.kpb_k, bvectors.kpb_G
+        nnkp = (; kpb_k, kpb_G)
+    end
+
+    unk = get(kwargs, :wannier_plot, false)
+    εF = auconvert(u"eV", scfres.εF).val
+
+    models = []
+    for (spin, prefix) in enumerate(prefixes)
+        @timing "Compute & save Wannier matrices" win, M, A, E = save_wannier(
+            basis, ψ, eigenvalues; fileprefix=prefix, nnkp, spin, unk,
+            fermi_energy=εF, kwargs...)
+
+        @timing "Run Wannierization" begin
+            # I simply call disentangle which works for both isolated and
+            # entangled cases. Also if no frozen window provided, just use
+            # 0.5 eV above Fermi.
+            dis_froz_max = get(win, :dis_froz_max, εF + 0.5)
+            dis_froz_min = get(win, :dis_froz_min, -Inf)
+            frozen_bands = Wannier.get_frozen_bands(E, dis_froz_max, dis_froz_min)
+
+            atom_labels = [first(i) for i in win.atoms_frac]
+            atom_positions = [last(i) for i in win.atoms_frac]
+            model = Wannier.Model(
+                win.unit_cell_cart, atom_positions, atom_labels,
+                bvectors, M, A, E, frozen_bands)
+            model.gauges .= Wannier.disentangle(model)
+
+            push!(models, model)
+        end
+    end
+    models
+end