stresses without unfold_bz

[jaemolihm]

Hi, regarding the new stresses functionality, isn't it possible to compute stresses using only irreducible k points - without unfold_bz - by symmetrizing the resulting stress tensor?

using Test
using DFTK
using ForwardDiff
import FiniteDiff
include("DFTK.jl/test/testcases.jl")

# Hellmann-Feynman stress
# via ForwardDiff & custom FFTW overloads on ForwardDiff.Dual

function make_basis(lattice, symmetries)
    Si = ElementPsp(silicon.atnum, psp=load_psp(silicon.psp))
    atoms = [Si => silicon.positions]
    model = model_DFT(lattice, atoms, [:lda_x, :lda_c_vwn]; symmetries)
    kgrid = [2, 2, 1]
    Ecut = 7
    PlaneWaveBasis(model; Ecut, kgrid)
end

function recompute_energy(lattice, symmetries)
    basis = make_basis(lattice, symmetries)
    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, lattice, symmetries)
    basis = make_basis(lattice, symmetries)
    ρ = DFTK.compute_density(basis, scfres.ψ, scfres.occupation)
    energies, H = energy_hamiltonian(basis, scfres.ψ, scfres.occupation; ρ=ρ)
    energies.total
end

a = 10.0  # slightly compressed
lattice = a / 2 * [[0 1 1.];
                    [1 0 1.];
                    [1 1 0.]]
is_converged = DFTK.ScfConvergenceDensity(1e-13)
scfres = self_consistent_field(make_basis(lattice, true); is_converged);
scfres_nosym = self_consistent_field(make_basis(lattice, false); is_converged);
stresses = compute_stresses(scfres)
stresses_nosym = compute_stresses(scfres_nosym)
@test isapprox(stresses, stresses_nosym, atol=1e-10)

# ===================================================
# New
function symmetrize_tensor(tensor, symmetries, lattice)
    tensor_symmetrized = zero(tensor)
    for (S, τ) in symmetries
        S_reduced = inv(lattice) * S * lattice
        tensor_symmetrized += S_reduced' * tensor * S_reduced
    end
    tensor_symmetrized /= length(symmetries)
    tensor_symmetrized
end

function compute_stresses_without_unfolding(scfres)
    # TODO optimize by only computing derivatives wrt 6 independent parameters
    # scfres = unfold_bz(scfres)
    # compute the Hellmann-Feynman energy (with fixed ψ/occ/ρ)
    function HF_energy(lattice)
        T = eltype(lattice)
        basis = scfres.basis
        model = basis.model
        new_model = Model(lattice;
                          model.n_electrons,
                          model.atoms,
                          magnetic_moments=[], # not used because we give symmetries explicitly
                          terms=model.term_types,
                          model.temperature,
                          model.smearing,
                          model.spin_polarization,
                          model.symmetries)
        new_basis = PlaneWaveBasis(new_model,
                                   basis.Ecut, basis.fft_size, basis.variational,
                                   basis.kcoords_global, basis.ksymops_global,
                                   basis.kgrid, basis.kshift, basis.symmetries,
                                   basis.comm_kpts)
        ρ = DFTK.compute_density(new_basis, scfres.ψ, scfres.occupation)
        energies, _ = energy_hamiltonian(new_basis, scfres.ψ, scfres.occupation;
                                         ρ, scfres.eigenvalues, scfres.εF)
        energies.total
    end
    L = scfres.basis.model.lattice
    Ω = scfres.basis.model.unit_cell_volume
    stresses_not_symmetrized = ForwardDiff.gradient(M -> HF_energy((I+M) * L), zero(L)) / Ω

    symmetrize_tensor(stresses_not_symmetrized, scfres.basis.symmetries, scfres.basis.model.lattice)
end
stresses_wo_unfold = compute_stresses_without_unfolding(scfres);
@test isapprox(stresses_wo_unfold, stresses_nosym, atol=1e-10)

The function compute_stresses_without_unfolding is taken from compute_stresses, and two changes are made: 1) unfold_bz is commented out, and 2) stress tensor is symmetrized based on basis.symmetries. Is there any problem with this method?

By the way, it took quite a while to compile compute_stresses_without_unfolding when developing. Is it natural?

P.S. This autodiff stress functionality is fascinating!

[antoine-levitt]

Nice, thanks for this! Yes I wondered about this and you're right, it should work. Also probably with forces, where we also unfold (see nonlocal.jl). I did the unfolding thing because it's more foolproof. I'm never been very confortable around symmetries, and we've had quite a few tricky bugs with them (and it's not the most pleasant thing in the world to debug) so I'm extra paranoid with it... Plus we needed the unfolding anyway for the wannier90 interface, and the cost is not great. But for MPI the symmetrization approach is definitely the way to go (since unfolding is tricky to make work with MPI), so we should do this. If you've got time to work on this feel free to submit a PR and we'll merge it gratefully!

By the way, it took quite a while to compile compute_stresses_without_unfolding when developing. Is it expected?

Unfortunately yes I've found the same thing, it takes ages to compile, very annoying... Haven't investigated it properly but it's forwarddiff-related. Hopefully we'll move to Diffractor at some point so I'm reluctant to investigate more.

P.S. This autodiff stress functionality is fascinating!

Yes, it works out very nicely! It's @niklasschmitz 's gsoc project, he's been doing a great job! Now the big excitement is seeing if reverse mode can be made to work...

[mfherbst]

Hey @jaemolihm. Many thanks for this snippet. We would indeed be very grateful if you (or someone else) turned this into a PR!

[jaemolihm]

Thanks for the comments! I'll make a PR in a few days.

[mfherbst]

Perfect. Thanks @jaemolihm !