Add initial exact exchange implementation

.github/workflows/CompatHelper.yaml

@@ -6,6 +6,7 @@ on:
 
 jobs:
   CompatHelper:
+    if: github.repository_owner = "JuliaMolSim"
     runs-on: ubuntu-latest
     steps:
       - uses: julia-actions/setup-julia@latest

Project.toml

@@ -89,7 +89,7 @@ IterativeSolvers = "0.9"
 JLD2 = "0.4"
 JSON3 = "1"
 LazyArtifacts = "1.3"
-Libxc = "0.3.17"
+Libxc = "0.3.18"
 LineSearches = "7"
 LinearAlgebra = "1"
 LinearMaps = "3"

examples/silicon_exx.jl

@@ -0,0 +1,31 @@
+# Very basic setup, useful for testing
+using DFTK
+
+a = 10.26  # Silicon lattice constant in Bohr
+lattice = a / 2 * [[0 1 1.];
+                   [1 0 1.];
+                   [1 1 0.]]
+Si = ElementPsp(:Si; psp=load_psp("hgh/lda/Si-q4"))
+atoms     = [Si, Si]
+positions = [ones(3)/8, -ones(3)/8]
+
+model  = model_PBE(lattice, atoms, positions)
+basis  = PlaneWaveBasis(model; Ecut=20, kgrid=[1, 1, 1])
+scfres = self_consistent_field(basis; tol=1e-6)
+
+# Run hybrid-DFT tuning some DFTK defaults
+# (Anderson does not work well right now as orbitals not taken into account)
+model  = model_PBE0(lattice, atoms, positions)
+basis  = PlaneWaveBasis(model; basis.Ecut, basis.kgrid)
+scfres = self_consistent_field(basis;
+                               solver=DFTK.scf_damping_solver(1.0),
+                               tol=1e-8, scfres.ρ, scfres.ψ,
+                               diagtolalg=DFTK.AdaptiveDiagtol(; ratio_ρdiff=5e-4))
+
+# Run Hartree-Fock
+model  = model_HF(lattice, atoms, positions)
+basis  = PlaneWaveBasis(model; basis.Ecut, basis.kgrid)
+scfres = self_consistent_field(basis;
+                               solver=DFTK.scf_damping_solver(1.0),
+                               tol=1e-8, scfres.ρ, scfres.ψ,
+                               diagtolalg=DFTK.AdaptiveDiagtol(; ratio_ρdiff=5e-4))

src/DFTK.jl

@@ -100,6 +100,7 @@ export Hamiltonian
 export HamiltonianBlock
 export energy_hamiltonian
 export Kinetic
+export ExactExchange
 export ExternalFromFourier
 export ExternalFromReal
 export AtomicLocal
@@ -143,7 +144,7 @@ include("eigen/preconditioners.jl")
 include("eigen/diag.jl")
 
 export model_atomic
-export model_DFT, model_PBE, model_LDA, model_SCAN
+export model_DFT, model_PBE, model_LDA, model_SCAN, model_PBE0, model_HF
 include("standard_models.jl")
 
 export KerkerMixing, KerkerDosMixing, SimpleMixing, DielectricMixing

src/DispatchFunctional.jl

@@ -12,12 +12,17 @@ function LibxcFunctional(identifier::Symbol)
     @assert fun.kind   in (:exchange, :correlation, :exchange_correlation)
     kind = Dict(:exchange => :x, :correlation => :c, :exchange_correlation => :xc)[fun.kind]
 
-    @assert fun.family in (:lda, :gga, :mgga)  # Hybrids not supported yet.
-    if fun.family == :mgga && Libxc.needs_laplacian(fun)
-        family = :mggal
+    @assert fun.family in (:lda, :gga, :mgga, :hyb_lda, :hyb_gga, :hyb_mgga)
+    # Libxc maintains the distinction between hybrid and non-hybrid equivalents,
+    # but this distinction is not relevant for the functional interface
+    if startswith(string(fun.family), "hyb")
+        family = Symbol(string(fun.family)[5:end])
     else
         family = fun.family
     end
+    if family == :mgga && Libxc.needs_laplacian(fun)
+        family = :mggal
+    end
     LibxcFunctional{family,kind}(identifier)
 end
 
@@ -154,3 +159,10 @@ for fun in (:potential_terms, :kernel_terms)
         end
     end
 end
+
+# TODO This is hackish for now until Libxc has fully picked up the DftFunctionals.jl interface
+exx_coefficient(::Functional{:lda})      = nothing
+exx_coefficient(::Functional{:gga})      = nothing
+exx_coefficient(::Functional{:mgga})     = nothing
+exx_coefficient(fun::DispatchFunctional) = exx_coefficient(fun.inner)
+exx_coefficient(fun::LibxcFunctional)    = Libxc.Functional(fun.identifier).exx_coefficient

src/standard_models.jl

@@ -33,8 +33,14 @@ function model_DFT(lattice::AbstractMatrix,
                    xc::Xc;
                    extra_terms=[], kwargs...)
     model_name = isempty(xc.functionals) ? "rHF" : join(string.(xc.functionals), "+")
+    exx = [ExactExchange(; scaling_factor=exx_coefficient(f))
+           for f in xc.functionals if !isnothing(exx_coefficient(f))]
+    if !isempty(exx)
+        @assert length(exx) == 1
+        @warn "Exact exchange in DFTK is hardly optimised and not yet production-ready."
+    end
     model_atomic(lattice, atoms, positions;
-                 extra_terms=[Hartree(), xc, extra_terms...], model_name, kwargs...)
+                 extra_terms=[Hartree(), xc, exx..., extra_terms...], model_name, kwargs...)
 end
 function model_DFT(lattice::AbstractMatrix,
                    atoms::Vector{<:Element},
@@ -44,6 +50,20 @@ function model_DFT(lattice::AbstractMatrix,
     model_DFT(lattice, atoms, positions, Xc(functionals); kwargs...)
 end
 
+"""
+Build an Hartree-Fock model from the specified atoms.
+"""
+function model_HF(lattice::AbstractMatrix, atoms::Vector{<:Element},
+                  positions::Vector{<:AbstractVector}; extra_terms=[], kwargs...)
+    @warn "Exact exchange in DFTK is hardly optimised and not yet production-ready."
+    model_atomic(lattice, atoms, positions;
+                 extra_terms=[Hartree(), ExactExchange(), extra_terms...],
+                 model_name="HF", kwargs...)
+end
+
+#
+# DFT shortcut models
+#
 
 """
 Build an LDA model (Perdew & Wang parametrization) from the specified atoms.
@@ -75,8 +95,19 @@ function model_SCAN(lattice::AbstractMatrix, atoms::Vector{<:Element},
 end
 
 
+"""
+Build an PBE0 model from the specified atoms.
+DOI:10.1103/PhysRevLett.77.3865
+"""
+function model_PBE0(lattice::AbstractMatrix, atoms::Vector{<:Element},
+                   positions::Vector{<:AbstractVector}; kwargs...)
+    model_DFT(lattice, atoms, positions, [:hyb_gga_xc_pbeh]; kwargs...)
+end
+
+
 # Generate equivalent functions for AtomsBase
-for fun in (:model_atomic, :model_DFT, :model_LDA, :model_PBE, :model_SCAN)
+for fun in (:model_atomic, :model_DFT, :model_HF,
+            :model_LDA, :model_PBE, :model_SCAN, :model_PBE0)
     @eval function $fun(system::AbstractSystem, args...; kwargs...)
         parsed = parse_system(system)
         $fun(parsed.lattice, parsed.atoms, parsed.positions, args...;

src/terms/exact_exchange.jl

@@ -0,0 +1,76 @@
+@doc raw"""
+Exact exchange term: the Hartree-Exact exchange energy of the orbitals
+```math
+-1/2 ∑_{nm} f_n f_m ∫∫ \frac{ψ_n^*(r)ψ_m^*(r')ψ_n(r')ψ_m(r)}{|r - r'|} dr dr'
+```
+"""
+struct ExactExchange
+    scaling_factor::Real  # to scale the term
+end
+ExactExchange(; scaling_factor=1) = ExactExchange(scaling_factor)
+(exchange::ExactExchange)(basis) = TermExactExchange(basis, exchange.scaling_factor)
+function Base.show(io::IO, exchange::ExactExchange)
+    fac = isone(exchange.scaling_factor) ? "" : "scaling_factor=$(exchange.scaling_factor)"
+    print(io, "ExactExchange($fac)")
+end
+struct TermExactExchange <: Term
+    scaling_factor::Real  # scaling factor, absorbed into poisson_green_coeffs
+    poisson_green_coeffs::AbstractArray
+end
+function TermExactExchange(basis::PlaneWaveBasis{T}, scaling_factor) where T
+    poisson_green_coeffs = 4T(π) ./ [sum(abs2, basis.model.recip_lattice * G)
+                                     for G in to_cpu(G_vectors(basis))]
+    poisson_green_coeffs[1] = 0
+    TermExactExchange(T(scaling_factor), scaling_factor .* poisson_green_coeffs)
+end
+
+# Note: Implementing exact exchange in a scalable and numerically stable way, such that it
+# rapidly converges with k-points is tricky. This implementation here is far too simple and
+# slow to be useful.
+#
+# ABINIT/src/66_wfs/m_getghc.F90
+# ABINIT/src/66_wfs/m_fock_getghc.F90
+# ABINIT/src/66_wfs/m_prep_kgb.F90
+# ABINIT/src/66_wfs/m_bandfft_kpt.F90
+#
+# For further information (in particular on regularising the Coulomb), consider the following
+#      https://www.vasp.at/wiki/index.php/Coulomb_singularity
+#      https://journals.aps.org/prb/pdf/10.1103/PhysRevB.34.4405   (QE default)
+#      https://journals.aps.org/prb/pdf/10.1103/PhysRevB.73.205119
+#      https://journals.aps.org/prb/pdf/10.1103/PhysRevB.77.193110
+#      https://docs.abinit.org/topics/documents/hybrids-2017.pdf (Abinit apparently
+#           uses a short-ranged Coulomb)
+
+@timing "ene_ops: ExactExchange" function ene_ops(term::TermExactExchange,
+                                                  basis::PlaneWaveBasis{T}, ψ, occupation;
+                                                  kwargs...) where {T}
+    if isnothing(ψ) || isnothing(occupation)
+        return (; E=T(0), ops=NoopOperator.(basis, basis.kpoints))
+    end
+
+    # TODO Occupation threshold
+    ψ, occupation = select_occupied_orbitals(basis, ψ, occupation; threshold=1e-8)
+
+    E = zero(T)
+
+    @assert length(ψ) == 1  # TODO: make it work for more kpoints
+    ik   = 1
+    kpt  = basis.kpoints[ik]
+    occk = occupation[ik]
+    ψk   = ψ[ik]
+
+    for (n, ψn) in enumerate(eachcol(ψk))
+        ψn_real = ifft(basis, kpt, ψn)
+        for (m, ψm) in enumerate(eachcol(ψk))
+            ψm_real  = ifft(basis, kpt, ψm)
+            ρnm_real = conj(ψn_real) .* ψm_real
+            ρnm_fourier = fft(basis, ρnm_real)
+
+            fac_mn = occk[n] * occk[m] / T(2)
+            E -= fac_mn * real(dot(ρnm_fourier .* term.poisson_green_coeffs, ρnm_fourier))
+        end
+    end
+
+    ops = [ExchangeOperator(basis, kpt, term.poisson_green_coeffs, occk, ψk)]
+    (; E, ops)
+end

src/terms/operators.jl

@@ -10,6 +10,22 @@ They also implement mul! and Matrix(op) for exploratory use.
 abstract type RealFourierOperator end
 # RealFourierOperator contain fields `basis` and `kpoint`
 
+Base.Array(op::RealFourierOperator) = Matrix(op)
+
+# Slow fallback for getting operator as a dense matrix
+function Matrix(op::RealFourierOperator)
+    T = complex(eltype(op.basis))
+    n_G = length(G_vectors(op.basis, op.kpoint))
+    H = zeros(T, n_G, n_G)
+    ψ = zeros(T, n_G)
+    @views for i in 1:n_G
+        ψ[i] = one(T)
+        mul!(H[:, i], op, ψ)
+        ψ[i] = zero(T)
+    end
+    H
+end
+
 # Unoptimized fallback, intended for exploratory use only.
 # For performance, call through Hamiltonian which saves FFTs.
 function LinearAlgebra.mul!(Hψ::AbstractVector, op::RealFourierOperator, ψ::AbstractVector)
@@ -147,7 +163,6 @@ function apply!(Hψ, op::DivAgradOperator, ψ,
         ∂αψ_real = ifft!(ψ_scratch, op.basis, op.kpoint, im .* G_plus_k[α] .* ψ.fourier)
         A∇ψ      = fft(op.basis, op.kpoint, ∂αψ_real .* op.A)
         Hψ.fourier .-= im .* G_plus_k[α] .* A∇ψ ./ 2
-
     end
 end
 # TODO Implement  Matrix(op::DivAgrad)
@@ -164,3 +179,28 @@ function optimize_operators_(ops)
                                                sum([op.potential for op in RSmults]))
     [nonRSmults..., combined_RSmults]
 end
+
+struct ExchangeOperator{T <: Real,Tocc,Tpsi} <: RealFourierOperator
+    basis::PlaneWaveBasis{T}
+    kpoint::Kpoint{T}
+    poisson_green_coeffs::Array{T, 3}
+    occk::Tocc
+    ψk::Tpsi
+end
+function apply!(Hψ, op::ExchangeOperator, ψ)
+    # Hψ = - ∑_n f_n ψ_n(r) ∫ (ψ_n)†(r') * ψ(r') / |r-r'| dr'
+    for (n, ψnk) in enumerate(eachcol(op.ψk))
+        ψnk_real = ifft(op.basis, op.kpoint, ψnk)
+        x_real   = conj(ψnk_real) .* ψ.real
+        # TODO Some symmetrisation of x_real might be needed here ...
+
+        # Compute integral by Poisson solve
+        x_four  = fft(op.basis, x_real)
+        Vx_four = x_four .* op.poisson_green_coeffs
+        Vx_real = ifft(op.basis, Vx_four)
+
+        # Real-space multiply and accumulate
+        Hψ.real .-= op.occk[n] .* ψnk_real .* Vx_real
+    end
+end
+# TODO Implement  Matrix(op::ExchangeOperator)

src/terms/terms.jl

@@ -51,6 +51,7 @@ include("ewald.jl")
 include("psp_correction.jl")
 include("entropy.jl")
 include("pairwise.jl")
+include("exact_exchange.jl")
 
 include("magnetic.jl")
 breaks_symmetries(::Magnetic) = true

test/exact_exchange.jl

@@ -0,0 +1,20 @@
+@testitem "Helium exchange energy" setup=[TestCases] begin
+    using DFTK
+    using LinearAlgebra
+    silicon = TestCases.silicon
+
+    lattice   = 10diagm(ones(3))
+    positions = [0.5ones(3)]
+    atoms     = [ElementCoulomb(:He)]
+    model     = model_atomic(lattice, atoms, positions)
+    basis     = PlaneWaveBasis(model; Ecut=15, kgrid=(1, 1, 1))
+    scfres    = self_consistent_field(basis)
+
+    Eh, oph = DFTK.ene_ops(Hartree()(basis), basis, scfres.ψ, scfres.occupation; scfres...)
+    Ex, opx = DFTK.ene_ops(ExactExchange()(basis), basis, scfres.ψ, scfres.occupation; scfres...)
+    @test abs(Ex + Eh) < 1e-12
+
+    mat_h = real.(Matrix(only(oph)))
+    mat_x = real.(Matrix(only(opx)))
+    @show maximum(abs, mat_x - mat_x')
+end

test/hamiltonian_consistency.jl

@@ -7,23 +7,9 @@ using LinearAlgebra
 using ..TestCases: silicon
 testcase = silicon
 
-function test_matrix_repr_operator(hamk, ψk; atol=1e-8)
-    for operator in hamk.operators
-        try
-            operator_matrix = Matrix(operator)
-            @test norm(operator_matrix*ψk - operator*ψk) < atol
-        catch e
-            allowed_missing_operators = Union{DFTK.DivAgradOperator,
-                                              DFTK.MagneticFieldOperator}
-            @assert operator isa allowed_missing_operators
-            @info "Matrix of operator $(nameof(typeof(operator))) not yet supported" maxlog=1
-        end
-    end
-end
-
 function test_consistency_term(term; rtol=1e-4, atol=1e-8, ε=1e-6, kgrid=[1, 2, 3],
-                               kshift=[0, 1, 0]/2, lattice=testcase.lattice, Ecut=10,
-                               spin_polarization=:none)
+                               kshift=[0, 1, 0]/2, lattice=testcase.lattice,
+                               Ecut=10, spin_polarization=:none)
     sspol = spin_polarization != :none ? " ($spin_polarization)" : ""
     xc    = term isa Xc ? "($(first(term.functionals)))" : ""
     @testset "$(typeof(term))$xc $sspol" begin
@@ -33,6 +19,7 @@ function test_consistency_term(term; rtol=1e-4, atol=1e-8, ε=1e-6, kgrid=[1, 2,
         model = Model(lattice, atoms, testcase.positions; terms=[term], spin_polarization,
                       symmetries=true)
         basis = PlaneWaveBasis(model; Ecut, kgrid, kshift)
+        @assert length(basis.terms) == 1
 
         n_electrons = testcase.n_electrons
         n_bands = div(n_electrons, 2, RoundUp)
@@ -48,8 +35,14 @@ function test_consistency_term(term; rtol=1e-4, atol=1e-8, ε=1e-6, kgrid=[1, 2,
         end
         E0, ham = energy_hamiltonian(basis, ψ, occupation; ρ)
 
-        @assert length(basis.terms) == 1
+        # Test operator is derivative of the energy
+        for ik = 1:length(basis.kpoints)
+            for operator in ham.blocks[ik].operators
+                @test norm(Matrix(operator) * ψ[ik] - operator * ψ[ik]) < atol
+            end
+        end
 
+        # Test operator is derivative of the energy
         δψ = [randn(ComplexF64, size(ψ[ik])) for ik = 1:length(basis.kpoints)]
         function compute_E(ε)
             ψ_trial = ψ .+ ε .* δψ
@@ -59,15 +52,13 @@ function test_consistency_term(term; rtol=1e-4, atol=1e-8, ε=1e-6, kgrid=[1, 2,
             E = energy_hamiltonian(basis, ψ_trial, occupation; ρ=ρ_trial).energies
             E.total
         end
-
         diff = (compute_E(ε) - compute_E(-ε)) / (2ε)
 
         diff_predicted = 0.0
         for ik = 1:length(basis.kpoints)
-            Hψk = ham.blocks[ik]*ψ[ik]
-            test_matrix_repr_operator(ham.blocks[ik], ψ[ik]; atol)
+            Hψk = ham.blocks[ik] * ψ[ik]
             δψkHψk = sum(occupation[ik][iband] * real(dot(δψ[ik][:, iband], Hψk[:, iband]))
-                         for iband=1:n_bands)
+                         for iband = 1:n_bands)
             diff_predicted += 2 * basis.kweights[ik] * δψkHψk
         end
         diff_predicted = mpi_sum(diff_predicted, basis.comm_kpts)
@@ -101,6 +92,7 @@ end
     test_consistency_term(Xc(:mgga_c_scan), spin_polarization=:collinear)
     test_consistency_term(Xc(:mgga_x_b00))
     test_consistency_term(Xc(:mgga_c_b94), spin_polarization=:collinear)
+    test_consistency_term(ExactExchange(); kgrid=(1, 1, 1), Ecut=7)
 
     let
         a = 6