Add basic Wannier.jl integration

Project.toml

@@ -127,8 +127,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,144 @@
+# # 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
+
+plot_bandstructure(scfres; kline_density=10)
+
+# ## 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
+
+## Helper functions to produce hydrogenic starting guesses
+s_guess(center) = DFTK.HydrogenicWannierProjection(center, 2, 0, 0, C.Z)
+pz_guess(center) = DFTK.HydrogenicWannierProjection(center, 2, 1, 0, C.Z)
+
+wannier_model = wannier_model(scfres;
+    fileprefix="wannier/graphene",
+    n_bands=scfres.n_bands_converge,
+    n_wannier=5,
+    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]),
+    ],
+    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)
+
+# 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=[
+                  ## 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]),
+              ],
+              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)

src/DFTK.jl

@@ -39,6 +39,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
@@ -184,9 +185,11 @@ include("pseudo/list_psp.jl")
 include("pseudo/attach_psp.jl")
 
 export atomic_system, periodic_system  # Reexport from AtomsBase
+export wannier_model
 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
@@ -240,6 +243,9 @@ function __init__()
     @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
 
     # 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

@@ -1,13 +1,27 @@
 # Stubs for conditionally defined functions
 
 """
-Wannerize the obtained bands using wannier90. By default all converged
+Build a Wannier.jl model for the obtained bands. 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`.
 
+!!! 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 wannier_model end
+
+"""
+Wannerize the obtained bands using wannier90.
+Same arguments as [`wannier_model`](@ref).
+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

src/external/wannier.jl

@@ -0,0 +1,36 @@
+"""
+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.
+"""
+@timing function wannier_model(scfres;
+        n_bands=scfres.n_bands_converge,
+        n_wannier=n_bands,
+        projections::AbstractVector{<:WannierProjection}=default_wannier_centers(n_wannier),
+        fileprefix=joinpath("wannierjl", "wannier"),
+        wannier_plot=false,
+        kwargs...)
+    # Write the files
+    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