From 55a43a1482f3a9cbe810b8501975a668cb63ce0c Mon Sep 17 00:00:00 2001
From: Junfeng Qiao <qiaojunfeng@outlook.com>
Date: Thu, 14 Dec 2023 12:54:58 +0100
Subject: [PATCH] Recursively search Fermi energy by refining kmesh (#36)

---
 Project.toml                       |   4 +
 src/Wannier.jl                     |   1 +
 src/interpolation/fermi_energy.jl  | 233 +++++++++++++++++++++++++++++
 test/interpolation/fermi_energy.jl |  28 ++++
 4 files changed, 266 insertions(+)
 create mode 100644 src/interpolation/fermi_energy.jl
 create mode 100644 test/interpolation/fermi_energy.jl

diff --git a/Project.toml b/Project.toml
index f24118a..92428d5 100644
--- a/Project.toml
+++ b/Project.toml
@@ -25,6 +25,8 @@ PeriodicTable = "7b2266bf-644c-5ea3-82d8-af4bbd25a884"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Spglib = "f761d5c5-86db-4880-b97f-9680a7cccfb5"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
@@ -56,6 +58,8 @@ PeriodicTable = "1.1"
 PlotlyJS = "0.18"
 ProgressMeter = "1.7"
 Reexport = "1.2"
+Roots = "2"
+SpecialFunctions = "2"
 Spglib = "0.6, 0.7, 0.8, 0.9"
 StaticArrays = "1"
 TOML = "1"
diff --git a/src/Wannier.jl b/src/Wannier.jl
index 53fbedf..92ccfd3 100644
--- a/src/Wannier.jl
+++ b/src/Wannier.jl
@@ -74,6 +74,7 @@ include("interpolation/hamiltonian_hessian.jl")
 include("interpolation/spin.jl")
 include("interpolation/berry_curvature.jl")
 include("interpolation/orbital_magnetization.jl")
+include("interpolation/fermi_energy.jl")
 
 # include("interpolation/magmom.jl")
 
diff --git a/src/interpolation/fermi_energy.jl b/src/interpolation/fermi_energy.jl
new file mode 100644
index 0000000..479e02f
--- /dev/null
+++ b/src/interpolation/fermi_energy.jl
@@ -0,0 +1,233 @@
+import Base.Broadcast.broadcastable
+using Roots
+using SpecialFunctions: erf
+
+abstract type SmearingFunction end
+
+Base.Broadcast.broadcastable(S::SmearingFunction) = Ref(S)
+
+struct FermiDiracSmearing <: SmearingFunction end
+
+occupation(x, ::FermiDiracSmearing) = 1 / (1 + exp(x))
+
+struct ColdSmearing <: SmearingFunction end
+
+function occupation(x::T, ::ColdSmearing) where {T}
+    return (
+        -erf(x + 1 / sqrt(T(2))) / 2 +
+        1 / sqrt(2 * T(π)) * exp(-(-x - 1 / sqrt(T(2)))^2) +
+        1 / T(2)
+    )
+end
+
+struct NoneSmearing <: SmearingFunction end
+
+occupation(x, ::NoneSmearing) = x > 0 ? zero(x) : one(x)
+
+@doc raw"""
+Compute occupation given eigenvalues and Fermi energy.
+
+# Arguments
+- `eigenvalues`: eigenvalues in eV
+- `εF`: Fermi energy in eV
+
+# Keyword arguments
+- `kBT`: temperature in the same unit as `E`, i.e., $k_B T$ in eV
+- `prefactor`: 1 for collinear calculation, 2 for spinless
+"""
+function occupation(
+    eigenvalues::AbstractVector, εF, kBT, smearing::SmearingFunction; prefactor
+)
+    T = promote_type(eltype(eltype(eigenvalues)), typeof(εF), typeof(kBT))
+    inv_kBT = iszero(kBT) ? T(Inf) : 1 / kBT
+
+    occ = map(eigenvalues) do εk
+        prefactor * occupation.((εk .- εF) .* inv_kBT, smearing)
+    end
+    return occ
+end
+
+"""
+    $(SIGNATURES)
+
+Compute the number of electrons with a given density of states.
+
+# Arguments
+- `energy`: Vector of energy, in eV unit
+- `dos`: density of states on energy grid, in states/eV unit
+- `εF`: Fermi energy, in eV unit
+
+# Returns
+- `n_electrons`: number of electrons
+"""
+function compute_n_electrons(energy::AbstractVector, dos::AbstractVector, εF::Number)
+    dE = energy[2] - energy[1]
+    cum_dos = cumsum(dos) * dE
+
+    idx = argmin(abs.(energy .- εF))
+    tol_energy = 1e-5
+    if abs(energy[idx] - εF) > tol_energy
+        error("Fermi energy not found in energy grid")
+    end
+
+    return cum_dos[idx]
+end
+
+default_kweights(eigenvalues) = 1 / length(eigenvalues)
+
+function compute_n_electrons(
+    occupation::AbstractVector, kweights=default_kweights(occupation)
+)
+    return sum(kweights .* sum.(occupation))
+end
+
+function compute_fermi_energy(
+    eigenvalues::AbstractVector,
+    n_electrons::Real,
+    kBT::Real,
+    smearing::SmearingFunction;
+    prefactor=2,
+    kweights=default_kweights(eigenvalues),
+    tol_n_electrons=1e-6,
+)
+    # Get rough bounds to bracket εF
+    min_ε = minimum(minimum, eigenvalues) - 1
+    max_ε = maximum(maximum, eigenvalues) + 1
+
+    excess(εF) = begin
+        occ = occupation(eigenvalues, εF, kBT, smearing; prefactor)
+        compute_n_electrons(occ, kweights) - n_electrons
+    end
+    @assert excess(min_ε) <= 0 <= excess(max_ε) "Fermi energy not bracketed $(excess(min_ε)) $(excess(max_ε))"
+
+    εF = Roots.find_zero(excess, (min_ε, max_ε), Roots.Bisection(); atol=tol_n_electrons)
+    Δn_elec = excess(εF)
+    abs(Δn_elec) > tol_n_electrons &&
+        error("Failed to find Fermi energy within tolerance, Δn_elec = $Δn_elec")
+
+    return εF
+end
+
+"""
+Compute Fermi energy by recursively refining the kgrid when interpolating the Hamiltonian.
+"""
+function compute_fermi_energy(
+    interp::HamiltonianInterpolator,
+    kgrid::AbstractVector,
+    n_electrons::Real,
+    kBT::Real,
+    smearing::SmearingFunction;
+    prefactor=2,
+    tol_n_electrons=1e-6,
+    tol_εF=1e-3,
+    max_refine=10,
+)
+    kpoints = get_kpoints(kgrid)
+    eigenvals, _ = interp(kpoints)
+    # the initial guessing Fermi energy
+    εF = compute_fermi_energy(
+        eigenvals, n_electrons, kBT, smearing; prefactor, tol_n_electrons
+    )
+    @printf("εF on input kgrid   : %15.9f eV, n_kpoints = %8d\n", εF, length(kpoints))
+
+    dv = Vec3(1 ./ kgrid)
+    kweight = default_kweights(eigenvals)
+    kvoxels = map(kpoints) do kpt
+        Kvoxel(kpt, dv, kweight)
+    end
+    adpt_kgrid = AdaptiveKgrid(kvoxels, eigenvals)
+
+    εF_prev = εF - 1
+    iter = 1
+    # search range
+    width_εF = 0.5
+    while abs(εF - εF_prev) > tol_εF && iter <= max_refine
+        refine_iks = filter(1:length(adpt_kgrid)) do ik
+            any(abs.(adpt_kgrid.vals[ik] .- εF) .<= width_εF)
+        end
+        # alternate between even and odd refinement, so it works for the
+        # K/K' point of graphene as well
+        # I should iterate odd grid 1st, otherwise it seems the graphene
+        # case could still stuck at wrong εF with [8, 8, 1] kgrid
+        n_subvoxels = iter % 2 == 0 ? 2 : 3
+        refine!(adpt_kgrid, refine_iks, x -> interp(x)[1]; n_subvoxels)
+
+        εF_prev = εF
+        εF = compute_fermi_energy(
+            adpt_kgrid.vals,
+            n_electrons,
+            kBT,
+            smearing;
+            prefactor,
+            kweights=[kv.weight for kv in adpt_kgrid.kvoxels],
+            tol_n_electrons,
+        )
+        # gradually reduce width_εF to save computation
+        ΔεF = εF - εF_prev
+        @printf(
+            "εF at iteration %3d : %15.9f eV, n_kpoints = %8d, ΔεF = %16.9e eV\n",
+            iter,
+            εF,
+            length(adpt_kgrid),
+            ΔεF,
+        )
+        iter += 1
+        # after 10 iters, the width is mutiplied by 0.8^10 ≈ 0.107
+        # width_εF *= 0.8
+        # set next search range according to ΔεF
+        width_εF = min(width_εF, abs(ΔεF) * 5)
+    end
+    return εF
+end
+
+struct Kvoxel{T,VT<:AbstractVector{T}}
+    """fractional coordinates of kpoint"""
+    point::VT
+
+    """length of the kvoxel along three dimensions"""
+    dv::VT
+
+    """weight of the kpoint"""
+    weight::T
+end
+
+struct AdaptiveKgrid{KV<:Kvoxel,VT}
+    kvoxels::Vector{KV}
+    vals::Vector{VT}
+end
+
+Base.length(ag::AdaptiveKgrid) = length(ag.kvoxels)
+
+"""
+
+# Arguments
+- `ag`: AdaptiveKgrid
+- `iks`: indices of kvoxels to be refined
+
+# Keyword arguments
+- `n_subvoxels`: number of subvoxels along each dimension. 2 -> split into 8 subvoxels
+"""
+function refine!(ag::AdaptiveKgrid, iks::AbstractVector, interp::Function; n_subvoxels=2)
+    new_kvoxels = eltype(ag.kvoxels)[]
+
+    # split the current kvoxel into 8 sub kvoxels, so 7 new kvoxels are added
+    range_subs = 0:(n_subvoxels - 1)
+    add_points = [Vec3(i, j, k) for i in range_subs for j in range_subs for k in range_subs]
+    deleteat!(add_points, 1)
+
+    for ik in iks
+        # split the current kvoxel into 8 sub kvoxels
+        vx0 = ag.kvoxels[ik]
+        voxel = Kvoxel(vx0.point, vx0.dv ./ n_subvoxels, vx0.weight / n_subvoxels^3)
+        ag.kvoxels[ik] = voxel
+        sub_voxels = map(add_points) do pt
+            Kvoxel(voxel.point + pt .* voxel.dv, voxel.dv, voxel.weight)
+        end
+        append!(new_kvoxels, sub_voxels)
+    end
+
+    new_vals = interp([v.point for v in new_kvoxels])
+    append!(ag.kvoxels, new_kvoxels)
+    append!(ag.vals, new_vals)
+    return nothing
+end
diff --git a/test/interpolation/fermi_energy.jl b/test/interpolation/fermi_energy.jl
new file mode 100644
index 0000000..6f40a69
--- /dev/null
+++ b/test/interpolation/fermi_energy.jl
@@ -0,0 +1,28 @@
+@testitem "compute_fermi_energy" begin
+    using Wannier.Datasets
+    model = read_w90_with_chk(
+        dataset"Si2_valence/Si2_valence", dataset"Si2_valence/reference/Si2_valence.chk.fmt"
+    )
+    hamiltonian = TBHamiltonian(model)
+    interp = HamiltonianInterpolator(hamiltonian)
+    kgrid = [12, 12, 12]
+
+    εF = Wannier.compute_fermi_energy(interp, kgrid, 7.1, 1e-2, Wannier.ColdSmearing())
+    εF_ref = 4.637512665199997
+    @test isapprox(εF, εF_ref; atol=1e-3)
+end
+
+@testitem "compute_fermi_energy graphene" begin
+    using Wannier.Datasets
+    model = load_dataset("graphene")
+    model.gauges .= read_amn(dataset"graphene/reference/graphene.dis.amn")
+    hamiltonian = TBHamiltonian(model)
+    interp = HamiltonianInterpolator(hamiltonian)
+    # on purposely choose 5x5x1 since this grid skips the K point, and
+    # a simple Fermi energy search would fail.
+    kgrid = [5, 5, 1]
+
+    εF = Wannier.compute_fermi_energy(interp, kgrid, 8, 0, Wannier.NoneSmearing())
+    εF_ref = -1.03673405699654
+    @test isapprox(εF, εF_ref; atol=1e-3)
+end