Reupload the projected density of states

examples/dos.jl

@@ -0,0 +1,32 @@
+# # Densities of states (DOS)
+# 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
+using Plots
+using LazyArtifacts
+
+## Define the geometry and pseudopotential
+a = 10.26  # Silicon lattice constant in Bohr
+lattice = a / 2 * [[0 1 1.0];
+    [1 0 1.0];
+    [1 1 0.0]]
+Si = ElementPsp(:Si; psp=load_psp(artifact"pd_nc_sr_lda_standard_0.4.1_upf", "Si.upf"))
+atoms = [Si, Si]
+positions = [ones(3) / 8, -ones(3) / 8]
+
+## Run SCF
+model = model_LDA(lattice, atoms, positions; temperature=5e-3)
+basis = PlaneWaveBasis(model; Ecut=15, kgrid=[4, 4, 4], symmetries_respect_rgrid=true)
+scfres = self_consistent_field(basis, tol=1e-8)
+
+## Plot the DOS
+plot_dos(scfres)
+
+## Plot the local DOS along one direction
+plot_ldos(scfres; n_points=100, ldos_xyz=[:, 10, 10])
+
+## Plot the projected DOS
+plot_pdos(scfres; εrange=(-0.3, 0.5))

ext/DFTKPlotsExt.jl

@@ -2,7 +2,7 @@
 using Brillouin: KPath
 using DFTK
 using DFTK: is_metal, data_for_plotting, spin_components, default_band_εrange
-import DFTK: plot_dos, plot_bandstructure
+import DFTK: plot_dos, plot_bandstructure, plot_ldos, plot_pdos
 using Plots
 using Unitful
 using UnitfulAtomic
@@ -108,4 +108,98 @@
 end
 plot_dos(scfres; kwargs...) = plot_dos(scfres.basis, scfres.eigenvalues; scfres.εF, kwargs...)
 
+function plot_ldos(basis, eigenvalues, ψ; εF=nothing, unit=u"hartree",
+                   temperature=basis.model.temperature,
+                   smearing=basis.model.smearing,
+                   εrange=default_band_εrange(eigenvalues; εF), 
+                   n_points=1000, ldos_xyz=[:, 1, 1], kwargs...)
+    eshift = something(εF, 0.0)
+    εs = range(austrip.(εrange)..., length=n_points)
+
+    # Constant to convert from AU to the desired unit
+    to_unit = ustrip(auconvert(unit, 1.0))
+
+    # LDε has three dimensions (x, y, z)
+    # map on a single axis to plot the variation with εs
+    LDεs = dropdims.(compute_ldos.(εs, Ref(basis), Ref(eigenvalues), Ref(ψ); smearing, temperature); dims=4)
+    LDεs_slice = similar(LDεs[1], n_points, length(LDεs[1][ldos_xyz...]))
+    for (i, LDε) in enumerate(LDεs)
+        LDεs_slice[i, :] = LDε[ldos_xyz...]
+    end
+    p = heatmap(1:size(LDεs_slice, 2), (εs .- eshift) .* to_unit, LDεs_slice; kwargs...)
+    if !isnothing(εF)
+        Plots.hline!(p, [0.0], label="εF", color=:green, lw=1.5)
+    end
+
+    if isnothing(εF)
+        Plots.ylabel!(p, "eigenvalues  ($(string(unit)))")
+    else
+        Plots.ylabel!(p, "eigenvalues -ε_F  ($(string(unit)))")
+    end
+    p
+end
+plot_ldos(scfres; kwargs...) = plot_ldos(scfres.basis, scfres.eigenvalues, scfres.ψ; scfres.εF, kwargs...)
+
+function plot_pdos(basis, eigenvalues, ψ, psp_group;
+                   εF=nothing, unit=u"hartree",
+                   temperature=basis.model.temperature,
+                   smearing=basis.model.smearing,
+                   εrange=default_band_εrange(eigenvalues; εF), 
+                   n_points=1000, p=nothing, kwargs...)
+    eshift = something(εF, 0.0)
+    εs = range(austrip.(εrange)..., length=n_points)
+
+    # Constant to convert from AU to the desired unit
+    to_unit = ustrip(auconvert(unit, 1.0))
+
+    # Calculate the projections of the first atom in the atom group
+    pdos = compute_pdos(εs, basis, eigenvalues, ψ, psp_group; 
+                        temperature, smearing)
+
+    # Summing all angular components for each atom orbital
+    psp = basis.model.atoms[first(psp_group)].psp
+    n_pswfc = DFTK.count_n_pswfc_radial(psp)
+    pdos_label = Vector{String}(undef, n_pswfc)
+    pdos_orbital = Vector{Vector{eltype(pdos)}}(undef, n_pswfc)
+    count = 1
+    for l = 0:psp.lmax
+        il_n = DFTK.count_n_pswfc_radial(psp, l)
+        mfet_count = sum(x -> DFTK.count_n_pswfc(psp,x), 0:l-1; init=1) 
+        for il = 1:il_n
+            mfet = mfet_count .+ il_n .* (collect(1:2l+1) .- 1)
+            pdos_label[count] = psp.pswfc_labels[l+1][il]
+            pdos_orbital[count] = dropdims(sum(pdos[:, mfet], dims=2); dims=2)
+            count += 1
+            mfet_count += 1
+        end
+    end
+    pdos_label = (String(basis.model.atoms[first(psp_group)].symbol) * "-") .* pdos_label
+
+    # Plot pdos for each orbital
+    p = something(p, Plots.plot(; kwargs...))
+    for i = 1:n_pswfc
+        Plots.plot!(p, (εs .- eshift) .* to_unit, pdos_orbital[i];
+                    label=pdos_label[i])
+    end
+    p
+end
+
+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).
+    scfres_unfold = DFTK.unfold_bz(scfres)
+    psp_groups = [group for group in scfres_unfold.basis.model.atom_groups
+                  if scfres_unfold.basis.model.atoms[first(group)] isa ElementPsp]
+
+    # Plot PDOS
+    for group in psp_groups
+        plot_pdos(scfres_unfold.basis, scfres_unfold.eigenvalues,
+                  scfres_unfold.ψ, group; scfres.εF, p, kwargs...)
+    end
+    p
+end
+
 end

src/DFTK.jl

@@ -212,7 +212,10 @@ export compute_stresses_cart
 include("postprocess/stresses.jl")
 export compute_dos
 export compute_ldos
+export compute_pdos
 export plot_dos
+export plot_ldos
+export plot_pdos
 include("postprocess/dos.jl")
 export compute_χ0
 export apply_χ0

src/postprocess/dos.jl

@@ -7,6 +7,10 @@
 #
 # 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
 
 """
 Total density of states at energy ε
@@ -63,7 +67,78 @@
     compute_ldos(ε, scfres.basis, scfres.eigenvalues, scfres.ψ; kwargs...)
 end
 
+"""
+Projected density of states at energy ε for the first atom in the atom group.
+"""
+function compute_pdos(ε, basis, eigenvalues, ψ, psp_group; 
+                      smearing=basis.model.smearing,
+                      temperature=basis.model.temperature)
+    if (temperature == 0) || smearing isa Smearing.None
+        error("compute_pdos only supports finite temperature")
+    end
+    filled_occ = filled_occupation(basis.model)
+
+    # Calculate the projections of the first atom in the atom group
+    psp = basis.model.atoms[first(psp_group)].psp
+    position = basis.model.positions[first(psp_group)]
+    projs = compute_pdos_projs(basis, eigenvalues, ψ, psp, position)
+
+    D = zeros(typeof(ε[1]), length(ε), size(projs[1], 2))
+    for (i, iε) in enumerate(ε)
+        for (ik, projk) in enumerate(projs)
+            @views for (iband, εnk) in enumerate(eigenvalues[ik])
+                enred = (εnk - iε) / temperature
+                D[i, :] .-= (filled_occ .* basis.kweights[ik] .* projk[iband, :]
+                             ./ temperature
+                             .* Smearing.occupation_derivative(smearing, enred))       
+            end
+        end
+    end
+    D = mpi_sum(D, basis.comm_kpts)
+end
+
+function compute_pdos(scfres::NamedTuple; ε=scfres.εF, kwargs...)
+    # TODO Require symmetrization with respect to kpoints and BZ symmetry 
+    #      (now achieved by unfolding all the quantities).
+    scfres = unfold_bz(scfres)
+    psp_groups = [group for group in scfres.basis.model.atom_groups
+                  if scfres.basis.model.atoms[first(group)] isa ElementPsp]
+    [compute_pdos(ε, scfres.basis, scfres.eigenvalues, scfres.ψ, group; kwargs...)
+     for group in psp_groups]
+end
+
+# Build projection matrix for a single atom
+# Stored as projs[K][nk,lmi] = |<ψnk, f_lmi>|^2,
+# where K is running over all k-points, l, m are AM quantum numbers, 
+# and i is running over all pseudo-atomic wavefunctions for a given l.
+function compute_pdos_projs(basis, eigenvalues, ψ, psp::PspUpf, position)
+    position = vector_red_to_cart(basis.model, position)
+    G_plus_k_all = [to_cpu(Gplusk_vectors_cart(basis, basis.kpoints[ik])) 
+                    for ik = 1:length(basis.kpoints)]
+
+    # Build Fourier transform factors of pseudo-wavefunctions centered at 0.
+    fourier_form = atomic_centered_function_form_factors(psp, eval_psp_pswfc_fourier,
+                                    count_n_pswfc_radial, count_n_pswfc, G_plus_k_all)
+
+    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]])
+        @views for iproj = 1:size(fourier_form_ik, 2)
+            # TODO Pseudo-atomic wave functions need to be orthogonalized.
+            fourier_form_ik[:, iproj] .= (structure_factor_ik .* fourier_form_ik[:, iproj]
+                                          ./ sqrt(basis.model.unit_cell_volume))
+        end
+        projs[ik] = abs2.(ψk' * fourier_form_ik)
+    end
+    projs
+end
+
 """
 Plot the density of states over a reasonable range. Requires to load `Plots.jl` beforehand.
 """
 function plot_dos end
+
+function plot_ldos end
+
+function plot_pdos end

src/pseudo/PspUpf.jl

@@ -227,3 +227,17 @@
         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
+function count_n_pswfc(psps, psp_positions)
+    sum(count_n_pswfc(psp) * length(positions)
+        for (psp, positions) in zip(psps, psp_positions))
+end

src/terms/nonlocal.jl

@@ -196,44 +196,64 @@ end
 Build form factors (Fourier transforms of projectors) for an atom centered at 0.
 """
 function build_form_factors(psp, G_plus_k::AbstractVector{Vec3{TT}}) where {TT}
+    atomic_centered_function_form_factors(psp, eval_psp_projector_fourier, 
+                                          count_n_proj_radial, count_n_proj, [G_plus_k])[1]
+end
+
+"""
+Build Fourier transform factors of a atomic function centered at 0.
+"""
+function atomic_centered_function_form_factors(psp, psp_fun::Function,                     
+                                               count_n_fun_radial::Function, 
+                                               count_n_fun::Function,
+                                               G_plus_ks::AbstractVector{<:AbstractVector{Vec3{TT}}}) where {TT}
     T = real(TT)
 
-    # Pre-compute the radial parts of the non-local projectors at unique |p| to speed up
+    # Pre-compute the radial parts of the non-local atomic functions at unique |p| to speed up
     # the form factor calculation (by a lot). Using a hash map gives O(1) lookup.
 
-    # Maximum number of projectors over angular momenta so that form factors
-    # for a given `p` can be stored in an `nproj x (lmax + 1)` matrix.
-    n_proj_max = maximum(l -> count_n_proj_radial(psp, l), 0:psp.lmax; init=0)
+    # Maximum number of atomic functions over angular momenta so that form factors
+    # for a given `p` can be stored in an `nfun x (lmax + 1)` matrix.
+    n_fun_per_l = map(l -> count_n_fun_radial(psp, l), 0:psp.lmax)
+    n_fun_max = maximum(n_fun_per_l)
 
     radials = IdDict{T,Matrix{T}}()  # IdDict for Dual compatibility
-    for p in G_plus_k
-        p_norm = norm(p)
-        if !haskey(radials, p_norm)
-            radials_p = Matrix{T}(undef, n_proj_max, psp.lmax + 1)
-            for l = 0:psp.lmax, iproj_l = 1:count_n_proj_radial(psp, l)
-                # TODO This might  be faster if we do this in batches of l
-                #      (i.e. make the inner loop run over k-points and G_plus_k)
-                #      and did recursion over l to compute the spherical bessels
-                radials_p[iproj_l, l+1] = eval_psp_projector_fourier(psp, iproj_l, l, p_norm)
+    for G_plus_k in G_plus_ks
+        for p in G_plus_k
+            p_norm = norm(p)
+            if !haskey(radials, p_norm)
+                radials_p = Matrix{T}(undef, n_fun_max, psp.lmax + 1)
+                for l = 0:psp.lmax, ifun_l = 1:count_n_fun_radial(psp, l)
+                    # TODO This might  be faster if we do this in batches of l
+                    #      (i.e. make the inner loop run over k-points and G_plus_k)
+                    #      and did recursion over l to compute the spherical bessels
+                    radials_p[ifun_l, l+1] = psp_fun(psp, ifun_l, l, p_norm)
+                end
+                radials[p_norm] = radials_p
             end
-            radials[p_norm] = radials_p
         end
     end
 
-    form_factors = Matrix{Complex{T}}(undef, length(G_plus_k), count_n_proj(psp))
-    for (ip, p) in enumerate(G_plus_k)
-        radials_p = radials[norm(p)]
-        count = 1
-        for l = 0:psp.lmax, m = -l:l
-            # see "Fourier transforms of centered functions" in the docs for the formula
-            angular = (-im)^l * ylm_real(l, m, p)
-            for iproj_l = 1:count_n_proj_radial(psp, l)
-                form_factors[ip, count] = radials_p[iproj_l, l+1] * angular
-                count += 1
+    form_factors = Vector{Matrix{Complex{T}}}(undef, length(G_plus_ks))
+    n_fun = count_n_fun(psp)
+    for (ik, G_plus_k) in enumerate(G_plus_ks)
+        form_factors_ik = Matrix{Complex{T}}(undef, length(G_plus_k), n_fun)
+        for (ip, p) in enumerate(G_plus_k)
+            radials_p = radials[norm(p)]
+            count = 1
+            for l = 0:psp.lmax, m = -l:l
+                # see "Fourier transforms of centered functions" in the docs for the formula
+                angular = (-im)^l * ylm_real(l, m, p)
+                for ifun_l = 1:n_fun_per_l[l+1]
+                    form_factors_ik[ip, count] = radials_p[ifun_l, l+1] * angular
+                    count += 1
+                end
             end
+            @assert count == n_fun + 1
         end
-        @assert count == count_n_proj(psp) + 1
+        form_factors[ik] = form_factors_ik
     end
+
     form_factors
 end