From ad3da5002417fc28cd9f7778953074a9105390ec Mon Sep 17 00:00:00 2001 From: xquan818 <840169780@qq.com> Date: Wed, 16 Oct 2024 18:35:38 +0200 Subject: [PATCH] refactor addition --- examples/dos.jl | 3 +-- ext/DFTKPlotsExt.jl | 3 +-- src/postprocess/dos.jl | 46 ++++++++++++++++----------------- src/pseudo/NormConservingPsp.jl | 39 +++++++++++++++++++--------- src/pseudo/PspUpf.jl | 12 ++------- src/terms/nonlocal.jl | 19 +++++++------- 6 files changed, 64 insertions(+), 58 deletions(-) diff --git a/examples/dos.jl b/examples/dos.jl index bca372055..4af8b06a3 100644 --- a/examples/dos.jl +++ b/examples/dos.jl @@ -1,7 +1,6 @@ # # Densities of states (DOS) -# In this example, we'll plot the DOS, local DOS, and projected DOS of Silicon. +# In this example, we'll plot the DOS, local DOS, and projected DOS of Silicon. # DOS computation only supports finite temperature. -# Projected DOS only supports PspUpf. using DFTK using Unitful diff --git a/ext/DFTKPlotsExt.jl b/ext/DFTKPlotsExt.jl index ed19e44b6..4afa2b361 100644 --- a/ext/DFTKPlotsExt.jl +++ b/ext/DFTKPlotsExt.jl @@ -170,8 +170,7 @@ function plot_pdos(scfres; kwargs...) # Plot DOS p = plot_dos(scfres; scfres.εF, kwargs...) - # TODO Require symmetrization with respect to kpoints and BZ symmetry - # (now achieved by unfolding all the quantities). + # TODO do the symmetrization instead of unfolding scfres_unfold = DFTK.unfold_bz(scfres) basis = scfres_unfold.basis psp_groups = [group for group in basis.model.atom_groups diff --git a/src/postprocess/dos.jl b/src/postprocess/dos.jl index 9f9a49da3..6e65f4f33 100644 --- a/src/postprocess/dos.jl +++ b/src/postprocess/dos.jl @@ -8,9 +8,9 @@ # LDOS (local density of states) # LD(ε) = sum_n f_n' |ψn|^2 = sum_n δ(ε - ε_n) |ψn|^2 # -# PDOS (projected density of states) -# PD(ε) = sum_n f_n' |<ψn,ϕlmi>|^2 -# ϕlmi = ∑_R ϕlmi(r-pos-R) is a periodic atomic wavefunction centered at pos +# PD(ε) = sum_n f_n' |<ψn,ϕ>|^2 +# ϕ = ∑_R ϕilm(r-pos-R) is the periodized atomic wavefunction, obtained from the pseudopotential +# This is computed for a given (i,l) (eg i=2,l=2 for the 3p) and summed over all m """ Total density of states at energy ε @@ -94,36 +94,36 @@ function compute_pdos(ε, basis, eigenvalues, ψ, i, l, psp, position; D = mpi_sum(D, basis.comm_kpts) end -function compute_pdos(scfres::NamedTuple; ε=scfres.εF, l=0, i=1, - psp=scfres.basis.model.atoms[1].psp, - position=basis.model.positions[1], kwargs...) - # TODO Require symmetrization with respect to kpoints and BZ symmetry - # (now achieved by unfolding all the quantities). +function compute_pdos(scfres::NamedTuple, iatom, i, l; ε=scfres.εF, kwargs...) + psp = scfres.basis.model.atoms[iatom].psp + position = scfres.basis.model.positions[iatom] + # TODO do the symmetrization instead of unfolding scfres = unfold_bz(scfres) compute_pdos(ε, scfres.basis, scfres.eigenvalues, scfres.ψ, i, l, psp, position; kwargs...) end -# Build projection vector for a single atom. -# Stored as projs[K][nk,lmi] = |<ψnk, f_m>|^2, -# where K is running over all k-points, -# m is running over -l:l for fixed l and i. -function compute_pdos_projs(basis, eigenvalues, ψ, i, l, psp::PspUpf, position) - position = vector_red_to_cart(basis.model, position) - G_plus_k_all = [to_cpu(Gplusk_vectors_cart(basis, basis.kpoints[ik])) +# Build atomic orbitals projections projs[ik][iband,m] = |<ψnk, ϕ>|^2 for a single atom. +# TODO optimization ? accept a range of positions, in case we want to compute the PDOS +# for all atoms of the same type (and reuse the computation of the atomic orbitals) +function compute_pdos_projs(basis, eigenvalues, ψ, i, l, psp::NormConservingPsp, position) + # Precompute the form factors on all kpoints at once so we can better use the cache (memory-compute tradeoff). + # Revisit this (pass the cache around explicitly) if RAM becomes an issue. + G_plus_k_all = [Gplusk_vectors(basis, basis.kpoints[ik]) for ik = 1:length(basis.kpoints)] + G_plus_k_all_cart = [map(recip_vector_red_to_cart(basis.model), gpk) + for gpk in G_plus_k_all] - # Build Fourier transform factors of pseudo-wavefunctions centered at 0. + # Build form factors of pseudo-wavefunctions centered at 0. fun(p) = eval_psp_pswfc_fourier(psp, i, l, p) - fourier_form = atomic_centered_fun_form_factors(fun, l, G_plus_k_all) + # form_factors_all[ik][p,m] + form_factors_all = build_form_factors(fun, l, G_plus_k_all_cart) projs = Vector{Matrix}(undef, length(basis.kpoints)) for (ik, ψk) in enumerate(ψ) - fourier_form_ik = fourier_form[ik] - structure_factor_ik = exp.(-im .* [dot(position, Gik) for Gik in G_plus_k_all[ik]]) - - # TODO Pseudo-atomic wave functions need to be orthogonalized. - proj_vectors_ik = structure_factor_ik .* fourier_form_ik ./ sqrt(basis.model.unit_cell_volume) - projs[ik] = abs2.(ψk' * proj_vectors_ik) + structure_factor = [cis2pi(-dot(position, p)) for p in G_plus_k_all[ik]] + # TODO orthogonalize pseudo-atomic wave functions? + proj_vectors = structure_factor .* form_factors_all[ik] ./ sqrt(basis.model.unit_cell_volume) + projs[ik] = abs2.(ψk' * proj_vectors) # contract on p -> projs[ik][iband,m] end projs diff --git a/src/pseudo/NormConservingPsp.jl b/src/pseudo/NormConservingPsp.jl index 6e8a1be1a..5746a070b 100644 --- a/src/pseudo/NormConservingPsp.jl +++ b/src/pseudo/NormConservingPsp.jl @@ -11,20 +11,24 @@ abstract type NormConservingPsp end # description::String # Descriptive string #### Methods: -# charge_ionic(psp::NormConservingPsp) -# has_valence_density(psp:NormConservingPsp) -# has_core_density(psp:NormConservingPsp) -# eval_psp_projector_real(psp::NormConservingPsp, i, l, r::Real) -# eval_psp_projector_fourier(psp::NormConservingPsp, i, l, p::Real) -# eval_psp_local_real(psp::NormConservingPsp, r::Real) -# eval_psp_local_fourier(psp::NormConservingPsp, p::Real) -# eval_psp_energy_correction(T::Type, psp::NormConservingPsp, n_electrons::Integer) +# charge_ionic(psp) +# has_valence_density(psp) +# has_core_density(psp) +# eval_psp_projector_real(psp, i, l, r::Real) +# eval_psp_projector_fourier(psp, i, l, p::Real) +# eval_psp_local_real(psp, r::Real) +# eval_psp_local_fourier(psp, p::Real) +# eval_psp_energy_correction(T::Type, psp, n_electrons::Integer) #### Optional methods: -# eval_psp_density_valence_real(psp::NormConservingPsp, r::Real) -# eval_psp_density_valence_fourier(psp::NormConservingPsp, p::Real) -# eval_psp_density_core_real(psp::NormConservingPsp, r::Real) -# eval_psp_density_core_fourier(psp::NormConservingPsp, p::Real) +# eval_psp_density_valence_real(psp, r::Real) +# eval_psp_density_valence_fourier(psp, p::Real) +# eval_psp_density_core_real(psp, r::Real) +# eval_psp_density_core_fourier(psp, p::Real) +# eval_psp_pswfc_real(psp, i::Int, l::Int, p::Real) +# eval_psp_pswfc_fourier(psp, i::Int, l::Int, p::Real) +# count_n_pswfc(psp, l::Integer) +# count_n_pswfc_radial(psp, l::Integer) """ eval_psp_projector_real(psp, i, l, r) @@ -203,3 +207,14 @@ function count_n_proj(psps, psp_positions) sum(count_n_proj(psp) * length(positions) for (psp, positions) in zip(psps, psp_positions)) end + +count_n_pswfc_radial(psp::NormConservingPsp, l) = error("Pseudopotential $psp does not implement atomic wavefunctions.") + +function count_n_pswfc_radial(psp::NormConservingPsp) + sum(l -> count_n_pswfc_radial(psp, l), 0:psp.lmax; init=0)::Int +end + +count_n_pswfc(psp::NormConservingPsp, l) = count_n_pswfc_radial(psp, l) * (2l + 1) +function count_n_pswfc(psp::NormConservingPsp) + sum(l -> count_n_pswfc(psp, l), 0:psp.lmax; init=0)::Int +end diff --git a/src/pseudo/PspUpf.jl b/src/pseudo/PspUpf.jl index 287625e93..6ebcc127f 100644 --- a/src/pseudo/PspUpf.jl +++ b/src/pseudo/PspUpf.jl @@ -171,6 +171,8 @@ function eval_psp_projector_fourier(psp::PspUpf, i, l, p::T)::T where {T<:Real} hankel(rgrid, r2_proj, l, p) end +count_n_pswfc_radial(psp::PspUpf, l) = length(psp.r2_pswfcs[l+1]) + function eval_psp_pswfc_real(psp::PspUpf, i, l, r::T)::T where {T<:Real} psp.r2_pswfcs_interp[l+1][i](r) / r^2 end @@ -227,13 +229,3 @@ function eval_psp_energy_correction(T, psp::PspUpf, n_electrons) r * (r * vloc[i] - -psp.Zion) end end - -count_n_pswfc_radial(psp::PspUpf, l::Integer) = length(psp.r2_pswfcs[l + 1]) -function count_n_pswfc_radial(psp::PspUpf) - sum(l -> count_n_pswfc_radial(psp, l), 0:psp.lmax; init=0)::Int -end - -count_n_pswfc(psp::PspUpf, l::Integer) = count_n_pswfc_radial(psp, l) * (2l + 1) -function count_n_pswfc(psp::PspUpf) - sum(l -> count_n_pswfc(psp, l), 0:psp.lmax; init=0)::Int -end diff --git a/src/terms/nonlocal.jl b/src/terms/nonlocal.jl index 4f1e81bfb..aba78c5f3 100644 --- a/src/terms/nonlocal.jl +++ b/src/terms/nonlocal.jl @@ -70,7 +70,7 @@ end # We compute the forces from the irreductible BZ; they are symmetrized later. G_plus_k = Gplusk_vectors(basis, kpt) G_plus_k_cart = to_cpu(Gplusk_vectors_cart(basis, kpt)) - form_factors = build_form_factors(element.psp, G_plus_k_cart) + form_factors = build_projector_form_factors(element.psp, G_plus_k_cart) for idx in group r = model.positions[idx] structure_factors = [cis2pi(-dot(p, r)) for p in G_plus_k] @@ -173,7 +173,7 @@ function build_projection_vectors(basis::PlaneWaveBasis{T}, kpt::Kpoint, for (psp, positions) in zip(psps, psp_positions) # Compute position-independent form factors G_plus_k_cart = to_cpu(Gplusk_vectors_cart(basis, kpt)) - form_factors = build_form_factors(psp, G_plus_k_cart) + form_factors = build_projector_form_factors(psp, G_plus_k_cart) # Combine with structure factors for r in positions @@ -195,8 +195,8 @@ end """ Build form factors (Fourier transforms of projectors) for all orbitals of an atom centered at 0. """ -function build_form_factors(psp::NormConservingPsp, - G_plus_k::AbstractVector{Vec3{TT}}) where {TT} +function build_projector_form_factors(psp::NormConservingPsp, + G_plus_k::AbstractVector{Vec3{TT}}) where {TT} G_plus_ks = [G_plus_k] n_proj = count_n_proj(psp) @@ -204,10 +204,10 @@ function build_form_factors(psp::NormConservingPsp, for l = 0:psp.lmax, n_proj_l = count_n_proj_radial(psp, l) offset = sum(x -> count_n_proj(psp, x), 0:l-1; init=0) .+ - n_proj_l .* (collect(1:2l+1) .- 1) # offset about m for given l and i + n_proj_l .* (collect(1:2l+1) .- 1) # offset about m for a given l for i = 1:n_proj_l proj_li(p) = eval_psp_projector_fourier(psp, i, l, p) - form_factors_li = atomic_centered_fun_form_factors(proj_li, l, G_plus_ks) + form_factors_li = build_form_factors(proj_li, l, G_plus_ks) @views form_factors[:, offset.+i] = form_factors_li[1] end end @@ -216,10 +216,11 @@ function build_form_factors(psp::NormConservingPsp, end """ -Build Fourier transform factors of a atomic function centered at 0 for a given l. +Build Fourier transform factors of an atomic function centered at 0 for a given l. """ -function atomic_centered_fun_form_factors(fun::Function, l::Int, - G_plus_ks::AbstractVector{<:AbstractVector{Vec3{TT}}}) where {TT} +function build_form_factors(fun::Function, l::Int, + G_plus_ks::AbstractVector{<:AbstractVector{Vec3{TT}}}) where {TT} + # TODO this function can be generally useful, should refactor to a separate file eventually T = real(TT) # Pre-compute the radial parts of the non-local atomic functions at unique |p| to speed up