Add basic Wannier.jl integration

Project.toml

@@ -50,10 +50,14 @@ UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
 [weakdeps]
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
+Wannier = "2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9"
+wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [extensions]
 DFTKJSON3Ext = "JSON3"
+DFTKWannierExt = "Wannier"
+DFTKWannier90Ext = "wannier90_jll"
 DFTKWriteVTK = "WriteVTK"
 
 [compat]
@@ -102,6 +106,8 @@ Statistics = "1"
 TimerOutputs = "0.5.12"
 Unitful = "1"
 UnitfulAtomic = "1"
+Wannier = "0.3.2"
+wannier90_jll = "3.1"
 WriteVTK = "1"
 julia = "1.9"
 
@@ -127,8 +133,9 @@ QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
+Wannier = "2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 
 [targets]
-test = ["Test", "TestItemRunner", "ASEconvert", "Aqua", "AtomsIO", "AtomsIOPython", "CUDA", "CUDA_Runtime_jll", "ComponentArrays", "DoubleFloats", "FiniteDiff", "FiniteDifferences", "GenericLinearAlgebra", "IntervalArithmetic", "JLD2", "JSON3", "Logging", "Plots", "QuadGK", "Random", "KrylovKit", "WriteVTK", "wannier90_jll"]
+test = ["Test", "TestItemRunner", "ASEconvert", "Aqua", "AtomsIO", "AtomsIOPython", "CUDA", "CUDA_Runtime_jll", "ComponentArrays", "DoubleFloats", "FiniteDiff", "FiniteDifferences", "GenericLinearAlgebra", "IntervalArithmetic", "JLD2", "JSON3", "Logging", "Plots", "QuadGK", "Random", "KrylovKit", "Wannier", "WriteVTK", "wannier90_jll"]

docs/Project.toml

@@ -24,6 +24,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
+Wannier = "2b19380a-1f7e-4d7d-b1b8-8aa60b3321c9"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 wannier90_jll = "c5400fa0-8d08-52c2-913f-1e3f656c1ce9"
 

docs/make.jl

@@ -47,7 +47,7 @@ PAGES = [
         # options we have
         "examples/atomsbase.jl",
         "examples/input_output.jl",
-        "examples/wannier90.jl",
+        "examples/wannier.jl",
     ],
     "Tipps and tricks" => [
         # Resolving convergence issues, what solver to use, improving performance or

docs/src/features.md

@@ -42,7 +42,7 @@
     - Integration with [ASE](https://wiki.fysik.dtu.dk/ase/) and
       [AtomsBase](https://github.com/JuliaMolSim/AtomsBase.jl) for passing
       atomic structures (see [AtomsBase integration](@ref)).
-    - [Wannierization using Wannier90](@ref)
+    - [Wannierization using Wannier.jl or Wannier90](@ref)
 
 
 * Runs out of the box on Linux, macOS and Windows

examples/wannier.jl

@@ -0,0 +1,161 @@
+# # Wannierization using Wannier.jl or Wannier90
+#
+# DFTK features an interface with the programs
+# [Wannier.jl](https://wannierjl.org) and [Wannier90](http://www.wannier.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 `wannier_model` (for Wannier.jl) or `run_wannier90` (for Wannier90).
+#
+# !!! warning "No guarantees on Wannier 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 Plots
+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
+
+bands = compute_bands(scfres; kline_density=10)
+plot_bandstructure(bands)
+
+# ## Wannierization with Wannier.jl
+#
+# Now we use the `wannier_model` routine to generate a Wannier.jl model
+# that can be used to perform the wannierization procedure.
+# For now, this model generation produces file in the Wannier90 convention,
+# where 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 "wannierjl" under the prefix "wannier" (i.e. "wannierjl/wannier.win",
+# "wannierjl/wannier.wout", etc.). A different file prefix can be given with the
+# keyword argument `fileprefix` as shown below.
+#
+# We now produce a simple Wannier model for 5 MLFWs.
+#
+# For a good MLWF, we need to provide initial projections that resemble the expected shape
+# of the Wannier functions.
+# Here we will use:
+# - 3 bond-centered 2s hydrogenic orbitals for the expected σ bonds
+# - 2 atom-centered 2pz hydrogenic orbitals for the expected π bands
+
+using Wannier # Needed to make Wannier.Model available
+
+# From chemical intuition, we know that the bonds with the lowest energy are:
+# - the 3 σ bonds,
+# - the π and π* bonds.
+# We provide relevant initial projections to help Wannierization
+# converge to functions with a similar shape.
+s_guess(center) = DFTK.HydrogenicWannierProjection(center, 2, 0, 0, C.Z)
+pz_guess(center) = DFTK.HydrogenicWannierProjection(center, 2, 1, 0, C.Z)
+projections = [
+    ## Note: fractional coordinates for the centers!
+    ## 3 bond-centered 2s hydrogenic orbitals to imitate σ bonds
+    s_guess((positions[1] + positions[2]) / 2),
+    s_guess((positions[1] + positions[2] + [0, -1, 0]) / 2),
+    s_guess((positions[1] + positions[2] + [-1, -1, 0]) / 2),
+    ## 2 atom-centered 2pz hydrogenic orbitals
+    pz_guess(positions[1]),
+    pz_guess(positions[2]),
+]
+
+# Wannierize:
+wannier_model = Wannier.Model(scfres;
+    fileprefix="wannier/graphene",
+    n_bands=scfres.n_bands_converge,
+    n_wannier=5,
+    projections,
+    dis_froz_max=ustrip(auconvert(u"eV", scfres.εF))+1) # maximum frozen window, for example 1 eV above Fermi level
+
+# Once we have the `wannier_model`, we can use the functions in the Wannier.jl package:
+#
+# Compute MLWF:
+U = disentangle(wannier_model, max_iter=200);
+
+# Inspect localization before and after Wannierization:
+omega(wannier_model)
+omega(wannier_model, U)
+
+# Build a Wannier interpolation model:
+kpath = irrfbz_path(model)
+interp_model = Wannier.InterpModel(wannier_model; kpath=kpath)
+
+# And so on...
+# Refer to the Wannier.jl documentation for further examples.
+
+# (Delete temporary files when done.)
+rm("wannier", recursive=true)
+
+# ### Custom initial guesses
+#
+# We can also provide custom initial guesses for Wannierization,
+# by passing a callable function in the `projections` array.
+# The function receives the basis and a list of points (fractional coordinates in reciprocal space),
+# and returns the Fourier transform of the initial guess function evaluated at each point.
+#
+# For example, we could use Gaussians for the σ and pz guesses with the following code:
+s_guess(center) = DFTK.GaussianWannierProjection(center)
+pz_guess(center) = begin
+    ## Approximate with two Gaussians offset by 0.5 Å from the center of the atom
+    offset = model.inv_lattice * [0, 0, austrip(0.5u"Å")]
+    center1 = center + offset
+    center2 = center - offset
+    ## Build the custom projector
+    (basis, qs) -> DFTK.GaussianWannierProjection(center1)(basis, qs) - DFTK.GaussianWannierProjection(center2)(basis, qs)
+end
+## Feed to Wannier via the `projections` as before...
+
+# This example assumes that Wannier.jl version 0.3.2 is used,
+# and may need to be updated to accomodate for changes in Wannier.jl.
+#
+# Note: Some parameters supported by Wannier90 may have to be passed to Wannier.jl differently,
+# for example the max number of iterations is passed to `disentangle` in Wannier.jl,
+# but as `num_iter` to `run_wannier90`.
+
+# ## Wannierization with Wannier90
+#
+# We can use the `run_wannier90` routine to generate all required files and perform the wannierization procedure:
+
+using wannier90_jll  # Needed to make run_wannier90 available
+run_wannier90(scfres;
+              fileprefix="wannier/graphene",
+              n_wannier=5,
+              projections,
+              num_print_cycles=25,
+              num_iter=200,
+              ##
+              dis_win_max=19.0,
+              dis_froz_max=ustrip(auconvert(u"eV", scfres.εF))+1, # 1 eV above Fermi level
+              dis_num_iter=300,
+              dis_mix_ratio=1.0,
+              ##
+              wannier_plot=true,
+              wannier_plot_format="cube",
+              wannier_plot_supercell=5,
+              write_xyz=true,
+              translate_home_cell=true,
+             );
+
+# As can be observed standard optional arguments for the disentanglement
+# can be passed directly to `run_wannier90` as keyword arguments.
+
+# (Delete temporary files.)
+rm("wannier", recursive=true)

ext/DFTKWannier90Ext.jl

@@ -0,0 +1,32 @@
+module DFTKWannier90Ext
+
+using DFTK
+import wannier90_jll
+
+"""
+Run the Wannierization procedure with Wannier90.
+"""
+@DFTK.timing function DFTK.run_wannier90(scfres;
+        n_bands=scfres.n_bands_converge,
+        n_wannier=n_bands,
+        projections=DFTK.default_wannier_centers(n_wannier),
+        fileprefix=joinpath("wannier90", "wannier"),
+        wannier_plot=false,
+        kwargs...)
+
+    prefix, dir = basename(fileprefix), dirname(fileprefix)
+
+    # Prepare files
+    DFTK.write_wannier90_files(scfres; n_bands, n_wannier, projections, fileprefix, wannier_plot, kwargs...) do
+        run(Cmd(`$(wannier90_jll.wannier90()) -pp $prefix`; dir))
+        DFTK.read_w90_nnkp(fileprefix)
+    end
+
+    # Run Wannierisation procedure
+    @DFTK.timing "Wannierization" begin
+        run(Cmd(`$(wannier90_jll.wannier90()) $prefix`; dir))
+    end
+    fileprefix
+end
+
+end

ext/DFTKWannierExt.jl

@@ -0,0 +1,43 @@
+module DFTKWannierExt
+
+using DFTK
+import Wannier
+
+"""
+Use Wannier.jl to read a .win file and produce the nnkp data (i.e. the b-vectors used in the Mmn matrix).
+"""
+function get_nnkpt_from_wannier(fileprefix)
+    model = Wannier.read_w90(fileprefix; amn=false, mmn=false, eig=false)
+    bvectors = model.bvectors
+
+    nnkpts = reduce(vcat,
+        map(1:model.n_kpts) do ik
+            zp = zip(bvectors.kpb_k[:, ik], eachcol(bvectors.kpb_b[:, :, ik]))
+            map(zp) do (ikpb, G_shift)
+                (ik, ikpb, G_shift=copy(G_shift))
+            end
+        end)
+
+    (nntot=model.n_bvecs, nnkpts)
+end
+
+"""
+Build a Wannier.jl model that can be used for Wannierization.
+"""
+@DFTK.timing function Wannier.Model(scfres;
+        n_bands=scfres.n_bands_converge,
+        n_wannier=n_bands,
+        projections=DFTK.default_wannier_centers(n_wannier),
+        fileprefix=joinpath("wannierjl", "wannier"),
+        wannier_plot=false,
+        kwargs...)
+    # Write the files
+    DFTK.write_wannier90_files(scfres; n_bands, n_wannier, projections, fileprefix, wannier_plot, kwargs...) do
+        get_nnkpt_from_wannier(fileprefix)
+    end
+
+    # Read Wannier.jl model
+    Wannier.read_w90(fileprefix)
+end
+
+end

src/DFTK.jl

@@ -43,6 +43,7 @@ include("common/zeros_like.jl")
 include("common/norm.jl")
 include("common/quadrature.jl")
 include("common/hankel.jl")
+include("common/hydrogenic.jl")
 
 export PspHgh
 export PspUpf
@@ -191,6 +192,7 @@ export atomic_system, periodic_system  # Reexport from AtomsBase
 export run_wannier90
 include("external/atomsbase.jl")
 include("external/stubs.jl")  # Function stubs for conditionally defined methods
+include("external/wannier_shared.jl")
 
 export compute_bands
 export plot_bandstructure
@@ -242,9 +244,6 @@ function __init__()
     # TODO Keep these requires for now as there are open PRs changing these files.
     @require JLD2="033835bb-8acc-5ee8-8aae-3f567f8a3819"     include("external/jld2io.jl")
     @require Plots="91a5bcdd-55d7-5caf-9e0b-520d859cae80"    include("plotting.jl")
-    @require wannier90_jll="c5400fa0-8d08-52c2-913f-1e3f656c1ce9" begin
-        include("external/wannier90.jl")
-    end
 
     # TODO Move out into extension module
     @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba"  begin

src/common/hydrogenic.jl

@@ -0,0 +1,23 @@
+@doc raw"""
+Radial functions from solutions of Hydrogenic Schrödinger equation.
+Same as Wannier90 user guide Table 3.3.
+# Arguments
+- `r`: radial grid
+- `n`: principal quantum number
+- `α`: diffusivity, ``\frac{Z}/{a}`` where ``Z`` is the atomic number and
+    ``a`` is the Bohr radius.
+"""
+function radial_hydrogenic(r::AbstractVector{T}, n::Integer, α::Real=one(T)) where {T<:Real}
+    # T(...) is used for some literals to ensure that the output is of type T
+    @assert n > 0
+    if n == 1
+        f = 2 * α^(3/T(2)) * exp.(-α * r)
+    elseif n == 2
+        f = 2^(-3/T(2)) * α^(3/T(2)) * (2 .- α * r) .* exp.(-α * r/2)
+    elseif n == 3
+        f = sqrt(4/T(27)) * α^(3/T(2)) * (1 .- 2/T(3) * α * r .+ 2/T(27) * α^2 * r.^2) .* exp.(-α * r/3)
+    else
+        error("n = $n is not supported")
+    end
+    f
+end

src/external/stubs.jl

@@ -3,8 +3,9 @@
 """
 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
+and `n_wannier = n_bands` wannier functions are computed.
+Random Gaussians are used as guesses by default, can be changed
+using the `projections` kwarg. All keyword arguments supported by
 Wannier90 for the disentanglement may be added as keyword arguments.
 The function returns the `fileprefix`.