From 329ff47f4ce1e8f53f204ebaea30af9df3e551d5 Mon Sep 17 00:00:00 2001 From: Antoine Levitt Date: Tue, 12 Apr 2022 15:57:37 +0200 Subject: [PATCH 1/6] support for isolated systems --- src/Model.jl | 11 +++- src/fft.jl | 8 +-- src/terms/ewald.jl | 69 +++++++++++++++--------- src/terms/hartree.jl | 81 +++++++++++++++++++++++++++- src/terms/local.jl | 72 ++++++++++++++----------- src/terms/pairwise.jl | 10 ++-- src/workarounds/forwarddiff_rules.jl | 1 + test/isolated.jl | 37 +++++++++++++ test/runtests.jl | 2 +- test/serialisation.jl | 1 + 10 files changed, 225 insertions(+), 67 deletions(-) create mode 100644 test/isolated.jl diff --git a/src/Model.jl b/src/Model.jl index 417f341765..57b1720b23 100644 --- a/src/Model.jl +++ b/src/Model.jl @@ -14,6 +14,10 @@ struct Model{T <: Real, VT <: Real} recip_lattice::Mat3{T} # Dimension of the system; 3 unless `lattice` has zero columns n_dim::Int + # Whether the problem is periodic along each direction + # When not periodic, the system is assumed to be centered + # at the middle of the cell (1/2) + periodic::Vec3{Bool} # Useful for conversions between cartesian and reduced coordinates inv_lattice::Mat3{T} inv_recip_lattice::Mat3{T} @@ -66,7 +70,7 @@ _is_well_conditioned(A; tol=1e5) = (cond(A) <= tol) """ Model(lattice, atoms, positions; n_electrons, magnetic_moments, terms, temperature, - smearing, spin_polarization, symmetries) + smearing, spin_polarization, symmetries, periodic) Creates the physical specification of a model (without any discretization information). @@ -109,6 +113,7 @@ function Model(lattice::AbstractMatrix{T}, spin_polarization=default_spin_polarization(magnetic_moments), symmetries=default_symmetries(lattice, atoms, positions, magnetic_moments, spin_polarization, terms), + periodic=Vec3(true, true, true) ) where {T <: Real} # Validate εF and n_electrons if !isnothing(εF) # fixed Fermi level @@ -176,7 +181,8 @@ function Model(lattice::AbstractMatrix{T}, @assert !isempty(symmetries) # Identity has to be always present. Model{T,value_type(T)}(model_name, - lattice, recip_lattice, n_dim, inv_lattice, inv_recip_lattice, + lattice, recip_lattice, n_dim, periodic, + inv_lattice, inv_recip_lattice, unit_cell_volume, recip_cell_volume, n_electrons, εF, spin_polarization, n_spin, temperature, smearing, atoms, positions, atom_groups, terms, symmetries) @@ -218,6 +224,7 @@ function Model{T}(model::Model; kwargs...) where {T <: Real} Model(T.(lattice), atoms, positions; model.model_name, + model.periodic, model.n_electrons, magnetic_moments=[], # not used because symmetries explicitly given terms=model.term_types, diff --git a/src/fft.jl b/src/fft.jl index 3adf785ccc..2a6c28f4b7 100644 --- a/src/fft.jl +++ b/src/fft.jl @@ -24,8 +24,8 @@ function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, f_fourier::Abstrac mul!(f_real, basis.opBFFT, f_fourier) f_real .*= basis.ifft_normalization end -function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, - kpt::Kpoint, f_fourier::AbstractVector; normalize=true) +function ifft!(f_real::AbstractArray3, basis::PlaneWaveBasis, kpt::Kpoint, + f_fourier::AbstractVector; normalize=true) @assert length(f_fourier) == length(kpt.mapping) @assert size(f_real) == basis.fft_size @@ -79,8 +79,8 @@ function fft!(f_fourier::AbstractArray3, basis::PlaneWaveBasis, f_real::Abstract mul!(f_fourier, basis.opFFT, f_real) f_fourier .*= basis.fft_normalization end -function fft!(f_fourier::AbstractVector, basis::PlaneWaveBasis, - kpt::Kpoint, f_real::AbstractArray3; normalize=true) +function fft!(f_fourier::AbstractVector, basis::PlaneWaveBasis, kpt::Kpoint, + f_real::AbstractArray3; normalize=true) @assert size(f_real) == basis.fft_size @assert length(f_fourier) == length(kpt.mapping) diff --git a/src/terms/ewald.jl b/src/terms/ewald.jl index 73aebbd824..80cc9104b7 100644 --- a/src/terms/ewald.jl +++ b/src/terms/ewald.jl @@ -20,8 +20,14 @@ end @timing "precomp: Ewald" function TermEwald(basis::PlaneWaveBasis{T}; η=default_η(basis.model.lattice)) where {T} model = basis.model + periodic = model.periodic + if all(.! periodic) + @assert all(periodic) || all(.! periodic) # only periodic or isolated atm + η = zero(T) # no Ewald splitting in this case: simple Coulomb for the unit cell only + end charges = charge_ionic.(model.atoms) - (; energy, forces) = energy_forces_ewald(model.lattice, charges, model.positions; η) + (; energy, forces) = energy_forces_ewald(model.lattice, charges, model.positions; η, + periodic) TermEwald(energy, forces, η) end @@ -34,17 +40,17 @@ compute_forces(term::TermEwald, ::PlaneWaveBasis, ψ, occupation; kwargs...) = t Standard computation of energy and forces. """ function energy_forces_ewald(lattice::AbstractArray{T}, charges::AbstractArray, - positions; kwargs...) where {T} - energy_forces_ewald(T, lattice, charges, positions, zero(Vec3{T}), nothing) + positions; periodic=Vec3(true, true, true), kwargs...) where {T} + energy_forces_ewald(T, lattice, charges, positions, zero(Vec3{T}), nothing; periodic) end """ Computation for phonons; required to build the dynamical matrix. """ function energy_forces_ewald(lattice::AbstractArray{T}, charges, positions, q, - ph_disp; kwargs...) where{T} + ph_disp; periodic=Vec3(true, true, true), kwargs...) where{T} S = promote_type(complex(T), eltype(ph_disp[1])) - energy_forces_ewald(S, lattice, charges, positions, q, ph_disp; kwargs...) + energy_forces_ewald(S, lattice, charges, positions, q, ph_disp; periodic, kwargs...) end # To compute the electrostatics of the system, we use the Ewald splitting method due to the @@ -80,7 +86,7 @@ For now this function returns zero energy and force on non-3D systems. Use a pai potential term if you want to customise this treatment. """ function energy_forces_ewald(S, lattice::AbstractArray{T}, charges, positions, q, ph_disp; - η=default_η(lattice)) where {T} + η=default_η(lattice), periodic=Vec3(true, true, true)) where {T} @assert length(charges) == length(positions) if isempty(charges) return (; energy=zero(T), forces=zero(positions)) @@ -115,27 +121,38 @@ function energy_forces_ewald(S, lattice::AbstractArray{T}, charges, positions, q poslims = [maximum(rj[i] - rk[i] for rj in positions for rk in positions) for i in 1:3] Rlims = estimate_integer_lattice_bounds(lattice, max_erfc_arg / η, poslims) + # Do we want shells in all directions? + is_dim_trivial = .!periodic + max_shell(n, trivial) = trivial ? 0 : n + Rlims = max_shell.(Rlims, is_dim_trivial) + # # Reciprocal space sum # - # Initialize reciprocal sum with correction term for charge neutrality - sum_recip::S = - (sum(charges)^2 / 4η^2) + if all(periodic) + # Initialize reciprocal sum with correction term for charge neutrality + sum_recip::S = - (sum(charges)^2 / 4η^2) + else + sum_recip = zero(T) + end forces_recip = zeros(Vec3{S}, length(positions)) - for G1 in -Glims[1]:Glims[1], G2 in -Glims[2]:Glims[2], G3 in -Glims[3]:Glims[3] - G = Vec3(G1, G2, G3) - iszero(G) && continue - Gsq = norm2(recip_lattice * G) - cos_strucfac = sum(Z * cos2pi(dot(r, G)) for (r, Z) in zip(positions, charges)) - sin_strucfac = sum(Z * sin2pi(dot(r, G)) for (r, Z) in zip(positions, charges)) - sum_strucfac = cos_strucfac^2 + sin_strucfac^2 - sum_recip += sum_strucfac * exp(-Gsq / 4η^2) / Gsq - for (ir, r) in enumerate(positions) - Z = charges[ir] - dc = -Z*2S(π)*G*sin2pi(dot(r, G)) - ds = +Z*2S(π)*G*cos2pi(dot(r, G)) - dsum = cos_strucfac*dc + sin_strucfac*ds - forces_recip[ir] -= dsum * exp(-Gsq / 4η^2)/Gsq + if all(periodic) # if non-periodic, no reciprocal-space contribution + for G1 in -Glims[1]:Glims[1], G2 in -Glims[2]:Glims[2], G3 in -Glims[3]:Glims[3] + G = Vec3(G1, G2, G3) + iszero(G) && continue + Gsq = norm2(recip_lattice * G) + cos_strucfac = sum(Z * cos2pi(dot(r, G)) for (r, Z) in zip(positions, charges)) + sin_strucfac = sum(Z * sin2pi(dot(r, G)) for (r, Z) in zip(positions, charges)) + sum_strucfac = cos_strucfac^2 + sin_strucfac^2 + sum_recip += sum_strucfac * exp(-Gsq / 4η^2) / Gsq + for (ir, r) in enumerate(positions) + Z = charges[ir] + dc = -Z*2S(π)*G*sin2pi(dot(r, G)) + ds = +Z*2S(π)*G*cos2pi(dot(r, G)) + dsum = cos_strucfac*dc + sin_strucfac*ds + forces_recip[ir] -= dsum * exp(-Gsq / 4η^2)/Gsq + end end end @@ -146,8 +163,12 @@ function energy_forces_ewald(S, lattice::AbstractArray{T}, charges, positions, q # # Real-space sum # - # Initialize real-space sum with correction term for uniform background - sum_real::S = -2η / sqrt(S(π)) * sum(Z -> Z^2, charges) + if all(periodic) + # Initialize real-space sum with correction term for uniform background + sum_real::S = -2η / sqrt(S(π)) * sum(Z -> Z^2, charges) + else + sum_real = zero(S) + end forces_real = zeros(Vec3{S}, length(positions)) for R1 in -Rlims[1]:Rlims[1], R2 in -Rlims[2]:Rlims[2], R3 in -Rlims[3]:Rlims[3] diff --git a/src/terms/hartree.jl b/src/terms/hartree.jl index 81372ce5f1..faca8f9f99 100644 --- a/src/terms/hartree.jl +++ b/src/terms/hartree.jl @@ -1,3 +1,4 @@ +using SpecialFunctions """ Hartree term: for a decaying potential V the energy would be @@ -39,11 +40,89 @@ function TermHartree(basis::PlaneWaveBasis{T}, scaling_factor) where {T} TermHartree(T(scaling_factor), T(scaling_factor) .* poisson_green_coeffs) end +# a gaussian of exponent α and integral M +ρref_real(r::T, M=1, α=1) where {T} = M * exp(-T(1)/2 * (α*r)^2) / ((2T(π))^(T(3)/2)) * α^3 +# solution of -ΔVref = 4π ρref +function Vref_real(r::T, M=1, α=1) where {T} + r == 0 && return M * 2 / sqrt(T(pi)) * α / sqrt(2) + M * erf(α/sqrt(2)*r)/r +end + +one_hot(i) = Vec3{Bool}(j == i for j=1:3) +∂f∂α(f, α, r) = ForwardDiff.derivative(ε -> f(r + ε * one_hot(α)), zero(eltype(r))) + +function get_center(basis, ρ) + sumρ = sum(ρ) + center = map(1:3) do α + rα = [r[α] for r in r_vectors_cart(basis)] + sum(rα .* ρ) / sumρ + end + Vec3(center...) +end +function get_integral(basis, ρ) + sum(ρ) * basis.dvol +end +function get_dipole(α, center, basis, ρ) + rα = [r[α]-center[α] for r in r_vectors_cart(basis)] + dot(rα, ρ) * basis.dvol +end +function get_variance(center, basis, ρ) + rr = [sum(abs2, r-center) for r in r_vectors_cart(basis)] + dot(rr, ρ) * basis.dvol +end + @timing "ene_ops: hartree" function ene_ops(term::TermHartree, basis::PlaneWaveBasis{T}, ψ, occupation; ρ, kwargs...) where {T} - ρtot_fourier = fft(basis, total_density(ρ)) + model = basis.model + ρtot_real = total_density(ρ) + ρtot_fourier = fft(basis, ρtot_real) pot_fourier = term.poisson_green_coeffs .* ρtot_fourier pot_real = irfft(basis, pot_fourier) + + # For isolated systems, the above does not compute a good potential (eg it assumes zero DC component) + # We correct it by solving -Δ V = 4πρ in two steps: we split ρ into ρref and ρ-ρref, + # where ρref is a gaussian has the same total charge as ρ. + # We compute the first potential in real space (explicitly since ρref is known), + # and the second (short-range) potential in Fourier space + # Compared to the usual scheme (ρref = 0), this results in a correction potential + # equal to Vref computed in real space minus Vref computed in Fourier space + if any(.!model.periodic) + @assert all(.!model.periodic) + # determine center and width from density + # Strictly speaking, these computations should result in extra terms to guarantee energy/ham consistency + center = get_center(basis, ρtot_real) + ρref_11 = [ρref_real(norm(r - center), model.n_electrons, 1) for r in r_vectors_cart(basis)] + spread_11 = get_variance(center, basis, ρref_11) + spread_ρ = get_variance(center, basis, ρtot_real) + α = 1/sqrt(spread_ρ/spread_11) + # α = 1 + + ρrad_fun(r) = ρref_real(norm(r - center), model.n_electrons, α) + ρrad = ρrad_fun.(r_vectors_cart(basis)) + + + # at this point we've determined a gaussian of same width and center as the original. + # Now we get the dipole moments of ρ and match them + ρders = [∂f∂α.(ρrad_fun, α, r_vectors_cart(basis)) for α=1:3] + + coeffs_ders = map(1:3) do α + get_dipole(α, center, basis, ρtot_real) / get_dipole(α, center, basis, ρders[α]) + end + ρref = ρrad + sum([coeffs_ders[α]*ρders[α] for α=1:3]) + + # compute corresponding solution of -ΔVref = 4π ρref + Vref_rad_fun(r) = Vref_real(norm(r - center), model.n_electrons, α) + Vref_rad = Vref_rad_fun.(r_vectors_cart(basis)) + Vref_ders = [∂f∂α.(Vref_rad_fun, α, r_vectors_cart(basis)) for α=1:3] + Vref = Vref_rad + sum([coeffs_ders[α]*Vref_ders[α] for α=1:3]) + + # TODO possibly optimize FFTs here + Vcorr_real = Vref - irfft(basis, term.poisson_green_coeffs .* fft(basis, ρref)) + Vcorr_fourier = fft(basis, Vcorr_real) + pot_real .+= Vcorr_real + pot_fourier .+= Vcorr_fourier + end + E = real(dot(pot_fourier, ρtot_fourier) / 2) ops = [RealSpaceMultiplication(basis, kpt, pot_real) for kpt in basis.kpoints] diff --git a/src/terms/local.jl b/src/terms/local.jl index d37df2c589..bea230e27e 100644 --- a/src/terms/local.jl +++ b/src/terms/local.jl @@ -83,43 +83,53 @@ Atomic local potential defined by `model.atoms`. """ struct AtomicLocal end function (::AtomicLocal)(basis::PlaneWaveBasis{T}) where {T} - # pot_fourier is expanded in a basis of e_{G-G'} - # Since V is a sum of radial functions located at atomic - # positions, this involves a form factor (`local_potential_fourier`) - # and a structure factor e^{-i G·r} model = basis.model - G_cart = to_cpu(G_vectors_cart(basis)) - # TODO Bring G_cart on the CPU for compatibility with the pseudopotentials which - # are not isbits ... might be able to solve this by restructuring the loop - - - # Pre-compute the form factors at unique values of |G| to speed up - # the potential Fourier transform (by a lot). Using a hash map gives O(1) - # lookup. - form_factors = IdDict{Tuple{Int,T},T}() # IdDict for Dual compatability - for G in G_cart - q = norm(G) - for (igroup, group) in enumerate(model.atom_groups) - if !haskey(form_factors, (igroup, q)) - element = model.atoms[first(group)] - form_factors[(igroup, q)] = local_potential_fourier(element, q) + + if all(model.periodic) + # pot_fourier is expanded in a basis of e_{G-G'} + # Since V is a sum of radial functions located at atomic + # positions, this involves a form factor (`local_potential_fourier`) + # and a structure factor e^{-i G·r} + G_cart = to_cpu(G_vectors_cart(basis)) + # TODO Bring G_cart on the CPU for compatibility with the pseudopotentials which + # are not isbits ... might be able to solve this by restructuring the loop + + + # Pre-compute the form factors at unique values of |G| to speed up + # the potential Fourier transform (by a lot). Using a hash map gives O(1) + # lookup. + form_factors = IdDict{Tuple{Int,T},T}() # IdDict for Dual compatability + for G in G_cart + q = norm(G) + for (igroup, group) in enumerate(model.atom_groups) + if !haskey(form_factors, (igroup, q)) + element = model.atoms[first(group)] + form_factors[(igroup, q)] = local_potential_fourier(element, q) + end end end - end - Gs = to_cpu(G_vectors(basis)) # TODO Again for GPU compatibility - pot_fourier = map(enumerate(Gs)) do (iG, G) - q = norm(G_cart[iG]) - pot = sum(enumerate(model.atom_groups)) do (igroup, group) - structure_factor = sum(r -> cis2pi(-dot(G, r)), @view model.positions[group]) - form_factors[(igroup, q)] * structure_factor + Gs = to_cpu(G_vectors(basis)) # TODO Again for GPU compatibility + pot_fourier = map(enumerate(Gs)) do (iG, G) + q = norm(G_cart[iG]) + pot = sum(enumerate(model.atom_groups)) do (igroup, group) + structure_factor = sum(r -> cis2pi(-dot(G, r)), @view model.positions[group]) + form_factors[(igroup, q)] * structure_factor + end + pot / sqrt(model.unit_cell_volume) end - pot / sqrt(model.unit_cell_volume) - end - - enforce_real!(basis, pot_fourier) # Symmetrize Fourier coeffs to have real iFFT - pot_real = irfft(basis, to_device(basis.architecture, pot_fourier)) + enforce_real!(basis, pot_fourier) # Symmetrize Fourier coeffs to have real iFFT + pot_real = irfft(basis, to_device(basis.architecture, pot_fourier)) + else + @assert all(.!basis.model.periodic) + # simple real-space sum + pot_real = map(r_vectors(basis)) do r + sum(zip(model.atoms, model.positions)) do (element, r_ion) + local_potential_real(element, norm(model.lattice * (r-r_ion))) + end + end + end TermAtomicLocal(pot_real) end diff --git a/src/terms/pairwise.jl b/src/terms/pairwise.jl index fd34e04725..f53f27652e 100644 --- a/src/terms/pairwise.jl +++ b/src/terms/pairwise.jl @@ -21,7 +21,8 @@ end model = basis.model symbols = Symbol.(atomic_symbol.(model.atoms)) (; energy, forces) = energy_forces_pairwise(model.lattice, symbols, model.positions, - P.V, P.params; P.max_radius) + P.V, P.params; P.max_radius, + basis.model.periodic) TermPairwisePotential(P.V, P.params, T(P.max_radius), energy, forces) end @@ -79,7 +80,8 @@ symbols). The potential is expected to decrease quickly at infinity. """ function energy_forces_pairwise(S, lattice::AbstractArray{T}, symbols, positions, V, params, - q, ph_disp; max_radius=100) where {T} + q, ph_disp; max_radius=100, + periodic=Vec3(true, true, true)) where {T} @assert length(symbols) == length(positions) if isempty(symbols) return (; energy=zero(T), forces=zero(positions)) @@ -97,8 +99,8 @@ function energy_forces_pairwise(S, lattice::AbstractArray{T}, symbols, positions poslims = [maximum(rj[i] - rk[i] for rj in positions for rk in positions) for i in 1:3] Rlims = estimate_integer_lattice_bounds(lattice, max_radius, poslims) - # Check if some coordinates are not used. - is_dim_trivial = iszero.(eachcol(lattice)) + # Do we want shells in all directions? + is_dim_trivial = iszero.(eachcol(lattice)) .|| (.!periodic)' max_shell(n, trivial) = trivial ? 0 : n Rlims = max_shell.(Rlims, is_dim_trivial) diff --git a/src/workarounds/forwarddiff_rules.jl b/src/workarounds/forwarddiff_rules.jl index 18f4b4207a..083eba91f9 100644 --- a/src/workarounds/forwarddiff_rules.jl +++ b/src/workarounds/forwarddiff_rules.jl @@ -143,6 +143,7 @@ function construct_value(model::Model{T}) where {T <: ForwardDiff.Dual} construct_value.(model.atoms), newpositions; model.model_name, + model.periodic, model.n_electrons, magnetic_moments=[], # Symmetries given explicitly terms=model.term_types, diff --git a/test/isolated.jl b/test/isolated.jl new file mode 100644 index 0000000000..9e3430eb01 --- /dev/null +++ b/test/isolated.jl @@ -0,0 +1,37 @@ +@testitem "Isolated systems" tags=[:off] begin + using DFTK + using LinearAlgebra + + RES = Dict() + for a in (10, 20, 30) + for per in (false, true) + @time begin + lattice = [a 0 0; + 0 a 0; + 0 0 a] + atoms = [ElementPsp(:Li, psp=load_psp("hgh/lda/Li-q3"))] + atoms = [ElementPsp(:He, psp=load_psp("hgh/lda/He-q2"))] + positions = [[1/2, 1/2, 1/2]] + + kgrid = [1, 1, 1] # no k-point sampling for an isolated system + Ecut = 40 + tol = 1e-6 + + model = model_LDA(lattice, atoms, positions, periodic=[per, per, per]) + basis = PlaneWaveBasis(model; Ecut, kgrid) + res = self_consistent_field(basis, is_converged=DFTK.ScfConvergenceDensity(tol)) + + rr = vec([norm(a * (r .- 1/2)) for r in r_vectors(basis)]) + + function quadrupole(basis, ρ) + rr = [norm(a * (r .- 1/2)) for r in r_vectors(basis)] + sum(rr .^ 2 .* ρ) * basis.dvol + end; + quad = quadrupole(basis, res.ρ) + println(quad) + RES[per, a] = (quad, res.energies.total) + end + end + end + display(sort(RES)) +end diff --git a/test/runtests.jl b/test/runtests.jl index 11c851b56b..4faf39e781 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -43,7 +43,7 @@ end const TAGS = isempty(DFTK_TEST_ARGS) ? [:all] : Symbol.(DFTK_TEST_ARGS) # Tags excluded from the first run of "all" tests -const EXCLUDED_FROM_ALL = Symbol[:example, :gpu] +const EXCLUDED_FROM_ALL = Symbol[:example, :gpu, :off] :fast ∈ TAGS && push!(EXCLUDED_FROM_ALL, :slow) mpi_nprocs() > 1 && push!(EXCLUDED_FROM_ALL, :dont_test_mpi) Sys.iswindows() && push!(EXCLUDED_FROM_ALL, :dont_test_windows) diff --git a/test/serialisation.jl b/test/serialisation.jl index c53b41f643..6c5e803373 100644 --- a/test/serialisation.jl +++ b/test/serialisation.jl @@ -4,6 +4,7 @@ using DFTK function test_scfres_agreement(tested, ref) @test tested.basis.model.lattice == ref.basis.model.lattice + @test tested.basis.model.periodic == ref.basis.model.periodic @test tested.basis.model.temperature == ref.basis.model.temperature @test tested.basis.model.smearing == ref.basis.model.smearing @test tested.basis.model.εF == ref.basis.model.εF From d5975af134988785cfbfde992711780128847366 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89tienne=20Polack?= Date: Tue, 14 Nov 2023 10:55:29 +0100 Subject: [PATCH 2/6] bilayer setup --- src/terms/local.jl | 2 +- test/isolated.jl | 38 +++++++++++++++++++++++++++++++++++++- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/src/terms/local.jl b/src/terms/local.jl index bea230e27e..1b934e369f 100644 --- a/src/terms/local.jl +++ b/src/terms/local.jl @@ -122,7 +122,7 @@ function (::AtomicLocal)(basis::PlaneWaveBasis{T}) where {T} enforce_real!(basis, pot_fourier) # Symmetrize Fourier coeffs to have real iFFT pot_real = irfft(basis, to_device(basis.architecture, pot_fourier)) else - @assert all(.!basis.model.periodic) + @assert all(.!model.periodic) # simple real-space sum pot_real = map(r_vectors(basis)) do r sum(zip(model.atoms, model.positions)) do (element, r_ion) diff --git a/test/isolated.jl b/test/isolated.jl index 9e3430eb01..68a609252d 100644 --- a/test/isolated.jl +++ b/test/isolated.jl @@ -1,4 +1,4 @@ -@testitem "Isolated systems" tags=[:off] begin +@testitem "Isolated systems" begin using DFTK using LinearAlgebra @@ -34,4 +34,40 @@ end end display(sort(RES)) + @test true +end + +@testitem "Graphene surface" tags=[:graphene] begin + using DFTK + using LinearAlgebra + + model_bilayer(; lz=30, periodic=[true for _ in 1:3]) = let + a = 4.66 + lattice = [ 1/2 1/2 0; + -√3/2 √3/2 0; + 0 0 lz/a] .* a + positions = [[1/3, -1/3, -6.45/lz/2], [-1/3, 1/3, -6.45/lz/2], + [1/3, -1/3, 6.45/lz/2], [-1/3, 1/3, 6.45/lz/2]] + atoms=fill(ElementPsp(6, psp=load_psp("hgh/pbe/c-q4")), 4) + model_PBE(lattice, atoms, positions) + end + + RES = Dict() + for lz in (30, 50, 100, 150, 300) + for per in (false, true) + @time begin + kgrid = [4, 4, 1] + Ecut = 20 + tol = 1e-5 + + model = model_bilayer(; lz, periodic=[true, true, per]) + basis = PlaneWaveBasis(model; Ecut, kgrid) + res = self_consistent_field(basis; tol) + + RES[per, lz] = res.energies.total + end + end + end + display(sort(RES)) + @test true end From caf7597c2fd146270160f59ada318eb7e98c199c Mon Sep 17 00:00:00 2001 From: Antoine Levitt Date: Mon, 4 Dec 2023 16:59:55 +0100 Subject: [PATCH 3/6] Extend reference functions to 1D --- src/terms/hartree.jl | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/terms/hartree.jl b/src/terms/hartree.jl index faca8f9f99..4978a66795 100644 --- a/src/terms/hartree.jl +++ b/src/terms/hartree.jl @@ -41,11 +41,19 @@ function TermHartree(basis::PlaneWaveBasis{T}, scaling_factor) where {T} end # a gaussian of exponent α and integral M -ρref_real(r::T, M=1, α=1) where {T} = M * exp(-T(1)/2 * (α*r)^2) / ((2T(π))^(T(3)/2)) * α^3 +ρref_real(r::T, M=1, α=1, d=3) where {T} = M * exp(-T(1)/2 * (α*r)^2) / ((2T(π))^(T(d)/2)) * α^d # solution of -ΔVref = 4π ρref -function Vref_real(r::T, M=1, α=1) where {T} - r == 0 && return M * 2 / sqrt(T(pi)) * α / sqrt(2) - M * erf(α/sqrt(2)*r)/r +function Vref_real(r::T, M=1, α=1, d=3) where {T} + if d == 3 + r == 0 && return M * 2 / sqrt(T(pi)) * α / sqrt(T(2)) + M * erf(α/sqrt(2)*r)/r + elseif d == 1 + ## TODO find a way to compute this more stably when r >> 1 + res = M * (-2π*r*erf(α*r/sqrt(2)) - 2sqrt(2π)*exp(-α^2 * r^2/2)/α) + res + else + error() + end end one_hot(i) = Vec3{Bool}(j == i for j=1:3) From 2ea2e60754109a2a36ed9af7b8648d03158404ec Mon Sep 17 00:00:00 2001 From: Antoine Levitt Date: Mon, 4 Dec 2023 17:04:23 +0100 Subject: [PATCH 4/6] Fixes --- src/terms/hartree.jl | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/src/terms/hartree.jl b/src/terms/hartree.jl index 4978a66795..b2a25759db 100644 --- a/src/terms/hartree.jl +++ b/src/terms/hartree.jl @@ -60,23 +60,20 @@ one_hot(i) = Vec3{Bool}(j == i for j=1:3) ∂f∂α(f, α, r) = ForwardDiff.derivative(ε -> f(r + ε * one_hot(α)), zero(eltype(r))) function get_center(basis, ρ) - sumρ = sum(ρ) + sumρ = sum(abs.(ρ)) center = map(1:3) do α rα = [r[α] for r in r_vectors_cart(basis)] - sum(rα .* ρ) / sumρ + sum(rα .* abs.(ρ)) / sumρ end Vec3(center...) end -function get_integral(basis, ρ) - sum(ρ) * basis.dvol -end function get_dipole(α, center, basis, ρ) rα = [r[α]-center[α] for r in r_vectors_cart(basis)] dot(rα, ρ) * basis.dvol end function get_variance(center, basis, ρ) rr = [sum(abs2, r-center) for r in r_vectors_cart(basis)] - dot(rr, ρ) * basis.dvol + dot(rr, abs.(ρ)) * basis.dvol end @timing "ene_ops: hartree" function ene_ops(term::TermHartree, basis::PlaneWaveBasis{T}, @@ -99,13 +96,13 @@ end # determine center and width from density # Strictly speaking, these computations should result in extra terms to guarantee energy/ham consistency center = get_center(basis, ρtot_real) - ρref_11 = [ρref_real(norm(r - center), model.n_electrons, 1) for r in r_vectors_cart(basis)] + ρref_11 = [ρref_real(norm(r - center), 1, 1, basis.model.n_dim) for r in r_vectors_cart(basis)] spread_11 = get_variance(center, basis, ρref_11) spread_ρ = get_variance(center, basis, ρtot_real) α = 1/sqrt(spread_ρ/spread_11) - # α = 1 - ρrad_fun(r) = ρref_real(norm(r - center), model.n_electrons, α) + total_charge = sum(ρtot_real) .* basis.dvol + ρrad_fun(r) = ρref_real(norm(r - center), 1, α, basis.model.n_dim) ρrad = ρrad_fun.(r_vectors_cart(basis)) @@ -116,13 +113,13 @@ end coeffs_ders = map(1:3) do α get_dipole(α, center, basis, ρtot_real) / get_dipole(α, center, basis, ρders[α]) end - ρref = ρrad + sum([coeffs_ders[α]*ρders[α] for α=1:3]) + ρref = total_charge * ρrad + sum([coeffs_ders[α]*ρders[α] for α=1:3]) # compute corresponding solution of -ΔVref = 4π ρref - Vref_rad_fun(r) = Vref_real(norm(r - center), model.n_electrons, α) + Vref_rad_fun(r) = Vref_real(norm(r - center), 1, α, basis.model.n_dim) Vref_rad = Vref_rad_fun.(r_vectors_cart(basis)) Vref_ders = [∂f∂α.(Vref_rad_fun, α, r_vectors_cart(basis)) for α=1:3] - Vref = Vref_rad + sum([coeffs_ders[α]*Vref_ders[α] for α=1:3]) + Vref = total_charge * Vref_rad + sum([coeffs_ders[α]*Vref_ders[α] for α=1:3]) # TODO possibly optimize FFTs here Vcorr_real = Vref - irfft(basis, term.poisson_green_coeffs .* fft(basis, ρref)) @@ -131,6 +128,8 @@ end pot_fourier .+= Vcorr_fourier end + pot_fourier *= term.scaling_factor + pot_real *= term.scaling_factor E = real(dot(pot_fourier, ρtot_fourier) / 2) ops = [RealSpaceMultiplication(basis, kpt, pot_real) for kpt in basis.kpoints] From 7d6658dfeacde4ef74aeb45ce1df772fbcff337d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89tienne=20Polack?= Date: Tue, 5 Dec 2023 10:28:30 +0100 Subject: [PATCH 5/6] structure for graphene slab --- src/terms/ewald.jl | 2 +- src/terms/hartree.jl | 24 ++++++++------- test/isolated.jl | 70 +++++++++++++++++++++++--------------------- 3 files changed, 50 insertions(+), 46 deletions(-) diff --git a/src/terms/ewald.jl b/src/terms/ewald.jl index 80cc9104b7..9e4b19a5aa 100644 --- a/src/terms/ewald.jl +++ b/src/terms/ewald.jl @@ -137,7 +137,7 @@ function energy_forces_ewald(S, lattice::AbstractArray{T}, charges, positions, q end forces_recip = zeros(Vec3{S}, length(positions)) - if all(periodic) # if non-periodic, no reciprocal-space contribution + if all(periodic) # if non-periodic, no reciprocal-space contribution for G1 in -Glims[1]:Glims[1], G2 in -Glims[2]:Glims[2], G3 in -Glims[3]:Glims[3] G = Vec3(G1, G2, G3) iszero(G) && continue diff --git a/src/terms/hartree.jl b/src/terms/hartree.jl index faca8f9f99..fdd6b5fe81 100644 --- a/src/terms/hartree.jl +++ b/src/terms/hartree.jl @@ -79,17 +79,19 @@ end pot_fourier = term.poisson_green_coeffs .* ρtot_fourier pot_real = irfft(basis, pot_fourier) - # For isolated systems, the above does not compute a good potential (eg it assumes zero DC component) + # For isolated systems, the above does not compute a good potential (e.g. it assumes + # zero DC component). # We correct it by solving -Δ V = 4πρ in two steps: we split ρ into ρref and ρ-ρref, - # where ρref is a gaussian has the same total charge as ρ. - # We compute the first potential in real space (explicitly since ρref is known), - # and the second (short-range) potential in Fourier space - # Compared to the usual scheme (ρref = 0), this results in a correction potential - # equal to Vref computed in real space minus Vref computed in Fourier space + # where ρref is a Gaussian has the same total charge as ρ. + # We compute the first potential in real space (explicitly since ρref is known), and the + # second (short-range) potential in Fourier space. + # Compared to the usual scheme (ρref = 0), this results in a correction potential equal + # to Vref computed in real space minus Vref computed in Fourier space if any(.!model.periodic) @assert all(.!model.periodic) - # determine center and width from density - # Strictly speaking, these computations should result in extra terms to guarantee energy/ham consistency + # Determine center and width from density. + # Strictly speaking, these computations should result in extra terms to guarantee + # energy/ham consistency. center = get_center(basis, ρtot_real) ρref_11 = [ρref_real(norm(r - center), model.n_electrons, 1) for r in r_vectors_cart(basis)] spread_11 = get_variance(center, basis, ρref_11) @@ -101,8 +103,8 @@ end ρrad = ρrad_fun.(r_vectors_cart(basis)) - # at this point we've determined a gaussian of same width and center as the original. - # Now we get the dipole moments of ρ and match them + # At this point we've determined a gaussian of same width and center as the original. + # Now we get the dipole moments of ρ and match them. ρders = [∂f∂α.(ρrad_fun, α, r_vectors_cart(basis)) for α=1:3] coeffs_ders = map(1:3) do α @@ -110,7 +112,7 @@ end end ρref = ρrad + sum([coeffs_ders[α]*ρders[α] for α=1:3]) - # compute corresponding solution of -ΔVref = 4π ρref + # Compute corresponding solution of -ΔVref = 4π ρref. Vref_rad_fun(r) = Vref_real(norm(r - center), model.n_electrons, α) Vref_rad = Vref_rad_fun.(r_vectors_cart(basis)) Vref_ders = [∂f∂α.(Vref_rad_fun, α, r_vectors_cart(basis)) for α=1:3] diff --git a/test/isolated.jl b/test/isolated.jl index 68a609252d..0358d99f6a 100644 --- a/test/isolated.jl +++ b/test/isolated.jl @@ -1,43 +1,44 @@ -@testitem "Isolated systems" begin +@testitem "Isolated systems" tags=[:iso1, :off] begin using DFTK using LinearAlgebra RES = Dict() for a in (10, 20, 30) for per in (false, true) - @time begin - lattice = [a 0 0; - 0 a 0; - 0 0 a] - atoms = [ElementPsp(:Li, psp=load_psp("hgh/lda/Li-q3"))] - atoms = [ElementPsp(:He, psp=load_psp("hgh/lda/He-q2"))] - positions = [[1/2, 1/2, 1/2]] + DFTK.reset_timer!(DFTK.timer) + lattice = [a 0 0; + 0 a 0; + 0 0 a] + atoms = [ElementPsp(:Li, psp=load_psp("hgh/lda/Li-q3"))] + atoms = [ElementPsp(:He, psp=load_psp("hgh/lda/He-q2"))] + positions = [[1/2, 1/2, 1/2]] - kgrid = [1, 1, 1] # no k-point sampling for an isolated system - Ecut = 40 - tol = 1e-6 + kgrid = [1, 1, 1] # no k-point sampling for an isolated system + Ecut = 40 + tol = 1e-6 - model = model_LDA(lattice, atoms, positions, periodic=[per, per, per]) - basis = PlaneWaveBasis(model; Ecut, kgrid) - res = self_consistent_field(basis, is_converged=DFTK.ScfConvergenceDensity(tol)) + model = model_LDA(lattice, atoms, positions, periodic=[per, per, per]) + basis = PlaneWaveBasis(model; Ecut, kgrid) + res = self_consistent_field(basis, is_converged=DFTK.ScfConvergenceDensity(tol)) - rr = vec([norm(a * (r .- 1/2)) for r in r_vectors(basis)]) + rr = vec([norm(a * (r .- 1/2)) for r in r_vectors(basis)]) - function quadrupole(basis, ρ) - rr = [norm(a * (r .- 1/2)) for r in r_vectors(basis)] - sum(rr .^ 2 .* ρ) * basis.dvol - end; - quad = quadrupole(basis, res.ρ) - println(quad) - RES[per, a] = (quad, res.energies.total) - end + function quadrupole(basis, ρ) + rr = [norm(a * (r .- 1/2)) for r in r_vectors(basis)] + sum(rr .^ 2 .* ρ) * basis.dvol + end; + quad = quadrupole(basis, res.ρ) + println(quad) + RES[per, a] = (quad, res.energies.total) + @show per, a + display(DFTK.timer) end end display(sort(RES)) @test true end -@testitem "Graphene surface" tags=[:graphene] begin +@testitem "Graphene surface" tags=[:iso2, :off] begin using DFTK using LinearAlgebra @@ -49,23 +50,24 @@ end positions = [[1/3, -1/3, -6.45/lz/2], [-1/3, 1/3, -6.45/lz/2], [1/3, -1/3, 6.45/lz/2], [-1/3, 1/3, 6.45/lz/2]] atoms=fill(ElementPsp(6, psp=load_psp("hgh/pbe/c-q4")), 4) - model_PBE(lattice, atoms, positions) + model_PBE(lattice, atoms, positions; periodic) end RES = Dict() for lz in (30, 50, 100, 150, 300) for per in (false, true) - @time begin - kgrid = [4, 4, 1] - Ecut = 20 - tol = 1e-5 + DFTK.reset_timer!(DFTK.timer) + kgrid = [4, 4, 1] + Ecut = 20 + tol = 1e-5 - model = model_bilayer(; lz, periodic=[true, true, per]) - basis = PlaneWaveBasis(model; Ecut, kgrid) - res = self_consistent_field(basis; tol) + model = model_bilayer(; lz, periodic=[true, true, per]) + basis = PlaneWaveBasis(model; Ecut, kgrid) + res = self_consistent_field(basis; tol) - RES[per, lz] = res.energies.total - end + RES[per, lz] = res.energies.total + @show per, lz + display(DFTK.timer) end end display(sort(RES)) From fe815c4b9d3445d3974898bcb58941ce999f7cc6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=89tienne=20Polack?= Date: Tue, 5 Dec 2023 10:35:31 +0100 Subject: [PATCH 6/6] adding examples --- src/elements.jl | 9 +++ src/terms/hartree.jl | 12 +++ test/isolated.jl | 186 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 207 insertions(+) diff --git a/src/elements.jl b/src/elements.jl index 839ad131e2..300812b5d4 100644 --- a/src/elements.jl +++ b/src/elements.jl @@ -205,3 +205,12 @@ end function local_potential_fourier(el::ElementGaussian, q::Real) -el.α * exp(- (q * el.L)^2 / 2) # = ∫_ℝ³ V(x) exp(-ix⋅q) dx end +# T@D@ { + +function charge_real(el::Element, r) + zero(r) + end +function charge_fourier(el::Element, q) + zero(q) +end +# } diff --git a/src/terms/hartree.jl b/src/terms/hartree.jl index 538aacc14c..40f423dcf0 100644 --- a/src/terms/hartree.jl +++ b/src/terms/hartree.jl @@ -80,6 +80,12 @@ end ψ, occupation; ρ, kwargs...) where {T} model = basis.model ρtot_real = total_density(ρ) + # T@D@ { + ρtot_real .-= [sum(DFTK.charge_real(el, norm(r - vector_red_to_cart(model, pos))) + for (el, pos) in zip(model.atoms, model.positions)) + for r in r_vectors_cart(basis)] + # println(sum(ρtot_real) * basis.dvol) + # } ρtot_fourier = fft(basis, ρtot_real) pot_fourier = term.poisson_green_coeffs .* ρtot_fourier pot_real = irfft(basis, pot_fourier) @@ -115,6 +121,12 @@ end coeffs_ders = map(1:3) do α get_dipole(α, center, basis, ρtot_real) / get_dipole(α, center, basis, ρders[α]) end + # T@D@ { + # println(coeffs_ders) + if basis.model.n_dim == 1 + coeffs_ders[2:3] .= 0 + end + # } ρref = total_charge * ρrad + sum([coeffs_ders[α]*ρders[α] for α=1:3]) # Compute corresponding solution of -ΔVref = 4π ρref. diff --git a/test/isolated.jl b/test/isolated.jl index 0358d99f6a..40ac53d70a 100644 --- a/test/isolated.jl +++ b/test/isolated.jl @@ -38,6 +38,7 @@ @test true end +# Broken while slab not implemented. @testitem "Graphene surface" tags=[:iso2, :off] begin using DFTK using LinearAlgebra @@ -73,3 +74,188 @@ end display(sort(RES)) @test true end + +@testitem "1d" tags=[:iso3, :off] begin + using DFTK + using LinearAlgebra + using Plots + Plots.scalefontsizes() # reset + Plots.scalefontsizes(2) + pyplot() + close() + + if false + x = 1.234 + M = 2.345 + α = 3.456 + d = 1 + ε = 1e-4 + f(x) = DFTK.Vref_real(x, M, α, d) + fpp(x) = -(f(x+ε)+f(x-ε)-2f(x))/(ε^2) / (4π) + xs = range(0, 5, 100) + p = plot(xs, fpp.(xs)) + plot!(p, xs, DFTK.ρref_real.(xs, M, α, d)) + display(p) + println(-(f(x+ε)+f(x-ε)-2f(x))/(ε^2) - 4π*DFTK.ρref_real(x, M, α, d)) + end + + + struct CustomPotential <: DFTK.Element + α # Prefactor + L # Width of the Gaussian nucleus + + α_charge + L_charge + end + + # Some default values + CustomPotential() = CustomPotential(1.0, 0.5, 1.0, 0.5); + + function DFTK.local_potential_real(el::CustomPotential, r::Real) + -el.α / (√(2π) * el.L) * exp(- (r / el.L)^2 / 2) + end + function DFTK.local_potential_fourier(el::CustomPotential, q::Real) + -el.α * exp(- (q * el.L)^2 / 2) + end + + function DFTK.charge_real(el::CustomPotential, r::Real) + el.α_charge / (√(2π) * el.L_charge) * exp(- (r / el.L_charge)^2 / 2) + end + function DFTK.charge_fourier(el::CustomPotential, q::Real) + el.α_charge * exp(- (q * el.L_charge)^2 / 2) + end + + + a = 10 + lattice = a .* [[1 0 0.]; [0 0 0]; [0 0 0]]; + + # x1 = 0.2 + # x2 = 0.8 + # positions = [[x1, 0, 0], [x2, 0, 0]] + gauss = CustomPotential(1, 0.5, 1, .5) + positions = [[0.5, 0, 0]] + atoms = [gauss]; + + + gauss = CustomPotential(1, 0.5, .333333, .5) + δ = 2 + positions = [[1/2+δ/a, 0, 0], + [1/2+δ/a, 0, 0], + [1/2-δ/a, 0, 0]] + atoms = [gauss, gauss, gauss]; + + # We setup a Gross-Pitaevskii model + n_electrons = 1 # Increase this for fun + terms = [Kinetic(), + AtomicLocal(), + Hartree(scaling_factor=.4), + # LocalNonlinearity(ρ -> C * ρ^α) + ] + + xlims = (-6, 6) + p = plot(; xlims, reuse=false, legend=:bottomleft) + for per in (true, false) + periodic = fill(per, 3) + model = Model(lattice, atoms, positions; n_electrons, terms, + spin_polarization=:spinless, periodic) + + # We discretize using a moderate Ecut and run a SCF algorithm to compute forces + # afterwards. As there is no ionic charge associated to `gauss` we have to specify + # a starting density and we choose to start from a zero density. + basis = PlaneWaveBasis(model; Ecut=500, kgrid=(1, 1, 1)) + ρ = ones(eltype(basis), basis.fft_size..., 1) + scfres = self_consistent_field(basis; tol=1e-8, ρ, maxiter=1000) + + ρ = scfres.ρ[:, 1, 1, 1] # converged density, first spin component + + x = vec(first.(DFTK.r_vectors_cart(basis))) + atomic_pot = (scfres.ham).blocks[1].operators[2].potential[:, 1, 1]; # use only dimension 1 + hartree_pot = (scfres.ham).blocks[1].operators[3].potential[:, 1, 1]; # use only dimension 1 + # plot(x, atomic_pot, label="Vat") + dash = per ? :dash : :dot + plot!(p, x.-a/2, ρ; ls=dash, label="ρ per=$per a=$a", xlims, lw=3) + plot!(p, x.-a/2, atomic_pot; ls=dash, label="Vnucl per=$per a=$a", lw=3) + plot!(p, x.-a/2, hartree_pot; ls=dash, label="Vh per=$per a=$a", lw=3) + end + display(p) +end + +@testitem "1D toy system" tags=[:iso4, :off] begin + using FFTW + using Plots + Plots.scalefontsizes() # reset + Plots.scalefontsizes(2) + pyplot() # for interactive plots + close() + + N = 1000 + L = 10 + x = range(-L, L, N) + ρ = @. exp(-(x) ^ 2) * (0.5 - (x)^2 + x) + Vconv = map(1:N) do i + -sum(abs(x[i] - x[j])*ρ[j] for j=1:N) * (2L/N) / 2 + end + + ks = map(1:N) do i + K = i < N/2 ? i-1 : i-N - 1 + K * (2π/(2L)) + end + + poisson = @. 1 / (ks^2) + poisson[1] = 0 + Vfft = real(ifft(poisson .* fft(ρ))) + + Vfft .-= Vfft[N÷2 + 1] + Vconv .-= Vconv[N÷2 + 1] + + p = plot() + p = plot!(p, x, ρ; label="ρ", lw=3, ls=:dot) + p = plot!(p, x, Vconv; label="V_{ref}", lw=3) + plot!(p, x, Vfft; label="V_{fft}", lw=3) + + # xlims = (-3,3) + # ylims = (-5, 2) + # plot!(p; xlims, ylims) + + display(p) +end + +@testitem "1D ElementGaussian" tags=[:iso5, :off] begin + using Plots + using LinearAlgebra + Plots.scalefontsizes() # reset + Plots.scalefontsizes(2) + pyplot() # for interactive plots + close() + p = plot() + + terms = [Kinetic(), AtomicLocal()] + + n_atoms = 1 + + if n_atoms == 1 + X = ElementGaussian(1.0, 0.5, :X) + atoms = [X] + positions = [[0.0, 0.0, 0.0]] + else + X = ElementGaussian(1.0, 0.5, :X) + Y = ElementGaussian(1.0, 0.2, :Y) + atoms = [X, Y] + positions = [[0.0, 0.0, 0.0], [1/4, 0, 0]] + end + n_atoms = length(atoms) + lattice = diagm([5.0 * n_atoms, 0, 0]) + model = Model(lattice, atoms, positions; n_electrons=2, terms) + basis = PlaneWaveBasis(model; Ecut=500, kgrid=[4,1,1]) + scfres = self_consistent_field(basis) + ham = scfres.ham + ρ = dropdims(scfres.ρ; dims=(2,3,4)) + + potential = dropdims(DFTK.total_local_potential(ham); dims=(2,3,4)) + xs = [r[1] for r in dropdims(r_vectors(basis); dims=(2,3))] + + plot!(p, xs, ρ; label="ρ", lw=3, ls=:dot) + plot!(p, xs, potential; label="V", lw=3) + + display(p) +end