Hellmann-Feynman stresses via ForwardDiff and custom rules

Project.toml

@@ -76,6 +76,7 @@ JLD2 = "0.4 - 0.4.7"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 DoubleFloats = "497a8b3b-efae-58df-a0af-a86822472b78"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 GenericLinearAlgebra = "14197337-ba66-59df-a3e3-ca00e7dcff7a"
 IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
@@ -86,4 +87,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [targets]
-test = ["Test", "Aqua", "DoubleFloats", "GenericLinearAlgebra", "IntervalArithmetic", "Plots", "Random", "KrylovKit", "JLD2", "WriteVTK"]
+test = ["Test", "Aqua", "DoubleFloats", "FiniteDiff", "GenericLinearAlgebra", "IntervalArithmetic", "Plots", "Random", "KrylovKit", "JLD2", "WriteVTK"]

src/DFTK.jl

@@ -177,13 +177,19 @@ function __init__()
     # DoubleFloats has been loaded (via a "using" or an "import").
     # See https://github.com/JuliaPackaging/Requires.jl for details.
     #
-    # The global variable GENERIC_FFT_LOADED makes sure that things are only
-    # included once.
+    # The global variables GENERIC_FFT_LOADED and DUMMY_INPLACE_LOADED
+    # make sure that things are only included once.
+    @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" begin
+        !isdefined(DFTK, :DUMMY_INPLACE_LOADED) && include("workarounds/dummy_inplace_fft.jl")
+        include("workarounds/forwarddiff_rules.jl")
+    end
     @require IntervalArithmetic="d1acc4aa-44c8-5952-acd4-ba5d80a2a253" begin
         include("workarounds/intervals.jl")
+        !isdefined(DFTK, :DUMMY_INPLACE_LOADED) && include("workarounds/dummy_inplace_fft.jl")
         !isdefined(DFTK, :GENERIC_FFT_LOADED) && include("workarounds/fft_generic.jl")
     end
     @require DoubleFloats="497a8b3b-efae-58df-a0af-a86822472b78" begin
+        !isdefined(DFTK, :DUMMY_INPLACE_LOADED) && include("workarounds/dummy_inplace_fft.jl")
         !isdefined(DFTK, :GENERIC_FFT_LOADED) && include("workarounds/fft_generic.jl")
     end
     @require Plots="91a5bcdd-55d7-5caf-9e0b-520d859cae80" include("plotting.jl")

src/Model.jl

@@ -103,8 +103,8 @@ function Model(lattice::AbstractMatrix{T};
         norm(lattice[:, i]) == norm(lattice[i, :]) == 0 || error(
             "For 1D and 2D systems, the non-empty dimensions must come first")
     end
-    cond(lattice[1:n_dim, 1:n_dim]) > 1e-5 || (
-        @warn "Your lattice is badly conditioned, the computation is likely to fail.")
+    _check_well_conditioned(lattice[1:n_dim, 1:n_dim]) || @warn (
+        "Your lattice is badly conditioned, the computation is likely to fail.")
 
     # Compute reciprocal lattice and volumes.
     # recall that the reciprocal lattice is the set of G vectors such
@@ -213,3 +213,5 @@ function spin_components(spin_polarization::Symbol)
     spin_polarization == :full      && return (:undefined, )
 end
 spin_components(model::Model) = spin_components(model.spin_polarization)
+
+_check_well_conditioned(A; tol=1e5) = (cond(A) <= tol)

src/densities.jl

@@ -60,7 +60,8 @@ is not collinear the spin density is `nothing`.
     @assert n_k > 0
 
     # Allocate an accumulator for ρ in each thread for each spin component
-    ρaccus = [similar(view(ψ[1], :, 1), (basis.fft_size..., n_spin))
+    T = promote_type(eltype(basis), eltype(ψ[1]))
+    ρaccus = [similar(ψ[1], T, (basis.fft_size..., n_spin))
               for ithread in 1:Threads.nthreads()]
 
     # TODO Better load balancing ... the workload per kpoint depends also on
@@ -79,7 +80,7 @@ is not collinear the spin density is `nothing`.
 
     Threads.@threads for (ikpts, ρaccu) in collect(zip(kpt_per_thread, ρaccus))
         ρaccu .= 0
-        ρ_k = similar(ψ[1][:, 1], basis.fft_size)
+        ρ_k = similar(ψ[1], T, basis.fft_size)
         for ik in ikpts
             kpt = basis.kpoints[ik]
             compute_partial_density!(ρ_k, basis, kpt, ψ[ik], occupation[ik])

src/workarounds/dummy_inplace_fft.jl

@@ -0,0 +1,15 @@
+# This is needed to flag that the dummy_inplace_fft.jl file has already been loaded
+const DUMMY_INPLACE_LOADED = true
+
+# A dummy wrapper around an out-of-place FFT plan to make it appear in-place
+# This is needed for some generic FFT implementations, which do not have in-place plans
+struct DummyInplace{opFFT}
+    fft::opFFT
+end
+LinearAlgebra.mul!(Y, p::DummyInplace, X) = (Y .= mul!(similar(X), p.fft, X))
+LinearAlgebra.ldiv!(Y, p::DummyInplace, X) = (Y .= ldiv!(similar(X), p.fft, X))
+
+import Base: *, \, length
+*(p::DummyInplace, X) = p.fft * X
+\(p::DummyInplace, X) = p.fft \ X
+length(p::DummyInplace) = length(p.fft)

src/workarounds/fft_generic.jl

@@ -90,17 +90,3 @@ function generic_plan_bfft(data::AbstractArray{T, 3}) where T
                     FourierTransforms.plan_bfft(data[1, :, 1]),
                     FourierTransforms.plan_bfft(data[1, 1, :])], T(1))
 end
-
-
-# A dummy wrapper around an out-of-place FFT plan to make it appear in-place
-# This is needed for some generic FFT implementations, which do not have in-place plans
-struct DummyInplace{opFFT}
-    fft::opFFT
-end
-LinearAlgebra.mul!(Y, p::DummyInplace, X) = (Y .= mul!(similar(X), p.fft, X))
-LinearAlgebra.ldiv!(Y, p::DummyInplace, X) = (Y .= ldiv!(similar(X), p.fft, X))
-
-import Base: *, \, length
-*(p::DummyInplace, X) = p.fft * X
-\(p::DummyInplace, X) = p.fft \ X
-length(p::DummyInplace) = length(p.fft)

src/workarounds/forwarddiff_rules.jl

@@ -0,0 +1,137 @@
+import ForwardDiff
+import AbstractFFTs
+
+# original PR by mcabbott: https://github.com/JuliaDiff/ForwardDiff.jl/pull/495
+
+ForwardDiff.value(x::Complex{<:ForwardDiff.Dual}) = Complex(x.re.value, x.im.value)
+
+ForwardDiff.partials(x::Complex{<:ForwardDiff.Dual}, n::Int) =
+    Complex(ForwardDiff.partials(x.re, n), ForwardDiff.partials(x.im, n))
+
+ForwardDiff.npartials(x::Complex{<:ForwardDiff.Dual{T,V,N}}) where {T,V,N} = N
+ForwardDiff.npartials(::Type{<:Complex{<:ForwardDiff.Dual{T,V,N}}}) where {T,V,N} = N
+
+ForwardDiff.tagtype(x::Complex{<:ForwardDiff.Dual{T,V,N}}) where {T,V,N} = T
+ForwardDiff.tagtype(::Type{<:Complex{<:ForwardDiff.Dual{T,V,N}}}) where {T,V,N} = T
+
+# AbstractFFTs.complexfloat(x::AbstractArray{<:ForwardDiff.Dual}) = float.(x .+ 0im)
+AbstractFFTs.complexfloat(x::AbstractArray{<:ForwardDiff.Dual}) = AbstractFFTs.complexfloat.(x)
+AbstractFFTs.complexfloat(d::ForwardDiff.Dual{T,V,N}) where {T,V,N} = convert(ForwardDiff.Dual{T,float(V),N}, d) + 0im
+
+AbstractFFTs.realfloat(x::AbstractArray{<:ForwardDiff.Dual}) = AbstractFFTs.realfloat.(x)
+AbstractFFTs.realfloat(d::ForwardDiff.Dual{T,V,N}) where {T,V,N} = convert(ForwardDiff.Dual{T,float(V),N}, d)
+
+for plan in [:plan_fft, :plan_ifft, :plan_bfft]
+    @eval begin
+        AbstractFFTs.$plan(x::AbstractArray{<:ForwardDiff.Dual}, region=1:ndims(x); kwargs...) =
+            AbstractFFTs.$plan(ForwardDiff.value.(x) .+ 0im, region; kwargs...)
+
+        AbstractFFTs.$plan(x::AbstractArray{<:Complex{<:ForwardDiff.Dual}}, region=1:ndims(x); kwargs...) =
+            AbstractFFTs.$plan(ForwardDiff.value.(x), region; kwargs...)
+    end
+end
+
+# rfft only accepts real arrays
+AbstractFFTs.plan_rfft(x::AbstractArray{<:ForwardDiff.Dual}, region=1:ndims(x); kwargs...) =
+    AbstractFFTs.plan_rfft(ForwardDiff.value.(x), region; kwargs...)
+
+for plan in [:plan_irfft, :plan_brfft]  # these take an extra argument, only when complex?
+    @eval begin
+        AbstractFFTs.$plan(x::AbstractArray{<:ForwardDiff.Dual}, region=1:ndims(x); kwargs...) =
+            AbstractFFTs.$plan(ForwardDiff.value.(x) .+ 0im, region; kwargs...)
+
+        AbstractFFTs.$plan(x::AbstractArray{<:Complex{<:ForwardDiff.Dual}}, d::Integer, region=1:ndims(x); kwargs...) =
+            AbstractFFTs.$plan(ForwardDiff.value.(x), d, region; kwargs...)
+    end
+end
+
+for P in [:Plan, :ScaledPlan]  # need ScaledPlan to avoid ambiguities
+    @eval begin
+        Base.:*(p::AbstractFFTs.$P, x::AbstractArray{<:ForwardDiff.Dual}) =
+            _apply_plan(p, x)
+
+        Base.:*(p::AbstractFFTs.$P, x::AbstractArray{<:Complex{<:ForwardDiff.Dual}}) =
+            _apply_plan(p, x)
+
+        LinearAlgebra.mul!(Y::AbstractArray, p::AbstractFFTs.$P, X::AbstractArray{<:ForwardDiff.Dual}) = 
+            (Y .= _apply_plan(p, X))
+
+        LinearAlgebra.mul!(Y::AbstractArray, p::AbstractFFTs.$P, X::AbstractArray{<:Complex{<:ForwardDiff.Dual}}) =
+            (Y .= _apply_plan(p, X))
+    end
+end
+
+LinearAlgebra.mul!(Y::AbstractArray{<:Complex{<:ForwardDiff.Dual}}, p::AbstractFFTs.ScaledPlan{T,P,<:ForwardDiff.Dual}, X::AbstractArray{<:ComplexF64}) where {T,P} =
+    (Y .= _apply_plan(p, X))
+
+function _apply_plan(p::AbstractFFTs.Plan, x::AbstractArray)
+    xtil = p * ForwardDiff.value.(x)
+    dxtils = ntuple(ForwardDiff.npartials(eltype(x))) do n
+        p * ForwardDiff.partials.(x, n)
+    end
+    T = ForwardDiff.tagtype(eltype(x))
+    map(xtil, dxtils...) do val, parts...
+        Complex(
+            ForwardDiff.Dual{T}(real(val), map(real, parts)),
+            ForwardDiff.Dual{T}(imag(val), map(imag, parts)),
+        )
+    end
+end
+
+function _apply_plan(p::AbstractFFTs.ScaledPlan{T,P,<:ForwardDiff.Dual}, x::AbstractArray) where {T,P}
+    _apply_plan(p.p, p.scale * x) # for when p.scale is Dual, need out-of-place
+end
+
+# DFTK setup specific
+
+next_working_fft_size(::Type{<:ForwardDiff.Dual}, size) = size
+
+_fftw_flags(::Type{<:ForwardDiff.Dual}) = FFTW.MEASURE | FFTW.UNALIGNED
+
+function build_fft_plans(T::Type{<:Union{ForwardDiff.Dual,Complex{<:ForwardDiff.Dual}}}, fft_size)
+    tmp = Array{Complex{T}}(undef, fft_size...)
+    opFFT  = FFTW.plan_fft(tmp, flags=_fftw_flags(T))
+    opBFFT = FFTW.plan_bfft(tmp, flags=_fftw_flags(T))
+
+    ipFFT  = DummyInplace{typeof(opFFT)}(opFFT)
+    ipBFFT = DummyInplace{typeof(opBFFT)}(opBFFT)
+    # backward by inverting and stripping off normalizations
+    ipFFT, opFFT, ipBFFT, opBFFT
+end
+
+# PlaneWaveBasis{<:Dual} contains dual-scaled fft, which means that the result f_fourier 
+# must be able to hold complex dual numbers even if f_real is not dual
+function r_to_G(basis::PlaneWaveBasis{T}, f_real::AbstractArray) where {T<:ForwardDiff.Dual}
+    f_fourier = similar(f_real, complex(T))
+    @assert length(size(f_real)) ∈ (3, 4)
+    # this exploits trailing index convention
+    for iσ = 1:size(f_real, 4)
+        @views r_to_G!(f_fourier[:, :, :, iσ], basis, f_real[:, :, :, iσ])
+    end
+    f_fourier
+end
+
+# determine symmetry operations only from primal lattice values
+function spglib_get_symmetry(lattice::Matrix{<:ForwardDiff.Dual}, atoms, magnetic_moments=[]; kwargs...)
+    spglib_get_symmetry(ForwardDiff.value.(lattice), atoms, magnetic_moments; kwargs...)
+end
+
+function _check_well_conditioned(A::AbstractArray{<:ForwardDiff.Dual}; kwargs...)
+    _check_well_conditioned(ForwardDiff.value.(A); kwargs...)
+end
+
+
+# other workarounds
+
+# problem: ForwardDiff of norm of SVector gives NaN derivative at zero
+# https://github.com/JuliaMolSim/DFTK.jl/issues/443#issuecomment-864930410
+# solution: follow ChainRules custom frule for norm
+# https://github.com/JuliaDiff/ChainRules.jl/blob/52a0eeadf8d19bff491f224517b7b064ce1ba378/src/rulesets/LinearAlgebra/norm.jl#L5
+# TODO delete, once forward diff AD tools use ChainRules natively
+function LinearAlgebra.norm(x::SVector{S,<:ForwardDiff.Dual}) where {S}
+    T = ForwardDiff.tagtype(eltype(x))
+    dx = ForwardDiff.partials.(x)
+    y = norm(ForwardDiff.value.(x))
+    dy = real(dot(ForwardDiff.value.(x), dx)) * pinv(y)
+    ForwardDiff.Dual{T}(y, dy)
+end

test/runtests.jl

@@ -107,5 +107,9 @@ Random.seed!(0)
         include("aqua.jl")
     end
 
+    if "all" in TAGS
+        include("stresses.jl")
+    end
+
     ("example" in TAGS) && include("runexamples.jl")
 end

test/stresses.jl

@@ -0,0 +1,47 @@
+using Test
+using DFTK
+using ForwardDiff
+import FiniteDiff
+include("testcases.jl")
+
+# Hellmann-Feynman stress
+# via ForwardDiff & custom FFTW overloads on ForwardDiff.Dual
+
+@testset "ForwardDiff stresses on silicon" begin
+    function make_basis(a)
+        lattice = a / 2 * [[0 1 1.];
+                           [1 0 1.];
+                           [1 1 0.]]
+        Si = ElementPsp(silicon.atnum, psp=load_psp(silicon.psp))
+        atoms = [Si => silicon.positions]
+        model = model_DFT(lattice, atoms, [:lda_x, :lda_c_vwn])
+        kgrid = [1, 1, 1]
+        Ecut = 7
+        PlaneWaveBasis(model, Ecut; kgrid=kgrid)
+    end
+
+    function recompute_energy(a)
+        basis = make_basis(a)
+        scfres = self_consistent_field(basis, is_converged=DFTK.ScfConvergenceDensity(1e-13))
+        energies, H = energy_hamiltonian(basis, scfres.ψ, scfres.occupation; ρ=scfres.ρ)
+        energies.total
+    end
+
+    function hellmann_feynman_energy(scfres_ref, a)
+        basis = make_basis(a)
+        ρ = DFTK.compute_density(basis, scfres_ref.ψ, scfres_ref.occupation)
+        energies, H = energy_hamiltonian(basis, scfres_ref.ψ, scfres_ref.occupation; ρ=ρ)
+        energies.total
+    end
+
+    a = 10.26
+    scfres = self_consistent_field(make_basis(a), is_converged=DFTK.ScfConvergenceDensity(1e-13))
+    hellmann_feynman_energy(a) = hellmann_feynman_energy(scfres, a)
+
+    ref_recompute = FiniteDiff.finite_difference_derivative(recompute_energy, a)
+    ref_hf = FiniteDiff.finite_difference_derivative(hellmann_feynman_energy, a)
+    s_hf = ForwardDiff.derivative(hellmann_feynman_energy, a)
+
+    @test isapprox(ref_hf, ref_recompute, atol=1e-4)
+    @test isapprox(s_hf, ref_hf, atol=1e-8)
+end