Fix ForwardDiff for lattice strain DFPT response

src/terms/Hamiltonian.jl

@@ -89,8 +89,10 @@ function LinearAlgebra.mul!(Hψ, H::Hamiltonian, ψ)
     end
     Hψ
 end
-# need `deepcopy` here to copy the elements of the array of arrays ψ (not just pointers)
-Base.:*(H::Hamiltonian, ψ) = mul!(deepcopy(ψ), H, ψ)
+function Base.:*(H::Hamiltonian, ψ)
+    result = ψ * one(eltype(H.basis))  # Includes type promotion
+    mul!(result, H, ψ)
+end
 
 # Loop through bands, IFFT to get ψ in real space, loop through terms, FFT and accumulate into Hψ
 # For the common DftHamiltonianBlock there is an optimized version below

src/workarounds/forwarddiff_rules.jl

@@ -157,35 +157,32 @@
     basis_primal = construct_value(basis_dual)
     scfres = self_consistent_field(basis_primal; kwargs...)
 
-    ## Compute external perturbation (contained in ham_dual) and from matvec with bands
+    # Compute explicit density perturbation (including strain)
+    ρ_basis = compute_density(basis_dual, scfres.ψ, scfres.occupation)
+
+    # Compute external perturbation (contained in ham_dual)
     Hψ_dual = let
-        occupation_dual = [T.(occk) for occk in scfres.occupation]
-        ψ_dual = [Complex.(T.(real(ψk)), T.(imag(ψk))) for ψk in scfres.ψ]
-        ρ_dual = compute_density(basis_dual, ψ_dual, occupation_dual)
-        εF_dual = T(scfres.εF)  # Only needed for entropy term
-        eigenvalues_dual = [T.(εk) for εk in scfres.eigenvalues]
-        ham_dual = energy_hamiltonian(basis_dual, ψ_dual, occupation_dual;
-                                      ρ=ρ_dual, eigenvalues=eigenvalues_dual,
-                                      εF=εF_dual).ham
-        ham_dual * ψ_dual
+        ham_dual = energy_hamiltonian(basis_dual, scfres.ψ, scfres.occupation;
+                                      ρ=ρ_basis, scfres.eigenvalues,
+                                      scfres.εF).ham
+        ham_dual * scfres.ψ
     end
 
-    ## Implicit differentiation
+    # Implicit differentiation
     response.verbose && println("Solving response problem")
     δresults = ntuple(ForwardDiff.npartials(T)) do α
         δHextψ = [ForwardDiff.partials.(δHextψk, α) for δHextψk in Hψ_dual]
         solve_ΩplusK_split(scfres, -δHextψ; tol=last(scfres.history_Δρ), response.verbose)
     end
 
-    ## Convert and combine
+    # Convert and combine
     DT = Dual{ForwardDiff.tagtype(T)}
     ψ = map(scfres.ψ, getfield.(δresults, :δψ)...) do ψk, δψk...
         map(ψk, δψk...) do ψnk, δψnk...
             Complex(DT(real(ψnk), real.(δψnk)),
                     DT(imag(ψnk), imag.(δψnk)))
         end
     end
-    ρ = map((ρi, δρi...) -> DT(ρi, δρi), scfres.ρ, getfield.(δresults, :δρ)...)
     eigenvalues = map(scfres.eigenvalues, getfield.(δresults, :δeigenvalues)...) do εk, δεk...
         map((εnk, δεnk...) -> DT(εnk, δεnk), εk, δεk...)
     end
@@ -194,6 +191,10 @@
     end
     εF = DT(scfres.εF, getfield.(δresults, :δεF)...)
 
+    # Add contributions from the basis (ρ_basis) and implicit diff (ρ_response)
+    ρ_response = map((ρi, δρi...) -> DT(ρi, δρi), zero(scfres.ρ), getfield.(δresults, :δρ)...)
+    ρ = ρ_basis + ρ_response
+
     # TODO Could add δresults[α].δVind the dual part of the total local potential in ham_dual
     # and in this way return a ham that represents also the total change in Hamiltonian
 

test/forwarddiff.jl

@@ -67,6 +67,48 @@
     end
 end
 
+@testitem "Strain sensitivity using ForwardDiff" #=
+    =#    tags=[:dont_test_mpi] setup=[TestCases] begin
+    using DFTK
+    using ForwardDiff
+    using LinearAlgebra
+    using ComponentArrays
+    using PseudoPotentialData
+    aluminium = TestCases.aluminium
+    Ecut = 5
+    kgrid = [2, 2, 2]
+    model = model_DFT(aluminium.lattice, aluminium.atoms, aluminium.positions;
+                      functionals=LDA(), temperature=1e-2, smearing=Smearing.Gaussian())
+    basis = PlaneWaveBasis(model; Ecut, kgrid)
+    nbandsalg = FixedBands(; n_bands_converge=10)
+    response = ResponseOptions(; verbose=true)
+
+    function compute_properties(ε)
+        model_strained = Model(model; lattice=(1 + ε) * model.lattice)
+        basis = PlaneWaveBasis(model_strained; Ecut, kgrid)
+        scfres = self_consistent_field(basis; tol=1e-10, nbandsalg, response)
+        ComponentArray(
+           eigenvalues=stack([ev[1:10] for ev in scfres.eigenvalues]),
+           ρ=scfres.ρ,
+           energies=collect(values(scfres.energies)),
+           εF=scfres.εF,
+           occupation=reduce(vcat, scfres.occupation),
+        )
+    end
+
+    dx = ForwardDiff.derivative(compute_properties, 0.)
+
+    h = 1e-4
+    x1 = compute_properties(-h)
+    x2 = compute_properties(+h)
+    dx_findiff = (x2 - x1) / 2h
+    @test norm(dx.ρ - dx_findiff.ρ) * sqrt(basis.dvol) < 1e-6
+    @test maximum(abs, dx.eigenvalues - dx_findiff.eigenvalues) < 1e-6
+    @test maximum(abs, dx.energies - dx_findiff.energies) < 1e-5
+    @test dx.εF - dx_findiff.εF < 1e-6
+    @test maximum(abs, dx.occupation - dx_findiff.occupation) < 1e-5
+end
+
 @testitem "scfres PSP sensitivity using ForwardDiff" #=
     =#    tags=[:dont_test_mpi] setup=[TestCases] begin
     using DFTK