Elastic constants

[antoine-levitt]

cc @LaurentVidal95

Elastic constants appear to work more or less out of the box (with one fix), we should probably fix and advertise them in an example. It does take ages to precompile the forwarddiff stuff though, I'm not sure what we can do about that...

@@ -192,10 +192,10 @@ function PlaneWaveBasis(model::Model{T}, Ecut::Number, fft_size, variational,
         @assert isnothing(kweights)
         kcoords, kweights, symmetries = bzmesh_ir_wedge(kgrid, symmetries; kshift)
     else
-        # Manual kpoint set based on kcoords/kweights
-        @assert length(kcoords) == length(kweights)
-        all_kcoords = unfold_kcoords(kcoords, symmetries)
-        symmetries  = symmetries_preserving_kgrid(symmetries, all_kcoords)
+        # # Manual kpoint set based on kcoords/kweights
+        # @assert length(kcoords) == length(kweights)
+        # all_kcoords = unfold_kcoords(kcoords, symmetries)
+        # symmetries  = symmetries_preserving_kgrid(symmetries, all_kcoords)
     end
 
     # Init MPI, and store MPI-global values for reference

# Very basic setup, useful for testing
using DFTK
using ForwardDiff

a = 10.26  # Silicon lattice constant in Bohr

Si = ElementPsp(:Si, psp=load_psp("hgh/lda/Si-q4"))
atoms     = [Si, Si]
positions = [ones(3)/8, -ones(3)/8]


function E(a)
    lattice = a / 2 * [[0 1 1.];
                       [1 0 1.];
                       [1 1 0.]]
    model = model_LDA(lattice, atoms, positions, symmetries=false)
    basis = PlaneWaveBasis(model; Ecut=15, kgrid=[2, 2, 2])
    scfres = self_consistent_field(basis, tol=1e-10)
    scfres.energies.total
end

# as = range(10, 11, 10)
# plot(as, E.(as))


function dEda(a)
    lattice = a / 2 * [[0 1 1.];
                       [1 0 1.];
                       [1 1 0.]]
    model = model_LDA(lattice, atoms, positions, symmetries=false)
    basis = PlaneWaveBasis(model; Ecut=15, kgrid=[2, 2, 2])
    scfres = self_consistent_field(basis, tol=1e-10)
    function HF_energy(new_a)
        basis = scfres.basis
        lattice = new_a / 2 * [[0 1 1.];
                           [1 0 1.];
                           [1 1 0.]]
        new_model = Model(basis.model; lattice)
        new_basis = PlaneWaveBasis(new_model,
                                   basis.Ecut, basis.fft_size, basis.variational,
                                   basis.kcoords_global, basis.kweights_global,
                                   basis.kgrid, basis.kshift, basis.symmetries_respect_rgrid,
                                   basis.comm_kpts, basis.architecture)
        ρ = compute_density(new_basis, scfres.ψ, scfres.occupation)
        energies = energy_hamiltonian(new_basis, scfres.ψ, scfres.occupation;
                                      ρ, scfres.eigenvalues, scfres.εF).energies
        energies.total
    end
    ForwardDiff.derivative(HF_energy, a)
end

a = 11
ε = 1e-4
println((E(a+ε) - E(a)) / ε)
println(dEda(a))

println((dEda(a+ε) - dEda(a))/ε)
println(ForwardDiff.derivative(dEda, a))

[antoine-levitt]

Actually I'm not sure if it's the precompilation or just the computation, should investigate (profile).

[LaurentVidal95]

Thanks for the basic setup ! I will play a bit with it. I'll let you know as soon as I've been able to do the profiling.

[LaurentVidal95]

From the profiling a lot of the time seems to be devoted to the computation of the Hessian involved in the ForwardDiff rule applying to the self_consistent_field function. But I'm not sure that I'm reading the graph correctly.

[antoine-levitt]

Can you maybe post a screenshot? If that's the case it's actually computation and not compilation and we should probably optimize it.

[mfherbst]

Cool, looks similar to what I did for the lattice constant sensitivities.

We can probably simplify a bit with proper update constructors or set lattice functions for planewavebasis, which is needed for lattice optimisations as well.

Regarding performance: The tolerances in solve omega+k are not yet tuned very well and parallelisation could be improved, but maybe there are more obvious things, too.

[LaurentVidal95]

Sorry for the delay ! I'm not sure a screen of the flame graph is readable but here is the file generated by ProfileView. Apparently you can just do "ProfileView.view(nothing)" and then open the file for an interactive flame graph !
HF_elastic_constant_pofiling.zip

[antoine-levitt]

Nothing weird here, just lots of ffts and matmuls, so yeah looks like it needs some finetuning