Updates for isbits functionals

Project.toml

@@ -15,6 +15,7 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
@@ -94,6 +95,7 @@ FFTW = "1.5"
 FiniteDiff = "2"
 FiniteDifferences = "0.12"
 ForwardDiff = "1"
+Functors = "0.5"
 GPUArraysCore = "0.1, 0.2"
 GenericLinearAlgebra = "0.3"
 GeometryOptimization = "0.1.3"

src/Model.jl

@@ -1,5 +1,12 @@
 # Contains the physical specification of the model
 
+# A builder for a term in the energy functional / Hamiltonian.
+# Term types are objects X that store the term parameters, and produce a
+# XTerm <: Term when instantiated with a `basis`.
+# See also `terms/terms.jl`
+# We have to define this here to avoid circular dependencies.
+abstract type TermType end
+
 # A physical specification of a model.
 # Contains the geometry information, but no discretization parameters.
 # The exact model used is defined by the list of terms.
@@ -53,10 +60,10 @@ struct Model{T <: Real, VT <: Real}
     positions::Vector{Vec3{T}}  # positions[i] is the location of atoms[i] in fract. coords
     atom_groups::Vector{Vector{Int}}  # atoms[i] == atoms[j] for all i, j in atom_group[α]
 
-    # each element t must implement t(basis), which instantiates a
-    # term in a given basis and gives back a term (<: Term)
+    # each element t must be <: TermType and implement t(basis), which
+    # instantiates a term in a given basis and gives back a term (<: Term)
     # see terms.jl for some default terms
-    term_types::Vector
+    term_types::Vector{TermType}
 
     # list of symmetries of the model
     symmetries::Vector{SymOp{VT}}

src/terms/anyonic.jl

@@ -58,7 +58,7 @@ function make_div_free(basis::PlaneWaveBasis{T}, A) where {T}
     [irfft(basis, out[α]) for α = 1:2]
 end
 
-struct Anyonic
+struct Anyonic <: TermType
     hbar
     β
 end

src/terms/entropy.jl

@@ -4,7 +4,7 @@ Turns the energy ``E`` into the free energy ``F=E-TS``.
 This is in particular useful because the free energy,
 not the energy, is minimized at self-consistency.
 """
-struct Entropy end
+struct Entropy <: TermType end
 (::Entropy)(basis) = TermEntropy()
 struct TermEntropy <: Term end
 

src/terms/ewald.jl

@@ -5,7 +5,7 @@ Ewald term: electrostatic energy per unit cell of the array of point
 charges defined by `model.atoms` in a uniform background of
 compensating charge yielding net neutrality.
 """
-Base.@kwdef struct Ewald
+Base.@kwdef struct Ewald <: TermType
     η = nothing  # Parameter used for the splitting 1/r ≡ erf(η·r)/r + erfc(η·r)/r
                  # (or nothing if autoselected)
 end

src/terms/hartree.jl

@@ -11,7 +11,7 @@ For the Coulomb potential with periodic boundary conditions, this is rather
 where G is the Green's function of the periodic Laplacian with zero
 mean (``-Δ G = ∑_R 4π δ_R``, integral of G zero on a unit cell).
 """
-struct Hartree
+struct Hartree <: TermType
     scaling_factor::Real  # to scale by an arbitrary factor (useful for exploration)
 end
 Hartree(; scaling_factor=1) = Hartree(scaling_factor)

src/terms/hubbard.jl

@@ -106,7 +106,7 @@ Hubbard energy, following the Dudarev et al. (1998) rotationally invariant forma
 1/2 Σ_{σI} U * Tr[hubbard_n[σ,I,I] * (1 - hubbard_n[σ,I,I])]
 ```
 """
-struct Hubbard{T}
+struct Hubbard{T} <: TermType
     manifold::OrbitalManifold
     U::T
     function Hubbard(manifold::OrbitalManifold, U)

src/terms/kinetic.jl

@@ -4,7 +4,7 @@ Kinetic energy:
 1/2 ∑_n f_n ∫ |∇ψ_n|^2 * {\rm blowup}(-i∇Ψ_n).
 ```
 """
-Base.@kwdef struct Kinetic{F}
+Base.@kwdef struct Kinetic{F} <: TermType
     scaling_factor::Real = 1
     blowup::F = BlowupIdentity()  # Blow-up to smooth energy bands.
 end

src/terms/local.jl

@@ -30,7 +30,7 @@ end
 """
 External potential given as values.
 """
-struct ExternalFromValues
+struct ExternalFromValues <: TermType
     potential_values::AbstractArray
 end
 function (external::ExternalFromValues)(basis::PlaneWaveBasis{T}) where {T}
@@ -43,7 +43,7 @@ end
 External potential from an analytic function `V` (in Cartesian coordinates).
 No low-pass filtering is performed.
 """
-struct ExternalFromReal
+struct ExternalFromReal <: TermType
     potential::Function
 end
 
@@ -56,7 +56,7 @@ end
 External potential from the (unnormalized) Fourier coefficients `V(G)`
 G is passed in Cartesian coordinates
 """
-struct ExternalFromFourier
+struct ExternalFromFourier <: TermType
     potential::Function
 end
 function (external::ExternalFromFourier)(basis::PlaneWaveBasis{T}) where {T}
@@ -103,7 +103,7 @@ end
 """
 Atomic local potential defined by `model.atoms`.
 """
-struct AtomicLocal end
+struct AtomicLocal <: TermType end
 function compute_local_potential(basis::PlaneWaveBasis{T}; positions=basis.model.positions,
                                  q=zero(Vec3{T})) where {T}
     # pot_fourier is <e_G|V|e_G'> expanded in a basis of e_{G-G'}

src/terms/local_nonlinearity.jl

@@ -1,7 +1,7 @@
 """
 Local nonlinearity, with energy ∫f(ρ) where ρ is the density
 """
-struct LocalNonlinearity
+struct LocalNonlinearity <: TermType
     f
 end
 struct TermLocalNonlinearity{TF} <: TermNonlinear

src/terms/magnetic.jl

@@ -4,13 +4,13 @@
 @doc raw"""
 Magnetic term ``A⋅(-i∇)``. It is assumed (but not checked) that ``∇⋅A = 0``.
 """
-struct Magnetic
+struct Magnetic <: TermType
     Afunction::Function  # A(x,y,z) returns [Ax,Ay,Az]
                          # both [x,y,z] and [Ax,Ay,Az] are in *Cartesian* coordinates
 end
 (M::Magnetic)(basis) = TermMagnetic(basis, M.Afunction)
 
-struct MagneticFromValues
+struct MagneticFromValues <: TermType
     # Apotential[α] is an array of size fft_size for α=1:3
     Apotential::Vector{<:AbstractArray}
 end

src/terms/nonlocal.jl

@@ -5,7 +5,7 @@ Nonlocal term coming from norm-conserving pseudopotentials in Kleinmann-Bylander
     ∑_a ∑_{ij} ∑_{n} f_n \braket{ψ_n}{{\rm proj}_{ai}} D_{ij} \braket{{\rm proj}_{aj}}{ψ_n}.
 ```
 """
-struct AtomicNonlocal end
+struct AtomicNonlocal <: TermType end
 function (::AtomicNonlocal)(basis::PlaneWaveBasis{T}) where {T}
     model = basis.model
 

src/terms/pairwise.jl

@@ -1,4 +1,4 @@
-struct PairwisePotential
+struct PairwisePotential <: TermType
     V
     params
     max_radius

src/terms/psp_correction.jl

@@ -1,7 +1,7 @@
 """
 Pseudopotential correction energy. TODO discuss the need for this.
 """
-struct PspCorrection end
+struct PspCorrection <: TermType end
 (::PspCorrection)(basis) = TermPspCorrection(basis)
 
 struct TermPspCorrection{T} <: TermLinear

src/terms/terms.jl

@@ -25,9 +25,6 @@ apply_kernel(term::TermLinear, basis::AbstractBasis, δρ; kwargs...) = nothing
 # contribution that is density-dependent or orbital-dependent as well)
 abstract type TermNonlinear <: Term end
 
-### Builders are objects X that store the term parameters, and produce a
-# XTerm <: Term when instantiated with a `basis`
-
 # TODO If needed improve this further by specialising energy() for certain terms
 function energy(term::Term, basis::AbstractBasis, ψ, occupation; kwargs...)
     ene_ops(term, basis, ψ, occupation; kwargs...).E
@@ -45,10 +42,10 @@ end
 
 include("Hamiltonian.jl")
 
-# breaks_symmetries on a term builder answers true if this term breaks
+# breaks_symmetries on a term type answers true if this term breaks
 # the symmetries of the lattice/atoms (in which case k-point reduction
 # is invalid)
-breaks_symmetries(::Any) = false
+breaks_symmetries(::TermType) = false
 
 include("kinetic.jl")
 

src/terms/xc.jl

@@ -1,7 +1,7 @@
 """
 Exchange-correlation term, defined by a list of functionals and usually evaluated through libxc.
 """
-struct Xc
+struct Xc <: TermType
     functionals::Vector{Functional}
     scaling_factor::Real  # Scales by an arbitrary factor (useful for exploration)
 
@@ -38,17 +38,8 @@ function (xc::Xc)(basis::PlaneWaveBasis{T}) where {T}
         ρcore = ρ_from_total(basis, atomic_total_density(basis, CoreDensity()))
         minimum(ρcore) < -sqrt(eps(T)) && @warn("Negative ρcore detected: $(minimum(ρcore))")
     end
-    functionals = map(xc.functionals) do fun
-        # Strip duals from functional parameters if needed
-        params = parameters(fun)
-        if !isempty(params)
-            newparams = convert_dual.(T, params)
-            fun = change_parameters(fun, newparams; keep_identifier=true)
-        end
-        fun
-    end
-    TermXc(convert(Vector{Functional}, functionals),
-           convert_dual(T, xc.scaling_factor),
+    TermXc(xc.functionals,
+           T(xc.scaling_factor),
            T(xc.potential_threshold), ρcore)
 end
 

src/workarounds/forwarddiff_rules.jl

@@ -1,5 +1,6 @@
 import ForwardDiff
 import ForwardDiff: Dual
+import Functors: fmap
 import AbstractFFTs
 
 # original PR by mcabbott: https://github.com/JuliaDiff/ForwardDiff.jl/pull/495
@@ -172,7 +173,7 @@ function construct_value(model::Model{T}) where {T <: Dual}
           model.model_name,
           model.n_electrons,
           magnetic_moments=value_type(T)[],  # Symmetries given explicitly
-          terms=model.term_types,
+          terms=map(construct_value, model.term_types),
           temperature=ForwardDiff.value(model.temperature),
           model.smearing,
           εF=ForwardDiff.value(model.εF),
@@ -203,6 +204,16 @@ function construct_value(psp::PspUpf{T,I}) where {T <: AbstractFloat, I <: Abstr
     psp
 end
 
+# by default, don't convert term types
+construct_value(x::TermType) = x
+function construct_value(x::Xc)
+    Xc(map(construct_value, x.functionals), ForwardDiff.value(x.scaling_factor),
+       x.potential_threshold, x.use_nlcc)
+end
+function construct_value(functional::DftFunctionals.Functional)
+    fmap(ForwardDiff.value, functional)
+end
+
 function construct_value(basis::PlaneWaveBasis{T}) where {T <: Dual}
     # NOTE: This is a pretty slow function as it *recomputes* basically
     #       everything instead of just extracting the primal data contained

test/forwarddiff.jl

@@ -214,7 +214,7 @@ end
         pos = [[1.01, 1.02, 1.03] / 8, -ones(3) / 8]
         pbec = DftFunctional(:gga_c_pbe)
         pbex = DftFunctional(:gga_x_pbe)
-        pbex = change_parameters(pbex, parameters(pbex) + ComponentArray(κ=0, μ=ε1))
+        pbex = DftFunctionals.PbeExchange(; κ=pbex.κ, μ=pbex.μ+ε1)
 
         model = model_DFT(Matrix{T}(silicon.lattice), silicon.atoms, pos;
                           functionals=[pbex, pbec])