Reupload the projected density of states

[xuequan818]

No description provided.

[antoine-levitt]

@mfherbst how is this going on your end? Can we merge this or are there conflicts with ongoing work?

[antoine-levitt]

Don't take i as input here

[xuequan818]

I have modified the inputs. Please check.

[antoine-levitt]

Just call fun(p_norm) here

[antoine-levitt]

Don't take the count functions or the psp_fun as input, they are implied by psp (just as it was before). Type psp as NormConservingPsp

[antoine-levitt]

I went through the code and did the following refactor (but I did not test it). Could you apply it, test it and commit it? After this is done I think we can merge, unless @mfherbst objects

diff --git a/src/postprocess/dos.jl b/src/postprocess/dos.jl
index 9f9a49da3..916df02d5 100644
--- a/src/postprocess/dos.jl
+++ b/src/postprocess/dos.jl
@@ -9,8 +9,9 @@
 # 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 +95,39 @@ 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, 
+function compute_pdos(scfres::NamedTuple, iatom, i, l; ε=scfres.εF, 
                       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).
+    psp = scfres.basis.model.atoms[iatom]
+    position=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)
+# 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)
     position = vector_red_to_cart(basis.model, position)
-    G_plus_k_all = [to_cpu(Gplusk_vectors_cart(basis, basis.kpoints[ik])) 
+    
+    # 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 = atomic_centered_fun_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_ik .* 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..c55bd03b4 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..688931773 100644
--- a/src/pseudo/PspUpf.jl
+++ b/src/pseudo/PspUpf.jl
@@ -171,6 +171,7 @@ 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 +228,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..23aa3d820 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)
@@ -207,7 +207,7 @@ function build_form_factors(psp::NormConservingPsp,
                  n_proj_l .* (collect(1:2l+1) .- 1) # offset about m for given l and i
         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, l,
+                            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

[xuequan818]

No problem.I will check it right now.

[antoine-levitt]

Thanks for all the work!