diff --git a/Project.toml b/Project.toml index b3cbee27a3..12298c2d2f 100644 --- a/Project.toml +++ b/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"] diff --git a/docs/Project.toml b/docs/Project.toml index 2fd13363cd..e5d4516dfa 100644 --- a/docs/Project.toml +++ b/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" diff --git a/docs/make.jl b/docs/make.jl index 3254a79900..ca700d3b4f 100644 --- a/docs/make.jl +++ b/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 diff --git a/docs/src/features.md b/docs/src/features.md index a2822c3483..a779cd57c5 100644 --- a/docs/src/features.md +++ b/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 diff --git a/examples/wannier.jl b/examples/wannier.jl new file mode 100644 index 0000000000..9bb0589cee --- /dev/null +++ b/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) +function pz_guess(center) + ## 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) diff --git a/examples/wannier90.jl b/examples/wannier90.jl deleted file mode 100644 index 5afb65837c..0000000000 --- a/examples/wannier90.jl +++ /dev/null @@ -1,73 +0,0 @@ -# # Wannierization using Wannier90 -# -# DFTK features an interface with the program -# [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 `run_wannier90`. -# -# !!! warning "No guarantees on Wannier90 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 - -bands = compute_bands(scfres; kline_density=10) -plot_bandstructure(bands) - -# Now we use the `run_wannier90` routine to generate all files needed by -# 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 -# keyword argument `fileprefix` as shown below. -# -# We now solve for 5 MLWF using wannier90: - -using wannier90_jll # Needed to make run_wannier90 available -run_wannier90(scfres; - fileprefix="wannier/graphene", - n_wannier=5, - num_print_cycles=25, - num_iter=200, - ## - dis_win_max=19.0, - dis_froz_max=0.1, - 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) diff --git a/ext/DFTKWannier90Ext.jl b/ext/DFTKWannier90Ext.jl new file mode 100644 index 0000000000..202baf4791 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/ext/DFTKWannierExt.jl b/ext/DFTKWannierExt.jl new file mode 100644 index 0000000000..305a0a6caa --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/DFTK.jl b/src/DFTK.jl index ed39db498e..a483372a7d 100644 --- a/src/DFTK.jl +++ b/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 diff --git a/src/common/hydrogenic.jl b/src/common/hydrogenic.jl new file mode 100644 index 0000000000..a5208ffd16 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/external/stubs.jl b/src/external/stubs.jl index 321587c63d..578b0a6726 100644 --- a/src/external/stubs.jl +++ b/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`. diff --git a/src/external/wannier90.jl b/src/external/wannier_shared.jl similarity index 68% rename from src/external/wannier90.jl rename to src/external/wannier_shared.jl index d0b7f5ca0b..0aa90bf2bf 100644 --- a/src/external/wannier90.jl +++ b/src/external/wannier_shared.jl @@ -1,5 +1,75 @@ +# Shared functionality to operate with Wannier programs (Wannier.jl and Wannier90). +# The communication happens with the Wannier90 file formats which Wannier.jl also understands. + using Dates +""" +A Gaussian-shaped initial guess. Can be used as an approximation of an s- or σ-like orbital. +""" +struct GaussianWannierProjection + center::AbstractVector +end + +function (proj::GaussianWannierProjection)(basis::PlaneWaveBasis, qs) + # associate a center with the fourier transform of the corresponding gaussian + # TODO: what is the normalization here? + + model = basis.model + map(qs) do q + q_cart = recip_vector_red_to_cart(model, q) + exp( 2π*(-im*dot(q, proj.center) - dot(q_cart, q_cart) / 4) ) + end +end + +@doc raw""" +A hydrogenic initial guess. + +`α` is the diffusivity, ``\frac{Z}/{a}`` where ``Z`` is the atomic number and + ``a`` is the Bohr radius. +""" +struct HydrogenicWannierProjection + center::AbstractVector + n::Integer + l::Integer + m::Integer + α::Real +end + +function (proj::HydrogenicWannierProjection)(basis::PlaneWaveBasis, qs) + # TODO: Performance can probably be improved a lot here. + # TODO: Some of this logic could be used for the pswfc from the UPF PSPs. + + # this is same as QE pw2wannier90, r = exp(x) / α + xmin = -6.0 + dx = 0.025 + rmax = 10.0 + n_r = round(Int, (log(rmax) - xmin) / dx) + 1 + x = range(; start=xmin, length=n_r, step=dx) + # rgrid depends on α + r = exp.(x) ./ proj.α + dr = r .* dx + R = radial_hydrogenic(r, proj.n, proj.α) + r2_R_dr = r.^2 .* R .* dr + + model = basis.model + map(qs) do q + q_cart = recip_vector_red_to_cart(model, q) + qnorm = norm(q_cart) + + radial_part = 0.0 + for (ir, rval) in enumerate(r) + @inbounds radial_part += r2_R_dr[ir] * sphericalbesselj_fast(proj.l, qnorm * rval) + end + + # both q and proj.center in reduced coordinates + center_offset = exp(-im*2π*dot(q, proj.center)) + # without the 4π normalization because the matrix will be orthonormalized anyway + angular_part = ylm_real(proj.l, proj.m, q_cart) * (-im)^proj.l + + center_offset * angular_part * radial_part + end +end + """ Write a win file at the indicated prefix. Parameters to Wannier90 can be added as kwargs: e.g. `num_iter=500`. @@ -61,7 +131,7 @@ function write_w90_win(fileprefix::String, basis::PlaneWaveBasis; "$(basis.kgrid.kgrid_size[3])\n") println(fp, "begin kpoints") for kpt in basis.kpoints - @printf fp "%10.6f %10.6f %10.6f\n" kpt.coordinate... + @printf fp "%10.10f %10.10f %10.10f\n" kpt.coordinate... end println(fp, "end kpoints") end @@ -128,7 +198,7 @@ end for iband = 1:n_bands ψnk_real = ifft(basis, basis.kpoints[ik], @view ψ[ik][:, iband]) for iz = 1:fft_size[3], iy = 1:fft_size[2], ix = 1:fft_size[1] - println(fp, real(ψnk_real[ix, iy, iz]), " ", imag(ψnk_real[ix, iy, iz])) + @printf fp "%25.18f %25.18f\n" real(ψnk_real[ix, iy, iz]) imag(ψnk_real[ix, iy, iz]) end end end @@ -184,39 +254,43 @@ end nothing end - @doc raw""" -Compute the matrix ``[A_k]_{m,n} = \langle ψ_m^k | g^{\text{per}}_n \rangle`` +Compute the starting matrix for Wannierization. -``g^{per}_n`` are periodized gaussians whose respective centers are given as an - (num_bands,1) array [ [center 1], ... ]. +Wannierization searches for a unitary matrix ``U_{m_n}``. +As a starting point for the search, we can provide an initial guess function ``g`` +for the shape of the Wannier functions, based on what we expect from knowledge of the problem or physical intuition. +This starting matrix is called ``[A_k]_{m,n}``, and is computed as follows: +``[A_k]_{m,n} = \langle ψ_m^k | g^{\text{per}}_n \rangle``. +The matrix will be orthonormalized by the chosen Wannier program, we don't need to do so ourselves. Centers are to be given in lattice coordinates and G_vectors in reduced coordinates. The dot product is computed in the Fourier space. Given an orbital ``g_n``, the periodized orbital is defined by : ``g^{per}_n = \sum\limits_{R \in {\rm lattice}} g_n( \cdot - R)``. -The Fourier coefficient of ``g^{per}_n`` at any G +The Fourier coefficient of ``g^{per}_n`` at any G is given by the value of the Fourier transform of ``g_n`` in G. + +Each projection is a callable object that accepts the basis and some qpoints as an argument, +and returns the Fourier transform of ``g_n`` at the qpoints. """ -function compute_Ak_gaussian_guess(basis::PlaneWaveBasis, ψk, kpt, centers, n_bands) - n_wannier = length(centers) +function compute_amn_kpoint(basis::PlaneWaveBasis, kpt, ψk, projections, n_bands) + n_wannier = length(projections) # TODO This function should be improved in performance - # associate a center with the fourier transform of the corresponding gaussian - fourier_gn(center, qs) = [exp( 2π*(-im*dot(q, center) - dot(q, q) / 4) ) for q in qs] qs = vec(map(G -> G .+ kpt.coordinate, G_vectors(basis))) # all q = k+G in reduced coordinates Ak = zeros(eltype(ψk), (n_bands, n_wannier)) # Compute Ak for n = 1:n_wannier - # Functions are l^2 normalized in Fourier, in DFTK conventions. - norm_gn_per = norm(fourier_gn(centers[n], qs), 2) + proj = projections[n] + gn_per = proj(basis, qs[kpt.mapping]) # Fourier coeffs of gn_per for k+G in common with ψk - coeffs_gn_per = fourier_gn(centers[n], qs[kpt.mapping]) ./ norm_gn_per + # Functions are l^2 normalized in Fourier, in DFTK conventions. + coeffs_gn_per = gn_per ./ norm(gn_per) # Compute overlap for m = 1:n_bands - # TODO Check the ordering of m and n here! Ak[m, n] = dot(ψk[:, m], coeffs_gn_per) end end @@ -224,13 +298,18 @@ function compute_Ak_gaussian_guess(basis::PlaneWaveBasis, ψk, kpt, centers, n_b end -@timing function write_w90_amn(fileprefix::String, basis::PlaneWaveBasis, ψ; n_bands, centers) +@timing function write_w90_amn( + fileprefix::String, + basis::PlaneWaveBasis, + projections, + ψ; + n_bands) open(fileprefix * ".amn", "w") do fp println(fp, "Generated by DFTK at $(now())") - println(fp, "$n_bands $(length(basis.kpoints)) $(length(centers))") + println(fp, "$n_bands $(length(basis.kpoints)) $(length(projections))") for (ik, (ψk, kpt)) in enumerate(zip(ψ, basis.kpoints)) - Ak = compute_Ak_gaussian_guess(basis, ψk, kpt, centers, n_bands) + Ak = compute_amn_kpoint(basis, kpt, ψk, projections, n_bands) for n = 1:size(Ak, 2) for m = 1:size(Ak, 1) @printf(fp, "%3i %3i %3i %22.18f %22.18f \n", @@ -246,17 +325,21 @@ end Default random Gaussian guess for maximally-localised wannier functions generated in reduced coordinates. """ -default_wannier_centres(n_wannier) = [rand(1, 3) for _ = 1:n_wannier] - -@timing function run_wannier90(scfres; - n_bands=scfres.n_bands_converge, - n_wannier=n_bands, - centers=default_wannier_centres(n_wannier), - fileprefix=joinpath("wannier90", "wannier"), - wannier_plot=false, kwargs...) +default_wannier_centers(n_wannier) = [GaussianWannierProjection(rand(3)) for _ = 1:n_wannier] + +""" +Shared file writing code for Wannier.jl and Wannier90. +""" +@timing function write_wannier90_files(preprocess_call, scfres; + n_bands, + n_wannier, + projections, + fileprefix, + wannier_plot, + kwargs...) # TODO None of the routines consider spin at the moment @assert scfres.basis.model.spin_polarization in (:none, :spinless) - @assert length(centers) == n_wannier + @assert length(projections) == n_wannier # TODO Use band_data_to_dict to get this easily MPI compatible. @@ -266,29 +349,20 @@ default_wannier_centres(n_wannier) = [rand(1, 3) for _ = 1:n_wannier] ψ = scfres_unfold.ψ # Make wannier directory ... - dir, prefix = dirname(fileprefix), basename(fileprefix) - mkpath(dir) + mkpath(dirname(fileprefix)) - # Write input file and launch Wannier90 preprocessing + # Write input file and launch preprocessing write_w90_win(fileprefix, basis; - num_wann=length(centers), num_bands=n_bands, wannier_plot, kwargs...) - run(Cmd(`$(wannier90_jll.wannier90()) -pp $prefix`; dir)) - - nnkp = read_w90_nnkp(fileprefix) + num_wann=n_wannier, num_bands=n_bands, wannier_plot, kwargs...) + nnkp = preprocess_call() # Files for main wannierization run write_w90_eig(fileprefix, scfres_unfold.eigenvalues; n_bands) - write_w90_amn(fileprefix, basis, ψ; n_bands, centers) + write_w90_amn(fileprefix, basis, projections, ψ; n_bands) write_w90_mmn(fileprefix, basis, ψ, nnkp; n_bands) # Writing the unk files is expensive (requires FFTs), so only do if needed. if wannier_plot write_w90_unk(fileprefix, basis, ψ; n_bands, spin=1) end - - # Run Wannierisation procedure - @timing "Wannierization" begin - run(Cmd(`$(wannier90_jll.wannier90()) $prefix`; dir)) - end - fileprefix end diff --git a/test/external/wannier.jl b/test/external/wannier.jl new file mode 100644 index 0000000000..ae08f76c1c --- /dev/null +++ b/test/external/wannier.jl @@ -0,0 +1,60 @@ +@testitem "Test Wannierization" tags=[:dont_test_mpi, :dont_test_windows] setup=[TestCases] begin + using DFTK + using Wannier + using wannier90_jll + silicon = TestCases.silicon + + model = model_LDA(silicon.lattice, silicon.atoms, silicon.positions) + basis = PlaneWaveBasis(model; Ecut=5, kgrid=[4, 4, 4], kshift=[1, 1, 1]/2) + nbandsalg = AdaptiveBands(model; n_bands_converge=12) + scfres = self_consistent_field(basis; nbandsalg, tol=1e-12) + + @testset begin + @testset "Using wannier90" begin + mktempdir() do tmpdir + run_wannier90(scfres; + fileprefix="$tmpdir/Si", + n_wannier=8, bands_plot=true, + num_print_cycles=50, num_iter=500, + dis_win_max=17.185257, + dis_froz_max=6.8318033, + dis_num_iter=120, + dis_mix_ratio=1.0, + wannier_plot=true) + + @test isfile("$tmpdir/Si.amn") + @test isfile("$tmpdir/Si.chk") + @test isfile("$tmpdir/Si.eig") + @test isfile("$tmpdir/Si.mmn") + @test isfile("$tmpdir/Si.nnkp") + @test isfile("$tmpdir/Si.win") + @test isfile("$tmpdir/Si.wout") + @test !isfile("$tmpdir/Si.werr") + end + end + + @testset "Using Wannier.jl" begin + mktempdir() do tmpdir + wannier_model = Wannier.Model(scfres; + fileprefix="$tmpdir/Si", + n_wannier=8, bands_plot=true, + dis_win_max=17.185257, + dis_froz_max=6.8318033, + dis_mix_ratio=1.0, + wannier_plot=true) + + wannier_model.U .= disentangle(wannier_model; max_iter=500) + + # for now the Wannier.jl compat writes the amn, eig, mmn and win files + @test isfile("$tmpdir/Si.amn") + @test !isfile("$tmpdir/Si.chk") + @test isfile("$tmpdir/Si.eig") + @test isfile("$tmpdir/Si.mmn") + @test !isfile("$tmpdir/Si.nnkp") + @test isfile("$tmpdir/Si.win") + @test !isfile("$tmpdir/Si.wout") # NO .wout file, the output is in the `wannier_model` + @test !isfile("$tmpdir/Si.werr") + end + end + end +end diff --git a/test/external/wannier90.jl b/test/external/wannier90.jl deleted file mode 100644 index b46094dd51..0000000000 --- a/test/external/wannier90.jl +++ /dev/null @@ -1,32 +0,0 @@ -@testitem "Test run_wannier90" tags=[:dont_test_mpi, :dont_test_windows] setup=[TestCases] begin - using DFTK - using wannier90_jll - silicon = TestCases.silicon - - model = model_LDA(silicon.lattice, silicon.atoms, silicon.positions) - basis = PlaneWaveBasis(model; Ecut=5, kgrid=[4, 4, 4], kshift=[1, 1, 1]/2) - nbandsalg = AdaptiveBands(model; n_bands_converge=12) - scfres = self_consistent_field(basis; nbandsalg, tol=1e-12) - - fileprefix = "wannier90_outputs/Si" - run_wannier90(scfres; fileprefix, - n_wannier=8, bands_plot=true, - num_print_cycles=50, num_iter=500, - dis_win_max=17.185257, - dis_froz_max=6.8318033, - dis_num_iter=120, - dis_mix_ratio=1.0, - wannier_plot=true) - - @test isfile("wannier90_outputs/Si.amn") - @test isfile("wannier90_outputs/Si.chk") - @test isfile("wannier90_outputs/Si.eig") - @test isfile("wannier90_outputs/Si.mmn") - @test isfile("wannier90_outputs/Si.nnkp") - @test isfile("wannier90_outputs/Si.win") - @test isfile("wannier90_outputs/Si.wout") - @test !isfile("wannier90_outputs/Si.werr") - - # remove produced files - rm("wannier90_outputs", recursive=true) -end diff --git a/test/hydrogenic.jl b/test/hydrogenic.jl new file mode 100644 index 0000000000..9ea629400f --- /dev/null +++ b/test/hydrogenic.jl @@ -0,0 +1,8 @@ +@testitem "Type-stability of radial_hydrogenic" begin + using DFTK: radial_hydrogenic + + for n = [1, 2, 3] + @inferred radial_hydrogenic([1.0, 1.0], n) + @inferred radial_hydrogenic([1f0, 1f0], n) + end +end \ No newline at end of file