Simple example for computing properties using (forward-mode) automatic differentiation. For a more classical approach and more details about computing polarizabilities, see Polarizability by linear response.
using DFTK
using LinearAlgebra
using ForwardDiff
# Construct PlaneWaveBasis given a particular electric field strength
# Again we take the example of a Helium atom.
function make_basis(ε::T; a=10., Ecut=30) where {T}
lattice=T(a) * I(3) # lattice is a cube of $a$ Bohrs
# Helium at the center of the box
atoms = [ElementPsp(:He, psp=load_psp("hgh/lda/He-q2"))]
positions = [[1/2, 1/2, 1/2]]
model = model_DFT(lattice, atoms, positions, [:lda_x, :lda_c_vwn];
extra_terms=[ExternalFromReal(r -> -ε * (r[1] - a/2))],
symmetries=false)
PlaneWaveBasis(model; Ecut, kgrid=[1, 1, 1]) # No k-point sampling on isolated system
end
# dipole moment of a given density (assuming the current geometry)
function dipole(basis, ρ)
@assert isdiag(basis.model.lattice)
a = basis.model.lattice[1, 1]
rr = [a * (r[1] - 1/2) for r in r_vectors(basis)]
sum(rr .* ρ) * basis.dvol
end
# Function to compute the dipole for a given field strength
function compute_dipole(ε; tol=1e-8, kwargs...)
scfres = self_consistent_field(make_basis(ε; kwargs...); tol)
dipole(scfres.basis, scfres.ρ)
end;
With this in place we can compute the polarizability from finite differences (just like in the previous example):
polarizability_fd = let
ε = 0.01
(compute_dipole(ε) - compute_dipole(0.0)) / ε
end
n Energy log10(ΔE) log10(Δρ) Diag Δtime --- --------------- --------- --------- ---- ------ 1 -2.770882126745 -0.52 9.0 2 -2.772147984284 -2.90 -1.32 1.0 158ms 3 -2.772170271702 -4.65 -2.46 1.0 138ms 4 -2.772170653526 -6.42 -3.17 1.0 158ms 5 -2.772170722681 -7.16 -4.14 2.0 159ms 6 -2.772170722846 -9.78 -4.31 1.0 143ms 7 -2.772170723011 -9.78 -5.32 1.0 156ms 8 -2.772170723015 -11.42 -5.95 2.0 165ms 9 -2.772170723015 -13.03 -6.57 1.0 160ms 10 -2.772170723015 + -14.21 -6.50 2.0 176ms 11 -2.772170723015 -13.97 -7.59 1.0 178ms 12 -2.772170723015 + -14.75 -7.82 2.0 167ms 13 -2.772170723015 + -14.65 -8.42 1.0 171ms n Energy log10(ΔE) log10(Δρ) Diag Δtime --- --------------- --------- --------- ---- ------ 1 -2.770804760980 -0.52 9.0 2 -2.772061412731 -2.90 -1.32 1.0 157ms 3 -2.772083168551 -4.66 -2.48 1.0 138ms 4 -2.772083351796 -6.74 -3.19 1.0 153ms 5 -2.772083415788 -7.19 -3.87 2.0 159ms 6 -2.772083417695 -8.72 -4.70 1.0 142ms 7 -2.772083417796 -10.00 -4.94 2.0 170ms 8 -2.772083417810 -10.85 -5.73 1.0 148ms 9 -2.772083417811 -12.23 -6.15 2.0 187ms 10 -2.772083417811 -13.17 -7.64 1.0 150ms 11 -2.772083417811 -14.57 -8.03 2.0 186ms
1.7735579529824077
We do the same thing using automatic differentiation. Under the hood this uses custom rules to implicitly differentiate through the self-consistent field fixed-point problem.
polarizability = ForwardDiff.derivative(compute_dipole, 0.0)
println()
println("Polarizability via ForwardDiff: $polarizability")
println("Polarizability via finite difference: $polarizability_fd")
n Energy log10(ΔE) log10(Δρ) Diag Δtime --- --------------- --------- --------- ---- ------ 1 -2.770738982964 -0.52 8.0 2 -2.772057858034 -2.88 -1.32 1.0 136ms 3 -2.772083163774 -4.60 -2.48 1.0 137ms 4 -2.772083350523 -6.73 -3.21 1.0 173ms 5 -2.772083414090 -7.20 -3.76 2.0 158ms 6 -2.772083417661 -8.45 -4.69 1.0 157ms 7 -2.772083417806 -9.84 -5.12 2.0 161ms 8 -2.772083417810 -11.33 -6.07 1.0 159ms 9 -2.772083417811 -12.61 -6.87 2.0 174ms 10 -2.772083417811 -14.10 -7.69 1.0 161ms 11 -2.772083417811 -13.94 -8.15 2.0 168ms Polarizability via ForwardDiff: 1.772534962791999 Polarizability via finite difference: 1.7735579529824077