Skip to content

Make DftFunctionals.jl types GPU compatible #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
+86 −36
Open

Make DftFunctionals.jl types GPU compatible #23

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,10 @@ authors = ["Michael F. Herbst <[email protected]>"]
version = "0.3.1"

[deps]
ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

[compat]
ComponentArrays = "0.15"
DiffResults = "1"
ForwardDiff = "1"
Libxc_jll = "5.1.0"
Expand Down
1 change: 0 additions & 1 deletion src/DftFunctionals.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
module DftFunctionals
using ForwardDiff
using ComponentArrays

include("interface.jl")
include("util.jl")
Expand Down
36 changes: 23 additions & 13 deletions src/functionals/gga_c_pbe.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,26 +1,36 @@
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))
function PbeCorrelation(parameters::NamedTuple, lda=DftFunctional(:lda_c_pw))
PbeCorrelation(parameters, lda, :gga_c_pbe_custom)
end
function PbeCorrelation(parameters::ComponentArray, identifier::Symbol)
function PbeCorrelation(parameters::NamedTuple, identifier::Symbol)
PbeCorrelation(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::NamedTuple;
keep_identifier=false)
if keep_identifier
PbeCorrelation(parameters, pbe.lda, pbe.identifier)
else
PbeCorrelation(parameters, pbe.lda)
end
end
# Change functional parameters based on an array of values. Assumes consistent ordering.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you mean by assumes consistent ordering ? I would not assume that the keys are in the same order ... as this cannot be enforced in practice.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also can we not use a ComponentArray here ?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In any case I would settle on one way of getting and setting the parameters and avoid the duplication of the interface. Otherwise this may derail. DftFunctionals is anyway already too complicated for what it's doing, so if we need to refactor parts of it to make GPU work and reshuffle functionality, that's ok.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also can we not use a ComponentArray here ?

Yes, we can. We could have a single definition of change_parameters that accepts anything that seamlessly converts to a NamedTuple (including a ComponentArray):

function change_parameters(pbe::PbeCorrelation, parameters; keep_identifier=false)                                                    
    if keep_identifier                                                                               
        PbeCorrelation(NamedTuple(parameters), pbe.lda, pbe.identifier)                              
    else                                                                                             
        PbeCorrelation(NamedTuple(parameters), pbe.lda)                                              
    end                                                                                              
end

That would be safer, and make derivatives wrt functional parameters cleaner.

function change_parameters(pbe::PbeCorrelation, parameter_vals::AbstractArray;
keep_identifier=false)
parameter_keys = keys(pbe.parameters)
parameters = NamedTuple{parameter_keys}(parameter_vals)
change_parameters(pbe, parameters; keep_identifier=keep_identifier)
end

function energy(pbe::PbeCorrelation, ρ::T, σ::U) where {T<:Number,U<:Number}
TT = arithmetic_type(pbe, T, U)
Expand Down Expand Up @@ -61,7 +71,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

"""
Expand All @@ -72,7 +82,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

"""
Expand All @@ -82,7 +92,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

"""
Expand All @@ -93,7 +103,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

"""
Expand All @@ -104,7 +114,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

"""
Expand All @@ -114,5 +124,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
35 changes: 23 additions & 12 deletions src/functionals/gga_x_pbe.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,21 +1,32 @@
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)
function PbeExchange(parameters::NamedTuple)
PbeExchange(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::NamedTuple;
keep_identifier=false)
if keep_identifier
PbeExchange(parameters, pbe.identifier)
else
PbeExchange(parameters)
end
end
# Change functional parameters based on an array of values. Assumes consistent ordering.
function change_parameters(pbe::PbeExchange, parameter_vals::AbstractArray;
keep_identifier=false)
parameter_keys = keys(pbe.parameters)
parameters = NamedTuple{parameter_keys}(parameter_vals)
change_parameters(pbe, parameters; keep_identifier=keep_identifier)
end

function energy(pbe::PbeExchange, ρ::T, σ::U) where {T<:Number,U<:Number}
TT = arithmetic_type(pbe, T, U)
Expand Down Expand Up @@ -48,23 +59,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

"""
Expand All @@ -73,7 +84,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

"""
Expand All @@ -82,7 +93,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

"""
Expand All @@ -91,13 +102,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
7 changes: 6 additions & 1 deletion src/interface.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
41 changes: 34 additions & 7 deletions test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
using DftFunctionals
using ComponentArrays
using Test
include("libxc.jl")

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -170,19 +169,19 @@ 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

pbe = DftFunctional(:gga_x_pbe)
Expand All @@ -192,7 +191,35 @@ end
ρ = reshape(ρ, 1, :)
σ = reshape(σ, 1, :)

θ = ComponentArray(; parameters(pbe)...)
# ForwardDiff expects functions that take arrays as arguments
θ = collect(values(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

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
θ = collect(values(parameters(pbe)))
egrad = ForwardDiff.jacobian(θ) do θ
potential_terms(change_parameters(pbe, θ), ρ, σ).e
end
Expand Down