using DFTK using Printf using LinearAlgebra using ForwardDiff using LinearMaps using IterativeSolvers Ti = ElementPsp(:Ti, psp=load_psp("hgh/lda/ti-q4.hgh")) O = ElementPsp(:O, psp=load_psp("hgh/lda/o-q6.hgh")) atoms = [Ti, Ti, O, O, O, O] positions = [[0.5, 0.5, 0.5], # Ti [0.0, 0.0, 0.0], # Ti [0.19542, 0.80458, 0.5], # O [0.80458, 0.19542, 0.5], # O [0.30458, 0.30458, 0.0], # O [0.69542, 0.69542, 0.0]] # O lattice = [[8.79341 0.0 0.0]; [0.0 8.79341 0.0]; [0.0 0.0 5.61098]]; positions[1] .+= [-0.022, 0.028, 0.035] model = model_LDA(lattice, atoms, positions) kgrid = [1, 1, 1] Ecut_ref = 35 basis_ref = PlaneWaveBasis(model; Ecut=Ecut_ref, kgrid) tol = 1e-5; scfres_ref = self_consistent_field(basis_ref; tol, callback=identity) ψ_ref, _ = DFTK.select_occupied_orbitals(basis_ref, scfres_ref.ψ, scfres_ref.occupation); Ecut = 15 basis = PlaneWaveBasis(model; Ecut, kgrid) scfres = self_consistent_field(basis; tol, callback=identity) ψr = DFTK.transfer_blochwave(scfres.ψ, basis, basis_ref) ρr = compute_density(basis_ref, ψr, scfres.occupation) Er, hamr = energy_hamiltonian(basis_ref, ψr, scfres.occupation; ρ=ρr); res = DFTK.compute_projected_gradient(basis_ref, ψr, scfres.occupation) res, occ = DFTK.select_occupied_orbitals(basis_ref, res, scfres.occupation) ψr, _ = DFTK.select_occupied_orbitals(basis_ref, ψr, scfres.occupation); function compute_error(basis, ϕ, ψ) map(zip(ϕ, ψ)) do (ϕk, ψk) S = ψk'ϕk U = S*(S'S)^(-1/2) ϕk - ψk*U end end err = compute_error(basis_ref, ψr, ψ_ref); P = [PreconditionerTPA(basis_ref, kpt) for kpt in basis_ref.kpoints] map(zip(P, ψr)) do (Pk, ψk) DFTK.precondprep!(Pk, ψk) end function apply_M(φk, Pk, δφnk, n) DFTK.proj_tangent_kpt!(δφnk, φk) δφnk = sqrt.(Pk.mean_kin[n] .+ Pk.kin) .* δφnk DFTK.proj_tangent_kpt!(δφnk, φk) δφnk = sqrt.(Pk.mean_kin[n] .+ Pk.kin) .* δφnk DFTK.proj_tangent_kpt!(δφnk, φk) end function apply_inv_M(φk, Pk, δφnk, n) DFTK.proj_tangent_kpt!(δφnk, φk) op(x) = apply_M(φk, Pk, x, n) function f_ldiv!(x, y) x .= DFTK.proj_tangent_kpt(y, φk) x ./= (Pk.mean_kin[n] .+ Pk.kin) DFTK.proj_tangent_kpt!(x, φk) end J = LinearMap{eltype(φk)}(op, size(δφnk, 1)) δφnk = cg(J, δφnk, Pl=DFTK.FunctionPreconditioner(f_ldiv!), verbose=false, reltol=0, abstol=1e-15) DFTK.proj_tangent_kpt!(δφnk, φk) end function apply_metric(φ, P, δφ, A::Function) map(enumerate(δφ)) do (ik, δφk) Aδφk = similar(δφk) φk = φ[ik] for n = 1:size(δφk,2) Aδφk[:,n] = A(φk, P[ik], δφk[:,n], n) end Aδφk end end Mres = apply_metric(ψr, P, res, apply_inv_M); resLF = DFTK.transfer_blochwave(res, basis_ref, basis) resHF = res - DFTK.transfer_blochwave(resLF, basis, basis_ref); e2 = apply_metric(ψr, P, resHF, apply_inv_M); # Rayleigh coefficients needed for `apply_Ω` Λ = map(enumerate(ψr)) do (ik, ψk) Hk = hamr.blocks[ik] Hψk = Hk * ψk ψk'Hψk end ΩpKe2 = DFTK.apply_Ω(e2, ψr, hamr, Λ) .+ DFTK.apply_K(basis_ref, e2, ψr, ρr, occ) ΩpKe2 = DFTK.transfer_blochwave(ΩpKe2, basis_ref, basis) rhs = resLF - ΩpKe2; ψ, _ = DFTK.select_occupied_orbitals(basis, scfres.ψ, scfres.occupation) e1 = DFTK.solve_ΩplusK(basis, ψ, rhs, occ; tol).δψ e1 = DFTK.transfer_blochwave(e1, basis, basis_ref) res_schur = e1 + Mres; f_ref = compute_forces(scfres_ref) forces = Dict("F(P_*)" => f_ref) relerror = Dict("F(P_*)" => 0.0) compute_relerror(f) = norm(f - f_ref) / norm(f_ref); f = compute_forces(scfres) forces["F(P)"] = f relerror["F(P)"] = compute_relerror(f); function df(basis, occupation, ψ, δψ, ρ) δρ = DFTK.compute_δρ(basis, ψ, δψ, occupation) ForwardDiff.derivative(ε -> compute_forces(basis, ψ.+ε.*δψ, occupation; ρ=ρ+ε.*δρ), 0) end; df_err = df(basis_ref, occ, ψr, DFTK.proj_tangent(err, ψr), ρr) forces["F(P) - df(P)⋅(P-P_*)"] = f - df_err relerror["F(P) - df(P)⋅(P-P_*)"] = compute_relerror(f - df_err); df_schur = df(basis_ref, occ, ψr, res_schur, ρr) forces["F(P) - df(P)⋅Rschur(P)"] = f - df_schur relerror["F(P) - df(P)⋅Rschur(P)"] = compute_relerror(f - df_schur); for (key, value) in pairs(forces) @printf("%30s = [%7.5f, %7.5f, %7.5f] (rel. error: %7.5f)\n", key, (value[1])..., relerror[key]) end