JuliaMolSim / DFTK.jl

Density-functional toolkit
https://docs.dftk.org
MIT License
426 stars 89 forks source link

AD wrong with symmetries #817

Open epolack opened 1 year ago

epolack commented 1 year ago
using DFTK
using StaticArrays
using ForwardDiff
using FiniteDifferences

include("test/testcases.jl")

for method in (ForwardDiff.derivative, FiniteDifferences.central_fdm(5, 1))
    text_meth = method == ForwardDiff.derivative ? "(ad)" : "(finite diff)"
    for symmetries in (true, false)
        text_sym = symmetries ? "with" : "without"
        cell = silicon
        Ecut = 5
        kgrid = [1, 1, 1]
        model = model_atomic(cell.lattice, cell.atoms, cell.positions; symmetries)
        basis = PlaneWaveBasis(model; Ecut, kgrid)

        T = eltype(basis)
        n_atoms = length(model.positions)
        n_dim = model.n_dim
        d2_term = zeros(ComplexF64, (n_dim, n_atoms, n_dim, n_atoms))
        for τ in 1:n_atoms
            for γ in 1:n_dim
                d2 = -method(zero(T)) do ε
                    displacement = zero.(model.positions)
                    displacement[τ] = StaticArrays.setindex(displacement[τ], one(T), γ)
                    cell_disp = (; lattice=eltype(ε).(cell.lattice),
                                 cell.atoms,
                                 positions=ε*displacement .+ cell.positions)
                    model_disp = model_atomic(cell_disp.lattice, cell_disp.atoms, cell_disp.positions; symmetries)
                    basis_disp = PlaneWaveBasis(model_disp; Ecut, kgrid)
                    scfres = self_consistent_field(basis_disp; callback=identity)
                    forces = compute_forces(scfres)
                    hcat(Array.(forces)...)
                end
                d2_term[:, :, γ, τ] = hcat(d2...)
            end
        end
        text = text_sym * " symmetries " * text_meth
        @info text reshape(d2_term, n_dim*n_atoms, n_dim*n_atoms)
    end
end

I expect the same results for all cases, especially because there is only one k-point. But AD with symmetries on shows wrong result (does not improve at all with bigger Ecut).

antoine-levitt commented 1 year ago

Huh, yes, there's no reason this should work actually. The place to fix this is in workarounds/forwarddiff, where the overload for Get_symmetry just ignores the variations of lattice and positions. It should instead return only those symmetries that are also preserved by the variations. Since I don't think spglib exposes this, we can either do it manually, or maybe deform by a finite amount and check. The use in stresses though is valid, so we need a flag to ignore this at model creation

antoine-levitt commented 1 year ago

https://github.com/JuliaMolSim/DFTK.jl/pull/819 is what I mean (untested and don't have time to polish it now)