Add PspLinComb

Author: mfherbst

src/DFTK.jl

@@ -52,6 +52,7 @@ export PspUpf
 include("pseudo/NormConservingPsp.jl")
 include("pseudo/PspHgh.jl")
 include("pseudo/PspUpf.jl")
+include("pseudo/PspLinComb.jl")
 
 export ElementPsp
 export ElementCohenBergstresser
@@ -60,6 +61,7 @@ export ElementGaussian
 export charge_nuclear, charge_ionic
 export n_elec_valence, n_elec_core
 export element_symbol, mass, species  # Note: Re-exported from AtomsBase
+export virtual_crystal_element
 include("elements.jl")
 
 export SymOp

src/elements.jl

@@ -52,10 +52,8 @@ function core_charge_density_fourier(::Element, ::T)::T where {T <: Real}
     error("Abstract elements do not necesesarily provide core charge density.")
 end
 
-# Fallback print function:
 Base.show(io::IO, el::Element) = print(io, "$(typeof(el))(:$(species(el)))")
 
-
 #
 # ElementCoulomb
 #
@@ -97,15 +95,19 @@ struct ElementPsp{P} <: Element
     family::Union{PseudoFamily,Nothing}  # PseudoFamily if known, else Nothing
     mass    # Atomic mass
 end
-function Base.show(io::IO, el::ElementPsp)
+function Base.show(io::IO, el::ElementPsp{P}) where {P}
     print(io, "ElementPsp(:$(el.species), ")
     if !isnothing(el.family)
-        print(io, el.family, ")")
+        print(io, el.family)
     elseif isempty(el.psp.identifier)
-        print(io, "custom psp)")
+        print(io, "custom psp")
     else
-        print(io, "\"$(el.psp.identifier)\")")
+        print(io, "\"$(el.psp.identifier)\"")
+    end
+    if P <: PspLinComb
+        print(io, ", ", el.psp.description)
     end
+    print(io, ")")
 end
 pseudofamily(el::ElementPsp) = el.family
 
@@ -275,3 +277,28 @@ function local_potential_fourier(el::ElementGaussian, p::Real)
     -el.α * exp(- (p * el.L)^2 / 2)  # = ∫_ℝ³ V(x) exp(-ix⋅p) dx
 end
 # TODO Strictly speaking needs a eval_psp_energy_correction
+
+#
+# Helper functions
+#
+
+# TODO Better name ?
+# TODO docstring
+function virtual_crystal_element(coefficients::Vector{<:Number}, elements::Vector{<:ElementPsp{<:NormConservingPsp}};
+                                 species=ChemicalSpecies(0))
+    length(coefficients) == length(elements) || throw(
+        ArgumentError("Expect coefficients and elements to have equal length."))
+    sum(coefficients) ≈ one(eltype(coefficients)) || throw(
+        ArgumentError("Expect coefficients to sum to 1"))
+    family = allequal(p -> p.family, elements) ? first(elements).family : nothing
+
+    description = "lin. comb. of " 
+    description *= join([(@sprintf "%.2f*%s" c "$(el.species)") for (c, el) in zip(coefficients, elements)], " ")
+    lincomb_psp  = PspLinComb(coefficients, [el.psp for el in elements]; description)
+    lincomb_mass = sum(c * mass(el) for (c, el) in zip(coefficients, elements))
+    ElementPsp(species, lincomb_psp, family, lincomb_mass)
+end
+function virtual_crystal_element(α::Number, αpsp::ElementPsp{<:NormConservingPsp},
+                                 β::Number, βpsp::ElementPsp{<:NormConservingPsp}; kwargs...)
+    virtual_crystal_element([α, β], [αpsp, βpsp]; kwargs...)
+end

src/pseudo/PspLinComb.jl

@@ -0,0 +1,83 @@
+# Linear combination of two pseudos
+#    Note that this data structure is deliberately not yet exported for now
+#    as we may want to find a different solution for representing virtual crystal elements
+#    in the future; users should use virtual_crystal_element to implicitly construct
+#    these objects.
+#
+struct PspLinComb{T, P <: NormConservingPsp} <: NormConservingPsp
+    lmax::Int
+    h::Vector{Matrix{T}}
+    hidx_to_psp_proj::Vector{Vector{Tuple{Int,Int}}}  # map per angular momentum and projector index of this object to the tuple (psp index, projector index)
+    coefficients::Vector{T}
+    pseudos::Vector{P}
+    identifier::String          # String identifying the PSP
+    description::String         # Descriptive string
+end
+
+function PspLinComb(coefficients::Vector{T}, pseudos::Vector{<:NormConservingPsp};
+                    description="psp linear combination") where {T}
+    @assert length(coefficients) == length(pseudos)
+    @assert sum(coefficients) ≈ one(eltype(coefficients))
+    @assert !any(p -> p isa PspLinComb, pseudos)
+
+    # These assumptions we could lift, but we make tham to make our lifes easier
+    @assert allequal(charge_ionic,        pseudos)
+    @assert allequal(has_valence_density, pseudos)
+    @assert allequal(has_core_density,    pseudos)
+    @assert allequal(p -> p.lmax,         pseudos)
+
+    TT   = promote_type(T, eltype(eltype(first(pseudos).h)))
+    lmax = first(pseudos).lmax
+    h = Matrix{TT}[]
+    hidx_to_psp_proj = Vector{Tuple{Int,Int}}[]  # (ipsp, iproj)
+    for l in 0:lmax
+        proj_l_total = sum(p -> count_n_proj_radial(p, l), pseudos)
+        hl = zeros(T, proj_l_total, proj_l_total)
+        splits = Tuple{Int,Int}[]  # (ipsp, iproj)
+
+        for (ipsp, p) in enumerate(pseudos)
+            iproj_offset = isempty(splits) ? 0 : last(splits)[2]
+            rnge = iproj_offset .+ (1:count_n_proj_radial(p, l))
+            # TODO Is this the correct scaling of the coefficient in h ?
+            hl[rnge, rnge] = p.h[l+1] * coefficients[ipsp]
+            append!(splits, [(ipsp, iproj) for iproj in 1:count_n_proj_radial(p, l)])
+        end
+        push!(h, hl)
+        push!(hidx_to_psp_proj, splits)
+    end
+    @assert length(h) == length(hidx_to_psp_proj) == lmax+1
+
+    identifier = ""
+    PspLinComb(lmax, h, hidx_to_psp_proj, coefficients, pseudos, identifier, description)
+end
+
+charge_ionic(psp::PspLinComb)        = charge_ionic(first(psp.pseudos))
+has_valence_density(psp::PspLinComb) = all(has_valence_density, psp.pseudos)
+has_core_density(psp::PspLinComb)    = any(has_core_density,    psp.pseudos)
+
+function eval_psp_projector_real(psp::PspLinComb, i, l, r::Real)
+    (ipsp, iproj) = psp.hidx_to_psp_proj[l+1][i]
+    eval_psp_projector_real(psp.pseudos[ipsp], iproj, l, r)
+end
+function eval_psp_projector_fourier(psp, i, l, p::Real)
+    (ipsp, iproj) = psp.hidx_to_psp_proj[l+1][i]
+    eval_psp_projector_fourier(psp.pseudos[ipsp], iproj, l, p)
+end
+
+function eval_psp_energy_correction(T, psp::PspLinComb)
+    sum(c * eval_psp_energy_correction(T, p) for (c, p) in zip(psp.coefficients, psp.pseudos))
+end
+
+macro make_psplincomb_call(fn)
+    quote
+        function $fn(psp::PspLinComb, arg::Real)
+            sum(c * $fn(pp, arg) for (c, pp) in zip(psp.coefficients, psp.pseudos))
+        end
+    end
+end
+@make_psplincomb_call DFTK.eval_psp_local_real
+@make_psplincomb_call DFTK.eval_psp_local_fourier
+@make_psplincomb_call DFTK.eval_psp_density_valence_real
+@make_psplincomb_call DFTK.eval_psp_density_valence_fourier
+@make_psplincomb_call DFTK.eval_psp_density_core_real
+@make_psplincomb_call DFTK.eval_psp_density_core_fourier

src/pseudo/pseudopotential_data.jl

@@ -29,7 +29,7 @@ Return the recommended kinetic energy cutoff, supersampling and density cutoff f
 Values may be `missing` if the data cannot be determined.
 """
 function PseudoPotentialData.recommended_cutoff(el::Element)
-    if isnothing(pseudofamily(el))
+    if isnothing(pseudofamily(el)) || :X == element_symbol(el)
         return (; Ecut=missing, supersampling=missing, Ecut_density=missing)
     else
         return recommended_cutoff(pseudofamily(el), element_symbol(el))
@@ -45,7 +45,7 @@ Effectively this returns `pseudometa(pseudofamily(element), element_symbol(eleme
 """
 function PseudoPotentialData.pseudometa(el::Element)
     family = pseudofamily(el)
-    if isnothing(family)
+    if isnothing(family) || :X == element_symbol(el)
         return Dict{String,Any}()
     else
         return pseudometa(family, element_symbol(el))

src/terms/operators.jl

@@ -1,6 +1,8 @@
 ### Linear operators operating on quantities in real or fourier space
 # This is the optimized low-level interface. Use the functions in Hamiltonian for high-level usage
 
+import Base: Matrix, Array
+
 """
 Linear operators that act on tuples (real, fourier)
 The main entry point is `apply!(out, op, in)` which performs the operation `out += op*in`
@@ -41,7 +43,7 @@ struct NoopOperator{T <: Real} <: RealFourierOperator
     kpoint::Kpoint{T}
 end
 apply!(Hψ, op::NoopOperator, ψ) = nothing
-function Base.Matrix(op::NoopOperator)
+function Matrix(op::NoopOperator)
     n_Gk = length(G_vectors(op.basis, op.kpoint))
     zeros_like(G_vectors(op.basis), eltype(op.basis), n_Gk, n_Gk)
 end
@@ -60,7 +62,7 @@ end
 function apply!(Hψ, op::RealSpaceMultiplication, ψ)
     Hψ.real .+= op.potential .* ψ.real
 end
-function Base.Matrix(op::RealSpaceMultiplication)
+function Matrix(op::RealSpaceMultiplication)
     # V(G, G') = <eG|V|eG'> = 1/sqrt(Ω) <e_{G-G'}|V>
     pot_fourier = fft(op.basis, op.potential)
     n_G = length(G_vectors(op.basis, op.kpoint))

src/workarounds/forwarddiff_rules.jl

@@ -202,6 +202,16 @@ function construct_value(psp::PspUpf{T,I}) where {T <: AbstractFloat, I <: Abstr
     #       but does not yet permit response derivatives w.r.t. UPF parameters.
     psp
 end
+construct_value(psp::PspLinComb) = psp
+function construct_value(psp::PspLinComb{<: Dual})
+    PspLinComb(psp.lmax,
+               [ForwardDiff.value.(hl) for hl in psp.h],
+               psp.hidx_to_psp_proj,
+               ForwardDiff.value.(psp.coefficients),
+               construct_value.(psp.pseudos),
+               psp.identifier,
+               psp.description)
+end
 
 function construct_value(basis::PlaneWaveBasis{T}) where {T <: Dual}
     # NOTE: This is a pretty slow function as it *recomputes* basically

test/PspLinComb.jl

@@ -0,0 +1,6 @@
+
+
+# Test whether doing a 0.5 of one element with itself yields the identical answer as not doing the lincomb at all
+
+
+# Test whether the linear combination pseudo does not break any important properities (like normalisation, being norm-conserving etc.)