diff --git a/gp_testing.jl b/gp_testing.jl new file mode 100644 index 0000000000..03b416f261 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/DFTK.jl b/src/DFTK.jl index 69aa2c5c14..5f5a464c1c 100644 --- a/src/DFTK.jl +++ b/src/DFTK.jl @@ -97,6 +97,7 @@ include("supercell.jl") export Energies include("Energies.jl") +export Densities export Hamiltonian export HamiltonianBlock export energy_hamiltonian diff --git a/src/PlaneWaveBasis.jl b/src/PlaneWaveBasis.jl index 1688e85ac2..19228da65f 100644 --- a/src/PlaneWaveBasis.jl +++ b/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 diff --git a/src/density_methods.jl b/src/density_methods.jl index 2c6bb7dde5..3747947349 100644 --- a/src/density_methods.jl +++ b/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) diff --git a/src/scf/scf_callbacks.jl b/src/scf/scf_callbacks.jl index c33c2ba384..36ad8b73b6 100644 --- a/src/scf/scf_callbacks.jl +++ b/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) diff --git a/src/scf/self_consistent_field.jl b/src/scf/self_consistent_field.jl index 6d04cd8435..540af36ffa 100644 --- a/src/scf/self_consistent_field.jl +++ b/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, diff --git a/src/terms/Hamiltonian.jl b/src/terms/Hamiltonian.jl index df120b566f..18a7bf0afe 100644 --- a/src/terms/Hamiltonian.jl +++ b/src/terms/Hamiltonian.jl @@ -183,37 +183,67 @@ 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 @@ -221,7 +251,10 @@ 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)) diff --git a/src/terms/kinetic.jl b/src/terms/kinetic.jl index ad71f596b5..f500f01906 100644 --- a/src/terms/kinetic.jl +++ b/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]) diff --git a/src/terms/local.jl b/src/terms/local.jl index 097a5bc88c..7136821d26 100644 --- a/src/terms/local.jl +++ b/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; diff --git a/src/terms/local_nonlinearity.jl b/src/terms/local_nonlinearity.jl index 94f80ae59b..10d842df2d 100644 --- a/src/terms/local_nonlinearity.jl +++ b/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, ρ) diff --git a/src/terms/terms.jl b/src/terms/terms.jl index aa7564ccc5..85afe2b353 100644 --- a/src/terms/terms.jl +++ b/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).