Make DftFunctionals.jl types GPU compatible

Project.toml

@@ -4,7 +4,6 @@ authors = ["Michael F. Herbst <info@michael-herbst.com>"]
 version = "0.3.1"
 
 [deps]
-ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 
@@ -16,8 +15,9 @@ Libxc_jll = "5.1.0"
 julia = "1.10"
 
 [extras]
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Libxc_jll = "a56a6d9d-ad03-58af-ab61-878bf78270d6"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Libxc_jll"]
+test = ["Test", "Libxc_jll", "ComponentArrays"]

src/DftFunctionals.jl

@@ -1,6 +1,5 @@
 module DftFunctionals
 using ForwardDiff
-using ComponentArrays
 
 include("interface.jl")
 include("util.jl")

src/functionals/gga_c_pbe.jl

@@ -1,24 +1,27 @@
-struct PbeCorrelation{Tlda,CA} <:
-       Functional{:gga,:c} where {Tlda,CA<:ComponentArray{<:Number}}
-    parameters::CA
+struct PbeCorrelation{NT,Tlda,Id} <:
+    Functional{:gga,:c} where {NT<:NamedTuple,Tlda,Id<:Union{Symbol,Val}}
+    parameters::NT
     lda::Tlda
-    identifier::Symbol
+    identifier::Id
 end
-function PbeCorrelation(parameters::ComponentArray, lda=DftFunctional(:lda_c_pw))
-    PbeCorrelation(parameters, lda, :gga_c_pbe_custom)
+function PbeCorrelation(parameters, lda=DftFunctional(:lda_c_pw))
+    PbeCorrelation(NamedTuple(parameters), lda, :gga_c_pbe_custom)
 end
-function PbeCorrelation(parameters::ComponentArray, identifier::Symbol)
-    PbeCorrelation(parameters, DftFunctional(:lda_c_pw), identifier)
+function PbeCorrelation(parameters, identifier::Symbol)
+    PbeCorrelation(NamedTuple(parameters), DftFunctional(:lda_c_pw), identifier)
 end
 
 identifier(pbe::PbeCorrelation) = pbe.identifier
 parameters(pbe::PbeCorrelation) = pbe.parameters
-function change_parameters(pbe::PbeCorrelation, parameters::ComponentArray;
+function to_isbits(pbe::PbeCorrelation)
+    PbeCorrelation(pbe.parameters, to_isbits(pbe.lda), Val{pbe.identifier}())
+end
+function change_parameters(pbe::PbeCorrelation, parameters;
                            keep_identifier=false)
     if keep_identifier
-        PbeCorrelation(parameters, pbe.lda, pbe.identifier)
+        PbeCorrelation(NamedTuple(parameters), pbe.lda, pbe.identifier)
     else
-        PbeCorrelation(parameters, pbe.lda)
+        PbeCorrelation(NamedTuple(parameters), pbe.lda)
     end
 end
 
@@ -61,7 +64,7 @@ Perdew, Burke, Ernzerhof 1996 (DOI: 10.1103/PhysRevLett.77.3865)
 function DftFunctional(::Val{:gga_c_pbe})
     β = 0.06672455060314922
     γ = (1 - log(2)) / π^2
-    PbeCorrelation(ComponentArray(; β, γ), :gga_c_pbe)
+    PbeCorrelation((; β, γ), :gga_c_pbe)
 end
 
 """
@@ -72,7 +75,7 @@ function DftFunctional(::Val{:gga_c_xpbe})
     β = 0.089809  # Fitted constants, Table I
     α = 0.197363  # Fitted constants, Table I
     γ = β^2 / 2α
-    PbeCorrelation(ComponentArray(; β, γ), :gga_c_xpbe)
+    PbeCorrelation((; β, γ), :gga_c_xpbe)
 end
 
 """
@@ -82,7 +85,7 @@ Perdew, Ruzsinszky, Csonka and others 2008 (DOI 10.1103/physrevlett.100.136406)
 function DftFunctional(::Val{:gga_c_pbe_sol})
     β = 0.046  # Page 3, left column below figure 1
     γ = (1 - log(2)) / π^2
-    PbeCorrelation(ComponentArray(; β, γ), :gga_c_pbe_sol)
+    PbeCorrelation((; β, γ), :gga_c_pbe_sol)
 end
 
 """
@@ -93,7 +96,7 @@ function DftFunctional(::Val{:gga_c_apbe})
     μ = 0.260   # p. 1, right column, bottom
     β = 3μ / π^2
     γ = (1 - log(2)) / π^2  # like in PBE
-    PbeCorrelation(ComponentArray(; β, γ), :gga_c_apbe)
+    PbeCorrelation((; β, γ), :gga_c_apbe)
 end
 
 """
@@ -104,7 +107,7 @@ function DftFunctional(::Val{:gga_c_pbe_mol})
     # β made to cancel self-interaction error in hydrogen
     β = 0.08384             # p. 4, right column, first paragraph
     γ = (1 - log(2)) / π^2  # like in PBE
-    PbeCorrelation(ComponentArray(; β, γ), :gga_c_pbe_mol)
+    PbeCorrelation((; β, γ), :gga_c_pbe_mol)
 end
 
 """
@@ -114,5 +117,5 @@ Sarmiento-Perez, Silvana, Marques 2015 (DOI 10.1021/acs.jctc.5b00529)
 function DftFunctional(::Val{:gga_c_pbefe})
     β = 0.043                    # Fitted constants, Table I
     γ = 0.031090690869654895034  # Fitted constants, Table I
-    PbeCorrelation(ComponentArray(; β, γ), :gga_c_pbefe)
+    PbeCorrelation((; β, γ), :gga_c_pbefe)
 end

src/functionals/gga_x_pbe.jl

@@ -1,19 +1,23 @@
-struct PbeExchange{CA} <: Functional{:gga,:x} where {CA<:ComponentArray{<:Number}}
-    parameters::CA
-    identifier::Symbol
+struct PbeExchange{NT, Id} <:
+    Functional{:gga,:x} where {NT<:NamedTuple, Id<:Union{Symbol,Val}}
+    parameters::NT
+    identifier::Id
 end
-function PbeExchange(parameters::ComponentArray)
-    PbeExchange(parameters, :gga_x_pbe_custom)
+function PbeExchange(parameters)
+    PbeExchange(NamedTuple(parameters), :gga_x_pbe_custom)
 end
 
 identifier(pbe::PbeExchange) = pbe.identifier
 parameters(pbe::PbeExchange) = pbe.parameters
-function change_parameters(pbe::PbeExchange, parameters::ComponentArray;
+function to_isbits(pbe::PbeExchange)
+    PbeExchange(pbe.parameters, Val{pbe.identifier}())
+end
+function change_parameters(pbe::PbeExchange, parameters;
                            keep_identifier=false)
     if keep_identifier
-        PbeExchange(parameters, pbe.identifier)
+        PbeExchange(NamedTuple(parameters), pbe.identifier)
     else
-        PbeExchange(parameters)
+        PbeExchange(NamedTuple(parameters))
     end
 end
 
@@ -48,23 +52,23 @@ Standard PBE exchange.
 Perdew, Burke, Ernzerhof 1996 (DOI: 10.1103/PhysRevLett.77.3865)
 """
 function DftFunctional(::Val{:gga_x_pbe})
-    PbeExchange(ComponentArray(κ=0.8040, μ=pbe_μ_from_β(0.06672455060314922)), :gga_x_pbe)
+    PbeExchange((; κ=0.8040, μ=pbe_μ_from_β(0.06672455060314922)), :gga_x_pbe)
 end
 
 """
 Revised PBE exchange.
 Zhang, Yang 1998 (DOI 10.1103/physrevlett.80.890)
 """
 function DftFunctional(::Val{:gga_x_pbe_r})
-    PbeExchange(ComponentArray(κ=1.245, μ=pbe_μ_from_β(0.06672455060314922)), :gga_x_pbe_r)
+    PbeExchange((; κ=1.245, μ=pbe_μ_from_β(0.06672455060314922)), :gga_x_pbe_r)
 end
 
 """
 XPBE exchange.
 Xu, Goddard 2004 (DOI 10.1063/1.1771632)
 """
 function DftFunctional(::Val{:gga_x_xpbe})
-    PbeExchange(ComponentArray(κ=0.91954, μ=0.23214), :gga_x_xpbe)  # Table 1
+    PbeExchange((; κ=0.91954, μ=0.23214), :gga_x_xpbe)  # Table 1
 end
 
 """
@@ -73,7 +77,7 @@ Perdew, Ruzsinszky, Csonka and others 2008 (DOI 10.1103/physrevlett.100.136406)
 """
 function DftFunctional(::Val{:gga_x_pbe_sol})
     # μ given below equation (2)
-    PbeExchange(ComponentArray(κ=0.8040, μ=10 / 81), :gga_x_pbe_sol)
+    PbeExchange((; κ=0.8040, μ=10 / 81), :gga_x_pbe_sol)
 end
 
 """
@@ -82,7 +86,7 @@ Constantin, Fabiano, Laricchia 2011 (DOI 10.1103/physrevlett.106.186406)
 """
 function DftFunctional(::Val{:gga_x_apbe})
     # p. 1, right column, bottom
-    PbeExchange(ComponentArray(κ=0.8040, μ=0.260), :gga_x_apbe)
+    PbeExchange((; κ=0.8040, μ=0.260), :gga_x_apbe)
 end
 
 """
@@ -91,13 +95,13 @@ del Campo, Gazqez, Trickey and others 2012 (DOI 10.1063/1.3691197)
 """
 function DftFunctional(::Val{:gga_x_pbe_mol})
     # p. 4, left column, bottom
-    PbeExchange(ComponentArray(κ=0.8040, μ=0.27583), :gga_x_pbe_mol)
+    PbeExchange((; κ=0.8040, μ=0.27583), :gga_x_pbe_mol)
 end
 
 """
 PBEfe exchange.
 Sarmiento-Perez, Silvana, Marques 2015 (DOI 10.1021/acs.jctc.5b00529)
 """
 function DftFunctional(::Val{:gga_x_pbefe})
-    PbeExchange(ComponentArray(κ=0.437, μ=0.346), :gga_x_pbefe)  # Table 1
+    PbeExchange((; κ=0.437, μ=0.346), :gga_x_pbefe)  # Table 1
 end

src/interface.jl

@@ -44,7 +44,12 @@ has_energy(::Functional) = true
 """
 Return adjustable parameters of the functional and their values.
 """
-parameters(::Functional) = ComponentArray{Bool}()
+parameters(::Functional) = NamedTuple()
+
+"""
+Return a isbits version of the functional (for GPU usage)
+"""
+to_isbits(func::Functional) = func
 
 """
 Return a new version of the passed functional with its parameters adjusted.

test/runtests.jl

@@ -1,5 +1,4 @@
 using DftFunctionals
-using ComponentArrays
 using Test
 include("libxc.jl")
 
@@ -47,8 +46,8 @@ end
     struct NewExchange <: Functional{:lda,:x}
     end
     f = NewExchange()
-    @test parameters(f) == ComponentArray{Bool}()
-    x = ComponentArray{Bool}()
+    @test parameters(f) == NamedTuple()
+    x = NamedTuple()
     @test_throws MethodError change_parameters(f, x)
 end
 
@@ -170,20 +169,21 @@ end
     @test :μ in keys(parameters(pbe))
     @test :κ in keys(parameters(pbe))
 
-    pbemod = change_parameters(pbe, ComponentArray(;μ=12, κ=1.2))
+    pbemod = change_parameters(pbe, (;μ=12, κ=1.2))
     @test parameters(pbemod).μ == 12
     @test parameters(pbemod).κ == 1.2
 
-    pbemod = change_parameters(DftFunctional(:gga_c_pbe), ComponentArray(;β=12, γ=1.2))
+    pbemod = change_parameters(DftFunctional(:gga_c_pbe), (;β=12, γ=1.2))
     @test parameters(pbemod).β == 12
     @test parameters(pbemod).γ == 1.2
 
     μ = rand()
     @test μ ≈ DftFunctionals.pbe_μ_from_β(DftFunctionals.pbe_β_from_μ(μ))
 end
 
-@testset "PBE parameter derivatives" begin
+@testset "PBE exchange parameter derivatives" begin
     using ForwardDiff
+    using ComponentArrays
 
     pbe = DftFunctional(:gga_x_pbe)
 
@@ -192,7 +192,36 @@ end
     ρ = reshape(ρ, 1, :)
     σ = reshape(σ, 1, :)
 
-    θ = ComponentArray(; parameters(pbe)...)
+    # ForwardDiff expects functions that take arrays as arguments
+    θ = ComponentArray(parameters(pbe))
+    egrad = ForwardDiff.jacobian(θ) do θ
+        potential_terms(change_parameters(pbe, θ), ρ, σ).e
+    end
+
+    egrad_fd = let ε=1e-5
+        δ = zero(θ)
+        δ[2] = ε
+
+        (  potential_terms(change_parameters(pbe, θ + δ), ρ, σ).e
+         - potential_terms(change_parameters(pbe, θ - δ), ρ, σ).e) / 2ε
+    end
+
+    @test maximum(abs, egrad[:, 2] - egrad_fd) < 1e-5
+end
+
+@testset "PBE correlation parameter derivatives" begin
+    using ForwardDiff
+    using ComponentArrays
+
+    pbe = DftFunctional(:gga_c_pbe)
+
+    ρ = [0.1, 0.2, 0.3, 0.4, 0.5]
+    σ = [0.2, 0.3, 0.4, 0.5, 0.6]
+    ρ = reshape(ρ, 1, :)
+    σ = reshape(σ, 1, :)
+
+    # ForwardDiff expects functions that take arrays as arguments
+    θ = ComponentArray(parameters(pbe))
     egrad = ForwardDiff.jacobian(θ) do θ
         potential_terms(change_parameters(pbe, θ), ρ, σ).e
     end