[WIP] State refactor

gp_testing.jl

@@ -0,0 +1,43 @@
+a = 10
+lattice = a .* [[1 0 0.]; [0 0 0]; [0 0 0]];
+pot(x) = (x - a/2)^2;
+C = 1.0
+α = 2;
+
+using DFTK
+using LinearAlgebra
+
+n_electrons = 1  # Increase this for fun
+terms = [Kinetic(),
+         ExternalFromReal(r -> pot(r[1])),
+         LocalNonlinearity(ρ -> C * ρ^α),
+]
+model = Model(lattice; n_electrons, terms, spin_polarization=:spinless);  # spinless electrons
+
+# We discretize using a moderate Ecut (For 1D values up to `5000` are completely fine)
+# and run a direct minimization algorithm:
+basis = PlaneWaveBasis(model, Ecut=500, kgrid=(1, 1, 1))
+scfres = self_consistent_field(basis; tol=1e-8, mixing=SimpleMixing()) # This is a constrained preconditioned LBFGS
+scfres.energies
+
+# ## Internals
+# We use the opportunity to explore some of DFTK internals.
+#
+# Extract the converged density and the obtained wave function:
+ρ = real(scfres.ρ)[:, 1, 1, 1]  # converged density, first spin component
+ψ_fourier = scfres.ψ[1][:, 1];    # first k-point, all G components, first eigenvector
+
+# Transform the wave function to real space and fix the phase:
+ψ = ifft(basis, basis.kpoints[1], ψ_fourier)[:, 1, 1]
+ψ /= (ψ[div(end, 2)] / abs(ψ[div(end, 2)]));
+
+# Check whether ``ψ`` is normalised:
+x = a * vec(first.(DFTK.r_vectors(basis)))
+N = length(x)
+dx = a / N  # real-space grid spacing
+@assert sum(abs2.(ψ)) * dx ≈ 1.0
+
+# The density is simply built from ψ:
+norm(scfres.ρ - abs2.(ψ))
+
+nothing

src/DFTK.jl

@@ -97,6 +97,7 @@ include("supercell.jl")
 export Energies
 include("Energies.jl")
 
+export Densities
 export Hamiltonian
 export HamiltonianBlock
 export energy_hamiltonian

src/PlaneWaveBasis.jl

@@ -92,7 +92,7 @@ struct PlaneWaveBasis{T,
     # not yet implement symmetry. See `unfold_bz` as a convenient way to use this.
     use_symmetries_for_kpoint_reduction::Bool
 
-    ## Instantiated terms (<: Term). See Hamiltonian for high-level usage
+    ## Instantiated terms (<: Union{OrbitalsTerm, DensitiesTerm}). See Hamiltonian for high-level usage
     terms::Vector{Any}
 end
 

src/density_methods.jl

@@ -29,6 +29,15 @@ function random_density(basis::PlaneWaveBasis{T}, n_electrons::Integer) where {T
     ρ_from_total_and_spin(ρtot, ρspin)
 end
 
+function guess_missing_densities(basis::PlaneWaveBasis, densities::Densities)
+    ρ = @something densities.ρ guess_density(basis)
+    τ = densities.τ
+    if isnothing(τ) && any(t -> t isa DensitiesTerm && :τ ∈ needed_densities(t), basis.terms)
+        τ = zero(ρ)
+    end
+    Densities(ρ, τ)
+end
+
 # Atomic density methods
 function guess_density(basis::PlaneWaveBasis, magnetic_moments=[],
                        n_electrons=basis.model.n_electrons)

src/scf/scf_callbacks.jl

@@ -198,7 +198,7 @@ end
 # as opposed to any of these more sophisticated criteria.
 
 function default_diagtolalg(basis; tol, kwargs...)
-    if any(t -> t isa TermNonlinear, basis.terms)
+    if any(t -> t isa NonlinearDensitiesTerm, basis.terms)
         AdaptiveDiagtol()
     else
         AdaptiveDiagtol(; diagtol_first=tol/5)

src/scf/self_consistent_field.jl

@@ -129,8 +129,7 @@ Overview of parameters:
 """
 @timing function self_consistent_field(
     basis::PlaneWaveBasis{T};
-    ρ=guess_density(basis),
-    τ=any(needs_τ, basis.terms) ? zero(ρ) : nothing,
+    densities=Densities(),
     ψ=nothing,
     tol=1e-6,
     is_converged=ScfConvergenceDensity(tol),
@@ -148,6 +147,7 @@ Overview of parameters:
     compute_consistent_energies=true,
     response=ResponseOptions(),  # Dummy here, only for AD
 ) where {T}
+    densities = guess_missing_densities(basis, densities)
     if !isnothing(ψ)
         @assert length(ψ) == length(basis.kpoints)
     end
@@ -162,7 +162,7 @@ Overview of parameters:
 
         # Note that ρin is not the density of ψ, and the eigenvalues
         # are not the self-consistent ones, which makes this energy non-variational
-        energies, ham = energy_hamiltonian(basis, ψ, occupation; ρ=ρin, τ, eigenvalues, εF)
+        energies, ham = energy_hamiltonian(basis, ψ, occupation, Densities(; ρ=ρin, τ); eigenvalues, εF)
 
         # Diagonalize `ham` to get the new state
         nextstate = next_density(ham, nbandsalg, fermialg; eigensolver, ψ, eigenvalues,
@@ -175,9 +175,10 @@ Overview of parameters:
         #      on the same footing. In the future such principles will also apply
         #      to other quantities. See discussion in
         #      https://github.com/JuliaMolSim/DFTK.jl/issues/1065
-        if any(needs_τ, basis.terms)
-            τ = compute_kinetic_energy_density(basis, ψ, occupation)
-        end
+        # TODO: remove τ
+        # if any(needs_τ, basis.terms)
+        #     τ = compute_kinetic_energy_density(basis, ψ, occupation)
+        # end
 
         # Update info with results gathered so far
         info_next = (; ham, basis, converged, stage=:iterate, algorithm="SCF",
@@ -187,7 +188,7 @@ Overview of parameters:
 
         # Compute the energy of the new state
         if compute_consistent_energies
-            (; energies) = energy(basis, ψ, occupation; ρ=ρout, τ, eigenvalues, εF)
+            (; energies) = energy(basis, ψ, occupation, Densities(; ρ=ρout, τ); eigenvalues, εF)
         end
         history_Etot = vcat(info.history_Etot, energies.total)
         history_Δρ = vcat(info.history_Δρ, norm(Δρ) * sqrt(basis.dvol))
@@ -209,18 +210,18 @@ Overview of parameters:
         ρnext, info_next
     end
 
-    info_init = (; ρin=ρ, τ, ψ, occupation=nothing, eigenvalues=nothing, εF=nothing,
+    info_init = (; ρin=densities.ρ, densities.τ, ψ, occupation=nothing, eigenvalues=nothing, εF=nothing,
                    n_iter=0, n_matvec=0, timedout=false, converged=false,
                    history_Etot=T[], history_Δρ=T[])
 
     # Convergence is flagged by is_converged inside the fixpoint_map.
-    _, info = solver(fixpoint_map, ρ, info_init; maxiter)
+    _, info = solver(fixpoint_map, densities.ρ, info_init; maxiter)
 
     # We do not use the return value of solver but rather the one that got updated by fixpoint_map
     # ψ is consistent with ρout, so we return that. We also perform a last energy computation
     # to return a correct variational energy
     (; ρin, ρout, τ, ψ, occupation, eigenvalues, εF, converged) = info
-    energies, ham = energy_hamiltonian(basis, ψ, occupation; ρ=ρout, τ, eigenvalues, εF)
+    energies, ham = energy_hamiltonian(basis, ψ, occupation, Densities(; ρ=ρout, τ); eigenvalues, εF)
 
     # Callback is run one last time with final state to allow callback to clean up
     scfres = (; ham, basis, energies, converged, nbandsalg.occupation_threshold,

src/terms/Hamiltonian.jl

@@ -183,45 +183,78 @@ end
     Hψ
 end
 
+function orbitals_term_energy(basis::PlaneWaveBasis{T}, ψ, occupation, ops) where {T}
+    E = zero(T)
+    for (ik, ψk) in enumerate(ψ)
+        # TODO: is this allocation a problem?
+        op_applied_storage = zeros_like(ψk)
+        # TODO: what about non-fourier terms?
+        apply!((; fourier=op_applied_storage), ops[ik], (; fourier=ψk))
+        # TODO: use columnwise_dots function here I think
+        E += basis.kweights[ik] *
+            sum(occupation[ik] .*
+                real(vec(sum(conj(ψk) .* op_applied_storage; dims=1))))
+    end
+    mpi_sum(E, basis.comm_kpts)
+end
 
 """
 Get energies and Hamiltonian
 kwargs is additional info that might be useful for the energy terms to precompute
 (eg the density ρ)
 """
-@timing function energy_hamiltonian(basis::PlaneWaveBasis, ψ, occupation; kwargs...)
-    # it: index into terms, ik: index into kpoints
-    @timing "ene_ops" ene_ops_arr = [ene_ops(term, basis, ψ, occupation; kwargs...)
-                                     for term in basis.terms]
+@timing function energy_hamiltonian(basis::PlaneWaveBasis{T}, ψ, occupation, densities;
+                                    eigenvalues=nothing, εF=nothing) where {T}
     term_names = [string(nameof(typeof(term))) for term in basis.model.term_types]
-    energy_values  = [eh.E for eh in ene_ops_arr]
-    operators = [eh.ops for eh in ene_ops_arr]         # operators[it][ik]
-
-    # flatten the inner arrays in case a term returns more than one operator
-    function flatten(arr)
-        ret = []
-        for a in arr
-            if a isa RealFourierOperator
-                push!(ret, a)
-            else
-                push!(ret, a...)
+    energy_values  = [zero(T) for _ in basis.model.term_types]
+
+    operators = [[] for _ in basis.kpoints]
+    total_potentials = Densities()
+
+    for (iterm, term) in enumerate(basis.terms)
+        if isa(term, OrbitalsTerm)
+            # TODO: is ops ambiguous without DFTK. prefix here?
+            ops = DFTK.ops(term, basis)
+            energy_values[iterm] = isnothing(ψ) ? T(Inf) : orbitals_term_energy(basis, ψ, occupation, ops)
+            for (ik, op) in enumerate(ops)
+                push!(operators[ik], op)
             end
+        else
+            (; E, potentials) = energy_potentials(term, basis, densities)
+            energy_values[iterm] = E
+            total_potentials = sum_densities(total_potentials, potentials)
+        end
+    end
+
+    if !isnothing(total_potentials.ρ)
+        for (ik, kpt) in enumerate(basis.kpoints)
+            push!(operators[ik], RealSpaceMultiplication(basis, kpt, total_potentials.ρ[:, :, :, kpt.spin]))
         end
-        ret
     end
+    if !isnothing(total_potentials.τ)
+        for (ik, kpt) in enumerate(basis.kpoints)
+            push!(operators[ik], DivAgradOperator(basis, kpt, total_potentials.τ[:, :, :, kpt.spin]))
+        end
+    end
+
     scratch = _ham_allocate_scratch(basis)
-    hks_per_k = [flatten([blocks[ik] for blocks in operators])
-                 for ik = 1:length(basis.kpoints)]      # hks_per_k[ik][it]
+    # TODO: Flattening should be restored
     ham = Hamiltonian(basis, [HamiltonianBlock(basis, kpt, hks; scratch)
-                              for (hks, kpt) in zip(hks_per_k, basis.kpoints)])
+                              for (hks, kpt) in zip(operators, basis.kpoints)])
+
+    # TODO: add Entropy contribution
+
     energies = Energies(term_names, energy_values)
     (; energies, ham)
 end
 
 """
 Faster version than energy_hamiltonian for cases where only the energy is needed.
 """
-@timing function energy(basis::PlaneWaveBasis, ψ, occupation; kwargs...)
+@timing function energy(basis::PlaneWaveBasis, ψ, occupation, densities; kwargs...)
+    # TODO: reimplement
+    return (; energies=energy_hamiltonian(basis, ψ, occupation, densities; kwargs...).energies)
+
     energy_values = [energy(term, basis, ψ, occupation; kwargs...) for term in basis.terms]
     term_names = [string(nameof(typeof(term))) for term in basis.model.term_types]
     (; energies=Energies(term_names, energy_values))

src/terms/kinetic.jl

@@ -16,7 +16,7 @@ function Base.show(io::IO, kin::Kinetic)
     print(io, "Kinetic($bup$fac)")
 end
 
-struct TermKinetic <: Term
+struct TermKinetic <: OrbitalsTerm
     scaling_factor::Real  # scaling factor, absorbed into kinetic_energies
     # kinetic energies 1/2(k+G)^2 *blowup(|k+G|, Ecut) for each k-point.
     kinetic_energies::Vector{<:AbstractVector}
@@ -37,6 +37,11 @@ function kinetic_energy(kin::Kinetic, Ecut, p)
     kinetic_energy(kin.blowup, kin.scaling_factor, Ecut, p)
 end
 
+function ops(term::TermKinetic, basis::PlaneWaveBasis{T}) where {T}
+    [FourierMultiplication(basis, kpoint, term.kinetic_energies[ik])
+     for (ik, kpoint) in enumerate(basis.kpoints)]
+end
+
 @timing "ene_ops: kinetic" function ene_ops(term::TermKinetic, basis::PlaneWaveBasis{T},
                                             ψ, occupation; kwargs...) where {T}
     ops = [FourierMultiplication(basis, kpoint, term.kinetic_energies[ik])

src/terms/local.jl

@@ -4,7 +4,15 @@
 # potential in real space on the grid. If the potential is different in the α and β
 # components then it should be a 4d-array with the last axis running over the
 # two spin components.
-abstract type TermLocalPotential <: Term end
+abstract type TermLocalPotential <: DensitiesTerm end
+
+function energy_potentials(term::TermLocalPotential, basis::PlaneWaveBasis{T}, densities::Densities) where {T}
+    E = sum(total_density(densities.ρ) .* term.potential_values) * basis.dvol
+    pot = ndims(term.potential_values) == 3 ? reshape(term.potential_values, basis.fft_size..., 1) :
+        term.potential_values
+    (; E, potentials=Densities(; ρ=pot))
+end
+needed_densities(::TermLocalPotential) = (:ρ,)
 
 @timing "ene_ops: local" function ene_ops(term::TermLocalPotential,
                                           basis::PlaneWaveBasis{T}, ψ, occupation;

src/terms/local_nonlinearity.jl

@@ -4,11 +4,20 @@ Local nonlinearity, with energy ∫f(ρ) where ρ is the density
 struct LocalNonlinearity
     f
 end
-struct TermLocalNonlinearity{TF} <: TermNonlinear
+struct TermLocalNonlinearity{TF} <: NonlinearDensitiesTerm
     f::TF
 end
 (L::LocalNonlinearity)(::AbstractBasis) = TermLocalNonlinearity(L.f)
 
+function energy_potentials(term::TermLocalNonlinearity, basis::PlaneWaveBasis{T}, densities::Densities) where {T}
+    fp(ρ) = ForwardDiff.derivative(term.f, ρ)
+    E = sum(fρ -> convert_dual(T, fρ), term.f.(densities.ρ)) * basis.dvol
+    potential = convert_dual.(T, fp.(densities.ρ))
+
+    (; E, potentials=Densities(; ρ=potential))
+end
+needed_densities(::TermLocalNonlinearity) = (:ρ,)
+
 function ene_ops(term::TermLocalNonlinearity, basis::PlaneWaveBasis{T}, ψ, occupation;
                  ρ, kwargs...) where {T}
     fp(ρ) = ForwardDiff.derivative(term.f, ρ)

src/terms/terms.jl

@@ -2,6 +2,36 @@ using DftFunctionals
 
 include("operators.jl")
 
+@kwdef struct Densities{Tρ, Tτ}
+    ρ::Tρ = nothing
+    τ::Tτ = nothing
+end
+
+function sum_densities(a::Densities, b::Densities)
+    sum_density(a, b) = isnothing(a) ? b : isnothing(b) ? a : a .+ b
+
+    Densities(sum_density(a.ρ, b.ρ),
+              sum_density(a.τ, b.τ))
+end
+
+# For an orbital term the ops are constant
+# and the energy is quadratic in the orbitals and computed from the ops
+# It must overload:
+# - `ops(term, basis)` returning `Vector{RealFourierOperator}` (for each kpoint)
+abstract type OrbitalsTerm end
+
+# A term that depends on the densities but not on the orbitals
+# It must overload:
+# - `needed_densities(term)` returning an iterable of symbols
+# - `energy_potentials(term, basis, densities::Densities)` returning `(; E, potentials::Densities)`
+abstract type DensitiesTerm end
+
+# TODO: linear subtype?
+
+# Terms that are non-linear in the density (i.e. which give rise to a Hamiltonian
+# contribution that is density-dependent as well)
+abstract type NonlinearDensitiesTerm <: DensitiesTerm end
+
 ### Terms
 # - A Term is something that, given a state, returns a named tuple (; E, hams) with an energy
 #   and a list of RealFourierOperator (for each kpoint).